resolve review comments

Author: shivaji-kharse

dgraph/cmd/alpha/run_test.go

@@ -1640,14 +1640,14 @@ func TestMain(m *testing.M) {
 
 	x.Panic(dg.Login(context.Background(), dgraphapi.DefaultUser, dgraphapi.DefaultPassword))
 
-	alphaGrpcPort, err := c.GetAlphaGrpcPublicPort()
+	alphaGrpcPort, err := c.GetAlphaGrpcPublicPort(0)
 	x.Panic(err)
 
 	alphaSockAdd = "0.0.0.0:" + alphaGrpcPort
 	alphaSockAddHttp, err := c.GetAlphaHttpPublicPort(0)
 	x.Panic(err)
 	addr = "http://0.0.0.0:" + alphaSockAddHttp
-	zeroSockAdd, err := c.GetZeroGrpcPublicPort()
+	zeroSockAdd, err := c.GetZeroGrpcPublicPort(0)
 	x.Panic(err)
 
 	// Increment lease, so that mutations work.

dgraph/cmd/bulk/loader.go

@@ -43,7 +43,7 @@ type options struct {
 	DataFormat       string
 	SchemaFile       string
 	GqlSchemaFile    string
-	OutDir           string
+	outDirNoMount    string
 	ReplaceOutDir    bool
 	TmpDir           string
 	NumGoroutines    int

dgraph/cmd/bulk/run.go

@@ -141,7 +141,7 @@ func run() {
 		GqlSchemaFile:    Bulk.Conf.GetString("graphql_schema"),
 		Encrypted:        Bulk.Conf.GetBool("encrypted"),
 		EncryptedOut:     Bulk.Conf.GetBool("encrypted_out"),
-		OutDir:           Bulk.Conf.GetString("out"),
+		outDirNoMount:    Bulk.Conf.GetString("out"),
 		ReplaceOutDir:    Bulk.Conf.GetBool("replace_out"),
 		TmpDir:           Bulk.Conf.GetString("tmp"),
 		NumGoroutines:    Bulk.Conf.GetInt("num_go_routines"),
@@ -259,8 +259,8 @@ func run() {
 
 	// Make sure it's OK to create or replace the directory specified with the --out option.
 	// It is always OK to create or replace the default output directory.
-	if opt.OutDir != defaultOutDir && !opt.ReplaceOutDir {
-		err := x.IsMissingOrEmptyDir(opt.OutDir)
+	if opt.outDirNoMount != defaultOutDir && !opt.ReplaceOutDir {
+		err := x.IsMissingOrEmptyDir(opt.outDirNoMount)
 		if err == nil {
 			fmt.Fprintf(os.Stderr, "Output directory exists and is not empty."+
 				" Use --replace_out to overwrite it.\n")
@@ -271,9 +271,9 @@ func run() {
 	}
 
 	// Delete and recreate the output dirs to ensure they are empty.
-	x.Check(os.RemoveAll(opt.OutDir))
+	x.Check(os.RemoveAll(opt.outDirNoMount))
 	for i := range opt.ReduceShards {
-		dir := filepath.Join(opt.OutDir, strconv.Itoa(i), "p")
+		dir := filepath.Join(opt.outDirNoMount, strconv.Itoa(i), "p")
 		x.Check(os.MkdirAll(dir, 0700))
 		opt.shardOutputDirs = append(opt.shardOutputDirs, dir)
 

dgraph/cmd/dgraphimport/import_client.go

@@ -20,28 +20,35 @@ import (
 	"google.golang.org/grpc"
 )
 
-// Client represents a Dgraph import client that handles snapshot streaming.
-type Client struct {
-	opts grpc.DialOption
-	dg   apiv25.DgraphClient
-}
-
 // NewClient creates a new import client with the specified endpoint and gRPC options.
-func NewClient(endpoint string, opts grpc.DialOption) (*Client, error) {
+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)
 	}
 
 	glog.Infof("Successfully connected to Dgraph endpoint: %s", endpoint)
-	return &Client{dg: apiv25.NewDgraphClient(conn), opts: opts}, nil
+	return apiv25.NewDgraphClient(conn), nil
+}
+
+func Import(ctx context.Context, endpoint string, opts grpc.DialOption, outDir string) error {
+	dg, err := newClient(endpoint, opts)
+	if err != nil {
+		return err
+	}
+	resp, err := startSnapshotStream(ctx, dg)
+	if err != nil {
+		return err
+	}
+
+	return sendSnapshot(ctx, dg, outDir, resp.Groups)
 }
 
 // StartSnapshotStream initiates a snapshot stream session with the Dgraph server.
-func (c *Client) StartSnapshotStream(ctx context.Context) (*apiv25.InitiateSnapshotStreamResponse, error) {
+func startSnapshotStream(ctx context.Context, dg apiv25.DgraphClient) (*apiv25.InitiateSnapshotStreamResponse, error) {
 	glog.V(2).Infof("Initiating snapshot stream")
 	req := &apiv25.InitiateSnapshotStreamRequest{}
-	resp, err := c.dg.InitiateSnapshotStream(ctx, req)
+	resp, err := dg.InitiateSnapshotStream(ctx, req)
 	if err != nil {
 		glog.Errorf("Failed to initiate snapshot stream: %v", err)
 		return nil, err
@@ -53,7 +60,7 @@ func (c *Client) StartSnapshotStream(ctx context.Context) (*apiv25.InitiateSnaps
 // SendSnapshot takes a p directory and a set of group IDs and streams the data from the
 // p directory to the corresponding group IDs. The function will skip any groups that do not
 // have a corresponding p directory.
-func (c *Client) SendSnapshot(ctx context.Context, pDir string, groups []uint32) error {
+func sendSnapshot(ctx context.Context, dg apiv25.DgraphClient, pDir string, groups []uint32) error {
 	glog.Infof("Starting to stream snapshots from directory: %s", pDir)
 
 	// Get mapping of group IDs to their respective p directories
@@ -77,7 +84,7 @@ func (c *Client) SendSnapshot(ctx context.Context, pDir string, groups []uint32)
 		}
 
 		glog.Infof("Streaming data for group %d from directory: %s", group, pDir)
-		err = streamData(ctx, c.dg, pDir, group)
+		err = streamData(ctx, dg, pDir, group)
 		if err != nil {
 			glog.Errorf("Failed to stream snapshot for group %d: %v", group, err)
 			return err

dgraph/cmd/dgraphimport/import_test.go

@@ -117,13 +117,14 @@ func TestStartSnapshotStreamForSingleNode(t *testing.T) {
 	require.NoError(t, gc.Login(context.Background(),
 		dgraphapi.DefaultUser, dgraphapi.DefaultPassword))
 
-	pubPort, err := c.GetAlphaGrpcPublicPort()
+	url, err := c.GetAlphaGrpcEndpoint(0)
+	require.NoError(t, err)
+	dg, err := newClient(url, grpc.WithTransportCredentials(insecure.NewCredentials()))
 	require.NoError(t, err)
-	client, err := NewClient(dgraphtest.GetLocalHostUrl(pubPort, ""), grpc.WithTransportCredentials(insecure.NewCredentials()))
 
 	require.NoError(t, err)
 
-	resp, err := client.StartSnapshotStream(context.Background())
+	resp, err := startSnapshotStream(context.Background(), dg)
 	require.NoError(t, err)
 
 	require.Equal(t, len(resp.Groups), 1)
@@ -141,12 +142,12 @@ func TestStartSnapshotStreamForHA(t *testing.T) {
 	defer func() { c.Cleanup(t.Failed()) }()
 	require.NoError(t, c.Start())
 
-	pubPort, err := c.GetAlphaGrpcPublicPort()
+	url, err := c.GetAlphaGrpcEndpoint(0)
 	require.NoError(t, err)
-	client, err := NewClient(dgraphtest.GetLocalHostUrl(pubPort, ""), grpc.WithTransportCredentials(insecure.NewCredentials()))
+	dg, err := newClient(url, grpc.WithTransportCredentials(insecure.NewCredentials()))
 	require.NoError(t, err)
 
-	resp, err := client.StartSnapshotStream(context.Background())
+	resp, err := startSnapshotStream(context.Background(), dg)
 	require.NoError(t, err)
 
 	require.Equal(t, len(resp.Groups), 1)
@@ -169,13 +170,12 @@ func TestStartSnapshotStreamForHASharded(t *testing.T) {
 	defer func() { c.Cleanup(t.Failed()) }()
 	require.NoError(t, c.Start())
 
-	pubPort, err := c.GetAlphaGrpcPublicPort()
+	url, err := c.GetAlphaGrpcEndpoint(0)
 	require.NoError(t, err)
-	client, err := NewClient(dgraphtest.GetLocalHostUrl(pubPort, ""),
-		grpc.WithTransportCredentials(insecure.NewCredentials()))
+	dg, err := newClient(url, grpc.WithTransportCredentials(insecure.NewCredentials()))
 	require.NoError(t, err)
 
-	resp, err := client.StartSnapshotStream(context.Background())
+	resp, err := startSnapshotStream(context.Background(), dg)
 	require.NoError(t, err)
 	require.Equal(t, len(resp.Groups), 2)
 
@@ -217,18 +217,11 @@ func runImportTest(t *testing.T, bulkAlphas, targetAlphas, replicasFactor int) {
 	defer func() { targetCluster.Cleanup(t.Failed()) }()
 	defer gcCleanup()
 
-	importClient, err := createImportClient(t, targetCluster)
-	require.NoError(t, err)
-
-	resp, err := importClient.StartSnapshotStream(context.Background())
+	url, err := targetCluster.GetAlphaGrpcEndpoint(0)
 	require.NoError(t, err)
-
-	_, err = gc.Query("schema{}")
-	require.Error(t, err)
-	require.ErrorContains(t, err, "the server is in draining mode")
-
 	outDir := filepath.Join(baseDir, "out")
-	err = importClient.SendSnapshot(context.Background(), outDir, resp.Groups)
+
+	err = Import(context.Background(), url, grpc.WithTransportCredentials(insecure.NewCredentials()), outDir)
 	require.NoError(t, err)
 
 	verifyImportResults(t, gc)
@@ -246,17 +239,13 @@ func setupBulkCluster(t *testing.T, numAlphas int) (*dgraphtest.LocalCluster, st
 	cluster, err := dgraphtest.NewLocalCluster(bulkConf)
 	require.NoError(t, err)
 
-	// Download and prepare data files
-	dataPath, err := dgraphtest.DownloadDataFiles()
-	require.NoError(t, err)
 	require.NoError(t, cluster.StartZero(0))
 
 	// Perform bulk load
-	schemaFile := filepath.Join(dataPath, "1million.schema")
-	rdfFile := filepath.Join(dataPath, "1million.rdf.gz")
+	oneMillion := dgraphtest.GetDataset(dgraphtest.OneMillionDataset)
 	opts := dgraphtest.BulkOpts{
-		DataFiles:   []string{rdfFile},
-		SchemaFiles: []string{schemaFile},
+		DataFiles:   []string{oneMillion.DataFilePath()},
+		SchemaFiles: []string{oneMillion.SchemaPath()},
 		OutDir:      filepath.Join(baseDir, "out"),
 	}
 	require.NoError(t, cluster.BulkLoad(opts))
@@ -282,15 +271,6 @@ func setupTargetCluster(t *testing.T, numAlphas, replicasFactor int) (*dgraphtes
 	return cluster, gc, cleanup
 }
 
-// createImportClient creates a new client for import operations
-func createImportClient(t *testing.T, cluster *dgraphtest.LocalCluster) (*Client, error) {
-	pubPort, err := cluster.GetAlphaGrpcPublicPort()
-	require.NoError(t, err)
-
-	return NewClient(dgraphtest.GetLocalHostUrl(pubPort, ""),
-		grpc.WithTransportCredentials(insecure.NewCredentials()))
-}
-
 // verifyImportResults validates the result of an import operation
 func verifyImportResults(t *testing.T, gc *dgraphapi.GrpcClient) {
 	// Check schema after streaming process

dgraph/cmd/zero/zero.go

@@ -860,18 +860,3 @@ func (s *Server) latestMembershipState(ctx context.Context) (*pb.MembershipState
 	}
 	return ms, nil
 }
-
-func (s *Server) ApplyDrainmode(ctx context.Context, req *pb.Drainmode) (*pb.Status, error) {
-	knownGroups := s.KnownGroups()
-
-	for _, grp := range knownGroups {
-		pl := s.Leader(grp)
-		wc := pb.NewWorkerClient(pl.Get())
-		in := &pb.Drainmode{State: req.State}
-
-		if status, err := wc.ApplyDrainmode(ctx, in); err != nil {
-			return status, errors.Wrapf(err, "while applying drainmode")
-		}
-	}
-	return nil, nil
-}

dgraphapi/cluster.go

@@ -17,9 +17,7 @@ import (
 	"strings"
 	"time"
 
-	"github.com/hypermodeinc/dgraph/v25/protos/pb"
 	"github.com/pkg/errors"
-	"google.golang.org/protobuf/encoding/protojson"
 
 	"github.com/dgraph-io/dgo/v250"
 	"github.com/dgraph-io/dgo/v250/protos/api"
@@ -650,27 +648,6 @@ func (hc *HTTPClient) GetZeroState() (*LicenseResponse, error) {
 	return &stateResponse, nil
 }
 
-func (hc *HTTPClient) GetAlphaState(url string) (*pb.MembershipState, error) {
-	if url == "" {
-		url = hc.alphaStateURL
-	}
-	resp, err := http.Get(url)
-	if err != nil {
-		return nil, err
-	}
-
-	body, err := io.ReadAll(resp.Body)
-	if err != nil {
-		return nil, errors.Wrapf(err, "error reading zero state response body")
-	}
-	var state pb.MembershipState
-	if err := protojson.Unmarshal(body, &state); err != nil {
-		return nil, err
-	}
-
-	return &state, err
-}
-
 func (hc *HTTPClient) PostDqlQuery(query string) ([]byte, error) {
 	req, err := http.NewRequest(http.MethodPost, hc.dqlURL, bytes.NewBufferString(query))
 	if err != nil {

dgraphtest/config.go

@@ -102,7 +102,7 @@ type ClusterConfig struct {
 	refillInterval        time.Duration
 	uidLease              int
 	portOffset            int // exposed port offset for grpc/http port for both alpha/zero
-	bulkOutDir            string
+	bulkOutDirForMount    string
 	lambdaURL             string
 	featureFlags          []string
 	customPlugins         bool
@@ -224,7 +224,7 @@ func (cc ClusterConfig) WithExposedPortOffset(offset uint64) ClusterConfig {
 // WithBulkLoadOutDir sets the out dir for the bulk loader. This ensures
 // that the same p directory is used while setting up alphas.
 func (cc ClusterConfig) WithBulkLoadOutDir(dir string) ClusterConfig {
-	cc.bulkOutDir = dir
+	cc.bulkOutDirForMount = dir
 	return cc
 }
 

dgraphtest/dgraph.go

@@ -311,8 +311,8 @@ func (a *alpha) mounts(c *LocalCluster) ([]mount.Mount, error) {
 		})
 	}
 
-	if c.conf.bulkOutDir != "" {
-		pDir := filepath.Join(c.conf.bulkOutDir, strconv.Itoa(a.id/c.conf.replicas), "p")
+	if c.conf.bulkOutDirForMount != "" {
+		pDir := filepath.Join(c.conf.bulkOutDirForMount, strconv.Itoa(a.id/c.conf.replicas), "p")
 		if err := os.MkdirAll(pDir, os.ModePerm); err != nil {
 			return nil, errors.Wrap(err, "erorr creating bulk dir")
 		}