add mongo support too

Author: jgao54

flow/activities/flowable.go

@@ -514,14 +514,15 @@ func (a *FlowableActivity) GetQRepPartitions(ctx context.Context,
 	}
 	defer srcClose(ctx)
 
-	if tableSizeEstimatorConn, ok := srcConn.(connectors.TableSizeEstimatorConnector); ok {
-		// expect estimate query execution to be fast, so set a short timeout
+	partitioned := config.WatermarkColumn != ""
+	if tableSizeEstimatorConn, ok := srcConn.(connectors.TableSizeEstimatorConnector); ok && partitioned {
+		// expect estimate query execution to be fast, set a short timeout defensively to avoid blocking workflow
 		timeoutCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
 		defer cancel()
 		if bytes, connErr := tableSizeEstimatorConn.GetTableSizeEstimatedBytes(timeoutCtx, config.WatermarkTable); connErr == nil {
-			if bytes > 1024*1024*1024*2024 {
+			if bytes > 1<<40 { // 1 TiB
 				msg := fmt.Sprintf("Large table detected: %s (%s). Partitioning query may take several hours to execute. "+
-					"This is normal for tables over 1 TB.", config.WatermarkTable, utils.FormatTableSize(bytes))
+					"This is normal for tables over 1 TiB.", config.WatermarkTable, utils.FormatTableSize(bytes))
 				a.Alerter.LogFlowInfo(ctx, config.FlowJobName, msg)
 			}
 		}

flow/connectors/core.go

@@ -682,4 +682,5 @@ var (
 
 	_ TableSizeEstimatorConnector = &connpostgres.PostgresConnector{}
 	_ TableSizeEstimatorConnector = &connmysql.MySqlConnector{}
+	_ TableSizeEstimatorConnector = &connmongo.MongoConnector{}
 )

flow/connectors/mongo/mongo.go

@@ -187,3 +187,15 @@ func (c *MongoConnector) GetLogRetentionHours(ctx context.Context) (float64, err
 
 	return float64(serverStatus.OplogTruncation.OplogMinRetentionHours), nil
 }
+
+func (c *MongoConnector) GetTableSizeEstimatedBytes(ctx context.Context, fullyQualifiedTable string) (int64, error) {
+	databaseTable, err := utils.ParseSchemaTable(fullyQualifiedTable)
+	if err != nil {
+		return 0, err
+	}
+	collStats, err := peerdb_mongo.GetCollStats(ctx, c.client, databaseTable.Schema, databaseTable.Table)
+	if err != nil {
+		return 0, err
+	}
+	return collStats.Size, nil
+}

flow/connectors/postgres/qrep.go

@@ -41,6 +41,10 @@ type QRepSyncSink interface {
 	CopyInto(context.Context, *PostgresConnector, pgx.Tx, pgx.Identifier) (int64, error)
 }
 
+func (c *PostgresConnector) SupportsTablePartitioning() bool {
+	return true
+}
+
 func (c *PostgresConnector) GetQRepPartitions(
 	ctx context.Context,
 	config *protos.QRepConfig,

flow/pkg/mongo/commands.go

@@ -60,6 +60,20 @@ func GetHelloResponse(ctx context.Context, client *mongo.Client) (HelloResponse,
 	return runCommand[HelloResponse](ctx, client, "hello")
 }
 
+type CollStats struct {
+	// uncompressed
+	Size int64 `bson:"size"`
+	// compressed
+	StorageSize int64 `bson:"storageSize"`
+}
+
+func GetCollStats(ctx context.Context, client *mongo.Client, database string, collection string) (CollStats, error) {
+	return runDatabaseCommand[CollStats](ctx, client, database, bson.D{
+		{Key: "collStats", Value: collection},
+		{Key: "scale", Value: 1},
+	})
+}
+
 func runCommand[T any](ctx context.Context, client *mongo.Client, command string) (T, error) {
 	var result T
 	singleResult := client.Database("admin").RunCommand(ctx, bson.D{
@@ -74,3 +88,16 @@ func runCommand[T any](ctx context.Context, client *mongo.Client, command string
 	}
 	return result, nil
 }
+
+func runDatabaseCommand[T any](ctx context.Context, client *mongo.Client, database string, commandDoc bson.D) (T, error) {
+	var result T
+	singleResult := client.Database(database).RunCommand(ctx, commandDoc)
+	if singleResult.Err() != nil {
+		return result, fmt.Errorf("command failed: %v", singleResult.Err())
+	}
+
+	if err := singleResult.Decode(&result); err != nil {
+		return result, fmt.Errorf("command decoding failed: %v", err)
+	}
+	return result, nil
+}