Proposal: Activity to fetch config from the DB

This uses a new helper in `internal.FetchConfigFromDB` to fetch a fully
hydrated `protos.FlowConnectionConfigs`. When passing this config to
Temporal we strip the `tableMappings` array element which can cause it
to go over the 2MB limit.

This contains a subset of the changes which were originally proposed in:
PeerDB-io#3407

Author: alexstoick

flow/activities/flowable.go

@@ -5,8 +5,10 @@ import (
 	"errors"
 	"fmt"
 	"log/slog"
+	"maps"
 	"net"
 	"os"
+	"slices"
 	"strconv"
 	"sync/atomic"
 	"time"
@@ -125,6 +127,13 @@ func (a *FlowableActivity) EnsurePullability(
 	}
 	defer connectors.CloseConnector(ctx, srcConn)
 
+	// We can fetch from the DB, as we are in the activity
+	cfg, err := internal.FetchConfigFromDB(config.FlowJobName, ctx)
+	if err != nil {
+		return nil, err
+	}
+	config.SourceTableIdentifiers = slices.Sorted(maps.Keys(internal.TableNameMapping(cfg.TableMappings, cfg.Resync)))
+
 	output, err := srcConn.EnsurePullability(ctx, config)
 	if err != nil {
 		return nil, a.Alerter.LogFlowError(ctx, config.FlowJobName, fmt.Errorf("failed to ensure pullability: %w", err))
@@ -145,6 +154,12 @@ func (a *FlowableActivity) CreateRawTable(
 	}
 	defer connectors.CloseConnector(ctx, dstConn)
 
+	cfg, err := internal.FetchConfigFromDB(config.FlowJobName, ctx)
+	if err != nil {
+		return nil, err
+	}
+	config.TableNameMapping = internal.TableNameMapping(cfg.TableMappings, cfg.Resync)
+
 	res, err := dstConn.CreateRawTable(ctx, config)
 	if err != nil {
 		return nil, a.Alerter.LogFlowError(ctx, config.FlowJobName, err)
@@ -174,11 +189,15 @@ func (a *FlowableActivity) SetupTableSchema(
 	}
 	defer connectors.CloseConnector(ctx, srcConn)
 
-	tableNameSchemaMapping, err := srcConn.GetTableSchema(ctx, config.Env, config.Version, config.System, config.TableMappings)
+	cfg, err := internal.FetchConfigFromDB(config.FlowName, ctx)
+	if err != nil {
+		return err
+	}
+	tableNameSchemaMapping, err := srcConn.GetTableSchema(ctx, config.Env, config.Version, config.System, cfg.TableMappings)
 	if err != nil {
 		return a.Alerter.LogFlowError(ctx, config.FlowName, fmt.Errorf("failed to get GetTableSchemaConnector: %w", err))
 	}
-	processed := internal.BuildProcessedSchemaMapping(config.TableMappings, tableNameSchemaMapping, logger)
+	processed := internal.BuildProcessedSchemaMapping(cfg.TableMappings, tableNameSchemaMapping, logger)
 
 	tx, err := a.CatalogPool.BeginTx(ctx, pgx.TxOptions{})
 	if err != nil {
@@ -243,9 +262,13 @@ func (a *FlowableActivity) CreateNormalizedTable(
 		return nil, err
 	}
 
+	cfg, err := internal.FetchConfigFromDB(config.FlowName, ctx)
+	if err != nil {
+		return nil, err
+	}
 	numTablesToSetup.Store(int32(len(tableNameSchemaMapping)))
 	tableExistsMapping := make(map[string]bool, len(tableNameSchemaMapping))
-	for _, tableMapping := range config.TableMappings {
+	for _, tableMapping := range cfg.TableMappings {
 		tableIdentifier := tableMapping.DestinationTableIdentifier
 		tableSchema := tableNameSchemaMapping[tableIdentifier]
 		existing, err := conn.SetupNormalizedTable(
@@ -292,6 +315,16 @@ func (a *FlowableActivity) SyncFlow(
 	var normalizeWaiting atomic.Bool
 	var syncingBatchID atomic.Int64
 	var syncState atomic.Pointer[string]
+
+	cfg, err := internal.FetchConfigFromDB(config.FlowJobName, ctx)
+	if err != nil {
+		return err
+	}
+
+	// Override config with DB values to deal with the large fields.
+	config = cfg
+	options.TableMappings = cfg.TableMappings
+
 	syncState.Store(shared.Ptr("setup"))
 	shutdown := heartbeatRoutine(ctx, func() string {
 		// Must load Waiting after BatchID to avoid race saying we're waiting on currently processing batch

flow/activities/snapshot_activity.go

@@ -65,6 +65,14 @@ func (a *SnapshotActivity) SetupReplication(
 		return nil, a.Alerter.LogFlowError(ctx, config.FlowJobName, fmt.Errorf("failed to get connector: %w", err))
 	}
 
+	configCtx := context.Background()
+	defer configCtx.Done()
+	cfg, err := internal.FetchConfigFromDB(config.FlowJobName, configCtx)
+	if err != nil {
+		return nil, err
+	}
+	config.TableNameMapping = internal.TableNameMapping(cfg.TableMappings, cfg.Resync)
+
 	logger.Info("waiting for slot to be created...")
 	slotInfo, err := conn.SetupReplication(ctx, config)
 
@@ -180,8 +188,12 @@ func (a *SnapshotActivity) GetDefaultPartitionKeyForTables(
 	}
 	defer connectors.CloseConnector(ctx, connector)
 
+	cfg, err := internal.FetchConfigFromDB(input.FlowJobName, ctx)
+	if err != nil {
+		return nil, err
+	}
 	output, err := connector.GetDefaultPartitionKeyForTables(ctx, &protos.GetDefaultPartitionKeyForTablesInput{
-		TableMappings: input.TableMappings,
+		TableMappings: cfg.TableMappings,
 	})
 	if err != nil {
 		return nil, a.Alerter.LogFlowError(ctx, input.FlowJobName, fmt.Errorf("failed to check if tables can parallel load: %w", err))

flow/internal/flow_configuration_helpers.go

@@ -0,0 +1,51 @@
+package internal
+
+import (
+	"context"
+	"database/sql"
+	"fmt"
+
+	"github.com/PeerDB-io/peerdb/flow/generated/protos"
+	"google.golang.org/protobuf/proto"
+)
+
+func TableNameMapping(tableMappings []*protos.TableMapping, resync bool) map[string]string {
+	tblNameMapping := make(map[string]string, len(tableMappings))
+	if resync {
+		for _, mapping := range tableMappings {
+			if mapping.Engine != protos.TableEngine_CH_ENGINE_NULL {
+				mapping.DestinationTableIdentifier += "_resync"
+			}
+		}
+	}
+	for _, v := range tableMappings {
+		tblNameMapping[v.SourceTableIdentifier] = v.DestinationTableIdentifier
+	}
+
+	return tblNameMapping
+}
+
+func MinimizeFlowConfiguration(cfg *protos.FlowConnectionConfigs) *protos.FlowConnectionConfigs {
+	if cfg == nil {
+		return nil
+	}
+	cfg.TableMappings = nil
+	return cfg
+}
+
+func FetchConfigFromDB(flowName string, ctx context.Context) (*protos.FlowConnectionConfigs, error) {
+	var configBytes sql.RawBytes
+	pool, _ := GetCatalogConnectionPoolFromEnv(ctx)
+	if err := pool.QueryRow(ctx,
+		"SELECT config_proto FROM flows WHERE name = $1 LIMIT 1", flowName,
+	).Scan(&configBytes); err != nil {
+		return nil, fmt.Errorf("unable to query flow config from catalog: %w", err)
+	}
+
+	var cfgFromDB protos.FlowConnectionConfigs
+	if err := proto.Unmarshal(configBytes, &cfgFromDB); err != nil {
+		return nil, fmt.Errorf("unable to unmarshal flow config: %w", err)
+	}
+
+	return &cfgFromDB, nil
+}

flow/workflows/cdc_flow.go

@@ -45,18 +45,14 @@ type CDCFlowWorkflowState struct {
 
 // returns a new empty PeerFlowState
 func NewCDCFlowWorkflowState(ctx workflow.Context, logger log.Logger, cfg *protos.FlowConnectionConfigs) *CDCFlowWorkflowState {
-	tableMappings := make([]*protos.TableMapping, 0, len(cfg.TableMappings))
-	for _, tableMapping := range cfg.TableMappings {
-		tableMappings = append(tableMappings, proto.CloneOf(tableMapping))
-	}
 	state := CDCFlowWorkflowState{
 		ActiveSignal:      model.NoopSignal,
 		CurrentFlowStatus: protos.FlowStatus_STATUS_SETUP,
 		FlowConfigUpdate:  nil,
 		SyncFlowOptions: &protos.SyncFlowOptions{
 			BatchSize:          cfg.MaxBatchSize,
 			IdleTimeoutSeconds: cfg.IdleTimeoutSeconds,
-			TableMappings:      tableMappings,
+			TableMappings:      []*protos.TableMapping{},
 		},
 		SnapshotNumRowsPerPartition:   cfg.SnapshotNumRowsPerPartition,
 		SnapshotNumPartitionsOverride: cfg.SnapshotNumPartitionsOverride,
@@ -129,6 +125,7 @@ func syncStateToConfigProtoInCatalog(
 	cfg *protos.FlowConnectionConfigs,
 	state *CDCFlowWorkflowState,
 ) *protos.FlowConnectionConfigs {
+	// TODO: make sure that `cfg` includes table mappings.
 	cloneCfg := updateFlowConfigWithLatestSettings(cfg, state)
 	uploadConfigToCatalog(ctx, cloneCfg)
 	return cloneCfg
@@ -278,7 +275,7 @@ func processTableAdditions(
 	alterPublicationAddAdditionalTablesFuture := workflow.ExecuteActivity(
 		alterPublicationAddAdditionalTablesCtx,
 		flowable.AddTablesToPublication,
-		cfg, flowConfigUpdate.AdditionalTables)
+		internal.MinimizeFlowConfiguration(cfg), flowConfigUpdate.AdditionalTables)
 
 	var res *CDCFlowWorkflowResult
 	var addTablesFlowErr error
@@ -320,6 +317,7 @@ func processTableAdditions(
 			childAddTablesCDCFlowFuture := workflow.ExecuteChildWorkflow(
 				childAddTablesCDCFlowCtx,
 				CDCFlowWorkflow,
+				// TODO: `additonalTableCfg` this cannot be minimized in the main branch; but the limitation is minimal.
 				additionalTablesCfg,
 				nil,
 			)
@@ -394,7 +392,7 @@ func processTableRemovals(
 		rawTableCleanupFuture := workflow.ExecuteActivity(
 			removeTablesCtx,
 			flowable.RemoveTablesFromRawTable,
-			cfg, state.FlowConfigUpdate.RemovedTables)
+			internal.MinimizeFlowConfiguration(cfg), state.FlowConfigUpdate.RemovedTables)
 		removeTablesSelector.AddFuture(rawTableCleanupFuture, func(f workflow.Future) {
 			if err := f.Get(ctx, nil); err != nil {
 				logger.Error("failed to clean up raw table for removed tables", slog.Any("error", err))
@@ -406,7 +404,7 @@ func processTableRemovals(
 			removeTablesFromCatalogFuture := workflow.ExecuteActivity(
 				removeTablesCtx,
 				flowable.RemoveTablesFromCatalog,
-				cfg, state.FlowConfigUpdate.RemovedTables)
+				internal.MinimizeFlowConfiguration(cfg), state.FlowConfigUpdate.RemovedTables)
 			removeTablesSelector.AddFuture(removeTablesFromCatalogFuture, func(f workflow.Future) {
 				if err := f.Get(ctx, nil); err != nil {
 					logger.Error("failed to clean up raw table for removed tables", slog.Any("error", err))
@@ -576,7 +574,7 @@ func CDCFlowWorkflow(
 
 		logger.Info("mirror resumed", slog.Duration("after", time.Since(startTime)))
 		state.updateStatus(ctx, logger, protos.FlowStatus_STATUS_RUNNING)
-		return state, workflow.NewContinueAsNewError(ctx, CDCFlowWorkflow, cfg, state)
+		return state, workflow.NewContinueAsNewError(ctx, CDCFlowWorkflow, internal.MinimizeFlowConfiguration(cfg), state)
 	}
 
 	originalRunID := workflow.GetInfo(ctx).OriginalRunID
@@ -609,22 +607,6 @@ func CDCFlowWorkflow(
 	// for safety, rely on the idempotency of SetupFlow instead
 	// also, no signals are being handled until the loop starts, so no PAUSE/DROP will take here.
 	if state.CurrentFlowStatus != protos.FlowStatus_STATUS_RUNNING {
-		originalTableMappings := make([]*protos.TableMapping, 0, len(cfg.TableMappings))
-		for _, tableMapping := range cfg.TableMappings {
-			originalTableMappings = append(originalTableMappings, proto.CloneOf(tableMapping))
-		}
-		// if resync is true, alter the table name schema mapping to temporarily add
-		// a suffix to the table names.
-		if cfg.Resync {
-			for _, mapping := range state.SyncFlowOptions.TableMappings {
-				if mapping.Engine != protos.TableEngine_CH_ENGINE_NULL {
-					mapping.DestinationTableIdentifier += "_resync"
-				}
-			}
-			// because we have renamed the tables.
-			cfg.TableMappings = state.SyncFlowOptions.TableMappings
-		}
-
 		// start the SetupFlow workflow as a child workflow, and wait for it to complete
 		// it should return the table schema for the source peer
 		setupFlowID := GetChildWorkflowID("setup-flow", cfg.FlowJobName, originalRunID)
@@ -648,7 +630,6 @@ func CDCFlowWorkflow(
 				state.ActiveSignal = model.ResyncSignal
 				cfg.Resync = true
 				cfg.DoInitialSnapshot = true
-				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
 				uploadConfigToCatalog(ctx, cfg)
@@ -672,7 +653,10 @@ func CDCFlowWorkflow(
 			WaitForCancellation:   true,
 		}
 		setupFlowCtx := workflow.WithChildOptions(ctx, childSetupFlowOpts)
-		setupFlowFuture := workflow.ExecuteChildWorkflow(setupFlowCtx, SetupFlowWorkflow, cfg)
+		// Resync will rely rely on the `cfg.Resync` flag to rename the tables
+		// during the snapshot process. This is how we're able to also remove the need
+		// to sync the config back into the DB / not rely on the `state.TableMappings`.
+		setupFlowFuture := workflow.ExecuteChildWorkflow(setupFlowCtx, SetupFlowWorkflow, internal.MinimizeFlowConfiguration(cfg))
 
 		var setupFlowOutput *protos.SetupFlowOutput
 		var setupFlowError error
@@ -808,7 +792,7 @@ func CDCFlowWorkflow(
 			logger.Info("executed setup flow and snapshot flow, start running")
 			state.updateStatus(ctx, logger, protos.FlowStatus_STATUS_RUNNING)
 		}
-		return state, workflow.NewContinueAsNewError(ctx, CDCFlowWorkflow, cfg, state)
+		return state, workflow.NewContinueAsNewError(ctx, CDCFlowWorkflow, internal.MinimizeFlowConfiguration(cfg), state)
 	}
 
 	var finished bool
@@ -819,7 +803,8 @@ func CDCFlowWorkflow(
 		WaitForCancellation: true,
 		RetryPolicy:         &temporal.RetryPolicy{MaximumAttempts: 1},
 	}))
-	syncFlowFuture := workflow.ExecuteActivity(syncCtx, flowable.SyncFlow, cfg, state.SyncFlowOptions)
+	state.SyncFlowOptions.TableMappings = []*protos.TableMapping{}
+	syncFlowFuture := workflow.ExecuteActivity(syncCtx, flowable.SyncFlow, internal.MinimizeFlowConfiguration(cfg), state.SyncFlowOptions)
 
 	mainLoopSelector := workflow.NewNamedSelector(ctx, "MainLoop")
 	mainLoopSelector.AddReceive(ctx.Done(), func(_ workflow.ReceiveChannel, _ bool) {
@@ -940,7 +925,7 @@ func CDCFlowWorkflow(
 			if state.ActiveSignal == model.TerminateSignal || state.ActiveSignal == model.ResyncSignal {
 				return state, workflow.NewContinueAsNewError(ctx, DropFlowWorkflow, state.DropFlowInput)
 			}
-			return state, workflow.NewContinueAsNewError(ctx, CDCFlowWorkflow, cfg, state)
+			return state, workflow.NewContinueAsNewError(ctx, CDCFlowWorkflow, internal.MinimizeFlowConfiguration(cfg), state)
 		}
 	}
 }

flow/workflows/setup_flow.go

@@ -36,15 +36,17 @@ type SetupFlowExecution struct {
 	tableNameMapping map[string]string
 	cdcFlowName      string
 	executionID      string
+	resync           bool
 }
 
 // NewSetupFlowExecution creates a new instance of SetupFlowExecution.
-func NewSetupFlowExecution(ctx workflow.Context, tableNameMapping map[string]string, cdcFlowName string) *SetupFlowExecution {
+func NewSetupFlowExecution(ctx workflow.Context, tableNameMapping map[string]string, cdcFlowName string, resync bool) *SetupFlowExecution {
 	return &SetupFlowExecution{
 		Logger:           log.With(workflow.GetLogger(ctx), slog.String(string(shared.FlowNameKey), cdcFlowName)),
 		tableNameMapping: tableNameMapping,
 		cdcFlowName:      cdcFlowName,
 		executionID:      workflow.GetInfo(ctx).WorkflowExecution.ID,
+		resync:           resync,
 	}
 }
 
@@ -119,10 +121,9 @@ func (s *SetupFlowExecution) ensurePullability(
 
 	// create EnsurePullabilityInput for the srcTableName
 	ensurePullabilityInput := &protos.EnsurePullabilityBatchInput{
-		PeerName:               config.SourceName,
-		FlowJobName:            s.cdcFlowName,
-		SourceTableIdentifiers: slices.Sorted(maps.Keys(s.tableNameMapping)),
-		CheckConstraints:       checkConstraints,
+		PeerName:         config.SourceName,
+		FlowJobName:      s.cdcFlowName,
+		CheckConstraints: checkConstraints,
 	}
 
 	future := workflow.ExecuteActivity(ctx, flowable.EnsurePullability, ensurePullabilityInput)
@@ -162,9 +163,8 @@ func (s *SetupFlowExecution) createRawTable(
 
 	// attempt to create the tables.
 	createRawTblInput := &protos.CreateRawTableInput{
-		PeerName:         config.DestinationName,
-		FlowJobName:      s.cdcFlowName,
-		TableNameMapping: s.tableNameMapping,
+		PeerName:    config.DestinationName,
+		FlowJobName: s.cdcFlowName,
 	}
 
 	rawTblFuture := workflow.ExecuteActivity(ctx, flowable.CreateRawTable, createRawTblInput)
@@ -191,12 +191,11 @@ func (s *SetupFlowExecution) setupNormalizedTables(
 	})
 
 	tableSchemaInput := &protos.SetupTableSchemaBatchInput{
-		PeerName:      flowConnectionConfigs.SourceName,
-		TableMappings: flowConnectionConfigs.TableMappings,
-		FlowName:      s.cdcFlowName,
-		System:        flowConnectionConfigs.System,
-		Env:           flowConnectionConfigs.Env,
-		Version:       flowConnectionConfigs.Version,
+		PeerName: flowConnectionConfigs.SourceName,
+		FlowName: s.cdcFlowName,
+		System:   flowConnectionConfigs.System,
+		Env:      flowConnectionConfigs.Env,
+		Version:  flowConnectionConfigs.Version,
 	}
 
 	if err := workflow.ExecuteActivity(ctx, flowable.SetupTableSchema, tableSchemaInput).Get(ctx, nil); err != nil {
@@ -207,7 +206,6 @@ func (s *SetupFlowExecution) setupNormalizedTables(
 	s.Info("setting up normalized tables on destination peer", slog.String("destination", flowConnectionConfigs.DestinationName))
 	setupConfig := &protos.SetupNormalizedTableBatchInput{
 		PeerName:          flowConnectionConfigs.DestinationName,
-		TableMappings:     flowConnectionConfigs.TableMappings,
 		SoftDeleteColName: flowConnectionConfigs.SoftDeleteColName,
 		SyncedAtColName:   flowConnectionConfigs.SyncedAtColName,
 		FlowName:          flowConnectionConfigs.FlowJobName,
@@ -261,13 +259,8 @@ func (s *SetupFlowExecution) executeSetupFlow(
 
 // SetupFlowWorkflow is the workflow that sets up the flow.
 func SetupFlowWorkflow(ctx workflow.Context, config *protos.FlowConnectionConfigs) (*protos.SetupFlowOutput, error) {
-	tblNameMapping := make(map[string]string, len(config.TableMappings))
-	for _, v := range config.TableMappings {
-		tblNameMapping[v.SourceTableIdentifier] = v.DestinationTableIdentifier
-	}
-
 	// create the setup flow execution
-	setupFlowExecution := NewSetupFlowExecution(ctx, tblNameMapping, config.FlowJobName)
+	setupFlowExecution := NewSetupFlowExecution(ctx, map[string]string{}, config.FlowJobName, config.Resync)
 
 	// execute the setup flow
 	setupFlowOutput, err := setupFlowExecution.executeSetupFlow(ctx, config)