From 3dd559a7b1239e749eaa07df7efdf694f53c0eb6 Mon Sep 17 00:00:00 2001
From: Ekaterina Lebedeva <ekaterina.lebedeva@yadro.com>
Date: Fri, 31 May 2024 21:20:23 +0300
Subject: [PATCH] [#139] Allow registry to export gRPC objects

Signed-off-by: Ekaterina Lebedeva <ekaterina.lebedeva@yadro.com>
---
 internal/registry/obj_exporter.go | 51 +++++++++++++++++++++++++++----
 1 file changed, 45 insertions(+), 6 deletions(-)

diff --git a/internal/registry/obj_exporter.go b/internal/registry/obj_exporter.go
index 73cda60..57e6d0d 100644
--- a/internal/registry/obj_exporter.go
+++ b/internal/registry/obj_exporter.go
@@ -37,6 +37,7 @@ func (o *ObjExporter) ExportJSONPreGen(fileName string) error {
 	}
 
 	bucketMap := make(map[string]struct{})
+	containerMap := make(map[string]struct{})
 
 	count, err := o.selector.Count()
 	if err != nil {
@@ -50,7 +51,7 @@ func (o *ObjExporter) ExportJSONPreGen(fileName string) error {
 			break
 		}
 
-		if _, err = f.WriteString(fmt.Sprintf(`%s{"bucket":"%s","object":"%s"}`, comma, info.S3Bucket, info.S3Key)); err != nil {
+		if err = writeObjectInfo(comma, info, f); err != nil {
 			return err
 		}
 
@@ -58,15 +59,54 @@ func (o *ObjExporter) ExportJSONPreGen(fileName string) error {
 			comma = ","
 		}
 
-		bucketMap[info.S3Bucket] = struct{}{}
+		if info.S3Bucket != "" {
+			bucketMap[info.S3Bucket] = struct{}{}
+		}
+		if info.CID != "" {
+			containerMap[info.CID] = struct{}{}
+		}
 	}
 
-	if _, err = f.WriteString(`],"buckets":[`); err != nil {
+	if _, err = f.WriteString(`]`); err != nil {
+		return err
+	}
+
+	if len(bucketMap) > 0 {
+		if err = writeContainerInfo("buckets", bucketMap, f); err != nil {
+			return err
+		}
+	}
+	if len(containerMap) > 0 {
+		if err = writeContainerInfo("containers", containerMap, f); err != nil {
+			return err
+		}
+	}
+
+	if _, err = f.WriteString(`}`); err != nil {
+		return err
+	}
+
+	return nil
+}
+
+func writeObjectInfo(comma string, info *ObjectInfo, f *os.File) (err error) {
+	var res string
+	if info.S3Bucket != "" || info.S3Key != "" {
+		res = fmt.Sprintf(`%s{"bucket":"%s","object":"%s"}`, comma, info.S3Bucket, info.S3Key)
+	} else {
+		res = fmt.Sprintf(`%s{"cid":"%s","oid":"%s"}`, comma, info.CID, info.OID)
+	}
+	_, err = f.WriteString(res)
+	return err
+}
+
+func writeContainerInfo(attrName string, bucketMap map[string]struct{}, f *os.File) (err error) {
+	if _, err = f.WriteString(fmt.Sprintf(`,"%s":[`, attrName)); err != nil {
 		return err
 	}
 
 	i := 0
-	comma = ""
+	comma := ""
 	for bucket := range bucketMap {
 		if _, err = f.WriteString(fmt.Sprintf(`%s"%s"`, comma, bucket)); err != nil {
 			return err
@@ -76,7 +116,6 @@ func (o *ObjExporter) ExportJSONPreGen(fileName string) error {
 		}
 		i++
 	}
-
-	_, err = f.WriteString(`]}`)
+	_, err = f.WriteString(`]`)
 	return err
 }