resolve review comments

Author: shivaji-kharse

dgraph/cmd/dgraphimport/import_client.go

@@ -7,6 +7,7 @@ package dgraphimport
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"io"
 	"math"
@@ -22,11 +23,11 @@ import (
 	"google.golang.org/grpc"
 )
 
-// NewClient creates a new import client with the specified endpoint and gRPC options.
+// newClient creates a new import client with the specified endpoint and gRPC options.
 func newClient(endpoint string, opts grpc.DialOption) (apiv25.DgraphClient, error) {
 	conn, err := grpc.NewClient(endpoint, opts)
 	if err != nil {
-		return nil, fmt.Errorf("failed to connect to endpoint [%s]: %w", endpoint, err)
+		return nil, fmt.Errorf("Failed to connect to endpoint [%s]: %w", endpoint, err)
 	}
 
 	glog.Infof("Successfully connected to Dgraph endpoint: %s", endpoint)
@@ -46,20 +47,20 @@ func Import(ctx context.Context, endpoint string, opts grpc.DialOption, bulkOutD
 	return sendSnapshot(ctx, dg, bulkOutDir, resp.Groups)
 }
 
-// StartSnapshotStream initiates a snapshot stream session with the Dgraph server.
+// startSnapshotStream initiates a snapshot stream session with the Dgraph server.
 func startSnapshotStream(ctx context.Context, dc apiv25.DgraphClient) (*apiv25.InitiateSnapshotStreamResponse, error) {
-	glog.V(2).Infof("Initiating snapshot stream")
+	glog.Info("Initiating snapshot stream")
 	req := &apiv25.InitiateSnapshotStreamRequest{}
 	resp, err := dc.InitiateSnapshotStream(ctx, req)
 	if err != nil {
 		glog.Errorf("Failed to initiate snapshot stream: %v", err)
-		return nil, err
+		return nil, fmt.Errorf("Failed to initiate snapshot stream: %v", err)
 	}
-	glog.Infof("Snapshot stream initiated successfully")
+	glog.Info("Snapshot stream initiated successfully")
 	return resp, nil
 }
 
-// SendSnapshot takes a p directory and a set of group IDs and streams the data from the
+// sendSnapshot 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 {
@@ -71,11 +72,12 @@ func sendSnapshot(ctx context.Context, dg apiv25.DgraphClient, baseDir string, g
 		errG.Go(func() error {
 			pDir := filepath.Join(baseDir, fmt.Sprintf("%d", group-1), "p")
 			if _, err := os.Stat(pDir); err != nil {
-				return fmt.Errorf("p directory does not exist for group %d: %s", group, pDir)
+				return fmt.Errorf("p directory does not exist for group [%d]: [%s]", group, pDir)
 			}
 
-			glog.Infof("Streaming data for group %d from directory: %s", group, pDir)
+			glog.Infof("Streaming data for group [%d] from directory: [%s]", group, pDir)
 			if err := streamData(ctx, dg, pDir, group); err != nil {
+				glog.Errorf("Failed to stream data for groups [%v] from directory: [%s]: %v", group, pDir, err)
 				return err
 			}
 
@@ -98,55 +100,59 @@ func streamData(ctx context.Context, dg apiv25.DgraphClient, pdir string, groupI
 	// Initialize stream with the server
 	out, err := dg.StreamSnapshot(ctx)
 	if err != nil {
-		return fmt.Errorf("failed to start snapshot stream: %w", err)
+		return fmt.Errorf("failed to start snapshot stream for group %d: %w", groupId, err)
 	}
 
 	// Open the BadgerDB instance at the specified directory
 	opt := badger.DefaultOptions(pdir)
 	ps, err := badger.OpenManaged(opt)
 	if err != nil {
-		return fmt.Errorf("failed to open BadgerDB at %s: %w", pdir, err)
+		return fmt.Errorf("failed to open BadgerDB at [%s]: %w", pdir, err)
 	}
-	defer ps.Close()
+
+	defer func() {
+		if err := ps.Close(); err != nil {
+			glog.Warningf("Error closing BadgerDB: %v", err)
+		}
+	}()
 
 	// Send group ID as the first message in the stream
-	glog.V(2).Infof("Sending group ID %d to server", groupId)
+	glog.V(2).Infof("Sending group ID [%d] to server", groupId)
 	groupReq := &apiv25.StreamSnapshotRequest{GroupId: groupId}
 	if err := out.Send(groupReq); err != nil {
-		return fmt.Errorf("failed to send group ID %d: %w", groupId, err)
+		return fmt.Errorf("failed to send group ID [%d]: %w", groupId, err)
 	}
 
 	// Configure and start the BadgerDB stream
-	glog.V(2).Infof("Starting BadgerDB stream for group %d", groupId)
+	glog.V(2).Infof("Starting BadgerDB stream for group [%d]", groupId)
 	stream := ps.NewStreamAt(math.MaxUint64)
-	stream.LogPrefix = fmt.Sprintf("Sending P dir (group %d)", groupId)
+	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 && err != io.EOF {
+		if err := out.Send(&apiv25.StreamSnapshotRequest{Pairs: kvs}); err != nil && !errors.Is(err, io.EOF) {
 			return fmt.Errorf("failed to send data chunk: %w", err)
 		}
 		return nil
 	}
 
 	// Execute the stream process
 	if err := stream.Orchestrate(ctx); err != nil {
-		return fmt.Errorf("stream orchestration failed for group %d: %w", groupId, err)
+		return fmt.Errorf("stream orchestration failed for group [%d]: %w", groupId, err)
 	}
 
 	// Send the final 'done' signal to mark completion
-	glog.V(2).Infof("Sending completion signal for group %d", groupId)
+	glog.V(2).Infof("Sending completion signal for group [%d]", groupId)
 	done := &apiv25.KVS{Done: true}
 
-	if err := out.Send(&apiv25.StreamSnapshotRequest{Pairs: done}); err != nil && err != io.EOF {
-		return fmt.Errorf("failed to send 'done' signal for group %d: %w", groupId, err)
+	if err := out.Send(&apiv25.StreamSnapshotRequest{Pairs: 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
-	ack, err := out.CloseAndRecv()
-	if err != nil {
-		return fmt.Errorf("failed to receive ACK for group %d: %w", groupId, err)
+	if _, err := out.CloseAndRecv(); err != nil {
+		return fmt.Errorf("failed to receive ACK for group [%d]: %w", groupId, err)
 	}
-	glog.Infof("Group %d: Received ACK with message: %v", groupId, ack.Done)
+	glog.Infof("Group [%d]: Received ACK ", groupId)
 
 	return nil
 }

dgraph/cmd/dgraphimport/import_test.go

@@ -11,11 +11,11 @@ import (
 	"encoding/json"
 	"path/filepath"
 	"testing"
-	"time"
 
 	"github.com/hypermodeinc/dgraph/v25/dgraphapi"
 	"github.com/hypermodeinc/dgraph/v25/dgraphtest"
 	"github.com/hypermodeinc/dgraph/v25/systest/1million/common"
+
 	"github.com/stretchr/testify/require"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/credentials/insecure"
@@ -191,13 +191,7 @@ func verifyImportResults(t *testing.T, gc *dgraphapi.GrpcClient) {
 	}
 
 	for _, tt := range common.OneMillionTCs {
-		ctx, cancel := context.WithTimeout(context.Background(), 10*time.Minute)
 		resp, err := gc.Query(tt.Query)
-		cancel()
-
-		if ctx.Err() == context.DeadlineExceeded {
-			t.Fatal("aborting test due to query timeout")
-		}
 		require.NoError(t, err)
 		require.NoError(t, dgraphapi.CompareJSON(tt.Resp, string(resp.Json)))
 	}

edgraph/server.go

@@ -1812,6 +1812,7 @@ func validateNamespace(ctx context.Context, tc *api.TxnContext) error {
 
 func (s *ServerV25) InitiateSnapshotStream(ctx context.Context,
 	c *apiv25.InitiateSnapshotStreamRequest) (v *apiv25.InitiateSnapshotStreamResponse, err error) {
+
 	drainMode := &pb.DrainModeRequest{State: true}
 	groups, err := worker.ProposeDrain(ctx, drainMode)
 	if err != nil {

protos/Makefile

@@ -41,6 +41,8 @@ copy-protos:
 	 ${TMPDIR}/src/${BADGER_PATH}/pb/pb.proto
 	@cp $(shell go list -m -f "{{.Dir}}" ${DGO_PATH})/protos/api.proto \
 	 ${TMPDIR}/src/${DGO_PATH}/protos/api.proto
+	@cp $(shell go list -m -f "{{.Dir}}" ${DGO_PATH})/protos/api.v25.proto \
+	 ${TMPDIR}/src/${DGO_PATH}/protos/api.v25.proto
 
 .PHONY: patch-pb
 patch-pb:
@@ -56,6 +58,7 @@ regenerate: tidy-deps copy-protos check clean
 		--go_opt=paths=source_relative \
 		--go-grpc_opt=paths=source_relative \
 		--go_opt=M{DGO_PATH}/protos/api.proto={DGO_PATH}/protos/api \
+		--go_opt=M{DGO_PATH}/protos/api.v25.proto={DGO_PATH}/protos/api.v25 \
 		pb.proto
 	@$(MAKE) patch-pb
 	@rm -rf ${TMPDIR}

protos/pb.proto

@@ -13,6 +13,7 @@ syntax = "proto3";
 package pb;
 
 import "github.com/dgraph-io/dgo/v250/protos/api.proto";
+import "github.com/dgraph-io/dgo/v250/protos/api.v25.proto";
 import "github.com/dgraph-io/badger/v4/pb/pb.proto";
 import "google/protobuf/descriptor.proto";
 
@@ -596,7 +597,7 @@ service Worker {
   rpc DeleteNamespace(DeleteNsRequest) returns (Status) {}
   rpc TaskStatus(TaskStatusRequest) returns (TaskStatusResponse) {}
   rpc ApplyDrainmode(DrainModeRequest) returns (Status) {}
-  rpc InternalStreamSnapshot(stream KVS) returns (ReceiveSnapshotKVRequest) {}
+  rpc InternalStreamSnapshot(stream api.v25.StreamSnapshotRequest) returns (api.v25.StreamSnapshotResponse) {}
 }
 
 message ReceiveSnapshotKVRequest {