From f45c14ed01254ed011982f3eb5e1c178e98a9526 Mon Sep 17 00:00:00 2001 From: Andrey Zhelnin Date: Tue, 23 Jun 2026 18:10:37 +0200 Subject: [PATCH 1/5] mongo: partition snapshots for numeric and string _id collections The MongoDB connector only parallelized initial snapshots when _id was a native ObjectID; numeric or string _id collections fell back to a single serial full-table partition. For large collections keyed by natural ids (e.g. app-store numeric ids or package-name strings) this makes the snapshot too slow to finish within the oplog retention window. Dispatch in GetQRepPartitions on the collection's min/max _id BSON type: - ObjectID: unchanged (uniform division of the 12-byte keyspace). - Int32/Int64: uniform min/max division, reusing utils.PartitionHelper to emit non-overlapping IntPartitionRange partitions. - String: sampling-based quantile boundaries ($sample, honouring the configured read preference) anchored by the real min/max, emitted as a new StringPartitionRange proto type with contiguous half-open ranges. Mixed-type, Double, or empty collections fall back to full-table. Add IntRange/StringRange cases to toRangeFilter and StringRange handling to monitoring.addPartitionToQRepRun. computeStringBoundaries is pure and unit-tested alongside the range-to-filter conversions. --- flow/connectors/mongo/qrep.go | 22 +- flow/connectors/mongo/qrep_partition.go | 305 +++++++++++++++--- flow/connectors/mongo/qrep_partition_test.go | 170 ++++++++++ .../connectors/utils/monitoring/monitoring.go | 2 + protos/flow.proto | 6 + 5 files changed, 462 insertions(+), 43 deletions(-) create mode 100644 flow/connectors/mongo/qrep_partition_test.go diff --git a/flow/connectors/mongo/qrep.go b/flow/connectors/mongo/qrep.go index 85c3b39d2..3e5c7eea5 100644 --- a/flow/connectors/mongo/qrep.go +++ b/flow/connectors/mongo/qrep.go @@ -73,7 +73,7 @@ func (c *MongoConnector) GetQRepPartitions( return fullTablePartition, nil } - return c.minMaxPartitions(ctx, collection, adjustedPartitions.AdjustedNumPartitions) + return c.buildPartitions(ctx, collection, adjustedPartitions.AdjustedNumPartitions) } func (c *MongoConnector) GetDefaultPartitionKeyForTables( @@ -234,6 +234,26 @@ func toRangeFilter(watermarkColumn string, partitionRange *protos.PartitionRange bson.E{Key: "$lte", Value: endObjectID}, }}, }, nil + case *protos.PartitionRange_IntRange: + // Numeric _id: inclusive range. PartitionHelper builds non-overlapping + // [start, end] integer ranges. Mongo compares int32/int64 numerically, so + // int64 bounds match int32-typed _ids. + return bson.D{ + bson.E{Key: watermarkColumn, Value: bson.D{ + bson.E{Key: "$gte", Value: r.IntRange.Start}, + bson.E{Key: "$lte", Value: r.IntRange.End}, + }}, + }, nil + case *protos.PartitionRange_StringRange: + // String _id: half-open [start, end) range. computeStringBoundaries makes + // ranges contiguous and sets the final end just past the real max, so $lt + // gives exact coverage with no gaps or overlap. + return bson.D{ + bson.E{Key: watermarkColumn, Value: bson.D{ + bson.E{Key: "$gte", Value: r.StringRange.Start}, + bson.E{Key: "$lt", Value: r.StringRange.End}, + }}, + }, nil default: return nil, fmt.Errorf("unsupported partition range type") } diff --git a/flow/connectors/mongo/qrep_partition.go b/flow/connectors/mongo/qrep_partition.go index 71e3e450b..d70631d9e 100644 --- a/flow/connectors/mongo/qrep_partition.go +++ b/flow/connectors/mongo/qrep_partition.go @@ -5,6 +5,7 @@ import ( "fmt" "log/slog" "math/big" + "slices" "github.com/google/uuid" "go.mongodb.org/mongo-driver/v2/bson" @@ -16,37 +17,95 @@ import ( "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{{ +// stringPartitionUpperSentinel is appended to the collection's max string _id to +// form the exclusive upper bound of the final string partition. It is the highest +// valid Unicode code point, so `max + sentinel` sorts strictly after `max` (and +// after every real _id, since `max` is the maximum). This lets every string +// partition use uniform half-open [start, end) ($gte/$lt) semantics while still +// capturing the maximum document in the last partition. +const stringPartitionUpperSentinel = "\U0010FFFF" + +const ( + // stringSampleOversample draws more samples than partitions so quantile + // boundaries are well distributed even with clustered keys. + stringSampleOversample = 20 + // stringSampleMaxSize caps sampling cost on very large collections. + stringSampleMaxSize = 100000 +) + +// fullTablePartitions returns the single-partition full-table fallback used when a +// collection cannot be partitioned (unsupported/mixed _id type, empty collection, +// or degenerate min/max range). +func fullTablePartitions() []*protos.QRepPartition { + return []*protos.QRepPartition{{ PartitionId: utils.FullTablePartitionID, Range: nil, FullTablePartition: true, }} +} - minID, maxID, found, err := findMinMaxObjectIDs(ctx, collection, c.config.ReadPreference) +// buildPartitions reads the collection's min and max _id (cheap, leverages the +// default _id index) and dispatches to a type-specific partitioner: +// - ObjectID -> uniform division of the 12-byte ObjectID keyspace +// - Int32/64 -> uniform division of the numeric range (reuses PartitionHelper) +// - String -> sampling-based quantile boundaries +// +// Mixed-type _id (min and max differ in type), Double, or any other type fall back +// to a full-table partition. The change stream resume token is captured before this +// runs, so any writes during partitioning are replayed by CDC. +func (c *MongoConnector) buildPartitions( + ctx context.Context, + collection *mongo.Collection, + numPartitions int64, +) ([]*protos.QRepPartition, error) { + minRaw, err := findBoundaryID(ctx, collection, Lower, c.config.ReadPreference) if err != nil { - return nil, err + return nil, fmt.Errorf("failed to find min _id: %w", err) + } + if minRaw.IsZero() { + c.logger.Info("[mongo] no documents found, falling back to full table partition") + return fullTablePartitions(), nil } - if !found { - c.logger.Info("[mongo] no valid min/max ObjectIDs found, falling back to full table partition") - return fullTablePartition, nil + maxRaw, err := findBoundaryID(ctx, collection, Upper, c.config.ReadPreference) + if err != nil { + return nil, fmt.Errorf("failed to find max _id: %w", err) + } + if maxRaw.IsZero() { + c.logger.Info("[mongo] no documents found, falling back to full table partition") + return fullTablePartitions(), nil } + + switch { + case minRaw.Type == bson.TypeObjectID && maxRaw.Type == bson.TypeObjectID: + return c.objectIDPartitions(minRaw.ObjectID(), maxRaw.ObjectID(), numPartitions) + case isNumericIDType(minRaw.Type) && isNumericIDType(maxRaw.Type): + minVal, _ := rawValueToInt64(minRaw) + maxVal, _ := rawValueToInt64(maxRaw) + return c.numericPartitions(minVal, maxVal, numPartitions) + case minRaw.Type == bson.TypeString && maxRaw.Type == bson.TypeString: + return c.stringPartitions(ctx, collection, minRaw.StringValue(), maxRaw.StringValue(), numPartitions) + default: + c.logger.Info("[mongo] _id type not supported for partitioning (or mixed types), falling back to full table partition", + slog.String("minType", minRaw.Type.String()), + slog.String("maxType", maxRaw.Type.String())) + return fullTablePartitions(), nil + } +} + +// objectIDPartitions divides the ObjectID keyspace uniformly using integer +// arithmetic on the full 12-byte value. Since the leading 4-byte timestamp is the +// most significant component, this naturally partitions by insertion time. +func (c *MongoConnector) objectIDPartitions( + minID bson.ObjectID, + maxID bson.ObjectID, + numPartitions int64, +) ([]*protos.QRepPartition, error) { 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 + return fullTablePartitions(), nil } c.logger.Info("[mongo] using min/max ObjectID partitioning", @@ -88,43 +147,205 @@ func (c *MongoConnector) minMaxPartitions( 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)") +// numericPartitions divides a numeric (int32/int64) _id range uniformly, reusing +// the shared PartitionHelper which builds non-overlapping IntPartitionRange +// partitions with tested +1 boundary semantics covering [min, max] inclusive. +func (c *MongoConnector) numericPartitions( + minVal int64, + maxVal int64, + numPartitions int64, +) ([]*protos.QRepPartition, error) { + if maxVal <= minVal { + c.logger.Info("[mongo] numeric min/max range is non-positive, falling back to full table partition") + return fullTablePartitions(), nil } - n.FillBytes(oid[:]) - return oid, nil + + c.logger.Info("[mongo] using numeric _id partitioning", + slog.Int64("minID", minVal), + slog.Int64("maxID", maxVal), + slog.Int64("numPartitions", numPartitions)) + + helper := utils.NewPartitionHelper(c.logger) + if err := helper.AddPartitionsWithRange(minVal, maxVal, numPartitions); err != nil { + return nil, fmt.Errorf("failed to build numeric partitions: %w", err) + } + return helper.GetPartitions(), 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( +// stringPartitions builds partitions for a string _id collection. Because string +// keys (e.g. package names) are not uniformly distributed, uniform division of the +// keyspace would be badly skewed; instead we sample the collection and pick quantile +// boundaries so each partition holds a roughly equal share of documents. +func (c *MongoConnector) stringPartitions( ctx context.Context, collection *mongo.Collection, - readPreference protos.ReadPreference, -) (bson.ObjectID, bson.ObjectID, bool, error) { - minRaw, err := findBoundaryID(ctx, collection, Lower, readPreference) + minVal string, + maxVal string, + numPartitions int64, +) ([]*protos.QRepPartition, error) { + if minVal >= maxVal { + c.logger.Info("[mongo] string min/max range is non-positive, falling back to full table partition") + return fullTablePartitions(), nil + } + + samples, err := c.sampleStringIDs(ctx, collection, numPartitions) if err != nil { - return bson.ObjectID{}, bson.ObjectID{}, false, fmt.Errorf("failed to find min _id: %w", err) + return nil, err + } + + ranges := computeStringBoundaries(minVal, maxVal, samples, numPartitions) + if len(ranges) < 2 { + c.logger.Info("[mongo] insufficient string boundaries from sampling, falling back to full table partition", + slog.Int("sampleCount", len(samples))) + return fullTablePartitions(), nil + } + + c.logger.Info("[mongo] using sampled string _id partitioning", + slog.String("minID", minVal), + slog.String("maxID", maxVal), + slog.Int("numPartitions", len(ranges)), + slog.Int("sampleCount", len(samples))) + + partitions := make([]*protos.QRepPartition, 0, len(ranges)) + for _, rng := range ranges { + partitions = append(partitions, &protos.QRepPartition{ + PartitionId: uuid.NewString(), + Range: &protos.PartitionRange{ + Range: &protos.PartitionRange_StringRange{ + StringRange: &protos.StringPartitionRange{ + Start: rng[0], + End: rng[1], + }, + }, + }, + FullTablePartition: false, + }) } - if minRaw.Type != bson.TypeObjectID { - return bson.ObjectID{}, bson.ObjectID{}, false, nil + return partitions, nil +} + +// sampleStringIDs draws a random sample of string _id values. $sample with a size +// below ~5% of the collection uses WiredTiger's random cursor, so this is cheap even +// on large collections; it honours the configured read preference (run on a +// secondary to avoid loading the primary). +func (c *MongoConnector) sampleStringIDs( + ctx context.Context, + collection *mongo.Collection, + numPartitions int64, +) ([]string, error) { + sampleSize := numPartitions * stringSampleOversample + if sampleSize > stringSampleMaxSize { + sampleSize = stringSampleMaxSize } - maxRaw, err := findBoundaryID(ctx, collection, Upper, readPreference) + aggCmd := bson.D{ + {Key: "aggregate", Value: collection.Name()}, + {Key: "pipeline", Value: bson.A{ + bson.D{{Key: "$sample", Value: bson.D{{Key: "size", Value: int32(sampleSize)}}}}, + bson.D{{Key: "$project", Value: bson.D{{Key: DefaultDocumentKeyColumnName, Value: 1}}}}, + }}, + {Key: "cursor", Value: bson.D{}}, + } + + cursor, err := collection.Database().RunCommandCursor(ctx, aggCmd, + options.RunCmd().SetReadPreference(protoToReadPref[c.config.ReadPreference])) if err != nil { - return bson.ObjectID{}, bson.ObjectID{}, false, fmt.Errorf("failed to find max _id: %w", err) + return nil, fmt.Errorf("failed to sample _id values: %w", err) + } + defer cursor.Close(ctx) + + samples := make([]string, 0, sampleSize) + for cursor.Next(ctx) { + rv := cursor.Current.Lookup(DefaultDocumentKeyColumnName) + if rv.Type == bson.TypeString { + samples = append(samples, rv.StringValue()) + } + } + if err := cursor.Err(); err != nil { + return nil, fmt.Errorf("cursor error while sampling _id values: %w", err) + } + return samples, nil +} + +// computeStringBoundaries turns a random sample of string _ids plus the real +// min/max into a contiguous set of half-open [start, end) ranges. Interior +// boundaries are quantiles of the (deduplicated, sorted) sample; the first range +// starts at the real min and the last range ends just past the real max (via the +// sentinel) so coverage of [min, max] is exact with no gaps or overlap. +// +// It is pure (no I/O) to keep the boundary math unit-testable. Returns fewer ranges +// than numPartitions when the sample yields too few distinct interior boundaries. +func computeStringBoundaries(minVal string, maxVal string, samples []string, numPartitions int64) [][2]string { + // Keep unique sampled values strictly inside (min, max). + seen := make(map[string]struct{}, len(samples)) + interior := make([]string, 0, len(samples)) + for _, s := range samples { + if s <= minVal || s >= maxVal { + continue + } + if _, ok := seen[s]; ok { + continue + } + seen[s] = struct{}{} + interior = append(interior, s) + } + slices.Sort(interior) + + // Pick up to numPartitions-1 evenly spaced interior boundaries (quantiles). + desiredBoundaries := numPartitions - 1 + var picked []string + if int64(len(interior)) <= desiredBoundaries { + picked = interior + } else { + picked = make([]string, 0, desiredBoundaries) + for i := int64(1); i <= desiredBoundaries; i++ { + idx := i * int64(len(interior)) / numPartitions + if idx >= int64(len(interior)) { + idx = int64(len(interior)) - 1 + } + picked = append(picked, interior[idx]) + } + picked = slices.Compact(picked) // drop consecutive duplicate quantiles } - if maxRaw.Type != bson.TypeObjectID { - return bson.ObjectID{}, bson.ObjectID{}, false, nil + + // Build contiguous half-open ranges: [min, p0), [p0, p1), ..., [pLast, max+sentinel). + starts := append([]string{minVal}, picked...) + ranges := make([][2]string, len(starts)) + for i := range starts { + var end string + if i+1 < len(starts) { + end = starts[i+1] + } else { + end = maxVal + stringPartitionUpperSentinel + } + ranges[i] = [2]string{starts[i], end} } + return ranges +} - return minRaw.ObjectID(), maxRaw.ObjectID(), true, 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 +} + +func isNumericIDType(t bson.Type) bool { + return t == bson.TypeInt32 || t == bson.TypeInt64 +} + +func rawValueToInt64(rv bson.RawValue) (int64, bool) { + switch rv.Type { + case bson.TypeInt32: + return int64(rv.Int32()), true + case bson.TypeInt64: + return rv.Int64(), true + default: + return 0, false + } } type Boundary int diff --git a/flow/connectors/mongo/qrep_partition_test.go b/flow/connectors/mongo/qrep_partition_test.go new file mode 100644 index 000000000..24daed131 --- /dev/null +++ b/flow/connectors/mongo/qrep_partition_test.go @@ -0,0 +1,170 @@ +package connmongo + +import ( + "testing" + + "github.com/stretchr/testify/require" + "go.mongodb.org/mongo-driver/v2/bson" + + "github.com/PeerDB-io/peerdb/flow/generated/protos" +) + +// rawValueOf marshals a single-field document and returns the _id RawValue, so we +// can exercise the BSON type-detection helpers with real driver values. +func rawValueOf(t *testing.T, id any) bson.RawValue { + t.Helper() + raw, err := bson.Marshal(bson.D{{Key: DefaultDocumentKeyColumnName, Value: id}}) + require.NoError(t, err) + return bson.Raw(raw).Lookup(DefaultDocumentKeyColumnName) +} + +func TestRawValueToInt64(t *testing.T) { + v, ok := rawValueToInt64(rawValueOf(t, int32(281704574))) + require.True(t, ok) + require.Equal(t, int64(281704574), v) + + v, ok = rawValueToInt64(rawValueOf(t, int64(9000000000))) + require.True(t, ok) + require.Equal(t, int64(9000000000), v) + + _, ok = rawValueToInt64(rawValueOf(t, "flipboard.app")) + require.False(t, ok) + + _, ok = rawValueToInt64(rawValueOf(t, 3.14)) + require.False(t, ok) +} + +func TestIsNumericIDType(t *testing.T) { + require.True(t, isNumericIDType(rawValueOf(t, int32(1)).Type)) + require.True(t, isNumericIDType(rawValueOf(t, int64(1)).Type)) + require.False(t, isNumericIDType(rawValueOf(t, "x").Type)) + require.False(t, isNumericIDType(rawValueOf(t, 1.5).Type)) + oid, _ := bson.ObjectIDFromHex("507f1f77bcf86cd799439011") + require.False(t, isNumericIDType(rawValueOf(t, oid).Type)) +} + +// assertContiguous verifies the ranges form a gap-free, non-overlapping cover of +// [min, max+sentinel): first start == min, last end == max+sentinel, and each +// range's end equals the next range's start. +func assertContiguous(t *testing.T, ranges [][2]string, minVal, maxVal string) { + t.Helper() + require.NotEmpty(t, ranges) + require.Equal(t, minVal, ranges[0][0], "first partition must start at min") + require.Equal(t, maxVal+stringPartitionUpperSentinel, ranges[len(ranges)-1][1], "last partition must end just past max") + for i := range ranges { + require.Less(t, ranges[i][0], ranges[i][1], "range start must be < end") + if i+1 < len(ranges) { + require.Equal(t, ranges[i][1], ranges[i+1][0], "ranges must be contiguous") + } + } +} + +func TestComputeStringBoundaries(t *testing.T) { + t.Run("even distribution", func(t *testing.T) { + samples := []string{"b", "c", "d", "e", "f", "g", "h", "i"} + ranges := computeStringBoundaries("a", "z", samples, 4) + require.Len(t, ranges, 4) + assertContiguous(t, ranges, "a", "z") + }) + + t.Run("clustered samples still contiguous", func(t *testing.T) { + samples := []string{"com.a", "com.b", "com.c", "com.d", "com.e", "com.f"} + ranges := computeStringBoundaries("com.a", "org.z", samples, 3) + require.LessOrEqual(t, len(ranges), 3) + assertContiguous(t, ranges, "com.a", "org.z") + }) + + t.Run("fewer distinct samples than partitions", func(t *testing.T) { + ranges := computeStringBoundaries("a", "z", []string{"m"}, 4) + require.Len(t, ranges, 2) // one interior boundary -> two partitions + assertContiguous(t, ranges, "a", "z") + require.Equal(t, [2]string{"a", "m"}, ranges[0]) + }) + + t.Run("duplicate samples deduplicated", func(t *testing.T) { + samples := []string{"m", "m", "m", "m"} + ranges := computeStringBoundaries("a", "z", samples, 4) + require.Len(t, ranges, 2) + assertContiguous(t, ranges, "a", "z") + }) + + t.Run("samples outside (min,max) ignored", func(t *testing.T) { + // "a" == min and "z" == max must be dropped; only "m" is a valid interior boundary + samples := []string{"a", "z", "m", "0", "zzzz"} + ranges := computeStringBoundaries("a", "z", samples, 4) + require.Len(t, ranges, 2) + assertContiguous(t, ranges, "a", "z") + require.Equal(t, [2]string{"a", "m"}, ranges[0]) + }) + + t.Run("no usable samples yields single range", func(t *testing.T) { + ranges := computeStringBoundaries("a", "z", nil, 4) + require.Len(t, ranges, 1) // caller treats <2 as full-table fallback + assertContiguous(t, ranges, "a", "z") + }) + + t.Run("many samples produce requested partition count", func(t *testing.T) { + samples := make([]string, 0, 100) + for i := 'b'; i < 'b'+20; i++ { + for j := 'a'; j < 'a'+5; j++ { + samples = append(samples, string([]rune{i, j})) + } + } + ranges := computeStringBoundaries("aa", "zz", samples, 8) + require.Len(t, ranges, 8) + assertContiguous(t, ranges, "aa", "zz") + }) +} + +func intPartition(start, end int64) *protos.PartitionRange { + return &protos.PartitionRange{Range: &protos.PartitionRange_IntRange{ + IntRange: &protos.IntPartitionRange{Start: start, End: end}, + }} +} + +func stringPartition(start, end string) *protos.PartitionRange { + return &protos.PartitionRange{Range: &protos.PartitionRange_StringRange{ + StringRange: &protos.StringPartitionRange{Start: start, End: end}, + }} +} + +func TestToRangeFilter(t *testing.T) { + t.Run("int range is inclusive on both ends", func(t *testing.T) { + filter, err := toRangeFilter(DefaultDocumentKeyColumnName, intPartition(10, 20)) + require.NoError(t, err) + require.Equal(t, bson.D{ + bson.E{Key: DefaultDocumentKeyColumnName, Value: bson.D{ + bson.E{Key: "$gte", Value: int64(10)}, + bson.E{Key: "$lte", Value: int64(20)}, + }}, + }, filter) + }) + + t.Run("string range is half-open", func(t *testing.T) { + filter, err := toRangeFilter(DefaultDocumentKeyColumnName, stringPartition("com.a", "com.m")) + require.NoError(t, err) + require.Equal(t, bson.D{ + bson.E{Key: DefaultDocumentKeyColumnName, Value: bson.D{ + bson.E{Key: "$gte", Value: "com.a"}, + bson.E{Key: "$lt", Value: "com.m"}, + }}, + }, filter) + }) + + t.Run("object id range still supported", func(t *testing.T) { + r := &protos.PartitionRange{Range: &protos.PartitionRange_ObjectIdRange{ + ObjectIdRange: &protos.ObjectIdPartitionRange{ + Start: "507f1f77bcf86cd799439011", + End: "507f1f77bcf86cd7994390ff", + }, + }} + _, err := toRangeFilter(DefaultDocumentKeyColumnName, r) + require.NoError(t, err) + }) + + t.Run("unsupported range type errors", func(t *testing.T) { + r := &protos.PartitionRange{Range: &protos.PartitionRange_NullRange{NullRange: &protos.NullPartitionRange{}}} + _, err := toRangeFilter(DefaultDocumentKeyColumnName, r) + require.Error(t, err) + }) +} diff --git a/flow/connectors/utils/monitoring/monitoring.go b/flow/connectors/utils/monitoring/monitoring.go index 147291bbf..98fa75e6a 100644 --- a/flow/connectors/utils/monitoring/monitoring.go +++ b/flow/connectors/utils/monitoring/monitoring.go @@ -362,6 +362,8 @@ func addPartitionToQRepRun(ctx context.Context, tx pgx.Tx, flowJobName string, rangeEnd = new(rangeEndValue.(string)) case *protos.PartitionRange_ObjectIdRange: rangeStart, rangeEnd = &x.ObjectIdRange.Start, &x.ObjectIdRange.End + case *protos.PartitionRange_StringRange: + rangeStart, rangeEnd = &x.StringRange.Start, &x.StringRange.End case *protos.PartitionRange_NullRange: // leave rangeStart and rangeEnd as nil default: diff --git a/protos/flow.proto b/protos/flow.proto index b9964bacc..717ed6892 100644 --- a/protos/flow.proto +++ b/protos/flow.proto @@ -324,6 +324,11 @@ message ObjectIdPartitionRange { string end = 2; } +message StringPartitionRange { + string start = 1; + string end = 2; +} + message PartitionRange { // can be a timestamp range or an integer range oneof range { @@ -333,6 +338,7 @@ message PartitionRange { UIntPartitionRange uint_range = 4; ObjectIdPartitionRange object_id_range = 5; NullPartitionRange null_range = 6; + StringPartitionRange string_range = 7; } } From 925889be164242daed25b1730ad1697efd2caeca Mon Sep 17 00:00:00 2001 From: Andrey Zhelnin Date: Thu, 25 Jun 2026 11:01:18 +0200 Subject: [PATCH 2/5] mongo: address PR review comments on _id partitioning - Add end_inclusive to StringPartitionRange proto (aligns with MySQL PR #4482 adding the same field; avoids future merge conflict) - Replace sentinel-string approach with end_inclusive=true on the last string partition, using $lte instead of $lt for its MongoDB filter - Sort string _id samples via MongoDB $sort stage instead of Go sort, preserving the collection's collation order - Filter boundary samples by exact equality instead of <=/>= comparison, trusting MongoDB's min/max guarantees - Add e2e tests for string and numeric _id partitioned snapshot flows - Document the double _id gap limitation on numericPartitions (follow-up) --- flow/connectors/mongo/qrep.go | 11 +- flow/connectors/mongo/qrep_partition.go | 48 ++++---- flow/connectors/mongo/qrep_partition_test.go | 30 +++-- flow/e2e/mongo_test.go | 114 +++++++++++++++++++ protos/flow.proto | 4 + 5 files changed, 173 insertions(+), 34 deletions(-) diff --git a/flow/connectors/mongo/qrep.go b/flow/connectors/mongo/qrep.go index 3e5c7eea5..690b68249 100644 --- a/flow/connectors/mongo/qrep.go +++ b/flow/connectors/mongo/qrep.go @@ -245,13 +245,16 @@ func toRangeFilter(watermarkColumn string, partitionRange *protos.PartitionRange }}, }, nil case *protos.PartitionRange_StringRange: - // String _id: half-open [start, end) range. computeStringBoundaries makes - // ranges contiguous and sets the final end just past the real max, so $lt - // gives exact coverage with no gaps or overlap. + // String _id: half-open [start, end) for all but the last partition; + // the last partition is closed [start, end] (EndInclusive == true). + endOp := "$lt" + if r.StringRange.EndInclusive { + endOp = "$lte" + } return bson.D{ bson.E{Key: watermarkColumn, Value: bson.D{ bson.E{Key: "$gte", Value: r.StringRange.Start}, - bson.E{Key: "$lt", Value: r.StringRange.End}, + bson.E{Key: endOp, Value: r.StringRange.End}, }}, }, nil default: diff --git a/flow/connectors/mongo/qrep_partition.go b/flow/connectors/mongo/qrep_partition.go index d70631d9e..34876d269 100644 --- a/flow/connectors/mongo/qrep_partition.go +++ b/flow/connectors/mongo/qrep_partition.go @@ -17,14 +17,6 @@ import ( "github.com/PeerDB-io/peerdb/flow/shared" ) -// stringPartitionUpperSentinel is appended to the collection's max string _id to -// form the exclusive upper bound of the final string partition. It is the highest -// valid Unicode code point, so `max + sentinel` sorts strictly after `max` (and -// after every real _id, since `max` is the maximum). This lets every string -// partition use uniform half-open [start, end) ($gte/$lt) semantics while still -// capturing the maximum document in the last partition. -const stringPartitionUpperSentinel = "\U0010FFFF" - const ( // stringSampleOversample draws more samples than partitions so quantile // boundaries are well distributed even with clustered keys. @@ -150,6 +142,12 @@ func (c *MongoConnector) objectIDPartitions( // numericPartitions divides a numeric (int32/int64) _id range uniformly, reusing // the shared PartitionHelper which builds non-overlapping IntPartitionRange // partitions with tested +1 boundary semantics covering [min, max] inclusive. +// +// Limitation: IntPartitionRange uses inclusive-inclusive [start, end] boundaries, +// so documents with double _id values that fall between two adjacent integer +// boundaries (e.g. _id=2.5 between [1,2] and [3,4]) would be missed. This is an +// edge case for collections with mixed int/double _ids; a follow-up will introduce +// a NumericPartitionRange with half-open [start, end) semantics to address it. func (c *MongoConnector) numericPartitions( minVal int64, maxVal int64, @@ -207,14 +205,15 @@ func (c *MongoConnector) stringPartitions( slog.Int("sampleCount", len(samples))) partitions := make([]*protos.QRepPartition, 0, len(ranges)) - for _, rng := range ranges { + for i, rng := range ranges { partitions = append(partitions, &protos.QRepPartition{ PartitionId: uuid.NewString(), Range: &protos.PartitionRange{ Range: &protos.PartitionRange_StringRange{ StringRange: &protos.StringPartitionRange{ - Start: rng[0], - End: rng[1], + Start: rng[0], + End: rng[1], + EndInclusive: i == len(ranges)-1, }, }, }, @@ -242,6 +241,7 @@ func (c *MongoConnector) sampleStringIDs( {Key: "aggregate", Value: collection.Name()}, {Key: "pipeline", Value: bson.A{ bson.D{{Key: "$sample", Value: bson.D{{Key: "size", Value: int32(sampleSize)}}}}, + bson.D{{Key: "$sort", Value: bson.D{{Key: DefaultDocumentKeyColumnName, Value: 1}}}}, bson.D{{Key: "$project", Value: bson.D{{Key: DefaultDocumentKeyColumnName, Value: 1}}}}, }}, {Key: "cursor", Value: bson.D{}}, @@ -267,20 +267,26 @@ func (c *MongoConnector) sampleStringIDs( return samples, nil } -// computeStringBoundaries turns a random sample of string _ids plus the real -// min/max into a contiguous set of half-open [start, end) ranges. Interior -// boundaries are quantiles of the (deduplicated, sorted) sample; the first range -// starts at the real min and the last range ends just past the real max (via the -// sentinel) so coverage of [min, max] is exact with no gaps or overlap. +// computeStringBoundaries turns a pre-sorted (by database collation) sample of +// string _ids plus the real min/max into a contiguous set of ranges. Interior +// boundaries are quantiles of the (deduplicated) sample; the first range starts +// at the real min and the last range ends at the real max. All ranges except the +// last are half-open [start, end); the last is closed [start, end]. +// +// samples must already be sorted in the database's collation order; this function +// does not re-sort them, to preserve collation semantics. // // It is pure (no I/O) to keep the boundary math unit-testable. Returns fewer ranges // than numPartitions when the sample yields too few distinct interior boundaries. func computeStringBoundaries(minVal string, maxVal string, samples []string, numPartitions int64) [][2]string { - // Keep unique sampled values strictly inside (min, max). + // Keep unique sampled values strictly inside (min, max), preserving order. + // We use exact equality to filter boundaries: MongoDB guarantees all sampled + // _ids are within [minVal, maxVal] by its own collation; only exact matches + // with the boundary values would create zero-width partitions. seen := make(map[string]struct{}, len(samples)) interior := make([]string, 0, len(samples)) for _, s := range samples { - if s <= minVal || s >= maxVal { + if s == minVal || s == maxVal { continue } if _, ok := seen[s]; ok { @@ -289,7 +295,6 @@ func computeStringBoundaries(minVal string, maxVal string, samples []string, num seen[s] = struct{}{} interior = append(interior, s) } - slices.Sort(interior) // Pick up to numPartitions-1 evenly spaced interior boundaries (quantiles). desiredBoundaries := numPartitions - 1 @@ -308,7 +313,8 @@ func computeStringBoundaries(minVal string, maxVal string, samples []string, num picked = slices.Compact(picked) // drop consecutive duplicate quantiles } - // Build contiguous half-open ranges: [min, p0), [p0, p1), ..., [pLast, max+sentinel). + // Build ranges: [min, p0), [p0, p1), ..., [pLast-1, pLast), [pLast, max]. + // The last range is closed (end_inclusive=true); all others are half-open. starts := append([]string{minVal}, picked...) ranges := make([][2]string, len(starts)) for i := range starts { @@ -316,7 +322,7 @@ func computeStringBoundaries(minVal string, maxVal string, samples []string, num if i+1 < len(starts) { end = starts[i+1] } else { - end = maxVal + stringPartitionUpperSentinel + end = maxVal } ranges[i] = [2]string{starts[i], end} } diff --git a/flow/connectors/mongo/qrep_partition_test.go b/flow/connectors/mongo/qrep_partition_test.go index 24daed131..0cc44a764 100644 --- a/flow/connectors/mongo/qrep_partition_test.go +++ b/flow/connectors/mongo/qrep_partition_test.go @@ -44,13 +44,13 @@ func TestIsNumericIDType(t *testing.T) { } // assertContiguous verifies the ranges form a gap-free, non-overlapping cover of -// [min, max+sentinel): first start == min, last end == max+sentinel, and each -// range's end equals the next range's start. +// [min, max]: first start == min, last end == max, and each range's end equals +// the next range's start (half-open except for the last which is closed). func assertContiguous(t *testing.T, ranges [][2]string, minVal, maxVal string) { t.Helper() require.NotEmpty(t, ranges) require.Equal(t, minVal, ranges[0][0], "first partition must start at min") - require.Equal(t, maxVal+stringPartitionUpperSentinel, ranges[len(ranges)-1][1], "last partition must end just past max") + require.Equal(t, maxVal, ranges[len(ranges)-1][1], "last partition must end at max") for i := range ranges { require.Less(t, ranges[i][0], ranges[i][1], "range start must be < end") if i+1 < len(ranges) { @@ -88,9 +88,10 @@ func TestComputeStringBoundaries(t *testing.T) { assertContiguous(t, ranges, "a", "z") }) - t.Run("samples outside (min,max) ignored", func(t *testing.T) { - // "a" == min and "z" == max must be dropped; only "m" is a valid interior boundary - samples := []string{"a", "z", "m", "0", "zzzz"} + t.Run("boundary samples ignored", func(t *testing.T) { + // "a" == min and "z" == max must be dropped to avoid zero-width partitions; + // only "m" is a valid interior boundary. Samples come pre-sorted from MongoDB. + samples := []string{"a", "m", "z"} ranges := computeStringBoundaries("a", "z", samples, 4) require.Len(t, ranges, 2) assertContiguous(t, ranges, "a", "z") @@ -122,9 +123,9 @@ func intPartition(start, end int64) *protos.PartitionRange { }} } -func stringPartition(start, end string) *protos.PartitionRange { +func stringPartition(start, end string, endInclusive bool) *protos.PartitionRange { return &protos.PartitionRange{Range: &protos.PartitionRange_StringRange{ - StringRange: &protos.StringPartitionRange{Start: start, End: end}, + StringRange: &protos.StringPartitionRange{Start: start, End: end, EndInclusive: endInclusive}, }} } @@ -141,7 +142,7 @@ func TestToRangeFilter(t *testing.T) { }) t.Run("string range is half-open", func(t *testing.T) { - filter, err := toRangeFilter(DefaultDocumentKeyColumnName, stringPartition("com.a", "com.m")) + filter, err := toRangeFilter(DefaultDocumentKeyColumnName, stringPartition("com.a", "com.m", false)) require.NoError(t, err) require.Equal(t, bson.D{ bson.E{Key: DefaultDocumentKeyColumnName, Value: bson.D{ @@ -151,6 +152,17 @@ func TestToRangeFilter(t *testing.T) { }, filter) }) + t.Run("last string range is closed", func(t *testing.T) { + filter, err := toRangeFilter(DefaultDocumentKeyColumnName, stringPartition("com.m", "org.z", true)) + require.NoError(t, err) + require.Equal(t, bson.D{ + bson.E{Key: DefaultDocumentKeyColumnName, Value: bson.D{ + bson.E{Key: "$gte", Value: "com.m"}, + bson.E{Key: "$lte", Value: "org.z"}, + }}, + }, filter) + }) + t.Run("object id range still supported", func(t *testing.T) { r := &protos.PartitionRange{Range: &protos.PartitionRange_ObjectIdRange{ ObjectIdRange: &protos.ObjectIdPartitionRange{ diff --git a/flow/e2e/mongo_test.go b/flow/e2e/mongo_test.go index 6ee223e3d..b35fead3a 100644 --- a/flow/e2e/mongo_test.go +++ b/flow/e2e/mongo_test.go @@ -157,6 +157,120 @@ func (s MongoClickhouseSuite) Test_Simple_Flow_Partitioned() { RequireEnvCanceled(t, env) } +func (s MongoClickhouseSuite) Test_Simple_Flow_Partitioned_StringID() { + t := s.T() + srcDatabase := GetTestDatabase(s.Suffix()) + srcTable := "test_simple_partitioned_string_id" + dstTable := "test_simple_dst_partitioned_string_id" + + connectionGen := FlowConnectionGenerationConfig{ + FlowJobName: AddSuffix(s, srcTable), + TableMappings: TableMappings(s, srcTable, dstTable), + Destination: s.Peer().Name, + } + flowConnConfig := s.generateFlowConnectionConfigsDefaultEnv(connectionGen) + flowConnConfig.DoInitialSnapshot = true + flowConnConfig.SnapshotNumRowsPerPartition = 10 + + adminClient := s.Source().(*MongoSource).AdminClient() + collection := adminClient.Database(srcDatabase).Collection(srcTable) + + docs := make([]any, 1000) + for i := range 1000 { + docs[i] = bson.D{ + {Key: "_id", Value: fmt.Sprintf("id-%05d", i)}, + {Key: fmt.Sprintf("init_key_%d", i), Value: fmt.Sprintf("init_value_%d", i)}, + } + } + _, err := collection.InsertMany(t.Context(), docs) + require.NoError(t, err) + + tc := NewTemporalClient(t) + env := ExecutePeerflow(t, tc, flowConnConfig) + + EnvWaitForEqualTablesWithNames(env, s, "initial load to match", srcTable, dstTable, "_id,doc") + + catalogPool, err := internal.GetCatalogConnectionPoolFromEnv(t.Context()) + require.NoError(t, err) + var partitionCount int + require.NoError(t, catalogPool.QueryRow(t.Context(), + `SELECT COUNT(*) FROM peerdb_stats.qrep_partitions WHERE parent_mirror_name = $1`, + flowConnConfig.FlowJobName).Scan(&partitionCount)) + require.Greater(t, partitionCount, 1, "expected multiple partitions for 1000 string-id rows") + + SetupCDCFlowStatusQuery(t, env, flowConnConfig) + cdcDocs := make([]any, 10) + for i := range 10 { + cdcDocs[i] = bson.D{ + {Key: "_id", Value: fmt.Sprintf("cdc-%05d", i)}, + {Key: fmt.Sprintf("cdc_key_%d", i), Value: fmt.Sprintf("cdc_value_%d", i)}, + } + } + _, err = collection.InsertMany(t.Context(), cdcDocs) + require.NoError(t, err) + + EnvWaitForEqualTablesWithNames(env, s, "cdc events to match", srcTable, dstTable, "_id,doc") + env.Cancel(t.Context()) + RequireEnvCanceled(t, env) +} + +func (s MongoClickhouseSuite) Test_Simple_Flow_Partitioned_NumericID() { + t := s.T() + srcDatabase := GetTestDatabase(s.Suffix()) + srcTable := "test_simple_partitioned_numeric_id" + dstTable := "test_simple_dst_partitioned_numeric_id" + + connectionGen := FlowConnectionGenerationConfig{ + FlowJobName: AddSuffix(s, srcTable), + TableMappings: TableMappings(s, srcTable, dstTable), + Destination: s.Peer().Name, + } + flowConnConfig := s.generateFlowConnectionConfigsDefaultEnv(connectionGen) + flowConnConfig.DoInitialSnapshot = true + flowConnConfig.SnapshotNumRowsPerPartition = 10 + + adminClient := s.Source().(*MongoSource).AdminClient() + collection := adminClient.Database(srcDatabase).Collection(srcTable) + + docs := make([]any, 1000) + for i := range 1000 { + docs[i] = bson.D{ + {Key: "_id", Value: int32(i + 1)}, + {Key: fmt.Sprintf("init_key_%d", i), Value: fmt.Sprintf("init_value_%d", i)}, + } + } + _, err := collection.InsertMany(t.Context(), docs) + require.NoError(t, err) + + tc := NewTemporalClient(t) + env := ExecutePeerflow(t, tc, flowConnConfig) + + EnvWaitForEqualTablesWithNames(env, s, "initial load to match", srcTable, dstTable, "_id,doc") + + catalogPool, err := internal.GetCatalogConnectionPoolFromEnv(t.Context()) + require.NoError(t, err) + var partitionCount int + require.NoError(t, catalogPool.QueryRow(t.Context(), + `SELECT COUNT(*) FROM peerdb_stats.qrep_partitions WHERE parent_mirror_name = $1`, + flowConnConfig.FlowJobName).Scan(&partitionCount)) + require.Greater(t, partitionCount, 1, "expected multiple partitions for 1000 numeric-id rows") + + SetupCDCFlowStatusQuery(t, env, flowConnConfig) + cdcDocs := make([]any, 10) + for i := range 10 { + cdcDocs[i] = bson.D{ + {Key: "_id", Value: int32(10000 + i)}, + {Key: fmt.Sprintf("cdc_key_%d", i), Value: fmt.Sprintf("cdc_value_%d", i)}, + } + } + _, err = collection.InsertMany(t.Context(), cdcDocs) + require.NoError(t, err) + + EnvWaitForEqualTablesWithNames(env, s, "cdc events to match", srcTable, dstTable, "_id,doc") + env.Cancel(t.Context()) + RequireEnvCanceled(t, env) +} + func (s MongoClickhouseSuite) Test_Snapshot_Collection_With_Single_Document() { t := s.T() srcDatabase := GetTestDatabase(s.Suffix()) diff --git a/protos/flow.proto b/protos/flow.proto index 717ed6892..a2ed13c83 100644 --- a/protos/flow.proto +++ b/protos/flow.proto @@ -327,6 +327,10 @@ message ObjectIdPartitionRange { message StringPartitionRange { string start = 1; string end = 2; + // unlike numeric/temporal ranges, string boundaries cannot be safely + // incremented; so the partition boundaries are [start, end) except + // for the last partition being [start, end]; hence this bool is needed. + bool end_inclusive = 3; } message PartitionRange { From 8dac19332baff440051874051b7f67a6fc5c27fe Mon Sep 17 00:00:00 2001 From: Andrey Zhelnin Date: Thu, 25 Jun 2026 14:51:28 +0200 Subject: [PATCH 3/5] mongo: reduce e2e test doc count to 100 for faster local runs --- flow/e2e/mongo_test.go | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/flow/e2e/mongo_test.go b/flow/e2e/mongo_test.go index b35fead3a..e4390eea9 100644 --- a/flow/e2e/mongo_test.go +++ b/flow/e2e/mongo_test.go @@ -175,8 +175,8 @@ func (s MongoClickhouseSuite) Test_Simple_Flow_Partitioned_StringID() { adminClient := s.Source().(*MongoSource).AdminClient() collection := adminClient.Database(srcDatabase).Collection(srcTable) - docs := make([]any, 1000) - for i := range 1000 { + docs := make([]any, 100) + for i := range 100 { docs[i] = bson.D{ {Key: "_id", Value: fmt.Sprintf("id-%05d", i)}, {Key: fmt.Sprintf("init_key_%d", i), Value: fmt.Sprintf("init_value_%d", i)}, @@ -196,7 +196,7 @@ func (s MongoClickhouseSuite) Test_Simple_Flow_Partitioned_StringID() { require.NoError(t, catalogPool.QueryRow(t.Context(), `SELECT COUNT(*) FROM peerdb_stats.qrep_partitions WHERE parent_mirror_name = $1`, flowConnConfig.FlowJobName).Scan(&partitionCount)) - require.Greater(t, partitionCount, 1, "expected multiple partitions for 1000 string-id rows") + require.Greater(t, partitionCount, 1, "expected multiple partitions for 100 string-id rows") SetupCDCFlowStatusQuery(t, env, flowConnConfig) cdcDocs := make([]any, 10) @@ -232,8 +232,8 @@ func (s MongoClickhouseSuite) Test_Simple_Flow_Partitioned_NumericID() { adminClient := s.Source().(*MongoSource).AdminClient() collection := adminClient.Database(srcDatabase).Collection(srcTable) - docs := make([]any, 1000) - for i := range 1000 { + docs := make([]any, 100) + for i := range 100 { docs[i] = bson.D{ {Key: "_id", Value: int32(i + 1)}, {Key: fmt.Sprintf("init_key_%d", i), Value: fmt.Sprintf("init_value_%d", i)}, @@ -253,7 +253,7 @@ func (s MongoClickhouseSuite) Test_Simple_Flow_Partitioned_NumericID() { require.NoError(t, catalogPool.QueryRow(t.Context(), `SELECT COUNT(*) FROM peerdb_stats.qrep_partitions WHERE parent_mirror_name = $1`, flowConnConfig.FlowJobName).Scan(&partitionCount)) - require.Greater(t, partitionCount, 1, "expected multiple partitions for 1000 numeric-id rows") + require.Greater(t, partitionCount, 1, "expected multiple partitions for 100 numeric-id rows") SetupCDCFlowStatusQuery(t, env, flowConnConfig) cdcDocs := make([]any, 10) From b89e2fc707693e9db11a06f1552bcec6dc9f5779 Mon Sep 17 00:00:00 2001 From: Andrey Zhelnin Date: Fri, 26 Jun 2026 16:54:46 +0200 Subject: [PATCH 4/5] mongo: address PR review comments (round 2) - Add utils.FullTablePartition() helper and use it in place of the local fullTablePartitions() in qrep_partition.go and the inline construction in qrep.go - Remove slices.Compact() from computeStringBoundaries: the seen map already deduplicates interior values, and the quantile index formula is strictly increasing when len(interior) >= numPartitions, so no consecutive duplicates can occur in the else branch - Tighten e2e partition count assertions to exact counts (10) instead of >1 for string and numeric _id tests - Delete Test_Snapshot_Non_ObjectID_Falls_Back_To_Single_Partition: string _id collections are now partitioned rather than falling back to full-table, so the test expectation no longer holds --- flow/connectors/mongo/qrep.go | 12 +----- flow/connectors/mongo/qrep_partition.go | 28 ++++---------- flow/connectors/utils/partition.go | 8 ++++ flow/e2e/mongo_test.go | 49 +------------------------ 4 files changed, 19 insertions(+), 78 deletions(-) diff --git a/flow/connectors/mongo/qrep.go b/flow/connectors/mongo/qrep.go index 690b68249..7290b08b3 100644 --- a/flow/connectors/mongo/qrep.go +++ b/flow/connectors/mongo/qrep.go @@ -26,17 +26,9 @@ func (c *MongoConnector) GetQRepPartitions( config *protos.QRepConfig, last *protos.QRepPartition, ) ([]*protos.QRepPartition, error) { - fullTablePartition := []*protos.QRepPartition{ - { - PartitionId: utils.FullTablePartitionID, - Range: nil, - FullTablePartition: true, - }, - } - if config.WatermarkColumn != DefaultDocumentKeyColumnName { c.logger.Warn("unexpected watermark column, falling back to full table partition") - return fullTablePartition, nil + return utils.FullTablePartition(), nil } if config.NumRowsPerPartition <= 0 { @@ -70,7 +62,7 @@ func (c *MongoConnector) GetQRepPartitions( if adjustedPartitions.AdjustedNumPartitions <= 1 { c.logger.Info("[mongo] insufficient partitions for parallel snapshot, falling back to full table partition") - return fullTablePartition, nil + return utils.FullTablePartition(), nil } return c.buildPartitions(ctx, collection, adjustedPartitions.AdjustedNumPartitions) diff --git a/flow/connectors/mongo/qrep_partition.go b/flow/connectors/mongo/qrep_partition.go index 34876d269..320ad5147 100644 --- a/flow/connectors/mongo/qrep_partition.go +++ b/flow/connectors/mongo/qrep_partition.go @@ -5,8 +5,6 @@ import ( "fmt" "log/slog" "math/big" - "slices" - "github.com/google/uuid" "go.mongodb.org/mongo-driver/v2/bson" "go.mongodb.org/mongo-driver/v2/mongo" @@ -25,17 +23,6 @@ const ( stringSampleMaxSize = 100000 ) -// fullTablePartitions returns the single-partition full-table fallback used when a -// collection cannot be partitioned (unsupported/mixed _id type, empty collection, -// or degenerate min/max range). -func fullTablePartitions() []*protos.QRepPartition { - return []*protos.QRepPartition{{ - PartitionId: utils.FullTablePartitionID, - Range: nil, - FullTablePartition: true, - }} -} - // buildPartitions reads the collection's min and max _id (cheap, leverages the // default _id index) and dispatches to a type-specific partitioner: // - ObjectID -> uniform division of the 12-byte ObjectID keyspace @@ -56,7 +43,7 @@ func (c *MongoConnector) buildPartitions( } if minRaw.IsZero() { c.logger.Info("[mongo] no documents found, falling back to full table partition") - return fullTablePartitions(), nil + return utils.FullTablePartition(), nil } maxRaw, err := findBoundaryID(ctx, collection, Upper, c.config.ReadPreference) if err != nil { @@ -64,7 +51,7 @@ func (c *MongoConnector) buildPartitions( } if maxRaw.IsZero() { c.logger.Info("[mongo] no documents found, falling back to full table partition") - return fullTablePartitions(), nil + return utils.FullTablePartition(), nil } switch { @@ -80,7 +67,7 @@ func (c *MongoConnector) buildPartitions( c.logger.Info("[mongo] _id type not supported for partitioning (or mixed types), falling back to full table partition", slog.String("minType", minRaw.Type.String()), slog.String("maxType", maxRaw.Type.String())) - return fullTablePartitions(), nil + return utils.FullTablePartition(), nil } } @@ -97,7 +84,7 @@ func (c *MongoConnector) objectIDPartitions( 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 fullTablePartitions(), nil + return utils.FullTablePartition(), nil } c.logger.Info("[mongo] using min/max ObjectID partitioning", @@ -155,7 +142,7 @@ func (c *MongoConnector) numericPartitions( ) ([]*protos.QRepPartition, error) { if maxVal <= minVal { c.logger.Info("[mongo] numeric min/max range is non-positive, falling back to full table partition") - return fullTablePartitions(), nil + return utils.FullTablePartition(), nil } c.logger.Info("[mongo] using numeric _id partitioning", @@ -183,7 +170,7 @@ func (c *MongoConnector) stringPartitions( ) ([]*protos.QRepPartition, error) { if minVal >= maxVal { c.logger.Info("[mongo] string min/max range is non-positive, falling back to full table partition") - return fullTablePartitions(), nil + return utils.FullTablePartition(), nil } samples, err := c.sampleStringIDs(ctx, collection, numPartitions) @@ -195,7 +182,7 @@ func (c *MongoConnector) stringPartitions( if len(ranges) < 2 { c.logger.Info("[mongo] insufficient string boundaries from sampling, falling back to full table partition", slog.Int("sampleCount", len(samples))) - return fullTablePartitions(), nil + return utils.FullTablePartition(), nil } c.logger.Info("[mongo] using sampled string _id partitioning", @@ -310,7 +297,6 @@ func computeStringBoundaries(minVal string, maxVal string, samples []string, num } picked = append(picked, interior[idx]) } - picked = slices.Compact(picked) // drop consecutive duplicate quantiles } // Build ranges: [min, p0), [p0, p1), ..., [pLast-1, pLast), [pLast, max]. diff --git a/flow/connectors/utils/partition.go b/flow/connectors/utils/partition.go index e1d1721a7..965aecd72 100644 --- a/flow/connectors/utils/partition.go +++ b/flow/connectors/utils/partition.go @@ -17,6 +17,14 @@ import ( const FullTablePartitionID = "full-table-partition-id" +func FullTablePartition() []*protos.QRepPartition { + return []*protos.QRepPartition{{ + PartitionId: FullTablePartitionID, + Range: nil, + FullTablePartition: true, + }} +} + type PartitionRangeType string const ( diff --git a/flow/e2e/mongo_test.go b/flow/e2e/mongo_test.go index e4390eea9..62a35f62a 100644 --- a/flow/e2e/mongo_test.go +++ b/flow/e2e/mongo_test.go @@ -196,7 +196,7 @@ func (s MongoClickhouseSuite) Test_Simple_Flow_Partitioned_StringID() { require.NoError(t, catalogPool.QueryRow(t.Context(), `SELECT COUNT(*) FROM peerdb_stats.qrep_partitions WHERE parent_mirror_name = $1`, flowConnConfig.FlowJobName).Scan(&partitionCount)) - require.Greater(t, partitionCount, 1, "expected multiple partitions for 100 string-id rows") + require.Equal(t, 10, partitionCount, "expected 10 partitions for 100 rows with 10 rows per partition") SetupCDCFlowStatusQuery(t, env, flowConnConfig) cdcDocs := make([]any, 10) @@ -253,7 +253,7 @@ func (s MongoClickhouseSuite) Test_Simple_Flow_Partitioned_NumericID() { require.NoError(t, catalogPool.QueryRow(t.Context(), `SELECT COUNT(*) FROM peerdb_stats.qrep_partitions WHERE parent_mirror_name = $1`, flowConnConfig.FlowJobName).Scan(&partitionCount)) - require.Greater(t, partitionCount, 1, "expected multiple partitions for 100 numeric-id rows") + require.Equal(t, 10, partitionCount, "expected 10 partitions for 100 rows with 10 rows per partition") SetupCDCFlowStatusQuery(t, env, flowConnConfig) cdcDocs := make([]any, 10) @@ -349,51 +349,6 @@ func (s MongoClickhouseSuite) Test_Snapshot_Empty_Collection() { RequireEnvCanceled(t, env) } -func (s MongoClickhouseSuite) Test_Snapshot_Non_ObjectID_Falls_Back_To_Single_Partition() { - t := s.T() - srcDatabase := GetTestDatabase(s.Suffix()) - srcTable := "test_non_objectid_snapshot" - dstTable := "test_non_objectid_snapshot_dst" - - connectionGen := FlowConnectionGenerationConfig{ - FlowJobName: AddSuffix(s, srcTable), - TableMappings: TableMappings(s, srcTable, dstTable), - Destination: s.Peer().Name, - } - flowConnConfig := s.generateFlowConnectionConfigsDefaultEnv(connectionGen) - flowConnConfig.DoInitialSnapshot = true - flowConnConfig.SnapshotNumRowsPerPartition = 5 - - adminClient := s.Source().(*MongoSource).AdminClient() - collection := adminClient.Database(srcDatabase).Collection(srcTable) - - docs := make([]any, 10) - for i := range 10 { - docs[i] = bson.D{ - {Key: "_id", Value: fmt.Sprintf("string_id_%d", i)}, - {Key: "value", Value: fmt.Sprintf("value_%d", i)}, - } - } - _, err := collection.InsertMany(t.Context(), docs) - require.NoError(t, err) - - tc := NewTemporalClient(t) - env := ExecutePeerflow(t, tc, flowConnConfig) - - EnvWaitForEqualTablesWithNames(env, s, "initial load", srcTable, dstTable, "_id,doc") - - catalogPool, err := internal.GetCatalogConnectionPoolFromEnv(t.Context()) - require.NoError(t, err) - var partitionCount int - require.NoError(t, catalogPool.QueryRow(t.Context(), - `SELECT COUNT(*) FROM peerdb_stats.qrep_partitions WHERE parent_mirror_name = $1`, - flowConnConfig.FlowJobName).Scan(&partitionCount)) - require.Equal(t, 1, partitionCount) - - env.Cancel(t.Context()) - RequireEnvCanceled(t, env) -} - func (s MongoClickhouseSuite) Test_Snapshot_Mixed_ObjectID_Falls_Back_To_Single_Partition() { t := s.T() srcDatabase := GetTestDatabase(s.Suffix()) From c21475e02ca191aad9695a06c6974f15f66f7297 Mon Sep 17 00:00:00 2001 From: Andrey Zhelnin Date: Fri, 26 Jun 2026 17:47:01 +0200 Subject: [PATCH 5/5] fix(lint): restore stdlib import group separator in qrep_partition.go --- flow/connectors/mongo/qrep_partition.go | 1 + 1 file changed, 1 insertion(+) diff --git a/flow/connectors/mongo/qrep_partition.go b/flow/connectors/mongo/qrep_partition.go index 320ad5147..c9bc55792 100644 --- a/flow/connectors/mongo/qrep_partition.go +++ b/flow/connectors/mongo/qrep_partition.go @@ -5,6 +5,7 @@ import ( "fmt" "log/slog" "math/big" + "github.com/google/uuid" "go.mongodb.org/mongo-driver/v2/bson" "go.mongodb.org/mongo-driver/v2/mongo"