feat(fetch): add flag '-t/--tmp-dir' (#19)

Fixes #19.
windvalley · Jan 11, 2022 · 8d3d6f6 · 8d3d6f6
1 parent a6a21c3
commit 8d3d6f6
Show file tree

Hide file tree

Showing 5 changed files with 33 additions and 10 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -7,6 +7,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+## [1.4.1]
+
+### Added
+
+- Add flag `-t/--tmp-dir` for subcommand `fetch`. For details at [#19](https://github.com/windvalley/gossh/issues/19).
+
 ## [1.4.0]
 
 ### Added

diff --git a/internal/cmd/fetch.go b/internal/cmd/fetch.go
@@ -32,6 +32,7 @@ import (
 var (
 	srcFiles    []string
 	localDstDir string
+	tmpDir      string
 )
 
 // fetchCmd represents the fetch command
@@ -53,7 +54,11 @@ Copy files/dirs from target hosts to local.`,
   #
   $ gossh fetch host1 host2 -f /path1/foo.txt -f /path2/bar/ -d /tmp/backup
     or
-  $ gossh fetch host1 host2 -f /path1/foo.txt,/path2/bar/ -d /tmp/backup`,
+  $ gossh fetch host1 host2 -f /path1/foo.txt,/path2/bar/ -d /tmp/backup
+
+  # Specify tmp dir on target host instead of default /tmp by flag '-t'.
+  # NOTE: If the tmp dir not exist, it will auto create it.
+  $ gossh fetch host1 -f /path/foo.txt -d ./backup/ -t /home/user/tmp/`,
 	PreRun: func(cmd *cobra.Command, args []string) {
 		if errs := config.Validate(); len(errs) != 0 {
 			util.CheckErr(errs)
@@ -64,7 +69,7 @@ Copy files/dirs from target hosts to local.`,
 
 		task.SetHosts(args)
 		task.SetFetchFiles(srcFiles)
-		task.SetFileOptions(localDstDir, true)
+		task.SetFetchOptions(localDstDir, tmpDir)
 
 		task.Start()
 	},
@@ -80,4 +85,8 @@ func init() {
 	fetchCmd.Flags().StringVarP(&localDstDir, "dest-path", "d", "",
 		"local directory that files/dirs from target hosts will be copied to",
 	)
+
+	fetchCmd.Flags().StringVarP(&tmpDir, "tmp-dir", "t", "/tmp",
+		"directory of target hosts for storing temporary zip file",
+	)
 }
diff --git a/internal/cmd/push.go b/internal/cmd/push.go
@@ -111,7 +111,7 @@ Copy local files/dirs to target hosts.`,
 
 		task.SetHosts(args)
 		task.SetPushfiles(files, zipFiles)
-		task.SetFileOptions(fileDstPath, allowOverwrite)
+		task.SetPushOptions(fileDstPath, allowOverwrite)
 
 		task.Start()
 	},

diff --git a/internal/pkg/sshtask/sshtask.go b/internal/pkg/sshtask/sshtask.go
@@ -97,11 +97,13 @@ type Task struct {
 	// hostnames or ips from command line arguments.
 	hosts []string
 
-	command        string
-	scriptFile     string
+	command    string
+	scriptFile string
+
 	pushFiles      *pushFiles
 	fetchFiles     []string
 	dstDir         string
+	tmpDir         string
 	remove         bool
 	allowOverwrite bool
 
@@ -184,12 +186,18 @@ func (t *Task) SetScriptOptions(destPath string, remove, allowOverwrite bool) {
 	t.allowOverwrite = allowOverwrite
 }
 
-// SetFileOptions ...
-func (t *Task) SetFileOptions(destPath string, allowOverwrite bool) {
+// SetPushOptions ...
+func (t *Task) SetPushOptions(destPath string, allowOverwrite bool) {
 	t.dstDir = destPath
 	t.allowOverwrite = allowOverwrite
 }
 
+// SetFetchOptions ...
+func (t *Task) SetFetchOptions(destPath, tmpDir string) {
+	t.dstDir = destPath
+	t.tmpDir = tmpDir
+}
+
 // RunSSH implements batchssh.Task
 func (t *Task) RunSSH(addr string) (string, error) {
 	lang := t.configFlags.Run.Lang
@@ -204,7 +212,7 @@ func (t *Task) RunSSH(addr string) (string, error) {
 	case PushTask:
 		return t.sshClient.PushFiles(addr, t.pushFiles.files, t.pushFiles.zipFiles, t.dstDir, t.allowOverwrite)
 	case FetchTask:
-		return t.sshClient.FetchFiles(addr, t.fetchFiles, t.dstDir)
+		return t.sshClient.FetchFiles(addr, t.fetchFiles, t.dstDir, t.tmpDir)
 	default:
 		return "", fmt.Errorf("unknown task type: %v", t.taskType)
 	}

diff --git a/pkg/batchssh/batchssh.go b/pkg/batchssh/batchssh.go
@@ -280,7 +280,7 @@ func (c *Client) PushFiles(
 func (c *Client) FetchFiles(
 	addr string,
 	srcFiles []string,
-	dstDir string,
+	dstDir, tmpDir string,
 ) (string, error) {
 	client, err := c.getClient(addr)
 	if err != nil {
@@ -337,7 +337,7 @@ func (c *Client) FetchFiles(
 	}
 	defer session.Close()
 
-	zippedFileTmpDir := fmt.Sprintf("/tmp/gossh-%s", addr)
+	zippedFileTmpDir := path.Join(tmpDir, "gossh-"+addr)
 	tmpZipFile := fmt.Sprintf("%s.%d", addr, time.Now().UnixMicro())
 	zippedFileFullpath := path.Join(zippedFileTmpDir, tmpZipFile)
 	_, err = c.executeCmd(