Dealing with the removed tableMappings - pt2; build is passing

Author: alexstoick

flow/activities/flowable.go

@@ -132,7 +132,12 @@ func (a *FlowableActivity) EnsurePullability(
 	if err != nil {
 		return nil, err
 	}
-	config.SourceTableIdentifiers = slices.Sorted(maps.Keys(internal.TableNameMapping(cfg.TableMappings, cfg.Resync)))
+	tableMappings, err := internal.FetchTableMappingsFromDB(ctx, config.FlowJobName, cfg.TableMappingVersion)
+	if err != nil {
+		return nil, err
+	}
+
+	config.SourceTableIdentifiers = slices.Sorted(maps.Keys(internal.TableNameMapping(tableMappings, cfg.Resync)))
 
 	output, err := srcConn.EnsurePullability(ctx, config)
 	if err != nil {
@@ -154,12 +159,6 @@ 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)
@@ -193,11 +192,15 @@ func (a *FlowableActivity) SetupTableSchema(
 	if err != nil {
 		return err
 	}
-	tableNameSchemaMapping, err := srcConn.GetTableSchema(ctx, config.Env, config.Version, config.System, cfg.TableMappings)
+	tableMappings, err := internal.FetchTableMappingsFromDB(ctx, cfg.FlowJobName, cfg.TableMappingVersion)
+	if err != nil {
+		return err
+	}
+	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(cfg.TableMappings, tableNameSchemaMapping, logger)
+	processed := internal.BuildProcessedSchemaMapping(tableMappings, tableNameSchemaMapping, logger)
 
 	tx, err := a.CatalogPool.BeginTx(ctx, pgx.TxOptions{})
 	if err != nil {
@@ -266,9 +269,13 @@ func (a *FlowableActivity) CreateNormalizedTable(
 	if err != nil {
 		return nil, err
 	}
+	tableMappings, err := internal.FetchTableMappingsFromDB(ctx, cfg.FlowJobName, cfg.TableMappingVersion)
+	if err != nil {
+		return nil, err
+	}
 	numTablesToSetup.Store(int32(len(tableNameSchemaMapping)))
 	tableExistsMapping := make(map[string]bool, len(tableNameSchemaMapping))
-	for _, tableMapping := range cfg.TableMappings {
+	for _, tableMapping := range tableMappings {
 		tableIdentifier := tableMapping.DestinationTableIdentifier
 		tableSchema := tableNameSchemaMapping[tableIdentifier]
 		existing, err := conn.SetupNormalizedTable(
@@ -321,9 +328,14 @@ func (a *FlowableActivity) SyncFlow(
 		return err
 	}
 
+	tableMappings, err := internal.FetchTableMappingsFromDB(ctx, cfg.FlowJobName, cfg.TableMappingVersion)
+	if err != nil {
+		return err
+	}
+
 	// Override config with DB values to deal with the large fields.
 	config = cfg
-	options.TableMappings = cfg.TableMappings
+	options.TableMappings = tableMappings
 
 	syncState.Store(shared.Ptr("setup"))
 	shutdown := heartbeatRoutine(ctx, func() string {
@@ -1073,7 +1085,8 @@ func (a *FlowableActivity) RecordMetricsCritical(ctx context.Context) error {
 		if isActive {
 			activeFlows = append(activeFlows, info)
 		}
-		a.OtelManager.Metrics.SyncedTablesGauge.Record(ctx, int64(len(info.config.TableMappings)))
+		//TODO: this will need a special query as we can extract this straight from the DB.
+		//a.OtelManager.Metrics.SyncedTablesGauge.Record(ctx, int64(len(info.config.TableMappings)))
 		a.OtelManager.Metrics.FlowStatusGauge.Record(ctx, 1, metric.WithAttributeSet(attribute.NewSet(
 			attribute.String(otel_metrics.FlowStatusKey, info.status.String()),
 			attribute.Bool(otel_metrics.IsFlowActiveKey, isActive),

flow/activities/flowable_core.go

@@ -675,13 +675,18 @@ func (a *FlowableActivity) startNormalize(
 		return fmt.Errorf("failed to get table name schema mapping: %w", err)
 	}
 
+	tableMappings, err := internal.FetchTableMappingsFromDB(ctx, config.FlowJobName, config.TableMappingVersion)
+	if err != nil {
+		return err
+	}
+
 	for {
 		logger.Info("normalizing batch", slog.Int64("syncBatchID", batchID))
 		res, err := dstConn.NormalizeRecords(ctx, &model.NormalizeRecordsRequest{
 			FlowJobName:            config.FlowJobName,
 			Env:                    config.Env,
 			TableNameSchemaMapping: tableNameSchemaMapping,
-			TableMappings:          config.TableMappings,
+			TableMappings:          tableMappings,
 			SoftDeleteColName:      config.SoftDeleteColName,
 			SyncedAtColName:        config.SyncedAtColName,
 			SyncBatchID:            batchID,

flow/activities/snapshot_activity.go

@@ -71,7 +71,11 @@ func (a *SnapshotActivity) SetupReplication(
 	if err != nil {
 		return nil, err
 	}
-	config.TableNameMapping = internal.TableNameMapping(cfg.TableMappings, cfg.Resync)
+	tableMappings, err := internal.FetchTableMappingsFromDB(ctx, cfg.FlowJobName, cfg.TableMappingVersion)
+	if err != nil {
+		return nil, err
+	}
+	config.TableNameMapping = internal.TableNameMapping(tableMappings, cfg.Resync)
 
 	logger.Info("waiting for slot to be created...")
 	slotInfo, err := conn.SetupReplication(ctx, config)
@@ -192,8 +196,12 @@ func (a *SnapshotActivity) GetDefaultPartitionKeyForTables(
 	if err != nil {
 		return nil, err
 	}
+	tableMappings, err := internal.FetchTableMappingsFromDB(ctx, cfg.FlowJobName, cfg.TableMappingVersion)
+	if err != nil {
+		return nil, err
+	}
 	output, err := connector.GetDefaultPartitionKeyForTables(ctx, &protos.GetDefaultPartitionKeyForTablesInput{
-		TableMappings: cfg.TableMappings,
+		TableMappings: 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/cmd/handler.go

@@ -469,8 +469,7 @@ func (h *FlowRequestHandler) FlowStateChange(
 				if err != nil {
 					return nil, NewInternalApiError(fmt.Errorf("unable to get flow config: %w", err))
 				}
-				tableMappings, err := internal.FetchTableMappingsFromDB(ctx, h.pool, req.FlowJobName)
-
+				tableMappings, err := internal.FetchTableMappingsFromDB(ctx, config.FlowJobName, config.TableMappingVersion)
 				if err != nil {
 					return nil, NewInternalApiError(fmt.Errorf("unable to get table mappings: %w", err))
 				}

flow/cmd/mirror_status.go

@@ -163,7 +163,8 @@ func (h *FlowRequestHandler) cdcFlowStatus(
 	if state.SyncFlowOptions != nil {
 		config.IdleTimeoutSeconds = state.SyncFlowOptions.IdleTimeoutSeconds
 		config.MaxBatchSize = state.SyncFlowOptions.BatchSize
-		config.TableMappings = state.SyncFlowOptions.TableMappings
+		//TODO: this will need to fetch from the DB?
+		//config.TableMappings = state.SyncFlowOptions.TableMappings
 	}
 
 	srcType, err := connectors.LoadPeerType(ctx, h.pool, config.SourceName)

flow/workflows/cdc_flow.go

@@ -327,7 +327,11 @@ func processTableAdditions(
 			additionalTablesCfg := proto.CloneOf(cfg)
 			additionalTablesCfg.DoInitialSnapshot = !flowConfigUpdate.SkipInitialSnapshotForTableAdditions
 			additionalTablesCfg.InitialSnapshotOnly = true
-			additionalTablesCfg.TableMappings = flowConfigUpdate.AdditionalTables
+			// TODO: thought - maybe we need to pass additionalTables and then in the
+			// CDCWorkflow we persist them to the DB and send an incremented `tableMappingVersion`
+			// to the new workflow?
+			//additionalTablesCfg.TableMappings = flowConfigUpdate.AdditionalTables
+
 			additionalTablesCfg.Resync = false
 			if state.SnapshotNumRowsPerPartition > 0 {
 				additionalTablesCfg.SnapshotNumRowsPerPartition = state.SnapshotNumRowsPerPartition

flow/workflows/snapshot_flow.go

@@ -287,7 +287,14 @@ func (s *SnapshotFlowExecution) cloneTables(
 		return err
 	}
 
-	for _, v := range cfg.TableMappings {
+	tableMappingsCtx := context.Background()
+	defer tableMappingsCtx.Done()
+	tableMappings, err := internal.FetchTableMappingsFromDB(tableMappingsCtx, cfg.FlowJobName, cfg.TableMappingVersion)
+	if err != nil {
+		return err
+	}
+
+	for _, v := range tableMappings {
 		source := v.SourceTableIdentifier
 		destination := v.DestinationTableIdentifier
 		s.logger.Info(