forked from TrueCloudLab/xk6-frostfs
121 lines
2.3 KiB
Go
121 lines
2.3 KiB
Go
package registry
|
|
|
|
import (
|
|
"fmt"
|
|
"os"
|
|
)
|
|
|
|
type ObjExporter struct {
|
|
selector *ObjSelector
|
|
}
|
|
|
|
type PreGenerateInfo struct {
|
|
Buckets []string `json:"buckets"`
|
|
Objects []ObjInfo `json:"objects"`
|
|
ObjSize string `json:"obj_size"`
|
|
}
|
|
|
|
type ObjInfo struct {
|
|
Bucket string `json:"bucket"`
|
|
Object string `json:"object"`
|
|
}
|
|
|
|
func NewObjExporter(selector *ObjSelector) *ObjExporter {
|
|
return &ObjExporter{selector: selector}
|
|
}
|
|
|
|
func (o *ObjExporter) ExportJSONPreGen(fileName string) error {
|
|
f, err := os.Create(fileName)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
defer f.Close()
|
|
|
|
// there can be a lot of object, so manually form json
|
|
if _, err = f.WriteString(`{"objects":[`); err != nil {
|
|
return err
|
|
}
|
|
|
|
bucketMap := make(map[string]struct{})
|
|
containerMap := make(map[string]struct{})
|
|
|
|
count, err := o.selector.Count()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
var comma string
|
|
for i := 0; i < count; i++ {
|
|
info := o.selector.NextObject()
|
|
if info == nil {
|
|
break
|
|
}
|
|
|
|
if err = writeObjectInfo(comma, info, f); err != nil {
|
|
return err
|
|
}
|
|
|
|
if i == 0 {
|
|
comma = ","
|
|
}
|
|
|
|
if info.S3Bucket != "" {
|
|
bucketMap[info.S3Bucket] = struct{}{}
|
|
}
|
|
if info.CID != "" {
|
|
containerMap[info.CID] = struct{}{}
|
|
}
|
|
}
|
|
|
|
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 := ""
|
|
for bucket := range bucketMap {
|
|
if _, err = f.WriteString(fmt.Sprintf(`%s"%s"`, comma, bucket)); err != nil {
|
|
return err
|
|
}
|
|
if i == 0 {
|
|
comma = ","
|
|
}
|
|
i++
|
|
}
|
|
_, err = f.WriteString(`]`)
|
|
return err
|
|
}
|