frostfs-node/cmd/frostfs-cli/modules/object/get.go
Evgenii Stratonikov 9562123c49 [#406] cli: Pass context to internal client
Signed-off-by: Evgenii Stratonikov <e.stratonikov@yadro.com>
2023-06-02 11:44:04 +00:00

152 lines
4.2 KiB
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
}