Add support for prestage async command

Author: bbockelm

client/handle_http.go

@@ -2519,7 +2519,7 @@ func fetchChecksum(ctx context.Context, types []ChecksumType, url *url.URL, toke
 // Verify that a file on disk matches the expected size. We ignore directories
 // and generic stat failures unless the file doesn't exist.
 func verifyFileSize(dest string, expectedSize int64, fields log.Fields) error {
-	if dest == os.DevNull {
+	if dest == os.DevNull || dest == "" {
 		log.WithFields(fields).Debugf("Skipping size check because destination is (%s)", os.DevNull)
 		return nil
 	}
@@ -3036,7 +3036,7 @@ Loop:
 
 		// Second sanity check to verify the file size as it appears on disk.
 		if err = verifyFileSize(dest, totalSize, fields); err != nil {
-			err = errors.Wrap(err, "failed to verify size of downloaded file on disk")
+			err = errors.Wrapf(err, "failed to verify size of downloaded file (%s) on disk", dest)
 			return
 		}
 	}

client_agent/cli_test.go

@@ -346,6 +346,163 @@ func TestCLIAsyncPut(t *testing.T) {
 	t.Logf("TestCLIAsyncPut complete, total time: %s", time.Since(startTime))
 }
 
+// TestCLIAsyncPrestage tests the pelican object prestage --async command
+func TestCLIAsyncPrestage(t *testing.T) {
+	// Reset test state
+	server_utils.ResetTestState()
+
+	// Create test federation
+	fed := fed_test_utils.NewFedTest(t, testOriginConfig)
+
+	// Create temporary directory
+	tempDir := t.TempDir()
+
+	// Create token
+	err := param.Set(param.IssuerKeysDirectory.GetName(), t.TempDir())
+	require.NoError(t, err)
+	issuer, err := config.GetServerIssuerURL()
+	require.NoError(t, err)
+
+	tokenConfig := token.NewWLCGToken()
+	tokenConfig.Lifetime = time.Minute * 5
+	tokenConfig.Issuer = issuer
+	tokenConfig.Subject = "test-cli-async-prestage"
+	tokenConfig.AddAudienceAny()
+
+	scopes := []token_scopes.TokenScope{}
+	readScope, err := token_scopes.Wlcg_Storage_Read.Path("/")
+	require.NoError(t, err)
+	scopes = append(scopes, readScope)
+	modScope, err := token_scopes.Wlcg_Storage_Modify.Path("/")
+	require.NoError(t, err)
+	scopes = append(scopes, modScope)
+	tokenConfig.AddScopes(scopes...)
+
+	tkn, err := tokenConfig.CreateToken()
+	require.NoError(t, err)
+
+	tokenFile := filepath.Join(tempDir, "token")
+	err = os.WriteFile(tokenFile, []byte(tkn), 0644)
+	require.NoError(t, err)
+
+	// Set up client API server
+	serverConfig, _ := client_agent.CreateTestServerConfig(t)
+
+	egrp, egrpCtx := errgroup.WithContext(context.Background())
+	ctx := context.WithValue(egrpCtx, config.EgrpKey, egrp)
+
+	server, err := client_agent.NewServer(ctx, serverConfig)
+	require.NoError(t, err)
+
+	err = server.Start()
+	require.NoError(t, err)
+
+	t.Cleanup(func() {
+		err := server.Shutdown()
+		assert.NoError(t, err)
+	})
+
+	// Build pelican binary
+	pelicanBin := buildPelicanBinary(t)
+
+	// Create test file and upload it first
+	testContent := []byte("Test file for async prestage\n")
+	uploadFile := filepath.Join(tempDir, "prestage-upload.txt")
+	err = os.WriteFile(uploadFile, testContent, 0644)
+	require.NoError(t, err)
+
+	federationPrefix := fed.Exports[0].FederationPrefix
+	discoveryUrl, err := url.Parse(param.Federation_DiscoveryUrl.GetString())
+	require.NoError(t, err)
+	uploadURL := fmt.Sprintf("pelican://%s%s/prestage-test.txt", discoveryUrl.Host, federationPrefix)
+
+	// Upload file first using async + wait
+	uploadCmd := exec.Command(pelicanBin, "object", "put", "--async", "--wait", uploadFile, uploadURL, "--token", tokenFile)
+	uploadCmd.Env = append(os.Environ(), fmt.Sprintf("PELICAN_CLIENTAGENT_SOCKET=%s", serverConfig.SocketPath))
+	output, err := uploadCmd.CombinedOutput()
+	require.NoError(t, err, "Failed to upload file: %s", output)
+
+	// Test async prestage without --wait
+	t.Run("AsyncPrestageWithoutWait", func(t *testing.T) {
+		cmd := exec.Command(pelicanBin, "object", "prestage", "--async", uploadURL, "--token", tokenFile)
+		cmd.Env = append(os.Environ(), fmt.Sprintf("PELICAN_CLIENTAGENT_SOCKET=%s", serverConfig.SocketPath))
+
+		output, err := cmd.CombinedOutput()
+		require.NoError(t, err, "Failed to run async prestage: %s", output)
+
+		outputStr := string(output)
+		t.Logf("Command output: %s", outputStr)
+
+		// Should contain job ID
+		assert.Contains(t, outputStr, "Job created:")
+		assert.Contains(t, outputStr, "Check status with: pelican job status")
+
+		// Extract job ID from output
+		re := regexp.MustCompile(`Job created: ([a-f0-9-]+)`)
+		matches := re.FindStringSubmatch(outputStr)
+		require.Len(t, matches, 2, "Could not extract job ID from output")
+		jobID := matches[1]
+		t.Logf("Created job ID: %s", jobID)
+	})
+
+	// Test async prestage with --wait
+	t.Run("AsyncPrestageWithWait", func(t *testing.T) {
+		cmd := exec.Command(pelicanBin, "object", "prestage", "--async", "--wait", uploadURL, "--token", tokenFile)
+		cmd.Env = append(os.Environ(), fmt.Sprintf("PELICAN_CLIENTAGENT_SOCKET=%s", serverConfig.SocketPath))
+
+		output, err := cmd.CombinedOutput()
+		outputStr := string(output)
+		t.Logf("Command output: %s", outputStr)
+
+		// Should contain job creation message
+		assert.Contains(t, outputStr, "Job created:")
+		assert.Contains(t, outputStr, "Waiting for job to complete")
+
+		// If prestage failed, log the client-agent log for debugging
+		if err != nil {
+			t.Logf("Prestage failed with error: %v", err)
+
+			// Extract job ID to query its status
+			re := regexp.MustCompile(`Job created: ([a-f0-9-]+)`)
+			matches := re.FindStringSubmatch(outputStr)
+			if len(matches) >= 2 {
+				jobID := matches[1]
+				t.Logf("Querying failed job ID: %s", jobID)
+
+				// Query the job status via API to get detailed error info
+				apiClient, apiErr := apiclient.NewAPIClient(serverConfig.SocketPath)
+				if apiErr != nil {
+					t.Logf("Failed to create API client: %v", apiErr)
+				} else {
+					ctx := context.Background()
+					if jobStatus, statusErr := apiClient.GetJobStatus(ctx, jobID); statusErr == nil {
+						t.Logf("Job Status: %s", jobStatus.Status)
+						if jobStatus.Error != "" {
+							t.Logf("Job Error: %s", jobStatus.Error)
+						}
+						if len(jobStatus.Transfers) > 0 {
+							for i, transfer := range jobStatus.Transfers {
+								t.Logf("Transfer %d: Status=%s, Operation=%s, Source=%s",
+									i, transfer.Status, transfer.Operation, transfer.Source)
+								if transfer.Error != "" {
+									t.Logf("  Transfer %d Error: %s", i, transfer.Error)
+								}
+							}
+						}
+					} else {
+						t.Logf("Failed to get job status: %v", statusErr)
+					}
+				}
+			}
+
+			require.NoError(t, err, "Prestage operation should not fail")
+		}
+
+		// Should contain completion message
+		assert.Contains(t, outputStr, "Job completed successfully")
+	})
+}
+
 // TestCLIJobCommands tests the pelican job subcommands
 func TestCLIJobCommands(t *testing.T) {
 	// Reset test state

client_agent/handlers.go

@@ -53,6 +53,17 @@ func (s *Server) CreateJobHandler(c *gin.Context) {
 		return
 	}
 
+	// Additional validation for specific operations
+	for _, transfer := range req.Transfers {
+		if transfer.Operation != "prestage" && transfer.Destination == "" {
+			c.JSON(http.StatusBadRequest, ErrorResponse{
+				Code:  ErrCodeInvalidRequest,
+				Error: "Destination is required for " + transfer.Operation + " operations",
+			})
+			return
+		}
+	}
+
 	// Build transfer options
 	options := buildTransferOptions(req.Options)
 

client_agent/models.go

@@ -24,9 +24,9 @@ import (
 
 // TransferRequest represents a single transfer operation within a job
 type TransferRequest struct {
-	Operation   string `json:"operation" binding:"required,oneof=get put copy"`
+	Operation   string `json:"operation" binding:"required,oneof=get put copy prestage"`
 	Source      string `json:"source" binding:"required"`
-	Destination string `json:"destination" binding:"required"`
+	Destination string `json:"destination"`
 	Recursive   bool   `json:"recursive"`
 }
 

client_agent/transfer_manager.go

@@ -475,6 +475,8 @@ func (tm *TransferManager) executeTransfer(transfer *Transfer, options []client.
 		results, err = client.DoPut(transfer.ctx, transfer.Source, transfer.Destination, transfer.Recursive, options...)
 	case "copy":
 		results, err = client.DoCopy(transfer.ctx, transfer.Source, transfer.Destination, transfer.Recursive, options...)
+	case "prestage":
+		results, err = client.DoPrestage(transfer.ctx, transfer.Source, options...)
 	default:
 		err = errors.Errorf("unknown operation: %s", transfer.Operation)
 	}

client_agent/types/types.go

@@ -38,7 +38,7 @@ type StoredJob struct {
 type StoredTransfer struct {
 	ID               string
 	JobID            string
-	Operation        string // "get", "put", "copy"
+	Operation        string // "get", "put", "copy", "prestage"
 	Source           string
 	Destination      string
 	Recursive        bool

cmd/object_prestage.go

@@ -19,13 +19,17 @@
 package main
 
 import (
+	"encoding/json"
+	"fmt"
 	"os"
+	"time"
 
 	"github.com/pkg/errors"
 	log "github.com/sirupsen/logrus"
 	"github.com/spf13/cobra"
 
 	"github.com/pelicanplatform/pelican/client"
+	"github.com/pelicanplatform/pelican/client_agent"
 	"github.com/pelicanplatform/pelican/config"
 	"github.com/pelicanplatform/pelican/error_codes"
 	"github.com/pelicanplatform/pelican/param"
@@ -48,6 +52,8 @@ func init() {
 	flagSet.StringP("cache", "c", "", `A comma-separated list of preferred caches to try for the transfer, where a "+" in the list indicates
 the client should fallback to discovered caches if all preferred caches fail.`)
 	flagSet.StringP("token", "t", "", "Token file to use for transfer")
+	flagSet.Bool("async", false, "Run the prestage asynchronously through the client API server and return a job ID")
+	flagSet.Bool("wait", false, "When used with --async, wait for the job to complete before returning")
 	objectCmd.AddCommand(prestageCmd)
 }
 
@@ -66,6 +72,131 @@ func prestageMain(cmd *cobra.Command, args []string) {
 		}
 	}
 
+	// Check for async mode
+	isAsync, _ := cmd.Flags().GetBool("async")
+	if isAsync {
+		// Validate arguments
+		if len(args) < 1 {
+			log.Errorln("Prefix(es) to prestage must be specified")
+			err = cmd.Help()
+			if err != nil {
+				log.Errorln("Failed to print out help:", err)
+			}
+			os.Exit(1)
+		}
+
+		// Ensure server is running, starting it if necessary
+		apiClient, err := ensureClientAgentRunning(cmd.Context(), 5)
+		if err != nil {
+			log.Errorln("Failed to ensure API server is running:", err)
+			log.Errorln("You can manually start it with 'pelican client-api serve --daemonize'")
+			os.Exit(1)
+		}
+
+		// Get flags for transfer options
+		tokenLocation, _ := cmd.Flags().GetString("token")
+
+		// Get preferred caches
+		caches, err := getPreferredCaches()
+		if err != nil {
+			log.Errorln("Failed to get preferred caches:", err)
+			os.Exit(1)
+		}
+
+		// Convert caches to strings
+		cacheStrings := make([]string, len(caches))
+		for i, cache := range caches {
+			cacheStrings[i] = cache.String()
+		}
+
+		// Build transfer options
+		options := client_agent.TransferOptions{
+			Token:  tokenLocation,
+			Caches: cacheStrings,
+		}
+
+		// Create transfers for each source prefix
+		transfers := make([]client_agent.TransferRequest, len(args))
+		for i, src := range args {
+			if !isPelicanUrl(src) {
+				log.Errorln("Provided URL is not a valid Pelican URL:", src)
+				os.Exit(1)
+			}
+			transfers[i] = client_agent.TransferRequest{
+				Operation:   "prestage",
+				Source:      src,
+				Destination: "", // Prestage doesn't have a destination
+				Recursive:   false,
+			}
+		}
+
+		// Create job
+		jobID, err := apiClient.CreateJob(ctx, transfers, options)
+		if err != nil {
+			log.Errorln("Failed to create job:", err)
+			os.Exit(1)
+		}
+
+		if outputJSON {
+			result := map[string]interface{}{
+				"job_id": jobID,
+				"status": "created",
+			}
+			jsonBytes, err := json.MarshalIndent(result, "", "  ")
+			if err != nil {
+				log.Errorln("Failed to marshal JSON:", err)
+				os.Exit(1)
+			}
+			fmt.Println(string(jsonBytes))
+		} else {
+			fmt.Printf("Job created: %s\n", jobID)
+		}
+
+		// Check if we should wait for completion
+		shouldWait, _ := cmd.Flags().GetBool("wait")
+		if shouldWait {
+			if !outputJSON {
+				fmt.Println("Waiting for job to complete...")
+			}
+
+			// Wait with a reasonable timeout (e.g., 1 hour)
+			err := apiClient.WaitForJob(ctx, jobID, 1*time.Hour)
+			if err != nil {
+				log.Errorln("Error waiting for job:", err)
+				os.Exit(1)
+			}
+
+			// Get final job status
+			status, err := apiClient.GetJobStatus(ctx, jobID)
+			if err != nil {
+				log.Errorln("Failed to get job status:", err)
+				os.Exit(1)
+			}
+
+			if outputJSON {
+				jsonBytes, err := json.MarshalIndent(status, "", "  ")
+				if err != nil {
+					log.Errorln("Failed to marshal JSON:", err)
+					os.Exit(1)
+				}
+				fmt.Println(string(jsonBytes))
+			} else {
+				fmt.Printf("Job completed successfully\n")
+			}
+
+			if status.Status != "completed" {
+				os.Exit(1)
+			}
+		} else {
+			if !outputJSON {
+				fmt.Printf("Check status with: pelican job status %s\n", jobID)
+			}
+		}
+
+		return
+	}
+
+	// Original synchronous behavior
 	tokenLocation, _ := cmd.Flags().GetString("token")
 
 	pb := newProgressBar()