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

package getp

import (
	"archive/tar"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"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:   "getp <bucket_name>/<package_name> <local_path>",
		Short: "download package all files to local disk",
		Args:  cobra.ExactArgs(2),
		RunE: func(c *cobra.Command, args []string) error {
			ctx := cmd.GetCmdCtx(c)
			return getp(c, ctx, opt, args)
		},
	}
	c.Flags().BoolVar(&opt.UseID, "id", false, "treat first argument as package id")
	c.Flags().StringVar(&opt.Prefix, "prefix", "", "download objects with this prefix")
	c.Flags().StringVar(&opt.NewPrefix, "new", "", "replace prefix specified by --prefix with this prefix")
	c.Flags().BoolVar(&opt.Zip, "zip", false, "download as zip file")
	c.Flags().StringVarP(&opt.Output, "output", "o", "", "output zip file name")
	cmd.RootCmd.AddCommand(c)
}

type option struct {
	UseID     bool
	Prefix    string
	NewPrefix string
	Zip       bool
	Output    string
}

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

		pkgID = clitypes.PackageID(id)
	} else {
		comps := strings.Split(args[0], "/")
		if len(comps) != 2 {
			return fmt.Errorf("invalid package name")
		}

		resp, err := ctx.Client.Package().GetByFullName(cliapi.PackageGetByFullName{
			BucketName:  comps[0],
			PackageName: comps[1],
		})
		if err != nil {
			return fmt.Errorf("get package by name: %w", err)
		}

		pkgID = resp.Package.PackageID
	}

	info, err := os.Stat(args[1])
	if err != nil {
		return err
	}

	if !info.IsDir() {
		return fmt.Errorf("local path should be a directory")
	}

	req := cliapi.PackageDownload{
		PackageID: pkgID,
		Prefix:    opt.Prefix,
	}

	if c.Flags().Changed("new") {
		req.NewPrefix = &opt.NewPrefix
	}
	if opt.Zip {
		req.Zip = true
	}

	downResp, err := ctx.Client.Package().Download(req)
	if err != nil {
		return fmt.Errorf("download package: %w", err)
	}
	defer downResp.File.Close()

	if opt.Zip {
		fileName := downResp.Name
		if opt.Output != "" {
			fileName = opt.Output
		}
		localFilePath := filepath.Join(args[1], fileName)

		file, err := os.OpenFile(localFilePath, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0644)
		if err != nil {
			return err
		}

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

		startTime := time.Now()
		n, err := io.Copy(file, downResp.File)
		if err != nil {
			return 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
	}

	startTime := time.Now()
	totalSize := int64(0)
	fileCnt := 0

	tr := tar.NewReader(downResp.File)
	for {
		header, err := tr.Next()
		if err == io.EOF {
			break
		}
		if err != nil {
			return err
		}

		localPath := filepath.Join(args[1], header.Name)

		fmt.Printf("%v", localPath)

		dir := filepath.Dir(localPath)
		err = os.MkdirAll(dir, 0755)
		if err != nil {
			fmt.Printf("\tx")
			return err
		}

		fileStartTime := time.Now()
		file, err := os.OpenFile(localPath, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0644)
		if err != nil {
			fmt.Printf("\tx")
			return err
		}

		_, err = io.Copy(file, tr)
		if err != nil {
			fmt.Printf("\tx")
			return err
		}

		dt := time.Since(fileStartTime)
		fmt.Printf("\t%v\t%v\n", bytesize.ByteSize(header.Size), dt)
		fileCnt++
		totalSize += header.Size
	}

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