JCS-pub/jcsctl/cmd/geto/geto.go

package geto

import (
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strconv"
	"time"

	"github.com/inhies/go-bytesize"
	"github.com/spf13/cobra"
	cliapi "gitlink.org.cn/cloudream/jcs-pub/client/sdk/api/v1"
	clitypes "gitlink.org.cn/cloudream/jcs-pub/client/types"
	"gitlink.org.cn/cloudream/jcs-pub/jcsctl/cmd"
)

func init() {
	var opt option
	c := &cobra.Command{
		Use:   "geto <bucket_name>/<package_name>:<object_path> <local_dir>",
		Short: "download object to local disk",
		Args:  cobra.ExactArgs(2),
		RunE: func(c *cobra.Command, args []string) error {
			ctx := cmd.GetCmdCtx(c)
			return geto(c, ctx, opt, args)
		},
	}
	c.Flags().BoolVar(&opt.UseID, "id", false, "treat first argument as object id")
	c.Flags().Int64Var(&opt.Offset, "offset", 0, "offset of object to download")
	c.Flags().Int64Var(&opt.Length, "length", 0, "length of object to download")
	c.Flags().Int64Var(&opt.Seek, "seek", 0, "seek position when save to local file, if set, will not truncate local file")
	c.Flags().StringVarP(&opt.Output, "output", "o", "", "output file name")
	cmd.RootCmd.AddCommand(c)
}

type option struct {
	UseID  bool
	Offset int64
	Length int64
	Seek   int64
	Output string
}

func geto(c *cobra.Command, ctx *cmd.CommandContext, opt option, args []string) error {
	var obj clitypes.Object
	if opt.UseID {
		id, err := strconv.ParseInt(args[0], 10, 64)
		if err != nil {
			return fmt.Errorf("invalid object id: %v", err)
		}

		resp, err := ctx.Client.Object().ListByIDs(cliapi.ObjectListByIDs{
			ObjectIDs: []clitypes.ObjectID{clitypes.ObjectID(id)},
		})
		if err != nil {
			return fmt.Errorf("list objects by ids: %v", err)
		}

		if resp.Objects[0] == nil {
			return fmt.Errorf("object not found")
		}

		obj = *resp.Objects[0]

	} else {
		bkt, pkg, objPath, ok := cmd.SplitObjectPath(args[0])
		if !ok {
			return fmt.Errorf("invalid object path")
		}

		pkgResp, err := ctx.Client.Package().GetByFullName(cliapi.PackageGetByFullName{
			BucketName:  bkt,
			PackageName: pkg,
		})
		if err != nil {
			return fmt.Errorf("get package by name: %v", err)
		}

		objResp, err := ctx.Client.Object().ListByPath(cliapi.ObjectListByPath{
			PackageID: pkgResp.Package.PackageID,
			Path:      objPath,
		})
		if err != nil {
			return fmt.Errorf("list objects by path: %v", err)
		}

		if len(objResp.Objects) != 1 {
			return fmt.Errorf("object not found")
		}

		obj = objResp.Objects[0]
	}

	filePath := args[1]
	if opt.Output != "" {
		filePath = filepath.Join(filePath, opt.Output)
	} else {
		filePath = filepath.Join(filePath, clitypes.BaseName(obj.Path))
	}

	flag := os.O_CREATE | os.O_WRONLY
	if !c.Flags().Changed("seek") {
		flag |= os.O_TRUNC
	}

	file, err := os.OpenFile(filePath, flag, 0644)
	if err != nil {
		return fmt.Errorf("open file: %v", err)
	}
	defer file.Close()

	if opt.Seek != 0 {
		if _, err := file.Seek(opt.Seek, 0); err != nil {
			return fmt.Errorf("seek file: %v", err)
		}
	}

	fmt.Printf("%v\n", filePath)

	var len *int64
	if c.Flags().Changed("length") {
		len = &opt.Length
	}
	resp, err := ctx.Client.Object().Download(cliapi.ObjectDownload{
		ObjectID: obj.ObjectID,
		Offset:   opt.Offset,
		Length:   len,
	})
	if err != nil {
		return fmt.Errorf("download object: %v", err)
	}

	startTime := time.Now()
	n, err := io.Copy(file, resp.File)
	if err != nil {
		return fmt.Errorf("copy object to file: %v", err)
	}

	dt := time.Since(startTime)
	fmt.Printf("size: %v, time: %v, speed: %v/s\n", bytesize.ByteSize(n), dt, bytesize.ByteSize(int64(float64(n)/dt.Seconds())))
	return nil
}