disable by default

Author: jgao54

flow/connectors/mongo/qrep.go

@@ -13,6 +13,7 @@ import (
 
 	"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,6 +36,16 @@ func (c *MongoConnector) GetQRepPartitions(
 	}
 
 	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
 	}
 

flow/e2e/mongo_test.go

@@ -109,6 +109,7 @@ func (s MongoClickhouseSuite) Test_Simple_Flow_Partitioned() {
 	flowConnConfig := s.generateFlowConnectionConfigsDefaultEnv(connectionGen)
 	flowConnConfig.DoInitialSnapshot = true
 	flowConnConfig.SnapshotNumRowsPerPartition = 10
+	flowConnConfig.Env["PEERDB_MONGODB_PARALLEL_SNAPSHOTTING"] = "true"
 
 	adminClient := s.Source().(*MongoSource).AdminClient()
 	collection := adminClient.Database(srcDatabase).Collection(srcTable)
@@ -172,6 +173,7 @@ func (s MongoClickhouseSuite) Test_Snapshot_Partition_Capped_To_Timestamp_Range(
 	// 1 row per partition would normally request 100 partitions for 100 docs,
 	// but the 5-second timestamp range should cap it to 5 partitions.
 	flowConnConfig.SnapshotNumRowsPerPartition = 1
+	flowConnConfig.Env["PEERDB_MONGODB_PARALLEL_SNAPSHOTTING"] = "true"
 
 	adminClient := s.Source().(*MongoSource).AdminClient()
 	collection := adminClient.Database(srcDatabase).Collection(srcTable)
@@ -221,6 +223,7 @@ func (s MongoClickhouseSuite) Test_Snapshot_Empty_Collection() {
 	flowConnConfig := s.generateFlowConnectionConfigsDefaultEnv(connectionGen)
 	flowConnConfig.DoInitialSnapshot = true
 	flowConnConfig.SnapshotNumRowsPerPartition = 1
+	flowConnConfig.Env["PEERDB_MONGODB_PARALLEL_SNAPSHOTTING"] = "true"
 
 	adminClient := s.Source().(*MongoSource).AdminClient()
 	err := adminClient.Database(srcDatabase).CreateCollection(t.Context(), srcTable)
@@ -255,6 +258,7 @@ func (s MongoClickhouseSuite) Test_Snapshot_Non_ObjectID_Falls_Back_To_Single_Pa
 	flowConnConfig := s.generateFlowConnectionConfigsDefaultEnv(connectionGen)
 	flowConnConfig.DoInitialSnapshot = true
 	flowConnConfig.SnapshotNumRowsPerPartition = 5
+	flowConnConfig.Env["PEERDB_MONGODB_PARALLEL_SNAPSHOTTING"] = "true"
 
 	adminClient := s.Source().(*MongoSource).AdminClient()
 	collection := adminClient.Database(srcDatabase).Collection(srcTable)
@@ -300,6 +304,7 @@ func (s MongoClickhouseSuite) Test_Snapshot_Mixed_ObjectID_Falls_Back_To_Single_
 	flowConnConfig := s.generateFlowConnectionConfigsDefaultEnv(connectionGen)
 	flowConnConfig.DoInitialSnapshot = true
 	flowConnConfig.SnapshotNumRowsPerPartition = 1
+	flowConnConfig.Env["PEERDB_MONGODB_PARALLEL_SNAPSHOTTING"] = "true"
 
 	adminClient := s.Source().(*MongoSource).AdminClient()
 	collection := adminClient.Database(srcDatabase).Collection(srcTable)

flow/internal/dynamicconf.go

@@ -435,6 +435,14 @@ var DynamicSettings = [...]*protos.DynamicSetting{
 		ApplyMode:        protos.DynconfApplyMode_APPLY_MODE_AFTER_RESUME,
 		TargetForSetting: protos.DynconfTarget_ALL,
 	},
+	{
+		Name:             "PEERDB_MONGODB_PARALLEL_SNAPSHOTTING",
+		Description:      "Use min/max ObjectID-based parallel partitioning for MongoDB initial snapshot",
+		DefaultValue:     "false",
+		ValueType:        protos.DynconfValueType_BOOL,
+		ApplyMode:        protos.DynconfApplyMode_APPLY_MODE_AFTER_RESUME,
+		TargetForSetting: protos.DynconfTarget_ALL,
+	},
 	{
 		Name: "PEERDB_MONGODB_DIRECT_BSON_CONVERTER",
 		Description: "Use direct BSON-to-JSON converter for MongoDB (faster, no intermediate deserialization). " +
@@ -797,6 +805,10 @@ func PeerDBPostgresApplyCtidBlockPartitioning(ctx context.Context, env map[strin
 	return dynamicConfBool(ctx, env, "PEERDB_POSTGRES_APPLY_CTID_BLOCK_PARTITIONING_OVERRIDE")
 }
 
+func PeerDBMongoDBParallelSnapshotting(ctx context.Context, env map[string]string) (bool, error) {
+	return dynamicConfBool(ctx, env, "PEERDB_MONGODB_PARALLEL_SNAPSHOTTING")
+}
+
 func PeerDBMongoDBDirectBsonConverter(ctx context.Context, env map[string]string) (bool, error) {
 	return dynamicConfBool(ctx, env, "PEERDB_MONGODB_DIRECT_BSON_CONVERTER")
 }