Skip to content

Commit

Permalink
cli: add dump-bin-put and dump-generate-index-file
Browse files Browse the repository at this point in the history
These commands can be useful for interaction with NeoFS.

Signed-off-by: Ekaterina Pavlova <[email protected]>
  • Loading branch information
AliceInHunterland committed Sep 11, 2024
1 parent 0b31a29 commit 8e1d6b3
Show file tree
Hide file tree
Showing 3 changed files with 384 additions and 0 deletions.
154 changes: 154 additions & 0 deletions cli/server/dump_bin_put.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,154 @@
package server

import (
"context"
"fmt"
"log"
"os"
"strconv"
"sync"
"time"

"github.com/google/uuid"
"github.com/nspcc-dev/neo-go/cli/cmdargs"
"github.com/nspcc-dev/neo-go/cli/options"
"github.com/nspcc-dev/neo-go/pkg/services/oracle/neofs"
"github.com/nspcc-dev/neo-go/pkg/wallet"
"github.com/nspcc-dev/neofs-sdk-go/client"
cid "github.com/nspcc-dev/neofs-sdk-go/container/id"
neofsecdsa "github.com/nspcc-dev/neofs-sdk-go/crypto/ecdsa"
"github.com/nspcc-dev/neofs-sdk-go/object"
"github.com/nspcc-dev/neofs-sdk-go/object/slicer"
"github.com/nspcc-dev/neofs-sdk-go/session"
"github.com/nspcc-dev/neofs-sdk-go/user"
"github.com/urfave/cli/v2"
)

const batchSize = 50

func dumpPut(ctx *cli.Context) error {
if err := cmdargs.EnsureNone(ctx); err != nil {
return err

Check warning on line 31 in cli/server/dump_bin_put.go

View check run for this annotation

Codecov / codecov/patch

cli/server/dump_bin_put.go#L29-L31

Added lines #L29 - L31 were not covered by tests
}
dir := ctx.String("dir")
batch := ctx.Uint("batch")
if batch == 0 {
batch = batchSize

Check warning on line 36 in cli/server/dump_bin_put.go

View check run for this annotation

Codecov / codecov/patch

cli/server/dump_bin_put.go#L33-L36

Added lines #L33 - L36 were not covered by tests
}
rpc := ctx.String("rpc-endpoint")
containerIDStr := ctx.String("container")
attribute := ctx.String("attribute")
start := ctx.Uint("start")
count := ctx.Uint("count")
acc, _, err := options.GetAccFromContext(ctx)
if err != nil {
return fmt.Errorf("failed to load wallet: %w", err)

Check warning on line 45 in cli/server/dump_bin_put.go

View check run for this annotation

Codecov / codecov/patch

cli/server/dump_bin_put.go#L38-L45

Added lines #L38 - L45 were not covered by tests
}

var containerID cid.ID
err = containerID.DecodeString(containerIDStr)
if err != nil {
return fmt.Errorf("failed to decode container ID: %w", err)

Check warning on line 51 in cli/server/dump_bin_put.go

View check run for this annotation

Codecov / codecov/patch

cli/server/dump_bin_put.go#L48-L51

Added lines #L48 - L51 were not covered by tests
}

clientSDK, err := neofs.GetSDKClient(context.Background(), rpc, 10*time.Minute)
if err != nil {
return fmt.Errorf("failed to create NeoFS client: %w", err)

Check warning on line 56 in cli/server/dump_bin_put.go

View check run for this annotation

Codecov / codecov/patch

cli/server/dump_bin_put.go#L54-L56

Added lines #L54 - L56 were not covered by tests
}
defer clientSDK.Close()

Check warning on line 58 in cli/server/dump_bin_put.go

View check run for this annotation

Codecov / codecov/patch

cli/server/dump_bin_put.go#L58

Added line #L58 was not covered by tests

sessionToken, err := createSessionToken(*acc, containerID)
if err != nil {
return fmt.Errorf("failed to create session token: %w", err)

Check warning on line 62 in cli/server/dump_bin_put.go

View check run for this annotation

Codecov / codecov/patch

cli/server/dump_bin_put.go#L60-L62

Added lines #L60 - L62 were not covered by tests
}

ctxCancel, cancel := context.WithCancel(context.Background())
defer cancel()

Check warning on line 66 in cli/server/dump_bin_put.go

View check run for this annotation

Codecov / codecov/patch

cli/server/dump_bin_put.go#L65-L66

Added lines #L65 - L66 were not covered by tests

for i := start; i < count; i += batch {
end := i + batch
if end > count {
end = count

Check warning on line 71 in cli/server/dump_bin_put.go

View check run for this annotation

Codecov / codecov/patch

cli/server/dump_bin_put.go#L68-L71

Added lines #L68 - L71 were not covered by tests
}

var wg sync.WaitGroup
errChan := make(chan error, batch)

Check warning on line 75 in cli/server/dump_bin_put.go

View check run for this annotation

Codecov / codecov/patch

cli/server/dump_bin_put.go#L74-L75

Added lines #L74 - L75 were not covered by tests

for j := i; j < end; j++ {
wg.Add(1)
go func(index uint) {
defer wg.Done()
path := fmt.Sprintf("%s/block-%d.bin", dir, index)
if err := uploadFileWithSlicer(ctxCancel, clientSDK, *acc, containerID, path, attribute, strconv.Itoa(int(index)), sessionToken); err != nil {
log.Printf("Failed to upload %s: %v", path, err)
errChan <- err
cancel()
} else {
log.Printf("Successfully uploaded: %s", path)

Check warning on line 87 in cli/server/dump_bin_put.go

View check run for this annotation

Codecov / codecov/patch

cli/server/dump_bin_put.go#L77-L87

Added lines #L77 - L87 were not covered by tests
}
}(j)
}

wg.Wait()
select {
case err := <-errChan:
return fmt.Errorf("batch upload failed: %w", err)
default:

Check warning on line 96 in cli/server/dump_bin_put.go

View check run for this annotation

Codecov / codecov/patch

cli/server/dump_bin_put.go#L92-L96

Added lines #L92 - L96 were not covered by tests
}
}

log.Println("Upload completed.")
return nil

Check warning on line 101 in cli/server/dump_bin_put.go

View check run for this annotation

Codecov / codecov/patch

cli/server/dump_bin_put.go#L100-L101

Added lines #L100 - L101 were not covered by tests
}

// uploadFileWithSlicer uploads a file to NeoFS using the Slicer with a session token.
func uploadFileWithSlicer(ctx context.Context, clientSDK *client.Client, account wallet.Account, containerID cid.ID, filePath, attributeKey, attributeValue string, sessionToken *session.Object) error {
file, err := os.Open(filePath)
if err != nil {
return fmt.Errorf("failed to open file: %w", err)

Check warning on line 108 in cli/server/dump_bin_put.go

View check run for this annotation

Codecov / codecov/patch

cli/server/dump_bin_put.go#L105-L108

Added lines #L105 - L108 were not covered by tests
}
defer file.Close()

Check warning on line 110 in cli/server/dump_bin_put.go

View check run for this annotation

Codecov / codecov/patch

cli/server/dump_bin_put.go#L110

Added line #L110 was not covered by tests

signer := user.NewAutoIDSignerRFC6979(account.PrivateKey().PrivateKey)
var ownerID user.ID
ownerID.SetScriptHash(account.PrivateKey().GetScriptHash())

Check warning on line 114 in cli/server/dump_bin_put.go

View check run for this annotation

Codecov / codecov/patch

cli/server/dump_bin_put.go#L112-L114

Added lines #L112 - L114 were not covered by tests

slc, err := slicer.New(ctx, clientSDK, signer, containerID, ownerID, sessionToken)
if err != nil {
return fmt.Errorf("failed to create slicer: %w", err)

Check warning on line 118 in cli/server/dump_bin_put.go

View check run for this annotation

Codecov / codecov/patch

cli/server/dump_bin_put.go#L116-L118

Added lines #L116 - L118 were not covered by tests
}

attrs := []object.Attribute{
*object.NewAttribute(attributeKey, attributeValue),

Check warning on line 122 in cli/server/dump_bin_put.go

View check run for this annotation

Codecov / codecov/patch

cli/server/dump_bin_put.go#L121-L122

Added lines #L121 - L122 were not covered by tests
}

_, err = slc.Put(ctx, file, attrs)
if err != nil {
return fmt.Errorf("failed to slice and upload file: %w", err)

Check warning on line 127 in cli/server/dump_bin_put.go

View check run for this annotation

Codecov / codecov/patch

cli/server/dump_bin_put.go#L125-L127

Added lines #L125 - L127 were not covered by tests
}

return nil

Check warning on line 130 in cli/server/dump_bin_put.go

View check run for this annotation

Codecov / codecov/patch

cli/server/dump_bin_put.go#L130

Added line #L130 was not covered by tests
}

func createSessionToken(account wallet.Account, containerID cid.ID) (*session.Object, error) {
sessionToken := session.Object{}

Check warning on line 134 in cli/server/dump_bin_put.go

View check run for this annotation

Codecov / codecov/patch

cli/server/dump_bin_put.go#L133-L134

Added lines #L133 - L134 were not covered by tests

sessionID := uuid.New()
sessionToken.SetID(sessionID)

Check warning on line 137 in cli/server/dump_bin_put.go

View check run for this annotation

Codecov / codecov/patch

cli/server/dump_bin_put.go#L136-L137

Added lines #L136 - L137 were not covered by tests

pubKey := account.PublicKey()
authKey := neofsecdsa.PublicKey(*pubKey)
sessionToken.SetAuthKey(&authKey)

Check warning on line 141 in cli/server/dump_bin_put.go

View check run for this annotation

Codecov / codecov/patch

cli/server/dump_bin_put.go#L139-L141

Added lines #L139 - L141 were not covered by tests

sessionToken.BindContainer(containerID)
sessionToken.ForVerb(session.VerbObjectPut)

Check warning on line 144 in cli/server/dump_bin_put.go

View check run for this annotation

Codecov / codecov/patch

cli/server/dump_bin_put.go#L143-L144

Added lines #L143 - L144 were not covered by tests

sessionToken.SetExp(uint64(time.Now().Add(100 * time.Minute).Unix()))

Check warning on line 146 in cli/server/dump_bin_put.go

View check run for this annotation

Codecov / codecov/patch

cli/server/dump_bin_put.go#L146

Added line #L146 was not covered by tests

signer := user.NewAutoIDSignerRFC6979(account.PrivateKey().PrivateKey)
if err := sessionToken.Sign(signer); err != nil {
return nil, fmt.Errorf("failed to sign session token: %w", err)

Check warning on line 150 in cli/server/dump_bin_put.go

View check run for this annotation

Codecov / codecov/patch

cli/server/dump_bin_put.go#L148-L150

Added lines #L148 - L150 were not covered by tests
}

return &sessionToken, nil

Check warning on line 153 in cli/server/dump_bin_put.go

View check run for this annotation

Codecov / codecov/patch

cli/server/dump_bin_put.go#L153

Added line #L153 was not covered by tests
}
161 changes: 161 additions & 0 deletions cli/server/generate_oid.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,161 @@
package server

import (
"context"
"fmt"
"log"
"os"
"sync"
"time"

"github.com/nspcc-dev/neo-go/pkg/crypto/keys"
"github.com/nspcc-dev/neo-go/pkg/services/oracle/neofs"
"github.com/nspcc-dev/neofs-sdk-go/client"
cid "github.com/nspcc-dev/neofs-sdk-go/container/id"
"github.com/nspcc-dev/neofs-sdk-go/object"
oid "github.com/nspcc-dev/neofs-sdk-go/object/id"
"github.com/urfave/cli/v2"
)

func generateOIDs(ctx *cli.Context) error {
rpc := ctx.String("rpc-endpoint")
containerIDStr := ctx.String("container")
attribute := ctx.String("attribute")
startIndex := uint32(ctx.Uint("start"))
count := uint32(ctx.Uint("count"))
outputDir := ctx.String("out")
fileSize := uint32(ctx.Uint("index-file-size"))

Check warning on line 27 in cli/server/generate_oid.go

View check run for this annotation

Codecov / codecov/patch

cli/server/generate_oid.go#L20-L27

Added lines #L20 - L27 were not covered by tests

if _, err := os.Stat(outputDir); os.IsNotExist(err) {
if err := os.MkdirAll(outputDir, os.ModePerm); err != nil {
return fmt.Errorf("failed to create directory %s: %w", outputDir, err)

Check warning on line 31 in cli/server/generate_oid.go

View check run for this annotation

Codecov / codecov/patch

cli/server/generate_oid.go#L29-L31

Added lines #L29 - L31 were not covered by tests
}
}

var containerID cid.ID
err := containerID.DecodeString(containerIDStr)
if err != nil {
return fmt.Errorf("failed to decode container ID: %w", err)

Check warning on line 38 in cli/server/generate_oid.go

View check run for this annotation

Codecov / codecov/patch

cli/server/generate_oid.go#L35-L38

Added lines #L35 - L38 were not covered by tests
}
pk, err := keys.NewPrivateKey()
if err != nil {
return fmt.Errorf("failed to create private key: %w", err)

Check warning on line 42 in cli/server/generate_oid.go

View check run for this annotation

Codecov / codecov/patch

cli/server/generate_oid.go#L40-L42

Added lines #L40 - L42 were not covered by tests
}

clientSDK, err := neofs.GetSDKClient(context.Background(), rpc, 10*time.Minute)
if err != nil {
return fmt.Errorf("failed to create NeoFS client: %w", err)

Check warning on line 47 in cli/server/generate_oid.go

View check run for this annotation

Codecov / codecov/patch

cli/server/generate_oid.go#L45-L47

Added lines #L45 - L47 were not covered by tests
}
defer clientSDK.Close()

Check warning on line 49 in cli/server/generate_oid.go

View check run for this annotation

Codecov / codecov/patch

cli/server/generate_oid.go#L49

Added line #L49 was not covered by tests

fileIndex := 1 + startIndex/fileSize

Check warning on line 51 in cli/server/generate_oid.go

View check run for this annotation

Codecov / codecov/patch

cli/server/generate_oid.go#L51

Added line #L51 was not covered by tests

type result struct {
index uint32
oidsBatch []oid.ID
err error

Check warning on line 56 in cli/server/generate_oid.go

View check run for this annotation

Codecov / codecov/patch

cli/server/generate_oid.go#L53-L56

Added lines #L53 - L56 were not covered by tests
}

for i := startIndex; i < startIndex+count; i += fileSize {
var wg sync.WaitGroup
results := make(chan result, fileSize)
workerPool := make(chan struct{}, 100)

Check warning on line 62 in cli/server/generate_oid.go

View check run for this annotation

Codecov / codecov/patch

cli/server/generate_oid.go#L59-L62

Added lines #L59 - L62 were not covered by tests

oids := make([]oid.ID, 0, fileSize)

Check warning on line 64 in cli/server/generate_oid.go

View check run for this annotation

Codecov / codecov/patch

cli/server/generate_oid.go#L64

Added line #L64 was not covered by tests

for j := i; j < i+fileSize && j < startIndex+count; j++ {
wg.Add(1)
workerPool <- struct{}{}

Check warning on line 68 in cli/server/generate_oid.go

View check run for this annotation

Codecov / codecov/patch

cli/server/generate_oid.go#L66-L68

Added lines #L66 - L68 were not covered by tests

go func(index uint32) {
defer wg.Done()
defer func() { <-workerPool }()

Check warning on line 72 in cli/server/generate_oid.go

View check run for this annotation

Codecov / codecov/patch

cli/server/generate_oid.go#L70-L72

Added lines #L70 - L72 were not covered by tests

prm := client.PrmObjectSearch{}
filters := object.NewSearchFilters()
filters.AddFilter(attribute, fmt.Sprintf("%d", index), object.MatchStringEqual)
prm.SetFilters(filters)
ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
defer cancel()

Check warning on line 79 in cli/server/generate_oid.go

View check run for this annotation

Codecov / codecov/patch

cli/server/generate_oid.go#L74-L79

Added lines #L74 - L79 were not covered by tests

oidsBatch, err := neofs.ObjectSearch(ctx, clientSDK, pk, containerID.String(), prm)
if err != nil || len(oidsBatch) == 0 {
results <- result{index, nil, fmt.Errorf("no OIDs found or failed for index %d: %w", index, err)}
return

Check warning on line 84 in cli/server/generate_oid.go

View check run for this annotation

Codecov / codecov/patch

cli/server/generate_oid.go#L81-L84

Added lines #L81 - L84 were not covered by tests
}
if len(oidsBatch) == 0 {
fmt.Println("Found OIDs:", oidsBatch, "for index", index, "total", len(oidsBatch))

Check warning on line 87 in cli/server/generate_oid.go

View check run for this annotation

Codecov / codecov/patch

cli/server/generate_oid.go#L86-L87

Added lines #L86 - L87 were not covered by tests
}
if index%2000 == 0 {
fmt.Println("Index:", index)

Check warning on line 90 in cli/server/generate_oid.go

View check run for this annotation

Codecov / codecov/patch

cli/server/generate_oid.go#L89-L90

Added lines #L89 - L90 were not covered by tests
}
results <- result{index, oidsBatch[:1], nil}

Check warning on line 92 in cli/server/generate_oid.go

View check run for this annotation

Codecov / codecov/patch

cli/server/generate_oid.go#L92

Added line #L92 was not covered by tests
}(j)
}

go func() {
wg.Wait()
close(results)
}()

Check warning on line 99 in cli/server/generate_oid.go

View check run for this annotation

Codecov / codecov/patch

cli/server/generate_oid.go#L96-L99

Added lines #L96 - L99 were not covered by tests

pendingResults := make(map[uint32][]oid.ID)
expectedIndex := i

Check warning on line 102 in cli/server/generate_oid.go

View check run for this annotation

Codecov / codecov/patch

cli/server/generate_oid.go#L101-L102

Added lines #L101 - L102 were not covered by tests

for res := range results {
if res.err != nil {
return res.err

Check warning on line 106 in cli/server/generate_oid.go

View check run for this annotation

Codecov / codecov/patch

cli/server/generate_oid.go#L104-L106

Added lines #L104 - L106 were not covered by tests
}

pendingResults[res.index] = res.oidsBatch

Check warning on line 109 in cli/server/generate_oid.go

View check run for this annotation

Codecov / codecov/patch

cli/server/generate_oid.go#L109

Added line #L109 was not covered by tests

for {
if batch, ok := pendingResults[expectedIndex]; ok {
delete(pendingResults, expectedIndex)
oids = append(oids, batch...)

Check warning on line 114 in cli/server/generate_oid.go

View check run for this annotation

Codecov / codecov/patch

cli/server/generate_oid.go#L111-L114

Added lines #L111 - L114 were not covered by tests

if len(oids) >= int(fileSize) {
if err := writeOIDFile(oids[:fileSize], outputDir, int(fileIndex)); err != nil {
return err

Check warning on line 118 in cli/server/generate_oid.go

View check run for this annotation

Codecov / codecov/patch

cli/server/generate_oid.go#L116-L118

Added lines #L116 - L118 were not covered by tests
}
oids = oids[fileSize:]
fileIndex++

Check warning on line 121 in cli/server/generate_oid.go

View check run for this annotation

Codecov / codecov/patch

cli/server/generate_oid.go#L120-L121

Added lines #L120 - L121 were not covered by tests
}
expectedIndex++
} else {
break

Check warning on line 125 in cli/server/generate_oid.go

View check run for this annotation

Codecov / codecov/patch

cli/server/generate_oid.go#L123-L125

Added lines #L123 - L125 were not covered by tests
}
}
}

if len(oids) > 0 {
if err := writeOIDFile(oids, outputDir, int(fileIndex)); err != nil {
return err

Check warning on line 132 in cli/server/generate_oid.go

View check run for this annotation

Codecov / codecov/patch

cli/server/generate_oid.go#L130-L132

Added lines #L130 - L132 were not covered by tests
}
fileIndex++

Check warning on line 134 in cli/server/generate_oid.go

View check run for this annotation

Codecov / codecov/patch

cli/server/generate_oid.go#L134

Added line #L134 was not covered by tests
}
}

log.Println("OID generation completed.")
return nil

Check warning on line 139 in cli/server/generate_oid.go

View check run for this annotation

Codecov / codecov/patch

cli/server/generate_oid.go#L138-L139

Added lines #L138 - L139 were not covered by tests
}

// writeOIDFile writes a list of OIDs to a file in the specified directory.
func writeOIDFile(oids []oid.ID, outputDir string, fileIndex int) error {
fileName := fmt.Sprintf("%s/oid-%d.bin", outputDir, fileIndex)
file, err := os.Create(fileName)
if err != nil {
return fmt.Errorf("failed to create file %s: %w", fileName, err)

Check warning on line 147 in cli/server/generate_oid.go

View check run for this annotation

Codecov / codecov/patch

cli/server/generate_oid.go#L143-L147

Added lines #L143 - L147 were not covered by tests
}
defer file.Close()

Check warning on line 149 in cli/server/generate_oid.go

View check run for this annotation

Codecov / codecov/patch

cli/server/generate_oid.go#L149

Added line #L149 was not covered by tests

for _, oid := range oids {
oidBytes := make([]byte, 32)
oid.Encode(oidBytes)
if _, err := file.Write(oidBytes); err != nil {
return fmt.Errorf("failed to write OID to file %s: %w", fileName, err)

Check warning on line 155 in cli/server/generate_oid.go

View check run for this annotation

Codecov / codecov/patch

cli/server/generate_oid.go#L151-L155

Added lines #L151 - L155 were not covered by tests
}
}

log.Printf("Successfully wrote %d OIDs to %s", len(oids), fileName)
return nil

Check warning on line 160 in cli/server/generate_oid.go

View check run for this annotation

Codecov / codecov/patch

cli/server/generate_oid.go#L159-L160

Added lines #L159 - L160 were not covered by tests
}
69 changes: 69 additions & 0 deletions cli/server/server.go
Original file line number Diff line number Diff line change
Expand Up @@ -83,6 +83,49 @@ func NewCommands() []*cli.Command {
Usage: "Height of the state to reset DB to",
Required: true,
})
var neofsFlags = []cli.Flag{
&cli.StringFlag{
Name: "rpc-endpoint",
Aliases: []string{"r"},
Usage: "NeoFS RPC address",
Required: true,
},
&cli.StringFlag{
Name: "container",
Aliases: []string{"cid"},
Usage: "NeoFS container ID",
Required: true,
},
&cli.StringFlag{
Name: "attribute",
Aliases: []string{"a"},
Usage: "Attribute key of the object",
Required: true,
},
&cli.UintFlag{
Name: "start",
Aliases: []string{"s"},
Usage: "Starting block index",
Required: true,
},
&cli.UintFlag{
Name: "count",
Aliases: []string{"n"},
Usage: "Number of blocks to be processed",
Required: true,
},
}
var putFlags = append(neofsFlags,
&cli.StringFlag{
Name: "dir",
Usage: "Directory containing the binary files",
Aliases: []string{"d"},
Required: true,
},
&cli.UintFlag{
Name: "batch",
Usage: "Number of blocks to be processed in parallel",
})
return []*cli.Command{
{
Name: "node",
Expand All @@ -109,6 +152,32 @@ func NewCommands() []*cli.Command {
Action: dumpBin,
Flags: cfgCountOutFlags,
},
{
Name: "dump-bin-put",
Usage: "Upload blocks from binary files in the specified directory to the NeoFS container",
UsageText: "neo-go db dump-bin-put --dir <directory> --rpc-endpoint <endpoint> --container <cid> --attribute <key> --start <block-index> --count <count> --wallet <wallet-path> ",
Action: dumpPut,
Flags: append(putFlags, options.Wallet...),
},
{
Name: "dump-generate-index-file",
Usage: "Generate index files of oids blocks from NeoFS container with the specified attribute",
UsageText: "neo-go db dump-generate-oid --start <block-index> --count <count> --out <output-dir>, --container <cid>, --index-file-size <file-size>, --rpc-endpoint <endpoint>, --attribute <key>",
Action: generateOIDs,
Flags: append([]cli.Flag{
&cli.StringFlag{
Name: "out",
Usage: "Output directory where OID files will be saved",
Aliases: []string{"o"},
Required: true,
},
&cli.UintFlag{
Name: "index-file-size",
Usage: "Number of blocks in oid file",
Required: true,
},
}, neofsFlags...),
},
{
Name: "restore",
Usage: "Restore blocks from the file",
Expand Down

0 comments on commit 8e1d6b3

Please sign in to comment.