Fetch configuration from DB

No longer passing the entire configuration object to the Temporal job as
we can hit the 2MB hard limit imposed by Temporal. Instead, pass
flowJobName and fetch the config from the DB in all the places required.

There are some additional changes that were required to be made, as some
of the jobs are used from multiple contexts - and I had to adapt some of
the information passed to ensure that they continue to work as expected.

Author: alexstoick

flow/activities/flowable.go

@@ -5,7 +5,9 @@ import (
 	"errors"
 	"fmt"
 	"log/slog"
+	"maps"
 	"os"
+	"slices"
 	"strconv"
 	"sync/atomic"
 	"time"
@@ -124,6 +126,12 @@ func (a *FlowableActivity) EnsurePullability(
 	}
 	defer connectors.CloseConnector(ctx, srcConn)
 
+	cfg, err := internal.FetchConfigFromDB(config.FlowJobName)
+	if err != nil {
+		return nil, a.Alerter.LogFlowError(ctx, config.FlowJobName, fmt.Errorf("failed to fetch config from DB: %w", err))
+	}
+
+	config.SourceTableIdentifiers = slices.Sorted(maps.Keys(internal.TableNameMapping(cfg.TableMappings)))
 	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))
@@ -165,6 +173,18 @@ func (a *FlowableActivity) SetupTableSchema(
 	})
 	defer shutdown()
 
+	// have to fetch config from the DB
+	cfg, err := internal.FetchConfigFromDB(config.FlowName)
+	if err != nil {
+		return a.Alerter.LogFlowError(ctx, config.FlowName, fmt.Errorf("failed to fetch config from DB: %w", err))
+	}
+	tableMappings := cfg.TableMappings
+	if len(config.FilteredTableMappings) > 0 {
+		// we use the filtered table mappings if provided. they are provided from
+		// the sync flow which includes changes to the schema.
+		tableMappings = config.FilteredTableMappings
+	}
+
 	logger := internal.LoggerFromCtx(ctx)
 	ctx = context.WithValue(ctx, shared.FlowNameKey, config.FlowName)
 	srcConn, err := connectors.GetByNameAs[connectors.GetTableSchemaConnector](ctx, config.Env, a.CatalogPool, config.PeerName)
@@ -173,11 +193,11 @@ func (a *FlowableActivity) SetupTableSchema(
 	}
 	defer connectors.CloseConnector(ctx, srcConn)
 
-	tableNameSchemaMapping, err := srcConn.GetTableSchema(ctx, config.Env, config.Version, config.System, config.TableMappings)
+	tableNameSchemaMapping, err := srcConn.GetTableSchema(ctx, config.Env, config.Version, config.System, 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(tableMappings, tableNameSchemaMapping, logger)
 
 	tx, err := a.CatalogPool.BeginTx(ctx, pgx.TxOptions{})
 	if err != nil {
@@ -213,6 +233,13 @@ func (a *FlowableActivity) CreateNormalizedTable(
 	numTablesSetup := atomic.Uint32{}
 	numTablesToSetup := atomic.Int32{}
 
+	cfg, err := internal.FetchConfigFromDB(config.FlowName)
+	if err != nil {
+		return nil, a.Alerter.LogFlowError(ctx, config.FlowName, fmt.Errorf("failed to fetch config from DB: %w", err))
+	}
+
+	tableMappings := cfg.TableMappings
+
 	shutdown := heartbeatRoutine(ctx, func() string {
 		return fmt.Sprintf("setting up normalized tables - %d of %d done", numTablesSetup.Load(), numTablesToSetup.Load())
 	})
@@ -244,7 +271,7 @@ func (a *FlowableActivity) CreateNormalizedTable(
 
 	numTablesToSetup.Store(int32(len(tableNameSchemaMapping)))
 	tableExistsMapping := make(map[string]bool, len(tableNameSchemaMapping))
-	for _, tableMapping := range config.TableMappings {
+	for _, tableMapping := range tableMappings {
 		tableIdentifier := tableMapping.DestinationTableIdentifier
 		tableSchema := tableNameSchemaMapping[tableIdentifier]
 		existing, err := conn.SetupNormalizedTable(
@@ -292,6 +319,12 @@ func (a *FlowableActivity) SyncFlow(
 	var syncingBatchID atomic.Int64
 	var syncState atomic.Pointer[string]
 	syncState.Store(shared.Ptr("setup"))
+
+	config, err := internal.FetchConfigFromDB(config.FlowJobName)
+	if err != nil {
+		return fmt.Errorf("unable to query flow config from catalog: %w", err)
+	}
+
 	shutdown := heartbeatRoutine(ctx, func() string {
 		// Must load Waiting after BatchID to avoid race saying we're waiting on currently processing batch
 		sBatchID := syncingBatchID.Load()
@@ -363,10 +396,10 @@ func (a *FlowableActivity) SyncFlow(
 		var syncResponse *model.SyncResponse
 		var syncErr error
 		if config.System == protos.TypeSystem_Q {
-			syncResponse, syncErr = a.syncRecords(groupCtx, config, options, srcConn.(connectors.CDCPullConnector),
+			syncResponse, syncErr = a.syncRecords(groupCtx, config, srcConn.(connectors.CDCPullConnector),
 				normRequests, normResponses, normBufferSize, &syncingBatchID, &syncState)
 		} else {
-			syncResponse, syncErr = a.syncPg(groupCtx, config, options, srcConn.(connectors.CDCPullPgConnector),
+			syncResponse, syncErr = a.syncPg(groupCtx, config, srcConn.(connectors.CDCPullPgConnector),
 				normRequests, normResponses, normBufferSize, &syncingBatchID, &syncState)
 		}
 
@@ -414,7 +447,6 @@ func (a *FlowableActivity) SyncFlow(
 func (a *FlowableActivity) syncRecords(
 	ctx context.Context,
 	config *protos.FlowConnectionConfigs,
-	options *protos.SyncFlowOptions,
 	srcConn connectors.CDCPullConnector,
 	normRequests *concurrency.LastChan,
 	normResponses *concurrency.LastChan,
@@ -451,7 +483,7 @@ func (a *FlowableActivity) syncRecords(
 			return stream, nil
 		}
 	}
-	return syncCore(ctx, a, config, options, srcConn,
+	return syncCore(ctx, a, config, srcConn,
 		normRequests, normResponses, normBufferSize,
 		syncingBatchID, syncWaiting, adaptStream,
 		connectors.CDCPullConnector.PullRecords,
@@ -461,15 +493,14 @@ func (a *FlowableActivity) syncRecords(
 func (a *FlowableActivity) syncPg(
 	ctx context.Context,
 	config *protos.FlowConnectionConfigs,
-	options *protos.SyncFlowOptions,
 	srcConn connectors.CDCPullPgConnector,
 	normRequests *concurrency.LastChan,
 	normResponses *concurrency.LastChan,
 	normBufferSize int64,
 	syncingBatchID *atomic.Int64,
 	syncWaiting *atomic.Pointer[string],
 ) (*model.SyncResponse, error) {
-	return syncCore(ctx, a, config, options, srcConn,
+	return syncCore(ctx, a, config, srcConn,
 		normRequests, normResponses, normBufferSize,
 		syncingBatchID, syncWaiting, nil,
 		connectors.CDCPullPgConnector.PullPg,
@@ -1230,9 +1261,13 @@ func (a *FlowableActivity) ReplicateXminPartition(ctx context.Context,
 	}
 }
 
-func (a *FlowableActivity) AddTablesToPublication(ctx context.Context, cfg *protos.FlowConnectionConfigs,
+func (a *FlowableActivity) AddTablesToPublication(ctx context.Context, flowJobName string,
 	additionalTableMappings []*protos.TableMapping,
 ) error {
+	cfg, err := internal.FetchConfigFromDB(flowJobName)
+	if err != nil {
+		return fmt.Errorf("unable to query flow config from catalog: %w", err)
+	}
 	ctx = context.WithValue(ctx, shared.FlowNameKey, cfg.FlowJobName)
 	srcConn, err := connectors.GetByNameAs[*connpostgres.PostgresConnector](ctx, cfg.Env, a.CatalogPool, cfg.SourceName)
 	if err != nil {

flow/activities/flowable_core.go

@@ -80,11 +80,10 @@ func (a *FlowableActivity) getTableNameSchemaMapping(ctx context.Context, flowNa
 func (a *FlowableActivity) applySchemaDeltas(
 	ctx context.Context,
 	config *protos.FlowConnectionConfigs,
-	options *protos.SyncFlowOptions,
 	schemaDeltas []*protos.TableSchemaDelta,
 ) error {
 	filteredTableMappings := make([]*protos.TableMapping, 0, len(schemaDeltas))
-	for _, tableMapping := range options.TableMappings {
+	for _, tableMapping := range config.TableMappings {
 		if slices.ContainsFunc(schemaDeltas, func(schemaDelta *protos.TableSchemaDelta) bool {
 			return schemaDelta.SrcTableName == tableMapping.SourceTableIdentifier &&
 				schemaDelta.DstTableName == tableMapping.DestinationTableIdentifier
@@ -95,12 +94,12 @@ func (a *FlowableActivity) applySchemaDeltas(
 
 	if len(schemaDeltas) > 0 {
 		if err := a.SetupTableSchema(ctx, &protos.SetupTableSchemaBatchInput{
-			PeerName:      config.SourceName,
-			TableMappings: filteredTableMappings,
-			FlowName:      config.FlowJobName,
-			System:        config.System,
-			Env:           config.Env,
-			Version:       config.Version,
+			PeerName:              config.SourceName,
+			FilteredTableMappings: filteredTableMappings,
+			FlowName:              config.FlowJobName,
+			System:                config.System,
+			Env:                   config.Env,
+			Version:               config.Version,
 		}); err != nil {
 			return a.Alerter.LogFlowError(ctx, config.FlowJobName, fmt.Errorf("failed to execute schema update at source: %w", err))
 		}
@@ -112,7 +111,6 @@ func syncCore[TPull connectors.CDCPullConnectorCore, TSync connectors.CDCSyncCon
 	ctx context.Context,
 	a *FlowableActivity,
 	config *protos.FlowConnectionConfigs,
-	options *protos.SyncFlowOptions,
 	srcConn TPull,
 	normRequests *concurrency.LastChan,
 	normResponses *concurrency.LastChan,
@@ -127,16 +125,17 @@ func syncCore[TPull connectors.CDCPullConnectorCore, TSync connectors.CDCSyncCon
 	ctx = context.WithValue(ctx, shared.FlowNameKey, flowName)
 	logger := internal.LoggerFromCtx(ctx)
 
-	tblNameMapping := make(map[string]model.NameAndExclude, len(options.TableMappings))
-	for _, v := range options.TableMappings {
+	// we should be able to rely on `config` here.
+	tblNameMapping := make(map[string]model.NameAndExclude, len(config.TableMappings))
+	for _, v := range config.TableMappings {
 		tblNameMapping[v.SourceTableIdentifier] = model.NewNameAndExclude(v.DestinationTableIdentifier, v.Exclude)
 	}
 
 	if err := srcConn.ConnectionActive(ctx); err != nil {
 		return nil, temporal.NewNonRetryableApplicationError("connection to source down", "disconnect", nil)
 	}
 
-	batchSize := options.BatchSize
+	batchSize := config.MaxBatchSize
 	if batchSize == 0 {
 		batchSize = 250_000
 	}
@@ -190,13 +189,13 @@ func syncCore[TPull connectors.CDCPullConnectorCore, TSync connectors.CDCSyncCon
 	errGroup.Go(func() error {
 		return pull(srcConn, errCtx, a.CatalogPool, a.OtelManager, &model.PullRecordsRequest[Items]{
 			FlowJobName:           flowName,
-			SrcTableIDNameMapping: options.SrcTableIdNameMapping,
+			SrcTableIDNameMapping: config.SrcTableIdNameMapping,
 			TableNameMapping:      tblNameMapping,
 			LastOffset:            lastOffset,
 			ConsumedOffset:        &consumedOffset,
 			MaxBatchSize:          batchSize,
 			IdleTimeout: internal.PeerDBCDCIdleTimeoutSeconds(
-				int(options.IdleTimeoutSeconds),
+				int(config.IdleTimeoutSeconds),
 			),
 			TableNameSchemaMapping:      tableNameSchemaMapping,
 			OverridePublicationName:     config.PublicationName,
@@ -234,11 +233,11 @@ func syncCore[TPull connectors.CDCPullConnectorCore, TSync connectors.CDCSyncCon
 		defer connectors.CloseConnector(ctx, dstConn)
 
 		syncState.Store(shared.Ptr("updating schema"))
-		if err := dstConn.ReplayTableSchemaDeltas(ctx, config.Env, flowName, options.TableMappings, recordBatchSync.SchemaDeltas); err != nil {
+		if err := dstConn.ReplayTableSchemaDeltas(ctx, config.Env, flowName, config.TableMappings, recordBatchSync.SchemaDeltas); err != nil {
 			return nil, fmt.Errorf("failed to sync schema: %w", err)
 		}
 
-		return nil, a.applySchemaDeltas(ctx, config, options, recordBatchSync.SchemaDeltas)
+		return nil, a.applySchemaDeltas(ctx, config, recordBatchSync.SchemaDeltas)
 	}
 
 	var res *model.SyncResponse
@@ -271,7 +270,7 @@ func syncCore[TPull connectors.CDCPullConnectorCore, TSync connectors.CDCSyncCon
 			Records:                recordBatchSync,
 			ConsumedOffset:         &consumedOffset,
 			FlowJobName:            flowName,
-			TableMappings:          options.TableMappings,
+			TableMappings:          config.TableMappings,
 			StagingPath:            config.CdcStagingPath,
 			Script:                 config.Script,
 			TableNameSchemaMapping: tableNameSchemaMapping,
@@ -334,7 +333,7 @@ func syncCore[TPull connectors.CDCPullConnectorCore, TSync connectors.CDCSyncCon
 	a.OtelManager.Metrics.CurrentBatchIdGauge.Record(ctx, res.CurrentSyncBatchID)
 
 	syncState.Store(shared.Ptr("updating schema"))
-	if err := a.applySchemaDeltas(ctx, config, options, res.TableSchemaDeltas); err != nil {
+	if err := a.applySchemaDeltas(ctx, config, res.TableSchemaDeltas); err != nil {
 		return nil, err
 	}
 

flow/cmd/handler.go

@@ -77,7 +77,6 @@ func (h *FlowRequestHandler) createCdcJobEntry(ctx context.Context,
 	if err != nil {
 		return fmt.Errorf("unable to marshal flow config: %w", err)
 	}
-
 	if _, err := h.pool.Exec(ctx,
 		`INSERT INTO flows (workflow_id, name, source_peer, destination_peer, config_proto, status,
 		description, source_table_identifier, destination_table_identifier) VALUES ($1,$2,$3,$4,$5,$6,'gRPC','','')`,

flow/cmd/mirror_status.go

@@ -473,7 +473,7 @@ func (h *FlowRequestHandler) getFlowConfigFromCatalog(
 ) (*protos.FlowConnectionConfigs, error) {
 	var configBytes sql.RawBytes
 	if err := h.pool.QueryRow(ctx,
-		"SELECT config_proto FROM flows WHERE name = $1", flowJobName,
+		"SELECT config_proto FROM flows WHERE name = $1 LIMIT 1", flowJobName,
 	).Scan(&configBytes); err != nil {
 		slog.ErrorContext(ctx, "unable to query flow config from catalog", slog.Any("error", err))
 		return nil, fmt.Errorf("unable to query flow config from catalog: %w", err)

flow/connectors/clickhouse/normalize.go

@@ -90,7 +90,12 @@ func (c *ClickHouseConnector) generateCreateTableSQLForNormalizedTable(
 	tmEngine := protos.TableEngine_CH_ENGINE_REPLACING_MERGE_TREE
 
 	var tableMapping *protos.TableMapping
-	for _, tm := range config.TableMappings {
+
+	cfg, err := internal.FetchConfigFromDB(config.FlowName)
+	if err != nil {
+		return nil, fmt.Errorf("failed to fetch config from DB: %w", err)
+	}
+	for _, tm := range cfg.TableMappings {
 		if tm.DestinationTableIdentifier == tableIdentifier {
 			tmEngine = tm.Engine
 			tableMapping = tm

flow/connectors/postgres/postgres.go

@@ -448,10 +448,15 @@ func pullCore[Items model.Items](
 		return fmt.Errorf("failed to get get setting for originMetaAsDestinationColumn: %w", err)
 	}
 
+	cfgFromDB, err := internal.FetchConfigFromDB(req.FlowJobName)
+	if err != nil {
+		return fmt.Errorf("unable to query flow config from catalog: %w", err)
+	}
+
 	cdc, err := c.NewPostgresCDCSource(ctx, &PostgresCDCConfig{
 		CatalogPool:                              catalogPool,
 		OtelManager:                              otelManager,
-		SrcTableIDNameMapping:                    req.SrcTableIDNameMapping,
+		SrcTableIDNameMapping:                    cfgFromDB.SrcTableIdNameMapping,
 		TableNameMapping:                         req.TableNameMapping,
 		TableNameSchemaMapping:                   req.TableNameSchemaMapping,
 		RelationMessageMapping:                   c.relationMessageMapping,

flow/e2e/postgres/peer_flow_pg_test.go

@@ -13,6 +13,7 @@ import (
 
 	"github.com/PeerDB-io/peerdb/flow/e2e"
 	"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/shared"
 	peerflow "github.com/PeerDB-io/peerdb/flow/workflows"
@@ -934,11 +935,12 @@ func (s PeerFlowE2ETestSuitePG) Test_Dynamic_Mirror_Config_Via_Signals() {
 		return s.comparePGTables(srcTable1Name, dstTable1Name, "id,t") == nil
 	})
 
-	workflowState := e2e.EnvGetWorkflowState(s.t, env)
-	assert.EqualValues(s.t, 7, workflowState.SyncFlowOptions.IdleTimeoutSeconds)
-	assert.EqualValues(s.t, 6, workflowState.SyncFlowOptions.BatchSize)
-	assert.Len(s.t, workflowState.SyncFlowOptions.TableMappings, 1)
-	assert.Len(s.t, workflowState.SyncFlowOptions.SrcTableIdNameMapping, 1)
+	updatedConfig, err := internal.FetchConfigFromDB(config.FlowJobName)
+	require.NoError(s.t, err)
+	assert.EqualValues(s.t, 7, updatedConfig.IdleTimeoutSeconds)
+	assert.EqualValues(s.t, 6, updatedConfig.MaxBatchSize)
+	assert.Len(s.t, updatedConfig.TableMappings, 1)
+	assert.Len(s.t, updatedConfig.SrcTableIdNameMapping, 1)
 
 	if !s.t.Failed() {
 		e2e.SignalWorkflow(s.t.Context(), env, model.FlowSignal, model.PauseSignal)
@@ -976,11 +978,12 @@ func (s PeerFlowE2ETestSuitePG) Test_Dynamic_Mirror_Config_Via_Signals() {
 			return s.comparePGTables(srcTable2Name, dstTable2Name, "id,t") == nil
 		})
 
-		workflowState = e2e.EnvGetWorkflowState(s.t, env)
-		assert.EqualValues(s.t, 14, workflowState.SyncFlowOptions.IdleTimeoutSeconds)
-		assert.EqualValues(s.t, 12, workflowState.SyncFlowOptions.BatchSize)
-		assert.Len(s.t, workflowState.SyncFlowOptions.TableMappings, 2)
-		assert.Len(s.t, workflowState.SyncFlowOptions.SrcTableIdNameMapping, 2)
+		updatedConfig, err = internal.FetchConfigFromDB(config.FlowJobName)
+		require.NoError(s.t, err)
+		assert.EqualValues(s.t, 14, updatedConfig.IdleTimeoutSeconds)
+		assert.EqualValues(s.t, 12, updatedConfig.MaxBatchSize)
+		assert.Len(s.t, updatedConfig.TableMappings, 2)
+		assert.Len(s.t, updatedConfig.SrcTableIdNameMapping, 2)
 	}
 
 	env.Cancel(s.t.Context())

flow/internal/flow_configuration_helpers.go

@@ -0,0 +1,38 @@
+package internal
+
+import (
+	"context"
+	"database/sql"
+	"fmt"
+
+	"google.golang.org/protobuf/proto"
+
+	"github.com/PeerDB-io/peerdb/flow/generated/protos"
+)
+
+func TableNameMapping(tableMappings []*protos.TableMapping) map[string]string {
+	tblNameMapping := make(map[string]string, len(tableMappings))
+	for _, v := range tableMappings {
+		tblNameMapping[v.SourceTableIdentifier] = v.DestinationTableIdentifier
+	}
+	return tblNameMapping
+}
+
+func FetchConfigFromDB(flowName string) (*protos.FlowConnectionConfigs, error) {
+	var configBytes sql.RawBytes
+	dbCtx := context.Background()
+	pool, _ := GetCatalogConnectionPoolFromEnv(dbCtx)
+	defer dbCtx.Done()
+	if err := pool.QueryRow(dbCtx,
+		"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
+}