mongo: validate each shard on sharded cluster

Author: jgao54

flow/connectors/mongo/validate.go

@@ -8,7 +8,13 @@ import (
 )
 
 func (c *MongoConnector) ValidateCheck(ctx context.Context) error {
-	if err := shared_mongo.ValidateServerCompatibility(ctx, c.client); err != nil {
+	if err := shared_mongo.ValidateServerCompatibility(ctx, c.client, shared_mongo.Credentials{
+		Username:   c.config.Username,
+		Password:   c.config.Password,
+		DisableTls: c.config.DisableTls,
+		RootCa:     c.config.RootCa,
+		TlsHost:    c.config.TlsHost,
+	}); err != nil {
 		return err
 	}
 
@@ -24,7 +30,13 @@ func (c *MongoConnector) ValidateMirrorSource(ctx context.Context, cfg *protos.F
 		return nil
 	}
 
-	if err := shared_mongo.ValidateOplogRetention(ctx, c.client); err != nil {
+	if err := shared_mongo.ValidateOplogRetention(ctx, c.client, shared_mongo.Credentials{
+		Username:   c.config.Username,
+		Password:   c.config.Password,
+		DisableTls: c.config.DisableTls,
+		RootCa:     c.config.RootCa,
+		TlsHost:    c.config.TlsHost,
+	}); err != nil {
 		return err
 	}
 

flow/shared/mongo/commands.go

@@ -68,6 +68,19 @@ func GetHelloResponse(ctx context.Context, client *mongo.Client) (*HelloResponse
 	return runCommand[HelloResponse](ctx, client, "hello")
 }
 
+type Shard struct {
+	Host string `bson:"host"`
+}
+
+type ListShards struct {
+	Shards []Shard `bson:"shards"`
+	Ok     int     `bson:"ok"`
+}
+
+func GetListShards(ctx context.Context, client *mongo.Client) (*ListShards, error) {
+	return runCommand[ListShards](ctx, client, "listShards")
+}
+
 func runCommand[T any](ctx context.Context, client *mongo.Client, command string) (*T, error) {
 	singleResult := client.Database("admin").RunCommand(ctx, bson.D{
 		bson.E{Key: command, Value: 1},

flow/shared/mongo/validation.go

@@ -2,10 +2,15 @@ package mongo
 
 import (
 	"context"
+	"crypto/tls"
 	"errors"
 	"fmt"
+	"strings"
+	"sync"
 
+	"github.com/PeerDB-io/peerdb/flow/shared"
 	"go.mongodb.org/mongo-driver/v2/mongo"
+	"go.mongodb.org/mongo-driver/v2/mongo/options"
 )
 
 const (
@@ -16,7 +21,15 @@ const (
 	ShardedCluster = "ShardedCluster"
 )
 
-func ValidateServerCompatibility(ctx context.Context, client *mongo.Client) error {
+type Credentials struct {
+	RootCa     *string
+	Username   string
+	Password   string
+	TlsHost    string
+	DisableTls bool
+}
+
+func ValidateServerCompatibility(ctx context.Context, client *mongo.Client, credentials Credentials) error {
 	buildInfo, err := GetBuildInfo(ctx, client)
 	if err != nil {
 		return err
@@ -48,8 +61,7 @@ func ValidateServerCompatibility(ctx context.Context, client *mongo.Client) erro
 	if topologyType == ReplicaSet {
 		return validateStorageEngine(ctx, client)
 	} else {
-		// TODO: run validation on shard
-		return nil
+		return runOnShardsInParallel(ctx, client, credentials, validateStorageEngine)
 	}
 }
 
@@ -79,7 +91,7 @@ func ValidateUserRoles(ctx context.Context, client *mongo.Client) error {
 	return nil
 }
 
-func ValidateOplogRetention(ctx context.Context, client *mongo.Client) error {
+func ValidateOplogRetention(ctx context.Context, client *mongo.Client, credentials Credentials) error {
 	validateOplogRetention := func(instanceCtx context.Context, instanceClient *mongo.Client) error {
 		ss, err := GetServerStatus(instanceCtx, instanceClient)
 		if err != nil {
@@ -100,8 +112,7 @@ func ValidateOplogRetention(ctx context.Context, client *mongo.Client) error {
 	if topology == ReplicaSet {
 		return validateOplogRetention(ctx, client)
 	} else {
-		// TODO: run validation on shard
-		return nil
+		return runOnShardsInParallel(ctx, client, credentials, validateOplogRetention)
 	}
 }
 
@@ -124,3 +135,91 @@ func GetTopologyType(ctx context.Context, client *mongo.Client) (string, error)
 	}
 	return "", errors.New("topology type must be ReplicaSet or ShardedCluster")
 }
+
+func runOnShardsInParallel(
+	ctx context.Context,
+	client *mongo.Client,
+	credentials Credentials,
+	runCommand func(ctx context.Context, client *mongo.Client) error,
+) error {
+	res, err := GetListShards(ctx, client)
+	if err != nil {
+		return err
+	}
+
+	if res.Ok != 1 || len(res.Shards) == 0 {
+		return errors.New("invalid shards")
+	}
+
+	hosts := getUniqueClusterHosts(res.Shards)
+	hostsErrors := make([]error, len(hosts))
+	var wg sync.WaitGroup
+	for idx, host := range hosts {
+		wg.Add(1)
+		go func(i int, h string) {
+			defer wg.Done()
+
+			shardOpts := options.Client().
+				ApplyURI("mongodb://" + h).
+				SetDirect(true).
+				SetAuth(options.Credential{
+					Username: credentials.Username,
+					Password: credentials.Password,
+				})
+
+			if !credentials.DisableTls {
+				tlsConfig, err := shared.CreateTlsConfig(tls.VersionTLS12, credentials.RootCa, "", credentials.TlsHost, false)
+				if err != nil {
+					hostsErrors[i] = fmt.Errorf("host %s TLS config error: %w", h, err)
+					return
+				}
+				shardOpts.SetTLSConfig(tlsConfig)
+			}
+
+			shardClient, err := mongo.Connect(shardOpts)
+			if err != nil {
+				hostsErrors[i] = fmt.Errorf("host %s connect error: %w", h, err)
+				return
+			}
+			defer shardClient.Disconnect(ctx) //nolint:errcheck
+
+			if err := runCommand(ctx, shardClient); err != nil {
+				hostsErrors[i] = fmt.Errorf("host %s command error: %w", h, err)
+				return
+			}
+		}(idx, host)
+	}
+	wg.Wait()
+
+	for _, err = range hostsErrors {
+		if err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
+func getUniqueClusterHosts(shards []Shard) []string {
+	hostSet := make(map[string]bool)
+	for _, shard := range shards {
+		hosts := shard.Host
+		if slashIdx := strings.Index(hosts, "/"); slashIdx != -1 {
+			hosts = hosts[slashIdx+1:]
+		}
+
+		for _, host := range strings.Split(hosts, ",") {
+			host = strings.TrimSpace(host)
+			if host != "" {
+				hostSet[host] = true
+			}
+		}
+	}
+
+	hosts := make([]string, 0, len(hostSet))
+	for host := range hostSet {
+		hosts = append(hosts, host)
+	}
+
+	return hosts
+}