feat(resync): Add support to update snapshot config (max_parallel_workers, num_tables, rows_per_partition) during resync signal

Author: masterashu

flow/cmd/handler.go

@@ -423,6 +423,7 @@ func (h *FlowRequestHandler) FlowStateChange(
 	ctx context.Context,
 	req *protos.FlowStateChangeRequest,
 ) (*protos.FlowStateChangeResponse, APIError) {
+	// Flow-API
 	logs := slog.String(string(shared.FlowNameKey), req.FlowJobName)
 	slog.InfoContext(ctx, "FlowStateChange called", logs, slog.Any("req", req))
 	if underMaintenance, err := internal.PeerDBMaintenanceModeEnabled(ctx, nil); err != nil {
@@ -505,6 +506,9 @@ func (h *FlowRequestHandler) FlowStateChange(
 
 				config.Resync = true
 				config.DoInitialSnapshot = true
+				// Override Snapshot Parameters
+				overrideSnapshotParametersInResync(req, config)
+
 				// validate mirror first because once the mirror is dropped, there's no going back
 				if _, err := h.ValidateCDCMirror(ctx, &protos.CreateCDCFlowRequest{
 					ConnectionConfigs: config,
@@ -551,6 +555,21 @@ func (h *FlowRequestHandler) FlowStateChange(
 	return &protos.FlowStateChangeResponse{}, nil
 }
 
+func overrideSnapshotParametersInResync(req *protos.FlowStateChangeRequest, configs *protos.FlowConnectionConfigs) {
+	if req.FlowConfigUpdate != nil && req.FlowConfigUpdate.GetCdcFlowConfigUpdate() != nil {
+		cdcConfigUpdate := req.FlowConfigUpdate.GetCdcFlowConfigUpdate()
+		if cdcConfigUpdate.SnapshotMaxParallelWorkers > 0 {
+			configs.SnapshotMaxParallelWorkers = cdcConfigUpdate.SnapshotMaxParallelWorkers
+		}
+		if cdcConfigUpdate.SnapshotNumTablesInParallel > 0 {
+			configs.SnapshotNumTablesInParallel = cdcConfigUpdate.SnapshotNumTablesInParallel
+		}
+		if cdcConfigUpdate.SnapshotNumRowsPerPartition > 0 {
+			configs.SnapshotNumRowsPerPartition = cdcConfigUpdate.SnapshotNumRowsPerPartition
+		}
+	}
+}
+
 func (h *FlowRequestHandler) handleCancelWorkflow(ctx context.Context, workflowID, runID string) error {
 	ctxWithTimeout, cancel := context.WithTimeout(ctx, 2*time.Minute)
 	defer cancel()

flow/shared/telemetry/activity_logging.go

@@ -93,6 +93,16 @@ func LogActivityStartFlowConfigUpdate(ctx context.Context, flowName string, upda
 		changes = append(changes, fmt.Sprintf("tables removed: %v", removedTables))
 	}
 
+	if update.SnapshotMaxParallelWorkers > 0 {
+		changes = append(changes, fmt.Sprintf("snapshotMaxParallelWorkers: %v", update.SnapshotMaxParallelWorkers))
+	}
+	if update.SnapshotNumTablesInParallel > 0 {
+		changes = append(changes, fmt.Sprintf("snapshotNumTablesInParallel: %v", update.SnapshotNumTablesInParallel))
+	}
+	if update.SnapshotNumRowsPerPartition > 0 {
+		changes = append(changes, fmt.Sprintf("snapshotNumRowsPerPartition: %v", update.SnapshotNumRowsPerPartition))
+	}
+
 	logActivity(ctx, ActionStartFlowConfigUpdate,
 		slog.String("flowName", flowName),
 		slog.String("activityDetails", strings.Join(changes, ", ")))

flow/workflows/cdc_flow.go

@@ -192,6 +192,9 @@ func handleFlowSignalStateChange(
 			// we should ContinueAsNew after the first signal in the selector, but just in case
 			cfg.Resync = true
 			cfg.DoInitialSnapshot = true
+			// This is just for in-memory
+			// Override Snapshot Parameters if request in State Change request
+			overrideSnapshotParametersInState(val, state)
 			state.DropFlowInput = &protos.DropFlowInput{
 				// to be filled in just before ContinueAsNew
 				FlowJobName:           cfg.FlowJobName,
@@ -206,6 +209,21 @@ func handleFlowSignalStateChange(
 	}
 }
 
+func overrideSnapshotParametersInState(req *protos.FlowStateChangeRequest, state *cdc_state.CDCFlowWorkflowState) {
+	if req.FlowConfigUpdate != nil && req.FlowConfigUpdate.GetCdcFlowConfigUpdate() != nil {
+		cdcConfigUpdate := req.FlowConfigUpdate.GetCdcFlowConfigUpdate()
+		if cdcConfigUpdate.SnapshotMaxParallelWorkers > 0 {
+			state.SnapshotMaxParallelWorkers = cdcConfigUpdate.SnapshotMaxParallelWorkers
+		}
+		if cdcConfigUpdate.SnapshotNumTablesInParallel > 0 {
+			state.SnapshotNumTablesInParallel = cdcConfigUpdate.SnapshotNumTablesInParallel
+		}
+		if cdcConfigUpdate.SnapshotNumRowsPerPartition > 0 {
+			state.SnapshotNumRowsPerPartition = cdcConfigUpdate.SnapshotNumRowsPerPartition
+		}
+	}
+}
+
 func processTableAdditions(
 	ctx workflow.Context,
 	logger log.Logger,
@@ -494,6 +512,8 @@ func CDCFlowWorkflow(
 				state.ActiveSignal = model.ResyncSignal
 				cfg.Resync = true
 				cfg.DoInitialSnapshot = true
+				// Update State with snapshot parameters
+				overrideSnapshotParametersInState(val, state)
 				resyncCfg := syncStateToConfigProtoInCatalog(ctx, cfg, state)
 				state.DropFlowInput = &protos.DropFlowInput{
 					FlowJobName:           resyncCfg.FlowJobName,
@@ -627,6 +647,9 @@ func CDCFlowWorkflow(
 				cfg.TableMappings = originalTableMappings
 				// this is the only place where we can have a resync during a resync
 				// so we need to NOT sync the tableMappings to catalog to preserve original names
+
+				// We still override the snapshot parameters (when resync with updated values)
+				overrideSnapshotParametersInState(val, state)
 				uploadConfigToCatalog(ctx, cfg)
 				state.DropFlowInput = &protos.DropFlowInput{
 					FlowJobName:           cfg.FlowJobName,
@@ -863,6 +886,7 @@ func CDCFlowWorkflow(
 			state.ActiveSignal = model.ResyncSignal
 			cfg.Resync = true
 			cfg.DoInitialSnapshot = true
+			overrideSnapshotParametersInState(val, state)
 			resyncCfg := syncStateToConfigProtoInCatalog(ctx, cfg, state)
 			state.DropFlowInput = &protos.DropFlowInput{
 				FlowJobName:           resyncCfg.FlowJobName,