Improve support for docdb

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)

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
+}

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
 			},