Add support for more properties to be passed with TableMapping (#500)

PeerDB-io · Oct 11, 2023 · 1fa5903 · 1fa5903
1 parent 1213a83
commit 1fa5903
Show file tree

Hide file tree

Showing 17 changed files with 1,210 additions and 954 deletions.
diff --git a/flow/activities/flowable.go b/flow/activities/flowable.go
@@ -188,11 +188,16 @@ func (a *FlowableActivity) StartFlow(ctx context.Context,
 		"flowName": input.FlowConnectionConfigs.FlowJobName,
 	}).Info("pulling records...")
 
+	tblNameMapping := make(map[string]string)
+	for _, v := range input.FlowConnectionConfigs.TableMappings {
+		tblNameMapping[v.SourceTableIdentifier] = v.DestinationTableIdentifier
+	}
+
 	startTime := time.Now()
 	recordsWithTableSchemaDelta, err := srcConn.PullRecords(&model.PullRecordsRequest{
 		FlowJobName:                 input.FlowConnectionConfigs.FlowJobName,
 		SrcTableIDNameMapping:       input.FlowConnectionConfigs.SrcTableIdNameMapping,
-		TableNameMapping:            input.FlowConnectionConfigs.TableNameMapping,
+		TableNameMapping:            tblNameMapping,
 		LastSyncState:               input.LastSyncState,
 		MaxBatchSize:                uint32(input.SyncFlowOptions.BatchSize),
 		IdleTimeout:                 10 * time.Second,

diff --git a/flow/cmd/handler.go b/flow/cmd/handler.go
@@ -65,18 +65,17 @@ func (h *FlowRequestHandler) createFlowJobEntry(ctx context.Context,
 			req.ConnectionConfigs.Destination.Name, srcErr)
 	}
 
-	for sourceTableIdentifier := range req.ConnectionConfigs.TableNameMapping {
-		destinationTableIdentifier := req.ConnectionConfigs.TableNameMapping[sourceTableIdentifier]
+	for _, v := range req.ConnectionConfigs.TableMappings {
 		_, err := h.pool.Exec(ctx, `
 		INSERT INTO flows (workflow_id, name, source_peer, destination_peer, description,
 		source_table_identifier, destination_table_identifier) VALUES ($1, $2, $3, $4, $5, $6, $7)
 		`, workflowID, req.ConnectionConfigs.FlowJobName, sourcePeerID, destinationPeerID,
 			"Mirror created via GRPC",
-			schemaForTableIdentifier(sourceTableIdentifier, sourePeerType),
-			schemaForTableIdentifier(destinationTableIdentifier, destinationPeerType))
+			schemaForTableIdentifier(v.SourceTableIdentifier, sourePeerType),
+			schemaForTableIdentifier(v.DestinationTableIdentifier, destinationPeerType))
 		if err != nil {
 			return fmt.Errorf("unable to insert into flows table for flow %s with source table %s: %w",
-				req.ConnectionConfigs.FlowJobName, sourceTableIdentifier, err)
+				req.ConnectionConfigs.FlowJobName, v.SourceTableIdentifier, err)
 		}
 	}
 

diff --git a/flow/e2e/congen.go b/flow/e2e/congen.go
@@ -145,9 +145,17 @@ func GenerateSnowflakePeer(snowflakeConfig *protos.SnowflakeConfig) (*protos.Pee
 }
 
 func (c *FlowConnectionGenerationConfig) GenerateFlowConnectionConfigs() (*protos.FlowConnectionConfigs, error) {
+	tblMappings := []*protos.TableMapping{}
+	for k, v := range c.TableNameMapping {
+		tblMappings = append(tblMappings, &protos.TableMapping{
+			SourceTableIdentifier:      k,
+			DestinationTableIdentifier: v,
+		})
+	}
+
 	ret := &protos.FlowConnectionConfigs{}
 	ret.FlowJobName = c.FlowJobName
-	ret.TableNameMapping = c.TableNameMapping
+	ret.TableMappings = tblMappings
 	ret.Source = GeneratePostgresPeer(c.PostgresPort)
 	ret.Destination = c.Destination
 	ret.CdcSyncMode = c.CDCSyncMode