resolve review comments

Author: shivaji-kharse

dgraph/cmd/dgraphimport/import_client.go

@@ -39,32 +39,32 @@ func Import(ctx context.Context, endpoint string, opts grpc.DialOption, bulkOutD
 	if err != nil {
 		return err
 	}
-	resp, err := startSnapshotStream(ctx, dg)
+	resp, err := startPDirStream(ctx, dg)
 	if err != nil {
 		return err
 	}
 
-	return sendSnapshot(ctx, dg, bulkOutDir, resp.Groups)
+	return sendPDir(ctx, dg, bulkOutDir, resp.Groups)
 }
 
-// startSnapshotStream initiates a snapshot stream session with the Dgraph server.
-func startSnapshotStream(ctx context.Context, dc apiv25.DgraphClient) (*apiv25.InitiateSnapshotStreamResponse, error) {
-	glog.Info("Initiating snapshot stream")
-	req := &apiv25.InitiateSnapshotStreamRequest{}
-	resp, err := dc.InitiateSnapshotStream(ctx, req)
+// startPDirStream initiates a snapshot stream session with the Dgraph server.
+func startPDirStream(ctx context.Context, dc apiv25.DgraphClient) (*apiv25.InitiatePDirStreamResponse, error) {
+	glog.Info("Initiating pdir stream")
+	req := &apiv25.InitiatePDirStreamRequest{}
+	resp, err := dc.InitiatePDirStream(ctx, req)
 	if err != nil {
-		glog.Errorf("failed to initiate snapshot stream: %v", err)
-		return nil, fmt.Errorf("failed to initiate snapshot stream: %v", err)
+		glog.Errorf("failed to initiate pdir stream: %v", err)
+		return nil, fmt.Errorf("failed to initiate pdir stream: %v", err)
 	}
-	glog.Info("Snapshot stream initiated successfully")
+	glog.Info("Pdir stream initiated successfully")
 	return resp, nil
 }
 
-// sendSnapshot takes a p directory and a set of group IDs and streams the data from the
+// sendPDir takes a p directory and a set of group IDs and streams the data from the
 // p directory to the corresponding group IDs. It first scans the provided directory for
 // subdirectories named with numeric group IDs.
-func sendSnapshot(ctx context.Context, dg apiv25.DgraphClient, baseDir string, groups []uint32) error {
-	glog.Infof("Starting to stream snapshots from directory: %s", baseDir)
+func sendPDir(ctx context.Context, dg apiv25.DgraphClient, baseDir string, groups []uint32) error {
+	glog.Infof("Starting to stream pdir from directory: %s", baseDir)
 
 	errG, ctx := errgroup.WithContext(ctx)
 	for _, group := range groups {
@@ -88,7 +88,7 @@ func sendSnapshot(ctx context.Context, dg apiv25.DgraphClient, baseDir string, g
 		return err
 	}
 
-	glog.Infof("Completed streaming all snapshots")
+	glog.Infof("Completed streaming all pdirs")
 	return nil
 }
 
@@ -98,9 +98,9 @@ func streamData(ctx context.Context, dg apiv25.DgraphClient, pdir string, groupI
 	glog.Infof("Opening stream for group %d from directory %s", groupId, pdir)
 
 	// Initialize stream with the server
-	out, err := dg.StreamSnapshot(ctx)
+	out, err := dg.StreamPDir(ctx)
 	if err != nil {
-		return fmt.Errorf("failed to start snapshot stream for group %d: %w", groupId, err)
+		return fmt.Errorf("failed to start pdir stream for group %d: %w", groupId, err)
 	}
 
 	// Open the BadgerDB instance at the specified directory
@@ -118,7 +118,7 @@ func streamData(ctx context.Context, dg apiv25.DgraphClient, pdir string, groupI
 
 	// Send group ID as the first message in the stream
 	glog.Infof("Sending group ID [%d] to server", groupId)
-	groupReq := &apiv25.StreamSnapshotRequest{GroupId: groupId}
+	groupReq := &apiv25.StreamPDirRequest{GroupId: groupId}
 	if err := out.Send(groupReq); err != nil {
 		return fmt.Errorf("failed to send group ID [%d]: %w", groupId, err)
 	}
@@ -129,8 +129,8 @@ func streamData(ctx context.Context, dg apiv25.DgraphClient, pdir string, groupI
 	stream.LogPrefix = fmt.Sprintf("Sending P dir for group [%d]", groupId)
 	stream.KeyToList = nil
 	stream.Send = func(buf *z.Buffer) error {
-		kvs := &apiv25.KVS{Data: buf.Bytes()}
-		if err := out.Send(&apiv25.StreamSnapshotRequest{Pairs: kvs}); err != nil && !errors.Is(err, io.EOF) {
+		p := &apiv25.StreamPacket{Data: buf.Bytes()}
+		if err := out.Send(&apiv25.StreamPDirRequest{StreamPacket: p}); err != nil && !errors.Is(err, io.EOF) {
 			return fmt.Errorf("failed to send data chunk: %w", err)
 		}
 		return nil
@@ -143,9 +143,9 @@ func streamData(ctx context.Context, dg apiv25.DgraphClient, pdir string, groupI
 
 	// Send the final 'done' signal to mark completion
 	glog.Infof("Sending completion signal for group [%d]", groupId)
-	done := &apiv25.KVS{Done: true}
+	done := &apiv25.StreamPacket{Done: true}
 
-	if err := out.Send(&apiv25.StreamSnapshotRequest{Pairs: done}); err != nil && !errors.Is(err, io.EOF) {
+	if err := out.Send(&apiv25.StreamPDirRequest{StreamPacket: done}); err != nil && !errors.Is(err, io.EOF) {
 		return fmt.Errorf("failed to send 'done' signal for group [%d]: %w", groupId, err)
 	}
 	// Wait for acknowledgment from the server

dgraph/cmd/dgraphimport/import_test.go

@@ -71,7 +71,7 @@ func TestDrainModeAfterStartSnapshotStream(t *testing.T) {
 			dc, err := newClient(url, grpc.WithTransportCredentials(insecure.NewCredentials()))
 			require.NoError(t, err)
 
-			resp, err := startSnapshotStream(context.Background(), dc)
+			resp, err := startPDirStream(context.Background(), dc)
 			require.NoError(t, err)
 
 			require.Equal(t, tt.expectedNum, len(resp.Groups))

edgraph/server.go

@@ -1810,21 +1810,21 @@ func validateNamespace(ctx context.Context, tc *api.TxnContext) error {
 	return nil
 }
 
-func (s *ServerV25) InitiateSnapshotStream(ctx context.Context,
-	c *apiv25.InitiateSnapshotStreamRequest) (v *apiv25.InitiateSnapshotStreamResponse, err error) {
+func (s *ServerV25) InitiatePDirStream(ctx context.Context,
+	c *apiv25.InitiatePDirStreamRequest) (v *apiv25.InitiatePDirStreamResponse, err error) {
 
 	drainMode := &pb.DrainModeRequest{State: true}
 	groups, err := worker.ProposeDrain(ctx, drainMode)
 	if err != nil {
 		return nil, err
 	}
 
-	resp := &apiv25.InitiateSnapshotStreamResponse{Groups: groups}
+	resp := &apiv25.InitiatePDirStreamResponse{Groups: groups}
 
 	return resp, nil
 }
 
-func (s *ServerV25) StreamSnapshot(stream apiv25.Dgraph_StreamSnapshotServer) error {
+func (s *ServerV25) StreamPDir(stream apiv25.Dgraph_StreamPDirServer) error {
 	if err := worker.InStream(stream); err != nil {
 		return err
 	}

go.mod

@@ -9,7 +9,7 @@ require (
 	github.com/Masterminds/semver/v3 v3.3.1
 	github.com/blevesearch/bleve/v2 v2.5.0
 	github.com/dgraph-io/badger/v4 v4.7.0
-	github.com/dgraph-io/dgo/v250 v250.0.0-20250422095606-337c8b78908a
+	github.com/dgraph-io/dgo/v250 v250.0.0-20250511125955-0567006c46cc
 	github.com/dgraph-io/gqlgen v0.13.2
 	github.com/dgraph-io/gqlparser/v2 v2.2.2
 	github.com/dgraph-io/graphql-transport-ws v0.0.0-20210511143556-2cef522f1f15

go.sum

@@ -122,8 +122,8 @@ github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc h1:U9qPSI2PIWSS1
 github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/dgraph-io/badger/v4 v4.7.0 h1:Q+J8HApYAY7UMpL8d9owqiB+odzEc0zn/aqOD9jhc6Y=
 github.com/dgraph-io/badger/v4 v4.7.0/go.mod h1:He7TzG3YBy3j4f5baj5B7Zl2XyfNe5bl4Udl0aPemVA=
-github.com/dgraph-io/dgo/v250 v250.0.0-20250422095606-337c8b78908a h1:RhE9o87JCsfvNnZmpLQ3KFQ2LRXX/93B32zue9KxHm0=
-github.com/dgraph-io/dgo/v250 v250.0.0-20250422095606-337c8b78908a/go.mod h1:h1cPCTmI0GZWc770q608FkWEdQfC4iiWYAgvQGnvQ9A=
+github.com/dgraph-io/dgo/v250 v250.0.0-20250511125955-0567006c46cc h1:TACB/WzfNGQxgxWQelzg6ur/NWoxOmeesrjJwVZ+Zjk=
+github.com/dgraph-io/dgo/v250 v250.0.0-20250511125955-0567006c46cc/go.mod h1:m5zwXW3OyxYK4asn97VxPRcQBn+dYj0uzT2MIu2R4is=
 github.com/dgraph-io/gqlgen v0.13.2 h1:TNhndk+eHKj5qE7BenKKSYdSIdOGhLqxR1rCiMso9KM=
 github.com/dgraph-io/gqlgen v0.13.2/go.mod h1:iCOrOv9lngN7KAo+jMgvUPVDlYHdf7qDwsTkQby2Sis=
 github.com/dgraph-io/gqlparser/v2 v2.1.1/go.mod h1:MYS4jppjyx8b9tuUtjV7jU1UFZK6P9fvO8TsIsQtRKU=

protos/pb.proto

@@ -597,11 +597,7 @@ service Worker {
   rpc DeleteNamespace(DeleteNsRequest) returns (Status) {}
   rpc TaskStatus(TaskStatusRequest) returns (TaskStatusResponse) {}
   rpc ApplyDrainmode(DrainModeRequest) returns (Status) {}
-  rpc InternalStreamSnapshot(stream api.v25.StreamSnapshotRequest) returns (api.v25.StreamSnapshotResponse) {}
-}
-
-message ReceiveSnapshotKVRequest {
-    bool done = 1;
+  rpc InternalStreamPDir(stream api.v25.StreamPDirRequest) returns (api.v25.StreamPDirResponse) {}
 }
 
 message DrainModeRequest {