From fa240302de730e02ccb7c0d1968b5d630858e820 Mon Sep 17 00:00:00 2001
From: Ekaterina Lebedeva <ekaterina.lebedeva@yadro.com>
Date: Wed, 7 Aug 2024 15:48:09 +0300
Subject: [PATCH] [#154] Add registry import cli utility

* Currently, objects created in preset are never deleted.
  k6 deletes only objects from registry, if registry file
  is not provided k6 delete load fails.
* Added cli utility to import objects created in preset
  into registry so k6 can delete them normally.

Signed-off-by: Ekaterina Lebedeva <ekaterina.lebedeva@yadro.com>
---
 cmd/xk6-registry/importer/root.go | 32 ++++++++++++++++++++++++++++++
 cmd/xk6-registry/main.go          | 18 +++++++++++++++++
 cmd/xk6-registry/root.go          | 33 +++++++++++++++++++++++++++++++
 internal/registry/obj_exporter.go |  9 +++++----
 internal/registry/obj_importer.go | 32 ++++++++++++++++++++++++++++++
 5 files changed, 120 insertions(+), 4 deletions(-)
 create mode 100644 cmd/xk6-registry/importer/root.go
 create mode 100644 cmd/xk6-registry/main.go
 create mode 100644 cmd/xk6-registry/root.go
 create mode 100644 internal/registry/obj_importer.go

diff --git a/cmd/xk6-registry/importer/root.go b/cmd/xk6-registry/importer/root.go
new file mode 100644
index 0000000..004ec5f
--- /dev/null
+++ b/cmd/xk6-registry/importer/root.go
@@ -0,0 +1,32 @@
+package importer
+
+import (
+	"fmt"
+
+	"git.frostfs.info/TrueCloudLab/xk6-frostfs/internal/registry"
+	"github.com/spf13/cobra"
+)
+
+// Cmd represents the import command.
+var Cmd = &cobra.Command{
+	Use:   "import",
+	Short: "Import objects into registry",
+	Long:  "Import objects into registry from pregenerated files",
+	Example: `xk6-registry import registry.bolt preset.json
+xk6-registry import --status created registry.bolt preset.json another_preset.json`,
+	RunE: runCmd,
+}
+
+func runCmd(cmd *cobra.Command, args []string) error {
+	if len(args) < 2 {
+		return fmt.Errorf("expected at least two non-flag argumets: paths to the registry and to the pregenerated file, got: %s", args)
+	}
+
+	objRegistry := registry.NewObjRegistry(cmd.Context(), args[0])
+	for i := 1; i < len(args); i++ {
+		if err := registry.ImportJSONPreGen(objRegistry, args[i]); err != nil {
+			return err
+		}
+	}
+	return nil
+}
diff --git a/cmd/xk6-registry/main.go b/cmd/xk6-registry/main.go
new file mode 100644
index 0000000..218416a
--- /dev/null
+++ b/cmd/xk6-registry/main.go
@@ -0,0 +1,18 @@
+package main
+
+import (
+	"context"
+	"os"
+	"os/signal"
+	"syscall"
+)
+
+func main() {
+	ctx, _ := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM, syscall.SIGHUP)
+
+	if cmd, err := rootCmd.ExecuteContextC(ctx); err != nil {
+		cmd.PrintErrln("Error:", err.Error())
+		cmd.PrintErrf("Run '%v --help' for usage.\n", cmd.CommandPath())
+		os.Exit(1)
+	}
+}
diff --git a/cmd/xk6-registry/root.go b/cmd/xk6-registry/root.go
new file mode 100644
index 0000000..ce2738d
--- /dev/null
+++ b/cmd/xk6-registry/root.go
@@ -0,0 +1,33 @@
+package main
+
+import (
+	"runtime"
+
+	"git.frostfs.info/TrueCloudLab/xk6-frostfs/cmd/xk6-registry/importer"
+	"git.frostfs.info/TrueCloudLab/xk6-frostfs/internal/version"
+	"github.com/spf13/cobra"
+)
+
+var rootCmd = &cobra.Command{
+	Use:     "xk6-registry",
+	Version: version.Version,
+	Short:   "Command Line Tool to work with Registry",
+	Long: `Registry provides tools to work with object registry for xk6.
+It contains command for importing objects in registry from preset`,
+	SilenceErrors: true,
+	SilenceUsage:  true,
+	Run:           rootCmdRun,
+}
+
+func init() {
+	cobra.AddTemplateFunc("runtimeVersion", runtime.Version)
+	rootCmd.SetVersionTemplate(`FrostFS xk6-registry
+{{printf "Version: %s" .Version }}
+GoVersion: {{ runtimeVersion }}
+`)
+	rootCmd.AddCommand(importer.Cmd)
+}
+
+func rootCmdRun(cmd *cobra.Command, _ []string) {
+	_ = cmd.Usage()
+}
diff --git a/internal/registry/obj_exporter.go b/internal/registry/obj_exporter.go
index e1e81e5..66110c1 100644
--- a/internal/registry/obj_exporter.go
+++ b/internal/registry/obj_exporter.go
@@ -17,10 +17,11 @@ type PreGenerateInfo struct {
 }
 
 type ObjInfo struct {
-	Bucket string `json:"bucket"`
-	Object string `json:"object"`
-	CID    string `json:"cid"`
-	OID    string `json:"oid"`
+	Bucket    string `json:"bucket"`
+	Object    string `json:"object"`
+	Container string `json:"container"`
+	CID       string `json:"cid"`
+	OID       string `json:"oid"`
 }
 
 func NewObjExporter(selector *ObjSelector) *ObjExporter {
diff --git a/internal/registry/obj_importer.go b/internal/registry/obj_importer.go
new file mode 100644
index 0000000..d0f0d4b
--- /dev/null
+++ b/internal/registry/obj_importer.go
@@ -0,0 +1,32 @@
+package registry
+
+import (
+	"encoding/json"
+	"os"
+)
+
+func ImportJSONPreGen(o *ObjRegistry, filename string) error {
+	f, err := os.ReadFile(filename)
+	if err != nil {
+		return err
+	}
+
+	var pregenInfo PreGenerateInfo
+	err = json.Unmarshal(f, &pregenInfo)
+	if err != nil {
+		return err
+	}
+
+	for _, obj := range pregenInfo.Objects {
+		if obj.Bucket != "" {
+			err = o.AddObject("", "", obj.Bucket, obj.Object, "")
+		} else {
+			err = o.AddObject(obj.Container, obj.Object, "", "", "")
+		}
+		if err != nil {
+			return err
+		}
+	}
+
+	return nil
+}