xk6-frostfs/internal/registry/obj_exporter.go

package registry

import (
	"fmt"
	"os"
)

type ObjExporter struct {
	selector *ObjSelector
}

type PreGenerateInfo struct {
	Buckets    []string  `json:"buckets"`
	Containers []string  `json:"containers"`
	Objects    []ObjInfo `json:"objects"`
	ObjSize    string    `json:"obj_size"`
}

type ObjInfo struct {
	Bucket string `json:"bucket"`
	Object string `json:"object"`
	CID    string `json:"cid"`
	OID    string `json:"oid"`
}

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
}