From 6d3ecb652889e82424d1ede2c27f3eda86a7a63d 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/import.go | 55 +++++++++++++++++++++++++++++
 cmd/xk6-registry/importer/root.go   | 27 ++++++++++++++
 cmd/xk6-registry/main.go            | 18 ++++++++++
 cmd/xk6-registry/root.go            | 33 +++++++++++++++++
 4 files changed, 133 insertions(+)
 create mode 100644 cmd/xk6-registry/importer/import.go
 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

diff --git a/cmd/xk6-registry/importer/import.go b/cmd/xk6-registry/importer/import.go
new file mode 100644
index 0000000..572b185
--- /dev/null
+++ b/cmd/xk6-registry/importer/import.go
@@ -0,0 +1,55 @@
+package importer
+
+import (
+	"encoding/json"
+	"os"
+
+	"git.frostfs.info/TrueCloudLab/xk6-frostfs/internal/registry"
+)
+
+type PreGenObj struct {
+	Bucket    string `json:"bucket"`
+	Object    string `json:"object"`
+	Container string `json:"container"`
+}
+
+type PreGenerateInfo struct {
+	Buckets    []string    `json:"buckets"`
+	Containers []string    `json:"containers"`
+	Objects    []PreGenObj `json:"objects"`
+	ObjSize    string      `json:"obj_size"`
+}
+
+// ImportJSONPreGen writes objects from pregenerated JSON file
+// to the registry.
+// Note that ImportJSONPreGen does not check if object already
+// exists in the registry so in case of re-entry the registry
+// will have two entities representing the same object.
+func ImportJSONPreGen(o *registry.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
+	}
+
+	// AddObject uses DB.Batch to combine concurrent Batch calls
+	// into a single Bolt transaction. DB.Batch is limited by
+	// DB.MaxBatchDelay which may affect perfomance.
+	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
+}
diff --git a/cmd/xk6-registry/importer/root.go b/cmd/xk6-registry/importer/root.go
new file mode 100644
index 0000000..2512fc3
--- /dev/null
+++ b/cmd/xk6-registry/importer/root.go
@@ -0,0 +1,27 @@
+package importer
+
+import (
+	"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,
+	Args: cobra.MinimumNArgs(2),
+}
+
+func runCmd(cmd *cobra.Command, args []string) error {
+	objRegistry := registry.NewObjRegistry(cmd.Context(), args[0])
+	for i := 1; i < len(args); i++ {
+		if err := 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()
+}