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

Author: andreyzhelnin-st

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)

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,15 +43,15 @@ 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 {
 		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
+		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].

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 (

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())