diff --git a/connectors/mongo/conn.go b/connectors/mongo/conn.go index eb39754..1b6b98f 100644 --- a/connectors/mongo/conn.go +++ b/connectors/mongo/conn.go @@ -37,6 +37,8 @@ type ConnectorSettings struct { TargetDocCountPerPartition int64 //target number of documents per partition (256k docs is 256MB with 1KB average doc size) SampleFactor int // a factor to determine how many extra samples per partition are used MaxPageSize int + NamespaceFanout int + DocumentDBSamplingFanout int PerNamespaceStreams bool SkipBatchOverwrite bool FullDocumentKey bool @@ -44,6 +46,11 @@ type ConnectorSettings struct { Query string // query filter, as a v2 Extended JSON string, e.g., '{\"x\":{\"$gt\":1}}'" } +const ( + defaultNamespaceFanoutLimit = 100 + defaultDocumentDBSamplingFanoutLimit = 100 +) + func setDefault[T comparable](field *T, defaultValue T) { if *field == *new(T) { *field = defaultValue @@ -239,6 +246,7 @@ func (c *conn) GeneratePlan(ctx context.Context, r *connect.Request[adiomv1.Gene done := make(chan struct{}) eg, ctx := errgroup.WithContext(ctx) + eg.SetLimit(c.planningFanoutLimit()) var finalPartitions []*adiomv1.Partition ch := make(chan *adiomv1.Partition) @@ -288,6 +296,7 @@ func (c *conn) GeneratePlan(ctx context.Context, r *connect.Request[adiomv1.Gene if numSamples > 1000000 { slog.Warn("More than 1000000 samples requested", "samples", numSamples) } + slog.Debug("Getting samples for namespace", "namespace", partition.GetNamespace(), "samples", numSamples) ids, err := c.sampleIDs(ctx, col, numSamples) if err != nil { return fmt.Errorf("error getting %v samples for partition: %w", numSamples, err) @@ -330,6 +339,13 @@ func (c *conn) GeneratePlan(ctx context.Context, r *connect.Request[adiomv1.Gene }), nil } +func (c *conn) planningFanoutLimit() int { + if c.settings.NamespaceFanout > 0 { + return c.settings.NamespaceFanout + } + return defaultNamespaceFanoutLimit +} + // GetInfo implements adiomv1connect.ConnectorServiceHandler. func (c *conn) GetInfo(ctx context.Context, r *connect.Request[adiomv1.GetInfoRequest]) (*connect.Response[adiomv1.GetInfoResponse], error) { // Get version of the MongoDB server @@ -691,11 +707,11 @@ func toTimestampPB(t bson.Timestamp) *timestamppb.Timestamp { } type MongoUpdate struct { - NS bson.M `bson:"ns"` + NS bson.M `bson:"ns"` ClusterTime *bson.Timestamp `bson:"clusterTime"` - DocumentKey bson.D `bson:"documentKey"` - FullDocument bson.Raw `bson:"fullDocument"` - OperationType string `bson:"operationType"` + DocumentKey bson.D `bson:"documentKey"` + FullDocument bson.Raw `bson:"fullDocument"` + OperationType string `bson:"operationType"` } // StreamUpdates implements adiomv1connect.ConnectorServiceHandler. @@ -1122,6 +1138,8 @@ func (c *conn) Teardown() { func NewConn(connSettings ConnectorSettings) (adiomv1connect.ConnectorServiceHandler, error) { setDefault(&connSettings.TargetDocCountPerPartition, 512*1000) setDefault(&connSettings.SampleFactor, 1) + setDefault(&connSettings.NamespaceFanout, defaultNamespaceFanoutLimit) + setDefault(&connSettings.DocumentDBSamplingFanout, defaultDocumentDBSamplingFanoutLimit) client, err := MongoClient(context.Background(), connSettings) if err != nil { slog.Error(fmt.Sprintf("unable to connect to mongo client: %v", err)) diff --git a/connectors/mongo/conn_unit_test.go b/connectors/mongo/conn_unit_test.go index 82ed573..37db63d 100644 --- a/connectors/mongo/conn_unit_test.go +++ b/connectors/mongo/conn_unit_test.go @@ -34,6 +34,17 @@ func bsonID(t *testing.T, name string, v interface{}) *adiomv1.BsonValue { return &adiomv1.BsonValue{Name: name, Type: uint32(typ), Data: data} } +func TestPlanningFanoutLimit(t *testing.T) { + assert.Equal(t, defaultNamespaceFanoutLimit, (&conn{flavor: FlavorMongoDB}).planningFanoutLimit()) + assert.Equal(t, defaultNamespaceFanoutLimit, (&conn{flavor: FlavorDocumentDB}).planningFanoutLimit()) + assert.Equal(t, 12, (&conn{settings: ConnectorSettings{NamespaceFanout: 12}, flavor: FlavorDocumentDB}).planningFanoutLimit()) +} + +func TestDocumentDBSamplingFanout(t *testing.T) { + assert.Equal(t, defaultDocumentDBSamplingFanoutLimit, (&conn{}).documentDBSamplingFanout()) + assert.Equal(t, 17, (&conn{settings: ConnectorSettings{DocumentDBSamplingFanout: 17}}).documentDBSamplingFanout()) +} + // buildIdFilter func TestBuildIdFilter_EmptyId(t *testing.T) { @@ -173,9 +184,9 @@ func deleteU(t *testing.T, id string) *adiomv1.Update { func partialU(t *testing.T, id string, data bson.M, unset ...string) *adiomv1.Update { u := &adiomv1.Update{ - Id: []*adiomv1.BsonValue{bsonID(t, "_id", id)}, - Type: adiomv1.UpdateType_UPDATE_TYPE_PARTIAL_UPDATE, - PartialUpdateUnset: unset, + Id: []*adiomv1.BsonValue{bsonID(t, "_id", id)}, + Type: adiomv1.UpdateType_UPDATE_TYPE_PARTIAL_UPDATE, + PartialUpdateUnset: unset, } if data != nil { u.Data = mustMarshal(t, data) diff --git a/connectors/mongo/docdb.go b/connectors/mongo/docdb.go index bd07308..3602fb9 100644 --- a/connectors/mongo/docdb.go +++ b/connectors/mongo/docdb.go @@ -71,6 +71,7 @@ func (c *conn) sampleIDs(ctx context.Context, col *mongo.Collection, numSamples results := make([]sampleResult, numSamples) eg, ctx := errgroup.WithContext(ctx) + eg.SetLimit(c.documentDBSamplingFanout()) for i := int64(0); i < numSamples; i++ { i := i eg.Go(func() error { @@ -122,3 +123,10 @@ func (c *conn) sampleIDs(ctx context.Context, col *mongo.Collection, numSamples return ids, nil } + +func (c *conn) documentDBSamplingFanout() int { + if c.settings.DocumentDBSamplingFanout > 0 { + return c.settings.DocumentDBSamplingFanout + } + return defaultDocumentDBSamplingFanoutLimit +} diff --git a/gen/adiom/commands/connectors/v1/cosmos_commandargs.go b/gen/adiom/commands/connectors/v1/cosmos_commandargs.go index b998379..941981c 100644 --- a/gen/adiom/commands/connectors/v1/cosmos_commandargs.go +++ b/gen/adiom/commands/connectors/v1/cosmos_commandargs.go @@ -38,6 +38,11 @@ func CosmosFlagsCommandFlags() []cli.Flag { Usage: "query filter for the initial data copy (v2 Extended JSON)", Aliases: []string{"q"}, }, + &cli.IntFlag{ + Name: "namespace-fanout", + Usage: "maximum number of namespaces to plan concurrently", + Value: 100, + }, &cli.IntFlag{ Name: "cosmos-reader-max-namespaces", Usage: "max namespaces to copy (keep under 15)", @@ -75,6 +80,7 @@ func ParseCosmosFlags(c *cli.Context) (*CosmosFlags, error) { cfg.Mongo.DocPartition = c.Int64("doc-partition") cfg.Mongo.MaxPageSize = int32(c.Int("max-page-size")) cfg.Mongo.InitialSyncQuery = c.String("initial-sync-query") + cfg.Mongo.NamespaceFanout = int32(c.Int("namespace-fanout")) cfg.CosmosReaderMaxNamespaces = int32(c.Int("cosmos-reader-max-namespaces")) cfg.CosmosDeletesCdc = c.String("cosmos-deletes-cdc") cfg.CosmosDeleteInterval = durationpb.New(c.Duration("cosmos-delete-interval")) diff --git a/gen/adiom/commands/connectors/v1/mongo.pb.go b/gen/adiom/commands/connectors/v1/mongo.pb.go index 7ded388..4e05ac4 100644 --- a/gen/adiom/commands/connectors/v1/mongo.pb.go +++ b/gen/adiom/commands/connectors/v1/mongo.pb.go @@ -34,6 +34,7 @@ type MongoBaseFlags struct { DocPartition int64 `protobuf:"varint,6,opt,name=doc_partition,json=docPartition,proto3" json:"doc_partition,omitempty"` MaxPageSize int32 `protobuf:"varint,7,opt,name=max_page_size,json=maxPageSize,proto3" json:"max_page_size,omitempty"` InitialSyncQuery string `protobuf:"bytes,8,opt,name=initial_sync_query,json=initialSyncQuery,proto3" json:"initial_sync_query,omitempty"` + NamespaceFanout int32 `protobuf:"varint,9,opt,name=namespace_fanout,json=namespaceFanout,proto3" json:"namespace_fanout,omitempty"` unknownFields protoimpl.UnknownFields sizeCache protoimpl.SizeCache } @@ -124,6 +125,13 @@ func (x *MongoBaseFlags) GetInitialSyncQuery() string { return "" } +func (x *MongoBaseFlags) GetNamespaceFanout() int32 { + if x != nil { + return x.NamespaceFanout + } + return 0 +} + // Full MongoDB flags = base + MongoDB-specific extras type MongoFlags struct { state protoimpl.MessageState `protogen:"open.v1"` @@ -189,7 +197,7 @@ var File_adiom_commands_connectors_v1_mongo_proto protoreflect.FileDescriptor const file_adiom_commands_connectors_v1_mongo_proto_rawDesc = "" + "\n" + - "(adiom/commands/connectors/v1/mongo.proto\x12\x1cadiom.commands.connectors.v1\x1a\x1ccommandargs/v1/options.proto\x1a\x1egoogle/protobuf/duration.proto\"\xc3\x05\n" + + "(adiom/commands/connectors/v1/mongo.proto\x12\x1cadiom.commands.connectors.v1\x1a\x1ccommandargs/v1/options.proto\x1a\x1egoogle/protobuf/duration.proto\"\xbe\x06\n" + "\x0eMongoBaseFlags\x120\n" + "\x14skip_batch_overwrite\x18\x01 \x01(\bR\x12skipBatchOverwrite\x12~\n" + "\x11full_document_key\x18\x02 \x01(\bBR\x82\xb5\x18N\x12Luses the full document key instead of just _id (except for batch overwrites)R\x0ffullDocumentKey\x12V\n" + @@ -204,7 +212,9 @@ const file_adiom_commands_connectors_v1_mongo_proto_rawDesc = "" + "\rmax_page_size\x18\a \x01(\x05B\x13\x82\xb5\x18\x0f\n" + "\rmax-page-sizeR\vmaxPageSize\x12\x84\x01\n" + "\x12initial_sync_query\x18\b \x01(\tBV\x82\xb5\x18R\n" + - "\x12initial-sync-query\x129query filter for the initial data copy (v2 Extended JSON):\x01qR\x10initialSyncQuery:\x0f\x82\xb5\x18\v\n" + + "\x12initial-sync-query\x129query filter for the initial data copy (v2 Extended JSON):\x01qR\x10initialSyncQuery\x12y\n" + + "\x10namespace_fanout\x18\t \x01(\x05BN\x82\xb5\x18J\n" + + "\x10namespace-fanout\x121maximum number of namespaces to plan concurrently\x1a\x03100R\x0fnamespaceFanout:\x0f\x82\xb5\x18\v\n" + "\tMongoBase\"\xe0\x02\n" + "\n" + "MongoFlags\x12@\n" + diff --git a/gen/adiom/commands/connectors/v1/mongo_commandargs.go b/gen/adiom/commands/connectors/v1/mongo_commandargs.go index cf290a9..f221032 100644 --- a/gen/adiom/commands/connectors/v1/mongo_commandargs.go +++ b/gen/adiom/commands/connectors/v1/mongo_commandargs.go @@ -38,6 +38,11 @@ func MongoBaseFlagsCommandFlags() []cli.Flag { Usage: "query filter for the initial data copy (v2 Extended JSON)", Aliases: []string{"q"}, }, + &cli.IntFlag{ + Name: "namespace-fanout", + Usage: "maximum number of namespaces to plan concurrently", + Value: 100, + }, } } @@ -51,6 +56,7 @@ func ParseMongoBaseFlags(c *cli.Context) (*MongoBaseFlags, error) { cfg.DocPartition = c.Int64("doc-partition") cfg.MaxPageSize = int32(c.Int("max-page-size")) cfg.InitialSyncQuery = c.String("initial-sync-query") + cfg.NamespaceFanout = int32(c.Int("namespace-fanout")) return cfg, nil } @@ -102,6 +108,11 @@ func MongoFlagsCommandFlags() []cli.Flag { Usage: "query filter for the initial data copy (v2 Extended JSON)", Aliases: []string{"q"}, }, + &cli.IntFlag{ + Name: "namespace-fanout", + Usage: "maximum number of namespaces to plan concurrently", + Value: 100, + }, &cli.IntFlag{ Name: "sample-factor", Usage: "Number of extra samples per partition", @@ -125,6 +136,7 @@ func ParseMongoFlags(c *cli.Context) (*MongoFlags, error) { cfg.Base.DocPartition = c.Int64("doc-partition") cfg.Base.MaxPageSize = int32(c.Int("max-page-size")) cfg.Base.InitialSyncQuery = c.String("initial-sync-query") + cfg.Base.NamespaceFanout = int32(c.Int("namespace-fanout")) cfg.SampleFactor = int32(c.Int("sample-factor")) cfg.PerNamespaceStreams = c.Bool("per-namespace-streams") return cfg, nil diff --git a/internal/app/options/connectorflags.go b/internal/app/options/connectorflags.go index 52f5be0..737eb2a 100644 --- a/internal/app/options/connectorflags.go +++ b/internal/app/options/connectorflags.go @@ -23,8 +23,8 @@ import ( "github.com/adiom-data/dsync/connectors/sqlbatch" "github.com/adiom-data/dsync/connectors/testconn" "github.com/adiom-data/dsync/connectors/vector" - adiomv1 "github.com/adiom-data/dsync/gen/adiom/v1" connectorsv1 "github.com/adiom-data/dsync/gen/adiom/commands/connectors/v1" + adiomv1 "github.com/adiom-data/dsync/gen/adiom/v1" "github.com/adiom-data/dsync/gen/adiom/v1/adiomv1connect" "github.com/urfave/cli/v2" "github.com/urfave/cli/v2/altsrc" @@ -359,12 +359,14 @@ func GetRegisteredConnectors() []RegisteredConnector { return false }, Create: func(args []string, as AdditionalSettings) (adiomv1connect.ConnectorServiceHandler, []string, error) { - return CreateHelper("MongoDB", "mongodb://connection-string [options]", connectorsv1.MongoFlagsCommandFlags(), func(c *cli.Context, args []string, _ AdditionalSettings) (adiomv1connect.ConnectorServiceHandler, error) { + return CreateHelper("MongoDB", "mongodb://connection-string [options]", mongoFlagsCommandFlags(), func(c *cli.Context, args []string, _ AdditionalSettings) (adiomv1connect.ConnectorServiceHandler, error) { flags, err := connectorsv1.ParseMongoFlags(c) if err != nil { return nil, err } - return mongo.NewConn(mongoSettingsFromFlags(flags, args[0])) + settings := mongoSettingsFromFlags(flags, args[0]) + settings.DocumentDBSamplingFanout = c.Int("documentdb-sampling-fanout") + return mongo.NewConn(settings) })(args, as) }, }, @@ -696,18 +698,18 @@ func fileSettingsFromFlags(f *connectorsv1.FileFlags, uri string) (fileconnector func s3SettingsFromFlags(f *connectorsv1.S3Flags, uri string) s3connector.ConnectorSettings { return s3connector.ConnectorSettings{ - Uri: uri, - PrettyJSON: f.PrettyJson, - Region: f.Region, - Prefix: f.Prefix, - OutputFormat: f.OutputFormat, - Profile: f.Profile, - Endpoint: f.Endpoint, - AccessKeyID: f.AccessKeyId, - SecretAccessKey: f.SecretAccessKey, - SessionToken: f.SessionToken, - UsePathStyle: f.UsePathStyle, - MaxFileSizeMB: f.MaxFileSize, + Uri: uri, + PrettyJSON: f.PrettyJson, + Region: f.Region, + Prefix: f.Prefix, + OutputFormat: f.OutputFormat, + Profile: f.Profile, + Endpoint: f.Endpoint, + AccessKeyID: f.AccessKeyId, + SecretAccessKey: f.SecretAccessKey, + SessionToken: f.SessionToken, + UsePathStyle: f.UsePathStyle, + MaxFileSizeMB: f.MaxFileSize, MaxTotalMemoryMB: f.MaxTotalMemory, } } @@ -724,10 +726,20 @@ func mongoBaseSettingsFromFlags(f *connectorsv1.MongoBaseFlags, connString strin s.WriterMaxBatchSize = int(f.WriterBatchSize) s.TargetDocCountPerPartition = f.DocPartition s.MaxPageSize = int(f.MaxPageSize) + s.NamespaceFanout = int(f.NamespaceFanout) s.Query = f.InitialSyncQuery return s } +func mongoFlagsCommandFlags() []cli.Flag { + flags := connectorsv1.MongoFlagsCommandFlags() + return append(flags, &cli.IntFlag{ + Name: "documentdb-sampling-fanout", + Value: 100, + Hidden: true, + }) +} + func mongoSettingsFromFlags(f *connectorsv1.MongoFlags, connString string) mongo.ConnectorSettings { s := mongoBaseSettingsFromFlags(f.Base, connString) s.SampleFactor = int(f.SampleFactor) diff --git a/proto/connectorcommands/adiom/commands/connectors/v1/mongo.proto b/proto/connectorcommands/adiom/commands/connectors/v1/mongo.proto index 1049771..8e2c543 100644 --- a/proto/connectorcommands/adiom/commands/connectors/v1/mongo.proto +++ b/proto/connectorcommands/adiom/commands/connectors/v1/mongo.proto @@ -36,6 +36,11 @@ message MongoBaseFlags { usage: "query filter for the initial data copy (v2 Extended JSON)" aliases: ["q"] }]; + int32 namespace_fanout = 9 [(commandargs.v1.flag) = { + name: "namespace-fanout" + usage: "maximum number of namespaces to plan concurrently" + default: "100" + }]; } // Full MongoDB flags = base + MongoDB-specific extras