Add temp_dir_path flag for bput to specify temp dir path for creating…

… bundle files
cyverse · Oct 11, 2022 · bfbb583 · bfbb583
1 parent b5cc5f2
commit bfbb583
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 7 deletions.
diff --git a/cmd/subcmd/gobput.go b/cmd/subcmd/gobput.go
@@ -26,6 +26,7 @@ func AddBputCommand(rootCmd *cobra.Command) {
 	bputCmd.Flags().IntP("max_file_num", "", commons.MaxBundleFileNum, "Specify max file number in a bundle file")
 	bputCmd.Flags().Int64P("max_file_size", "", commons.MaxBundleFileSize, "Specify max file size of a bundle file")
 	bputCmd.Flags().BoolP("progress", "", false, "Display progress bar")
+	bputCmd.Flags().StringP("temp_dir_path", "", os.TempDir(), "Specify a local temp directory path to create bundle files")
 
 	rootCmd.AddCommand(bputCmd)
 }
@@ -90,6 +91,12 @@ func processBputCommand(command *cobra.Command, args []string) error {
 		}
 	}
 
+	tempDirPath := os.TempDir()
+	tempDirPathFlag := command.Flags().Lookup("temp_dir_path")
+	if tempDirPathFlag != nil {
+		tempDirPath = tempDirPathFlag.Value.String()
+	}
+
 	// Create a file system
 	account := commons.GetAccount()
 	filesystem, err := commons.GetIRODSFSClient(account)
@@ -101,7 +108,7 @@ func processBputCommand(command *cobra.Command, args []string) error {
 
 	defer filesystem.Release()
 
-	bundleTransferManager := commons.NewBundleTransferManager(maxFileNum, maxFileSize)
+	bundleTransferManager := commons.NewBundleTransferManager(maxFileNum, maxFileSize, tempDirPath)
 
 	targetPath := ""
 

diff --git a/commons/bundle.go b/commons/bundle.go
@@ -22,15 +22,17 @@ const (
 )
 
 type Bundle struct {
+	tempDirPath     string
 	name            string
 	files           []string
 	size            int64
 	localBundlePath string
 	irodsBundlePath string
 }
 
-func NewBundle() *Bundle {
+func NewBundle(tempDirPath string) *Bundle {
 	return &Bundle{
+		tempDirPath:     tempDirPath,
 		files:           []string{},
 		size:            0,
 		localBundlePath: "",
@@ -58,26 +60,27 @@ func (bundle *Bundle) Seal(name string) error {
 	bundle.name = name
 
 	// set local bundle path
-	tempDir := os.TempDir()
-	bundle.localBundlePath = filepath.Join(tempDir, fmt.Sprintf("%s_%s.tar", bundle.name, bundleID))
+	bundle.localBundlePath = filepath.Join(bundle.tempDirPath, fmt.Sprintf("%s_%s.tar", bundle.name, bundleID))
 	return nil
 }
 
 type BundleTransferManager struct {
 	pendingBundles          *list.List // *Bundle
 	maxBundleFileNum        int
 	maxBundleFileSize       int64
+	tempDirPath             string
 	errors                  *list.List // error
 	progressTrackerCallback ProgressTrackerCallback
 	mutex                   sync.Mutex
 }
 
 // NewBundleTransferManager creates a new BundleTransferManager
-func NewBundleTransferManager(maxBundleFileNum int, maxBundleFileSize int64) *BundleTransferManager {
+func NewBundleTransferManager(maxBundleFileNum int, maxBundleFileSize int64, tempDirPath string) *BundleTransferManager {
 	manager := &BundleTransferManager{
 		pendingBundles:          list.New(),
 		maxBundleFileNum:        maxBundleFileNum,
 		maxBundleFileSize:       maxBundleFileSize,
+		tempDirPath:             tempDirPath,
 		errors:                  list.New(),
 		progressTrackerCallback: nil,
 	}
@@ -109,7 +112,7 @@ func (manager *BundleTransferManager) ScheduleBundleUpload(source string, size i
 			if lastBundle.size >= manager.maxBundleFileSize || len(lastBundle.files) >= manager.maxBundleFileNum {
 				// exceed bundle size or file num
 				// create a new
-				currentBundle = NewBundle()
+				currentBundle = NewBundle(manager.tempDirPath)
 
 				logger.Debugf("assigning a new bundle %d", manager.pendingBundles.Len())
 				manager.pendingBundles.PushBack(currentBundle)
@@ -122,7 +125,7 @@ func (manager *BundleTransferManager) ScheduleBundleUpload(source string, size i
 		}
 	} else {
 		// add new
-		currentBundle = NewBundle()
+		currentBundle = NewBundle(manager.tempDirPath)
 
 		logger.Debugf("assigning a new bundle %d", manager.pendingBundles.Len())
 		manager.pendingBundles.PushBack(currentBundle)