Make dgraph import work over the internet

This PR fixes the timeout issue when cloudflare is not happy
for just one way data send. It adds application level ACKs
to work around the 120s timeout for a HTTP response.

Additionally, it adds an argument to take p directory as
input for the dgraph import command.

Author: mangalaman93

Dockerfile

@@ -4,11 +4,13 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
     bzip2=1.0.8-5+b1 \
     git=1:2.39.5-0+deb12u2 \
     && rm -rf /var/lib/apt/lists/*
+ARG TARGETARCH=amd64
+ARG TARGETOS=linux
 WORKDIR /go/src/repo
 COPY go.mod go.sum ./
 RUN go mod download && go mod verify
 COPY . .
-RUN CGO_ENABLED=0 make
+RUN CGO_ENABLED=0 GOOS=${TARGETOS} GOARCH=${TARGETARCH} make
 
 ###################### Stage II ######################
 FROM ubuntu:24.04

dgraph/cmd/dgraphimport/import_client.go

@@ -127,23 +127,18 @@ func streamSnapshotForGroup(ctx context.Context, dc apiv2.DgraphClient, pdir str
 	if err != nil {
 		return fmt.Errorf("failed to start external snapshot stream for group %d: %w", groupId, err)
 	}
-
 	defer func() {
-		if _, err := out.CloseAndRecv(); err != nil {
-			glog.Errorf("failed to close the stream for group [%v]: %v", groupId, err)
-		}
-
-		glog.Infof("[import] Group [%v]: Received ACK ", groupId)
+		_ = out.CloseSend()
 	}()
 
 	// Open the BadgerDB instance at the specified directory
 	opt := badger.DefaultOptions(pdir)
+	opt.ReadOnly = true
 	ps, err := badger.OpenManaged(opt)
 	if err != nil {
 		glog.Errorf("failed to open BadgerDB at [%s]: %v", pdir, err)
 		return fmt.Errorf("failed to open BadgerDB at [%v]: %v", pdir, err)
 	}
-
 	defer func() {
 		if err := ps.Close(); err != nil {
 			glog.Warningf("[import] Error closing BadgerDB: %v", err)
@@ -154,17 +149,18 @@ func streamSnapshotForGroup(ctx context.Context, dc apiv2.DgraphClient, pdir str
 	glog.Infof("[import] Sending request for streaming external snapshot for group ID [%v]", groupId)
 	groupReq := &apiv2.StreamExtSnapshotRequest{GroupId: groupId}
 	if err := out.Send(groupReq); err != nil {
-		return fmt.Errorf("failed to send request for streaming external snapshot for group ID [%v] to the server: %w",
-			groupId, err)
+		return fmt.Errorf("failed to send request for group ID [%v] to the server: %w", groupId, err)
+	}
+	if _, err := out.Recv(); err != nil {
+		return fmt.Errorf("failed to receive response for group ID [%v] from the server: %w", groupId, err)
 	}
+	glog.Infof("[import] Group [%v]: Received ACK for sending group request", groupId)
 
 	// Configure and start the BadgerDB stream
 	glog.Infof("[import] Starting BadgerDB stream for group [%v]", groupId)
-
 	if err := streamBadger(ctx, ps, out, groupId); err != nil {
 		return fmt.Errorf("badger streaming failed for group [%v]: %v", groupId, err)
 	}
-
 	return nil
 }
 
@@ -180,6 +176,11 @@ func streamBadger(ctx context.Context, ps *badger.DB, out apiv2.Dgraph_StreamExt
 		if err := out.Send(&apiv2.StreamExtSnapshotRequest{Pkt: p}); err != nil && !errors.Is(err, io.EOF) {
 			return fmt.Errorf("failed to send data chunk: %w", err)
 		}
+		if _, err := out.Recv(); err != nil {
+			return fmt.Errorf("failed to receive response for group ID [%v] from the server: %w", groupId, err)
+		}
+		glog.Infof("[import] Group [%v]: Received ACK for sending data chunk", groupId)
+
 		return nil
 	}
 
@@ -196,5 +197,10 @@ func streamBadger(ctx context.Context, ps *badger.DB, out apiv2.Dgraph_StreamExt
 		return fmt.Errorf("failed to send 'done' signal for group [%d]: %w", groupId, err)
 	}
 
+	if _, err := out.Recv(); err != nil {
+		return fmt.Errorf("failed to receive response for group ID [%v] from the server: %w", groupId, err)
+	}
+	glog.Infof("[import] Group [%v]: Received ACK for sending completion signal", groupId)
+
 	return nil
 }

dgraph/cmd/dgraphimport/import_test.go

@@ -1,4 +1,4 @@
-//go:build integration
+//go:build integration2
 
 /*
  * SPDX-FileCopyrightText: © Hypermode Inc. <[email protected]>
@@ -75,7 +75,8 @@ func TestDrainModeAfterStartSnapshotStream(t *testing.T) {
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			conf := dgraphtest.NewClusterConfig().WithNumAlphas(tt.numAlphas).WithNumZeros(tt.numZeros).WithReplicas(tt.replicas)
+			conf := dgraphtest.NewClusterConfig().WithNumAlphas(tt.numAlphas).
+				WithNumZeros(tt.numZeros).WithReplicas(tt.replicas)
 			c, err := dgraphtest.NewLocalCluster(conf)
 			require.NoError(t, err)
 			defer func() { c.Cleanup(t.Failed()) }()
@@ -268,14 +269,12 @@ func runImportTest(t *testing.T, tt testcase) {
 			require.NoError(t, targetCluster.StopAlpha(alphaID))
 		}
 	}
-
 	if tt.err != "" {
 		err := Import(context.Background(), connectionString, outDir)
 		require.Error(t, err)
 		require.ErrorContains(t, err, tt.err)
 		return
 	}
-
 	require.NoError(t, Import(context.Background(), connectionString, outDir))
 
 	for group, alphas := range alphaGroups {
@@ -287,7 +286,6 @@ func runImportTest(t *testing.T, tt testcase) {
 	}
 
 	require.NoError(t, targetCluster.HealthCheck(false))
-
 	t.Log("Import completed")
 
 	for i := 0; i < tt.targetAlphas; i++ {
@@ -330,7 +328,9 @@ func setupBulkCluster(t *testing.T, numAlphas int, encrypted bool) (*dgraphtest.
 }
 
 // setupTargetCluster creates and starts a cluster that will receive the imported data
-func setupTargetCluster(t *testing.T, numAlphas, replicasFactor int) (*dgraphtest.LocalCluster, *dgraphapi.GrpcClient, func()) {
+func setupTargetCluster(t *testing.T, numAlphas, replicasFactor int) (
+	*dgraphtest.LocalCluster, *dgraphapi.GrpcClient, func()) {
+
 	conf := dgraphtest.NewClusterConfig().
 		WithNumAlphas(numAlphas).
 		WithNumZeros(3).

dgraph/cmd/dgraphimport/run.go

@@ -34,6 +34,7 @@ func init() {
 
 	flag := ImportCmd.Cmd.Flags()
 	flag.StringP("files", "f", "", "Location of *.rdf(.gz) or *.json(.gz) file(s) to load.")
+	flag.StringP("snapshot-dir", "p", "", "Location of p directory")
 	flag.StringP("schema", "s", "", "Location of DQL schema file.")
 	flag.StringP("graphql_schema", "g", "", "Location of the GraphQL schema file.")
 	flag.StringP("graphql-schema", "", "", "Location of the GraphQL schema file.")
@@ -72,6 +73,17 @@ func run() {
 		os.Exit(1)
 	}
 
+	// if snapshot p directory is already provided, there is no need to run bulk loader
+	if ImportCmd.Conf.GetString("snapshot-dir") != "" {
+		connStr := ImportCmd.Conf.GetString("conn-str")
+		snapshotDir := ImportCmd.Conf.GetString("snapshot-dir")
+		if err := Import(context.Background(), connStr, snapshotDir); err != nil {
+			fmt.Println("Failed to import data:", err)
+			os.Exit(1)
+		}
+		return
+	}
+
 	cacheSize := 64 << 20 // These are the default values. User can overwrite them using --badger.
 	cacheDefaults := fmt.Sprintf("indexcachesize=%d; blockcachesize=%d; ",
 		(70*cacheSize)/100, (30*cacheSize)/100)

dgraphtest/load.go

@@ -503,8 +503,12 @@ func (c *LocalCluster) BulkLoad(opts BulkOpts) error {
 		args = append(args, "-g", strings.Join(opts.GQLSchemaFiles, ","))
 	}
 
+	dgraphCmdPath := os.Getenv("DGRAPH_CMD_PATH")
+	if dgraphCmdPath == "" {
+		dgraphCmdPath = filepath.Join(c.tempBinDir, "dgraph")
+	}
 	log.Printf("[INFO] running bulk loader with args: [%v]", strings.Join(args, " "))
-	cmd := exec.Command(filepath.Join(c.tempBinDir, "dgraph"), args...)
+	cmd := exec.Command(dgraphCmdPath, args...)
 	if out, err := cmd.CombinedOutput(); err != nil {
 		return errors.Wrapf(err, "error running bulk loader: %v", string(out))
 	} else {

edgraph/server.go

@@ -1829,18 +1829,20 @@ func (s *ServerV25) UpdateExtSnapshotStreamingState(ctx context.Context,
 
 	groups, err := worker.ProposeDrain(ctx, req)
 	if err != nil {
+		glog.Errorf("[import] failed to propose drain mode: %v", err)
 		return nil, err
 	}
 
-	resp := &apiv2.UpdateExtSnapshotStreamingStateResponse{Groups: groups}
-
-	return resp, nil
+	return &apiv2.UpdateExtSnapshotStreamingStateResponse{Groups: groups}, nil
 }
 
 func (s *ServerV25) StreamExtSnapshot(stream apiv2.Dgraph_StreamExtSnapshotServer) error {
 	defer x.ExtSnapshotStreamingState(false)
-
-	return worker.InStream(stream)
+	if err := worker.InStream(stream); err != nil {
+		glog.Errorf("[import] failed to stream external snapshot: %v", err)
+		return err
+	}
+	return nil
 }
 
 // CommitOrAbort commits or aborts a transaction.

go.mod

@@ -9,7 +9,7 @@ require (
 	github.com/Masterminds/semver/v3 v3.4.0
 	github.com/blevesearch/bleve/v2 v2.5.2
 	github.com/dgraph-io/badger/v4 v4.7.0
-	github.com/dgraph-io/dgo/v250 v250.0.0-preview4.0.20250619041351-4a519e53fb9d
+	github.com/dgraph-io/dgo/v250 v250.0.0-preview4.0.20250709182152-32901102e0d0
 	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

@@ -128,8 +128,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-preview4.0.20250619041351-4a519e53fb9d h1:9PLyvZY1Nih05g+2womk+kNnX3Gb20kx5BsK3foA5a8=
-github.com/dgraph-io/dgo/v250 v250.0.0-preview4.0.20250619041351-4a519e53fb9d/go.mod h1:gLr7uM+x/8PjSQJ4Ca9kfQF15uBzruDzRK3bnELt3vE=
+github.com/dgraph-io/dgo/v250 v250.0.0-preview4.0.20250709182152-32901102e0d0 h1:gvNB/M+LrjkX9c4QJCPdODVoxgMt3zlIBW9i8xVfoYo=
+github.com/dgraph-io/dgo/v250 v250.0.0-preview4.0.20250709182152-32901102e0d0/go.mod h1:gLr7uM+x/8PjSQJ4Ca9kfQF15uBzruDzRK3bnELt3vE=
 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/Makefile

@@ -53,6 +53,7 @@ regenerate: tidy-deps copy-protos check clean
 	@protoc \
 		--proto_path=/usr/local/include \
 		--proto_path=/usr/include \
+		--proto_path=/opt/homebrew/include/google/protobuf \
 		--proto_path=${PROTO_PATH} \
 		--go_out=pb --go-grpc_out=pb \
 		--go_opt=paths=source_relative \

protos/depcheck.sh

@@ -37,7 +37,7 @@ function CompareSemVer() {
 
 function CheckProtobufIncludes() {
 	echo -n "Checking for directory /usr/include/google/protobuf or /usr/local/include/google/protobuf... "
-	if !([[ -d /usr/include/google/protobuf ]] || [[ -d /usr/local/include/google/protobuf ]]); then
+	if !([[ -d /usr/include/google/protobuf ]] || [[ -d /usr/local/include/google/protobuf ]] || [[ -d /opt/homebrew/include/google/protobuf ]]); then
 		echo "FAIL" >&2
 		echo "Missing protobuf headers in /usr/include/google/protobuf or /usr/local/include/google/protobuf:" \
 			"directory not found." >&2