Add prestage support to the client

client/acquire_token.go

@@ -123,6 +123,30 @@
 	tg.Token.Store(&info)
 }
 
+// Copy the contents
+func (tg *tokenGenerator) Copy() *tokenGenerator {
+	return &tokenGenerator{
+		DirResp:       tg.DirResp,
+		Destination:   tg.Destination,
+		IsWrite:       tg.IsWrite,
+		EnableAcquire: tg.EnableAcquire,
+		Sync:          new(singleflight.Group),
+	}
+}
+
+// Determine the token name if it is embedded in the scheme, Condor-style
+func getTokenName(destination *url.URL) (scheme, tokenName string) {
+	schemePieces := strings.Split(destination.Scheme, "+")
+	tokenName = ""
+	// Scheme is always the last piece
+	scheme = schemePieces[len(schemePieces)-1]
+	// If there are 2 or more pieces, token name is everything but the last item, joined with a +
+	if len(schemePieces) > 1 {
+		tokenName = strings.Join(schemePieces[:len(schemePieces)-1], "+")
+	}
+	return
+}
+
 // Read a token from a file; ensure
 func getTokenFromFile(tokenLocation string) (string, error) {
 	//Read in the JSON

client/fed_test.go

@@ -940,6 +940,84 @@ func TestClientUnpack(t *testing.T) {
 	assert.Equal(t, int64(11), fi.Size())
 }
 
+// A test that spins up a federation, and tests object get and put
+func TestPrestage(t *testing.T) {
+	viper.Reset()
+	server_utils.ResetOriginExports()
+	fed := fed_test_utils.NewFedTest(t, bothAuthOriginCfg)
+
+	te, err := client.NewTransferEngine(fed.Ctx)
+	require.NoError(t, err)
+
+	// Other set-up items:
+	// The cache will open the file to stat it, downloading the first block.
+	// Make sure we are greater than 64kb in size.
+	testFileContent := strings.Repeat("test file content", 10000)
+	// Create the temporary file to upload
+	tempFile, err := os.CreateTemp(t.TempDir(), "test")
+	assert.NoError(t, err, "Error creating temp file")
+	defer os.Remove(tempFile.Name())
+	_, err = tempFile.WriteString(testFileContent)
+	assert.NoError(t, err, "Error writing to temp file")
+	tempFile.Close()
+
+	tempToken, _ := getTempToken(t)
+	defer tempToken.Close()
+	defer os.Remove(tempToken.Name())
+	// Disable progress bars to not reuse the same mpb instance
+	viper.Set("Logging.DisableProgressBars", true)
+
+	oldPref, err := config.SetPreferredPrefix(config.PelicanPrefix)
+	assert.NoError(t, err)
+	defer func() {
+		_, err := config.SetPreferredPrefix(oldPref)
+		require.NoError(t, err)
+	}()
+
+	// Set path for object to upload/download
+	for _, export := range fed.Exports {
+		tempPath := tempFile.Name()
+		fileName := filepath.Base(tempPath)
+		uploadURL := fmt.Sprintf("pelican://%s:%s%s/prestage/%s", param.Server_Hostname.GetString(), strconv.Itoa(param.Server_WebPort.GetInt()),
+			export.FederationPrefix, fileName)
+
+		// Upload the file with COPY
+		transferResultsUpload, err := client.DoCopy(fed.Ctx, tempFile.Name(), uploadURL, false, client.WithTokenLocation(tempToken.Name()))
+		assert.NoError(t, err)
+		assert.Equal(t, int64(len(testFileContent)), transferResultsUpload[0].TransferredBytes)
+
+		// Check the cache info twice, make sure it's not cached.
+		tc, err := te.NewClient(client.WithTokenLocation(tempToken.Name()))
+		require.NoError(t, err)
+		innerFileUrl, err := url.Parse(uploadURL)
+		require.NoError(t, err)
+		age, size, err := tc.CacheInfo(fed.Ctx, innerFileUrl)
+		require.NoError(t, err)
+		assert.Equal(t, int64(len(testFileContent)), size)
+		assert.Equal(t, -1, age)
+
+		age, size, err = tc.CacheInfo(fed.Ctx, innerFileUrl)
+		require.NoError(t, err)
+		assert.Equal(t, int64(len(testFileContent)), size)
+		assert.Equal(t, -1, age)
+
+		// Prestage the object
+		tj, err := tc.NewPrestageJob(fed.Ctx, innerFileUrl)
+		require.NoError(t, err)
+		err = tc.Submit(tj)
+		require.NoError(t, err)
+		results, err := tc.Shutdown()
+		require.NoError(t, err)
+		assert.Equal(t, 1, len(results))
+
+		// Check if object is cached.
+		age, size, err = tc.CacheInfo(fed.Ctx, innerFileUrl)
+		require.NoError(t, err)
+		assert.Equal(t, int64(len(testFileContent)), size)
+		require.NotEqual(t, -1, age)
+	}
+}
+
 // A test that generates a token locally from the private key
 func TestTokenGenerate(t *testing.T) {
 	viper.Reset()