diff --git a/connectors/mongo/conn.go b/connectors/mongo/conn.go index ee5879e5..266872d6 100644 --- a/connectors/mongo/conn.go +++ b/connectors/mongo/conn.go @@ -70,6 +70,7 @@ type conn struct { client *mongo.Client settings ConnectorSettings + flavor MongoFlavor nextCursorID atomic.Int64 ctx context.Context @@ -265,19 +266,19 @@ func (c *conn) GeneratePlan(ctx context.Context, r *connect.Request[adiomv1.Gene if numSamples > 1000000 { slog.Warn("More than 1000000 samples requested", "samples", numSamples) } - res, err := col.Aggregate(ctx, mongo.Pipeline{{{"$sample", bson.D{{"size", numSamples}}}}, {{"$project", bson.D{{"_id", 1}}}}, {{"$sort", bson.D{{"_id", 1}}}}}) + ids, err := c.sampleIDs(ctx, col, numSamples) if err != nil { return fmt.Errorf("error getting %v samples for partition: %w", numSamples, err) } var factorCount = c.settings.SampleFactor / 2 var low bson.RawValue - for res.Next(ctx) { + for _, id := range ids { if factorCount > 0 { factorCount -= 1 continue } factorCount = c.settings.SampleFactor - 1 - high := res.Current.Lookup("_id") + high := id ch <- &adiomv1.Partition{ Namespace: partition.GetNamespace(), EstimatedCount: uint64(c.settings.TargetDocCountPerPartition), @@ -566,7 +567,7 @@ func (c *conn) StreamLSN(ctx context.Context, r *connect.Request[adiomv1.StreamL {{"$match", nsFilter}}, } } - opts := moptions.ChangeStream().SetStartAfter(bson.Raw(r.Msg.GetCursor())) + opts := c.changeStreamOpts(r.Msg.GetCursor()) changeStream, err := watcher.Watch(ctx, pipeline, opts) if err != nil { @@ -707,7 +708,7 @@ func (c *conn) StreamUpdates(ctx context.Context, r *connect.Request[adiomv1.Str {{"$match", nsFilter}}, } } - opts := moptions.ChangeStream().SetStartAfter(bson.Raw(r.Msg.GetCursor())).SetFullDocument("updateLookup") + opts := c.changeStreamOpts(r.Msg.GetCursor()).SetFullDocument("updateLookup") changeStream, err := watcher.Watch(ctx, pipeline, opts) if err != nil { @@ -995,6 +996,7 @@ func NewConnWithClient(client *mongo.Client, settings ConnectorSettings) adiomv1 return &conn{ client: client, settings: settings, + flavor: GetMongoFlavor(settings.ConnectionString), ctx: ctx, cancel: cancel, buffers: map[int64]buffer{}, @@ -1003,6 +1005,16 @@ func NewConnWithClient(client *mongo.Client, settings ConnectorSettings) adiomv1 } } +func (c *conn) changeStreamOpts(cursor []byte) *moptions.ChangeStreamOptions { + opts := moptions.ChangeStream() + if c.flavor == FlavorDocumentDB { + opts.SetResumeAfter(bson.Raw(cursor)) + } else { + opts.SetStartAfter(bson.Raw(cursor)) + } + return opts +} + func maybeUnavailableError(err error) error { if mongo.IsNetworkError(err) { return connect.NewError(connect.CodeUnavailable, err) diff --git a/connectors/mongo/docdb.go b/connectors/mongo/docdb.go new file mode 100644 index 00000000..5e6d4d96 --- /dev/null +++ b/connectors/mongo/docdb.go @@ -0,0 +1,125 @@ +package mongo + +import ( + "bytes" + "cmp" + "encoding/binary" + "encoding/hex" + "fmt" + "sort" + + "context" + + "go.mongodb.org/mongo-driver/bson" + "go.mongodb.org/mongo-driver/bson/bsontype" + "go.mongodb.org/mongo-driver/mongo" + "golang.org/x/sync/errgroup" +) + +var supportedIDTypes = map[bsontype.Type]bool{ + bson.TypeObjectID: true, + bson.TypeString: true, + bson.TypeInt32: true, + bson.TypeInt64: true, + bson.TypeBinary: true, +} + +func compareBSONRawValues(a, b bson.RawValue) int { + switch a.Type { + case bson.TypeObjectID: + return bytes.Compare(a.Value, b.Value) + case bson.TypeString: + return bytes.Compare(a.Value[4:len(a.Value)-1], b.Value[4:len(b.Value)-1]) + case bson.TypeInt32: + ai := int32(binary.LittleEndian.Uint32(a.Value)) + bi := int32(binary.LittleEndian.Uint32(b.Value)) + return cmp.Compare(ai, bi) + case bson.TypeInt64: + ai := int64(binary.LittleEndian.Uint64(a.Value)) + bi := int64(binary.LittleEndian.Uint64(b.Value)) + return cmp.Compare(ai, bi) + case bson.TypeBinary: + return bytes.Compare(a.Value[5:], b.Value[5:]) + } + panic("compareBSONRawValues called with unsupported type") +} + +func (c *conn) sampleIDs(ctx context.Context, col *mongo.Collection, numSamples int64) ([]bson.RawValue, error) { + if c.flavor != FlavorDocumentDB { + res, err := col.Aggregate(ctx, mongo.Pipeline{ + {{"$sample", bson.D{{"size", numSamples}}}}, + {{"$project", bson.D{{"_id", 1}}}}, + {{"$sort", bson.D{{"_id", 1}}}}, + }) + if err != nil { + return nil, fmt.Errorf("error getting %v samples: %w", numSamples, err) + } + defer res.Close(ctx) + var ids []bson.RawValue + for res.Next(ctx) { + ids = append(ids, res.Current.Lookup("_id")) + } + if res.Err() != nil { + return nil, fmt.Errorf("err iterating through samples: %w", res.Err()) + } + return ids, nil + } + + type sampleResult struct { + id bson.RawValue + key string + } + + results := make([]sampleResult, numSamples) + eg, ctx := errgroup.WithContext(ctx) + for i := int64(0); i < numSamples; i++ { + i := i + eg.Go(func() error { + res, err := col.Aggregate(ctx, mongo.Pipeline{ + {{"$sample", bson.D{{"size", 1}}}}, + {{"$project", bson.D{{"_id", 1}}}}, + }) + if err != nil { + return fmt.Errorf("error getting sample %v: %w", i, err) + } + defer res.Close(ctx) + if res.Next(ctx) { + id := res.Current.Lookup("_id") + results[i] = sampleResult{id: id, key: hex.EncodeToString(id.Value)} + } + return nil + }) + } + if err := eg.Wait(); err != nil { + return nil, err + } + + seen := map[string]struct{}{} + var ids []bson.RawValue + for _, r := range results { + if r.key == "" { + continue + } + if _, dup := seen[r.key]; !dup { + seen[r.key] = struct{}{} + ids = append(ids, r.id) + } + } + + if len(ids) > 0 { + t := ids[0].Type + for _, id := range ids[1:] { + if id.Type != t { + return nil, fmt.Errorf("mixed _id types not supported for DocumentDB sampling") + } + } + if !supportedIDTypes[t] { + return nil, fmt.Errorf("unsupported _id type for DocumentDB sampling: %v", t) + } + sort.Slice(ids, func(i, j int) bool { + return compareBSONRawValues(ids[i], ids[j]) < 0 + }) + } + + return ids, nil +} diff --git a/internal/app/options/connectorflags.go b/internal/app/options/connectorflags.go index 3ffd5198..61fcb478 100644 --- a/internal/app/options/connectorflags.go +++ b/internal/app/options/connectorflags.go @@ -479,7 +479,7 @@ func GetRegisteredConnectors() []RegisteredConnector { IsConnector: func(s string) bool { if strings.HasPrefix(s, "mongodb://") || strings.HasPrefix(s, "mongodb+srv://") { flavor := mongo.GetMongoFlavor(s) - return flavor == mongo.FlavorMongoDB || flavor == mongo.FlavorCosmosDB_VCORE + return flavor == mongo.FlavorMongoDB || flavor == mongo.FlavorCosmosDB_VCORE || flavor == mongo.FlavorDocumentDB } return false },