stream bulk #9308
+1,803
−1,072
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,49 @@ | ||
| package dgraphimport | ||
|
|
||
| import ( | ||
| "context" | ||
| "fmt" | ||
| "os" | ||
|
|
||
| "github.com/hypermodeinc/dgraph/v24/edgraph" | ||
| "github.com/hypermodeinc/dgraph/v24/x" | ||
| "github.com/spf13/cobra" | ||
| "google.golang.org/grpc" | ||
| "google.golang.org/grpc/credentials/insecure" | ||
| ) | ||
|
|
||
| var ( | ||
| Import x.SubCommand | ||
| ) | ||
|
|
||
| func init() { | ||
| Import.Cmd = &cobra.Command{ | ||
| Use: "import", | ||
| Short: "Run the import tool", | ||
| Run: func(cmd *cobra.Command, args []string) { | ||
| run() | ||
| }, | ||
| Annotations: map[string]string{"group": "tool"}, | ||
| } | ||
| Import.Cmd.SetHelpTemplate(x.NonRootTemplate) | ||
| } | ||
|
|
||
| func run() { | ||
| if err := importP(); err != nil { | ||
| fmt.Fprintln(os.Stderr, err) | ||
| } | ||
| } | ||
|
|
||
| func importP() error { | ||
| client, err := edgraph.NewImportClient("localhost:9080", grpc.WithTransportCredentials(insecure.NewCredentials())) | ||
| if err != nil { | ||
| return err | ||
| } | ||
|
|
||
| alphas, err := client.InitiateSnapShotStream(context.Background()) | ||
| if err := client.StreamSnapshot(context.Background(), "/home/shiva/workspace/dgraph-work/stream_data/out", alphas.LeaderAlphas); err != nil { | ||
| fmt.Println("error is---------", err) | ||
| } | ||
|
|
||
| return nil | ||
| } | ||
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,138 @@ | ||
| package edgraph | ||
|
|
||
| import ( | ||
| "context" | ||
| "fmt" | ||
| "math" | ||
| "os" | ||
| "path/filepath" | ||
| "strconv" | ||
|
|
||
| "github.com/dgraph-io/badger/v4" | ||
| apiv25 "github.com/dgraph-io/dgo/v240/protos/api.v25" | ||
| "github.com/dgraph-io/ristretto/v2/z" | ||
| "github.com/golang/glog" | ||
| "google.golang.org/grpc" | ||
| ) | ||
|
|
||
| type Client struct { | ||
| opts grpc.DialOption | ||
| dg apiv25.DgraphHMClient | ||
| } | ||
|
|
||
| func NewImportClient(endpoint string, opts grpc.DialOption) (*Client, error) { | ||
| conn, err := grpc.NewClient(endpoint, opts) | ||
| if err != nil { | ||
| return nil, fmt.Errorf("failed to connect to endpoint [%s]: %w", endpoint, err) | ||
| } | ||
|
|
||
| return &Client{dg: apiv25.NewDgraphHMClient(conn), opts: opts}, nil | ||
| } | ||
|
|
||
| func (c *Client) InitiateSnapShotStream(ctx context.Context) (*apiv25.InitiateSnapShotStreamResponse, error) { | ||
| req := &apiv25.InitiateSnapShotStreamRequest{} | ||
| return c.dg.InitiateSnapShotStream(ctx, req) | ||
| } | ||
| func (c *Client) StreamSnapshot(ctx context.Context, pDir string, alphas map[uint32]string) error { | ||
| groupDirs, err := getPDiPdirrectories(pDir) | ||
| if err != nil { | ||
| return fmt.Errorf("Error getting p directories: %v", err) | ||
| } | ||
|
|
||
| for key, leader := range alphas { | ||
| pDir, exists := groupDirs[key-1] | ||
| if !exists { | ||
| fmt.Printf("No p directory found for group %d, skipping...\n", key) | ||
| continue | ||
| } | ||
|
|
||
| if _, err := os.Stat(pDir); os.IsNotExist(err) { | ||
| fmt.Printf("p directory does not exist: %s, skipping...\n", pDir) | ||
| continue | ||
| } | ||
|
|
||
| conn, err := grpc.NewClient(leader, c.opts) | ||
| if err != nil { | ||
| return fmt.Errorf("Failed to connect to leader %s: %v", leader, err) | ||
| } | ||
| defer conn.Close() | ||
|
|
||
| dg := apiv25.NewDgraphHMClient(conn) | ||
| err = stream(ctx, dg, pDir) | ||
| if err != nil { | ||
| return err | ||
| } | ||
| } | ||
|
|
||
| return nil | ||
| } | ||
|
|
||
| func stream(ctx context.Context, dg apiv25.DgraphHMClient, pdir string) error { | ||
| out, err := dg.StreamPSnapshot(ctx) | ||
| if err != nil { | ||
| return fmt.Errorf("failed to start snapshot stream: %w", err) | ||
| } | ||
|
|
||
| opt := badger.DefaultOptions(pdir) | ||
| ps, err := badger.OpenManaged(opt) | ||
| if err != nil { | ||
| return fmt.Errorf("failed to open BadgerDB: %w", err) | ||
| } | ||
| defer ps.Close() | ||
|
|
||
| stream := ps.NewStreamAt(math.MaxUint64) | ||
| stream.LogPrefix = "Sending P dir" | ||
| stream.KeyToList = nil | ||
| stream.Send = func(buf *z.Buffer) error { | ||
| kvs := &apiv25.KVS{Data: buf.Bytes()} | ||
| if err := out.Send(kvs); err != nil { | ||
| return fmt.Errorf("failed to send data: %w", err) | ||
| } | ||
| return nil | ||
| } | ||
|
|
||
| if err := stream.Orchestrate(ctx); err != nil { | ||
| return fmt.Errorf("stream orchestration failed: %w", err) | ||
| } | ||
|
|
||
| done := &apiv25.KVS{ | ||
| Done: true, | ||
| Predicates: []string{}, | ||
| Types: []string{}, | ||
| } | ||
| if err := out.Send(done); err != nil { | ||
| return fmt.Errorf("failed to send 'done' signal: %w", err) | ||
| } | ||
|
|
||
| fmt.Println("Snapshot writes DONE. Sending ACK") | ||
| ack, err := out.CloseAndRecv() | ||
| if err != nil { | ||
| return fmt.Errorf("failed to receive ACK: %w", err) | ||
| } | ||
| glog.Infof("Received ACK with message: %v\n", ack.Message) | ||
|
|
||
| return nil | ||
| } | ||
|
|
||
| func getPDiPdirrectories(basePath string) (map[uint32]string, error) { | ||
| groupDirs := make(map[uint32]string) | ||
|
|
||
| entries, err := os.ReadDir(basePath) | ||
| if err != nil { | ||
| return nil, err | ||
| } | ||
|
|
||
| for _, entry := range entries { | ||
| if entry.IsDir() { | ||
| groupID, err := strconv.Atoi(entry.Name()) | ||
| if err == nil { | ||
| pDir := filepath.Join(basePath, entry.Name(), "p") | ||
| if _, err := os.Stat(pDir); err == nil { | ||
| groupDirs[uint32(groupID)] = pDir | ||
| } | ||
| } | ||
| } | ||
| } | ||
|
|
||
| return groupDirs, nil | ||
| } | ||
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Check failure
Code scanning / CodeQL
Incorrect conversion between integer types High
Copilot Autofix AI 7 days ago
To fix the problem, we need to ensure that the value parsed from the directory name is within the valid range for
uint32before performing the conversion. We can achieve this by usingstrconv.ParseUintwith a bit size of 32, which directly parses the string into auint32value. This approach avoids the need for an additional bounds check.