mongodb parallel snapshotting support

flow/connectors/mongo/qrep.go

@@ -8,12 +8,12 @@ import (
 	"math"
 	"time"
 
-	"github.com/google/uuid"
 	"go.mongodb.org/mongo-driver/v2/bson"
 	"go.mongodb.org/mongo-driver/v2/mongo/options"
 
 	"github.com/PeerDB-io/peerdb/flow/connectors/utils"
 	"github.com/PeerDB-io/peerdb/flow/generated/protos"
+	"github.com/PeerDB-io/peerdb/flow/internal"
 	"github.com/PeerDB-io/peerdb/flow/model"
 	"github.com/PeerDB-io/peerdb/flow/otel_metrics"
 	"github.com/PeerDB-io/peerdb/flow/pkg/common"
@@ -35,7 +35,17 @@ func (c *MongoConnector) GetQRepPartitions(
 		},
 	}
 
-	if config.WatermarkColumn == "" {
+	if config.WatermarkColumn != DefaultDocumentKeyColumnName {
+		c.logger.Warn("unexpected watermark column, falling back to full table partition")
+		return fullTablePartition, nil
+	}
+
+	parallelSnapshotting, err := internal.PeerDBMongoDBParallelSnapshotting(ctx, config.Env)
+	if err != nil {
+		c.logger.Warn("failed to get parallel snapshotting config", slog.Any("error", err))
+	}
+	if !parallelSnapshotting {
+		c.logger.Info("parallel snapshotting disabled, falling back to full table partition")
 		return fullTablePartition, nil
 	}
 
@@ -59,9 +69,6 @@ func (c *MongoConnector) GetQRepPartitions(
 	if err != nil {
 		return nil, fmt.Errorf("failed to count documents in collection %s: %w", parseWatermarkTable.Table, err)
 	}
-	if totalRows == 0 {
-		return []*protos.QRepPartition{}, nil
-	}
 
 	// Calculate the number of partitions
 	adjustedPartitions := shared.AdjustNumPartitions(totalRows, numRowsPerPartition)
@@ -71,77 +78,24 @@ func (c *MongoConnector) GetQRepPartitions(
 		slog.Int64("adjustedNumPartitions", adjustedPartitions.AdjustedNumPartitions),
 		slog.Int64("adjustedNumRowsPerPartition", adjustedPartitions.AdjustedNumRowsPerPartition))
 
-	// no need to bother with bucketAuto if we have only one partition
-	if adjustedPartitions.AdjustedNumPartitions == 1 {
+	if adjustedPartitions.AdjustedNumPartitions <= 1 {
+		c.logger.Info("[mongo] insufficient partitions for parallel snapshot, falling back to full table partition")
 		return fullTablePartition, nil
 	}
 
-	if config.WatermarkColumn != DefaultDocumentKeyColumnName {
-		return nil, fmt.Errorf("only %s is currently supported as watermark column for MongoDB connector", DefaultDocumentKeyColumnName)
-	}
-
-	// Use bucketAuto to create partitions based on _id field
-	bucketAutoPipeline := []bson.D{
-		{
-			{Key: "$bucketAuto", Value: bson.D{
-				{Key: "groupBy", Value: "$" + config.WatermarkColumn},
-				{Key: "buckets", Value: adjustedPartitions.AdjustedNumPartitions},
-			}},
-		},
-	}
-
-	cursor, err := collection.Aggregate(ctx, bucketAutoPipeline)
-	if err != nil {
-		return nil, fmt.Errorf("failed to aggregate for bucket partitions: %w", err)
-	}
-	defer cursor.Close(ctx)
-
-	partitions := make([]*protos.QRepPartition, 0, adjustedPartitions.AdjustedNumPartitions)
-	for cursor.Next(ctx) {
-		var bucket struct {
-			ID struct {
-				Min bson.ObjectID `bson:"min"`
-				Max bson.ObjectID `bson:"max"`
-			} `bson:"_id"`
-		}
-		if err := bson.Unmarshal(cursor.Current, &bucket); err != nil {
-			return nil, fmt.Errorf("failed to unmarshal bucket: %w", err)
-		}
-
-		partitions = append(partitions, &protos.QRepPartition{
-			PartitionId: uuid.NewString(),
-			Range: &protos.PartitionRange{
-				Range: &protos.PartitionRange_ObjectIdRange{
-					ObjectIdRange: &protos.ObjectIdPartitionRange{
-						Start: bucket.ID.Min.Hex(),
-						End:   bucket.ID.Max.Hex(),
-					},
-				},
-			},
-			FullTablePartition: false,
-		})
-	}
-	if err := cursor.Err(); err != nil {
-		if errors.Is(err, context.Canceled) {
-			c.logger.Warn("context canceled while performing bucketAuto aggregation",
-				slog.String("watermark_table", config.WatermarkTable))
-		} else {
-			c.logger.Error("error while performing bucketAuto aggregation",
-				slog.String("watermark_table", config.WatermarkTable),
-				slog.String("error", err.Error()))
-		}
-		return nil, fmt.Errorf("cursor error during bucketAuto aggregation: %w", err)
-	}
-
-	return partitions, nil
+	return c.minMaxPartitions(ctx, collection, adjustedPartitions.AdjustedNumPartitions)
 }
 
 func (c *MongoConnector) GetDefaultPartitionKeyForTables(
 	ctx context.Context,
 	input *protos.GetDefaultPartitionKeyForTablesInput,
 ) (*protos.GetDefaultPartitionKeyForTablesOutput, error) {
+	mapping := make(map[string]string, len(input.TableMappings))
+	for _, tm := range input.TableMappings {
+		mapping[tm.SourceTableIdentifier] = DefaultDocumentKeyColumnName
+	}
 	return &protos.GetDefaultPartitionKeyForTablesOutput{
-		TableDefaultPartitionKeyMapping: nil,
+		TableDefaultPartitionKeyMapping: mapping,
 	}, nil
 }
 
@@ -275,8 +229,6 @@ func GetDefaultSchema(internalVersion uint32) types.QRecordSchema {
 	return types.QRecordSchema{Fields: schema}
 }
 
-// with $bucketAuto, buckets except the last bucket treat their max value as exclusive
-// we can't tell what bucket is the "last" bucket without additional tracking, so we accept bounday records being inserted twice
 func toRangeFilter(watermarkColumn string, partitionRange *protos.PartitionRange) (bson.D, error) {
 	switch r := partitionRange.Range.(type) {
 	case *protos.PartitionRange_ObjectIdRange:

flow/connectors/mongo/qrep_partition.go

@@ -0,0 +1,165 @@
+package connmongo
+
+import (
+	"context"
+	"fmt"
+	"log/slog"
+	"math/big"
+
+	"github.com/google/uuid"
+	"go.mongodb.org/mongo-driver/v2/bson"
+	"go.mongodb.org/mongo-driver/v2/mongo"
+	"go.mongodb.org/mongo-driver/v2/mongo/options"
+
+	"github.com/PeerDB-io/peerdb/flow/connectors/utils"
+	"github.com/PeerDB-io/peerdb/flow/generated/protos"
+	"github.com/PeerDB-io/peerdb/flow/shared"
+)
+
+// minMaxPartitions creates partitions by querying only the min and max _id
+// from the collection (leveraging the default _id index), then uniformly
+// dividing the ObjectID range using integer arithmetic on the full 12-byte
+// value. Since the first 4-byte timestamp is the most significant component,
+// this naturally partitions by time while also provides best-effort handling
+// of edge cases where records are all inserted in the same second.
+func (c *MongoConnector) minMaxPartitions(
+	ctx context.Context,
+	collection *mongo.Collection,
+	numPartitions int64,
+) ([]*protos.QRepPartition, error) {
+	fullTablePartition := []*protos.QRepPartition{{
+		PartitionId:        utils.FullTablePartitionID,
+		Range:              nil,
+		FullTablePartition: true,
+	}}
+
+	minID, maxID, found, err := findMinMaxObjectIDs(ctx, collection, c.config.ReadPreference)
+	if err != nil {
+		return nil, err
+	}
+	if !found {
+		c.logger.Info("[mongo] no valid min/max ObjectIDs found, falling back to full table partition")
+		return fullTablePartition, nil
+	}
+	minInt := new(big.Int).SetBytes(minID[:])
+	maxInt := new(big.Int).SetBytes(maxID[:])
+	intRange := new(big.Int).Sub(maxInt, minInt)
+	if intRange.Sign() <= 0 {
+		c.logger.Info("[mongo] min/max ObjectID range is non-positive, falling back to full table partition")
+		return fullTablePartition, nil
+	}
+
+	c.logger.Info("[mongo] using min/max ObjectID partitioning",
+		slog.String("minID", minID.Hex()),
+		slog.String("maxID", maxID.Hex()),
+		slog.Int64("numPartitions", numPartitions))
+
+	step := shared.BigIntDivCeil(intRange, big.NewInt(numPartitions))
+
+	partitions := make([]*protos.QRepPartition, 0, numPartitions)
+	for start := new(big.Int).Set(minInt); start.Cmp(maxInt) <= 0; start.Add(start, step) {
+		end := new(big.Int).Add(start, step)
+		end.Sub(end, big.NewInt(1))
+		if end.Cmp(maxInt) > 0 {
+			end.Set(maxInt)
+		}
+		startObjectID, err := bigIntToObjectID(start)
+		if err != nil {
+			return nil, err
+		}
+		endObjectID, err := bigIntToObjectID(end)
+		if err != nil {
+			return nil, err
+		}
+		partitions = append(partitions, &protos.QRepPartition{
+			PartitionId: uuid.NewString(),
+			Range: &protos.PartitionRange{
+				Range: &protos.PartitionRange_ObjectIdRange{
+					ObjectIdRange: &protos.ObjectIdPartitionRange{
+						Start: startObjectID.Hex(),
+						End:   endObjectID.Hex(),
+					},
+				},
+			},
+			FullTablePartition: false,
+		})
+	}
+
+	return partitions, nil
+}
+
+// bigIntToObjectID converts a big.Int back to a 12-byte ObjectID.
+func bigIntToObjectID(n *big.Int) (bson.ObjectID, error) {
+	var oid bson.ObjectID
+	if n.BitLen() > 96 {
+		return oid, fmt.Errorf("big.Int value exceeds 96 bits (ObjectID size)")
+	}
+	n.FillBytes(oid[:])
+	return oid, nil
+}
+
+// findMinMaxObjectIDs returns the min and max object IDs from the collection.
+// The two boundary queries are not wrapped in a snapshot session: a concurrent
+// insert between them may shift the true min/max, but this is safe because the
+// change stream resume token is captured before either query runs. Any writes
+// that occur during or after partitioning will be replayed by the change stream.
+func findMinMaxObjectIDs(
+	ctx context.Context,
+	collection *mongo.Collection,
+	readPreference protos.ReadPreference,
+) (bson.ObjectID, bson.ObjectID, bool, error) {
+	minRaw, err := findBoundaryID(ctx, collection, Lower, readPreference)
+	if err != nil {
+		return bson.ObjectID{}, bson.ObjectID{}, false, fmt.Errorf("failed to find min _id: %w", err)
+	}
+	if minRaw.Type != bson.TypeObjectID {
+		return bson.ObjectID{}, bson.ObjectID{}, false, nil
+	}
+
+	maxRaw, err := findBoundaryID(ctx, collection, Upper, readPreference)
+	if err != nil {
+		return bson.ObjectID{}, bson.ObjectID{}, false, fmt.Errorf("failed to find max _id: %w", err)
+	}
+	if maxRaw.Type != bson.TypeObjectID {
+		return bson.ObjectID{}, bson.ObjectID{}, false, nil
+	}
+
+	return minRaw.ObjectID(), maxRaw.ObjectID(), true, nil
+}
+
+type Boundary int
+
+const (
+	Lower Boundary = 1
+	Upper Boundary = -1
+)
+
+// findBoundaryID returns the lower- or upper-bound _id of a collection
+func findBoundaryID(
+	ctx context.Context,
+	collection *mongo.Collection,
+	boundary Boundary,
+	readPreference protos.ReadPreference,
+) (bson.RawValue, error) {
+	findCmd := bson.D{
+		{Key: "find", Value: collection.Name()},
+		{Key: "sort", Value: bson.D{{Key: DefaultDocumentKeyColumnName, Value: boundary}}},
+		{Key: "limit", Value: 1},
+		{Key: "projection", Value: bson.D{{Key: DefaultDocumentKeyColumnName, Value: 1}}},
+	}
+
+	cursor, err := collection.Database().RunCommandCursor(ctx, findCmd, options.RunCmd().SetReadPreference(protoToReadPref[readPreference]))
+	if err != nil {
+		return bson.RawValue{}, fmt.Errorf("failed to run find command: %w", err)
+	}
+	defer cursor.Close(ctx)
+
+	if !cursor.Next(ctx) {
+		if err := cursor.Err(); err != nil {
+			return bson.RawValue{}, fmt.Errorf("cursor error: %w", err)
+		}
+		// no results
+		return bson.RawValue{}, nil
+	}
+	return cursor.Current.Lookup(DefaultDocumentKeyColumnName), nil
+}

flow/e2e/mongo.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"errors"
 	"os"
+	"sort"
 	"testing"
 
 	"github.com/stretchr/testify/require"
@@ -83,6 +84,11 @@ func (s *MongoSource) GetRows(ctx context.Context, suffix, table, cols string) (
 		recordBatch.Records = append(recordBatch.Records, record)
 	}
 
+	// sort by _id string value to match ClickHouse's `ORDER BY 1`
+	sort.Slice(recordBatch.Records, func(i, j int) bool {
+		return recordBatch.Records[i][0].Value().(string) < recordBatch.Records[j][0].Value().(string)
+	})
+
 	return recordBatch, nil
 }