增加任务取回功能

cmd/saastool/task.go

@@ -25,6 +25,8 @@ func RunTask(args ...string) error {
 		return RunTaskInfo(args...)
 	case "upload":
 		return RunTaskUpload(args...)
+	case "download":
+		return RunTaskDownload(args...)
 	case "run":
 		return RunTaskRun(args...)
 	default:
@@ -51,6 +53,7 @@ Commands:
     delete               Delete a task on server
     info                 Get a task info on server
     upload               Upload task's file block to server
+    download             Download task's file block to local
 
 "help" is the default command.
 

cmd/saastool/task_download.go

@@ -0,0 +1,134 @@
+package main
+
+import (
+	"flag"
+	"fmt"
+	"net/http"
+	"os"
+	"path/filepath"
+
+	"e.coding.net/rta/public/saasapi"
+	"e.coding.net/rta/public/saasapi/pkg/saashttp"
+)
+
+type downloadTaskParams struct {
+	taskSha256 string
+	destPath   string
+	saasHttp   *saashttp.SaasClient
+}
+
+func RunTaskDownload(args ...string) error {
+	fs := flag.NewFlagSet("download", flag.ExitOnError)
+	cfgFile := paramConfig(fs)
+	sha256 := paramSha256(fs)
+	destPath := paramDestPath(fs)
+
+	if err := fs.Parse(args); err != nil {
+		fmt.Fprintln(os.Stderr, "command line parse error", "err", err)
+		return err
+	}
+
+	if fs.NArg() > 0 || len(*sha256) == 0 || len(*destPath) == 0 {
+		fs.PrintDefaults()
+		return nil
+	}
+
+	cfg, err := LoadConfigFile(*cfgFile)
+	if err != nil {
+		fmt.Fprintln(os.Stderr, "LoadConfigFile error", "err", err)
+		return err
+	}
+
+	downloadTaskParams := downloadTaskParams{
+		taskSha256: *sha256,
+		destPath:   *destPath,
+		saasHttp: &saashttp.SaasClient{
+			Client:  &http.Client{},
+			ApiUrls: &cfg.ApiUrls,
+			Auth:    &cfg.Auth,
+		},
+	}
+
+	return doTaskDownload(downloadTaskParams)
+
+}
+func doTaskDownload(downloadTaskParams downloadTaskParams) error {
+	infoTaskParams := infoTaskParams{
+		taskSha256: downloadTaskParams.taskSha256,
+		saasHttp:   downloadTaskParams.saasHttp,
+	}
+	taskInfo, err := doTaskInfo(infoTaskParams)
+	if err != nil {
+		return err
+	}
+
+	if len(taskInfo.GetSourcePath()) == 0 {
+		err = fmt.Errorf("task download failed. source path is empty")
+		fmt.Fprintln(os.Stderr, err)
+		return err
+	}
+
+	totalFiles := len(taskInfo.GetTaskFileInfos())
+	fi := 0
+	for _, finfo := range taskInfo.GetTaskFileInfos() {
+		fi++
+		var f *os.File
+		offset := int64(0)
+		totalBlocks := len(finfo.GetFileBlocks())
+		bi := 0
+		for _, binfo := range finfo.GetFileBlocks() {
+			bi++
+			if binfo.GetUploaded() {
+				if f == nil {
+					fname := finfo.GetFileName()
+					if len(downloadTaskParams.destPath) > 0 {
+						fname = filepath.Join(downloadTaskParams.destPath, fname)
+					}
+
+					os.MkdirAll(filepath.Dir(fname), os.ModePerm)
+
+					f, err = os.Create(fname)
+					if err != nil {
+						return err
+					}
+				}
+
+				blockRes, err := downloadTaskParams.saasHttp.TaskDownload(
+					binfo.GetBlockSha256(),
+					f,
+					offset,
+					int(binfo.GetBlockLength()),
+				)
+
+				if err != nil {
+					return err
+				}
+
+				if blockRes.GetCode() != saasapi.ErrorCode_SUCC {
+					err = fmt.Errorf("download block error, code %d, msg %s", blockRes.GetCode(), blockRes.GetStatus())
+					fmt.Fprintln(os.Stderr, err)
+					return err
+				} else {
+					fmt.Printf("download block success. file: %v, sha256 %v. block %v/%v, file %v/%v\n",
+						finfo.GetFileName(), binfo.GetBlockSha256(),
+						bi, totalBlocks, fi, totalFiles,
+					)
+				}
+
+			} else {
+				fmt.Printf("download block. file: %v, sha256 %v. block %v/%v, file %v/%v\n",
+					finfo.GetFileName(), binfo.GetBlockSha256(),
+					bi, totalBlocks, fi, totalFiles,
+				)
+			}
+			offset += int64(binfo.GetBlockLength())
+
+		}
+
+		if f != nil {
+			f.Close()
+		}
+	}
+
+	return nil
+}