frostfs-node/cmd/frostfs-cli/modules/object/get.go

package object

import (
	"bytes"
	"fmt"
	"io"
	"os"

	internalclient "git.frostfs.info/TrueCloudLab/frostfs-node/cmd/frostfs-cli/internal/client"
	"git.frostfs.info/TrueCloudLab/frostfs-node/cmd/frostfs-cli/internal/commonflags"
	"git.frostfs.info/TrueCloudLab/frostfs-node/cmd/frostfs-cli/internal/key"
	commonCmd "git.frostfs.info/TrueCloudLab/frostfs-node/cmd/internal/common"
	cid "git.frostfs.info/TrueCloudLab/frostfs-sdk-go/container/id"
	"git.frostfs.info/TrueCloudLab/frostfs-sdk-go/object"
	oid "git.frostfs.info/TrueCloudLab/frostfs-sdk-go/object/id"
	"github.com/cheggaaa/pb"
	"github.com/spf13/cobra"
)

var objectGetCmd = &cobra.Command{
	Use:   "get",
	Short: "Get object from FrostFS",
	Long:  "Get object from FrostFS",
	Run:   getObject,
}

func initObjectGetCmd() {
	commonflags.Init(objectGetCmd)
	initFlagSession(objectGetCmd, "GET")

	flags := objectGetCmd.Flags()

	flags.String(commonflags.CIDFlag, "", commonflags.CIDFlagUsage)
	_ = objectGetCmd.MarkFlagRequired(commonflags.CIDFlag)

	flags.String(commonflags.OIDFlag, "", commonflags.OIDFlagUsage)
	_ = objectGetCmd.MarkFlagRequired(commonflags.OIDFlag)

	flags.String(fileFlag, "", "File to write object payload to(with -b together with signature and header). Default: stdout.")
	flags.Bool(rawFlag, false, rawFlagDesc)
	flags.Bool(noProgressFlag, false, "Do not show progress bar")
	flags.Bool(binaryFlag, false, "Serialize whole object structure into given file(id + signature + header + payload).")
}

func getObject(cmd *cobra.Command, _ []string) {
	var cnr cid.ID
	var obj oid.ID

	objAddr := readObjectAddress(cmd, &cnr, &obj)

	filename := cmd.Flag(fileFlag).Value.String()
	out, closer := createOutWriter(cmd, filename)
	defer closer()

	pk := key.GetOrGenerate(cmd)

	cli := internalclient.GetSDKClientByFlag(cmd, pk, commonflags.RPC)

	var prm internalclient.GetObjectPrm
	prm.SetClient(cli)
	Prepare(cmd, &prm)
	readSession(cmd, &prm, pk, cnr, obj)

	raw, _ := cmd.Flags().GetBool(rawFlag)
	prm.SetRawFlag(raw)
	prm.SetAddress(objAddr)

	var p *pb.ProgressBar
	noProgress, _ := cmd.Flags().GetBool(noProgressFlag)

	var payloadWriter io.Writer
	var payloadBuffer *bytes.Buffer
	binary, _ := cmd.Flags().GetBool(binaryFlag)
	if binary {
		payloadBuffer = new(bytes.Buffer)
		payloadWriter = payloadBuffer
	} else {
		payloadWriter = out
	}

	if filename == "" || noProgress {
		prm.SetPayloadWriter(payloadWriter)
	} else {
		p = pb.New64(0)
		p.Output = cmd.OutOrStdout()
		prm.SetPayloadWriter(p.NewProxyWriter(payloadWriter))
		prm.SetHeaderCallback(func(o *object.Object) {
			p.SetTotal64(int64(o.PayloadSize()))
			p.Start()
		})
	}

	res, err := internalclient.GetObject(cmd.Context(), prm)
	if p != nil {
		p.Finish()
	}
	if err != nil {
		if ok := printSplitInfoErr(cmd, err); ok {
			return
		}

		commonCmd.ExitOnErr(cmd, "rpc error: %w", err)
	}

	processResult(cmd, res, binary, payloadBuffer, out, filename)
}

func processResult(cmd *cobra.Command, res *internalclient.GetObjectRes, binary bool, payloadBuffer *bytes.Buffer, out io.Writer, filename string) {
	if binary {
		objToStore := res.Header()
		// TODO(@acid-ant): #1932 Use streams to marshal/unmarshal payload
		objToStore.SetPayload(payloadBuffer.Bytes())
		objBytes, err := objToStore.Marshal()
		commonCmd.ExitOnErr(cmd, "", err)
		_, err = out.Write(objBytes)
		commonCmd.ExitOnErr(cmd, "unable to write binary object in out: %w ", err)
	}

	if filename != "" && !strictOutput(cmd) {
		cmd.Printf("[%s] Object successfully saved\n", filename)
	}

	// Print header only if file is not streamed to stdout.
	if filename != "" {
		err := printHeader(cmd, res.Header())
		commonCmd.ExitOnErr(cmd, "", err)
	}
}

func createOutWriter(cmd *cobra.Command, filename string) (out io.Writer, closer func()) {
	if filename == "" {
		out = os.Stdout
		closer = func() {}
	} else {
		f, err := os.OpenFile(filename, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0644)
		if err != nil {
			commonCmd.ExitOnErr(cmd, "", fmt.Errorf("can't open file '%s': %w", filename, err))
		}

		out = f
		closer = func() {
			f.Close()
		}
	}
	return
}

func strictOutput(cmd *cobra.Command) bool {
	toJSON, _ := cmd.Flags().GetBool(commonflags.JSON)
	toProto, _ := cmd.Flags().GetBool("proto")
	return toJSON || toProto
}