Fix subsetting table with json column (#1982)

nucleuscloud · May 21, 2024 · 34ceca3 · 34ceca3
1 parent 97d8f15
commit 34ceca3
Show file tree

Hide file tree

Showing 4 changed files with 76 additions and 34 deletions.
diff --git a/worker/internal/benthos/utils.go b/worker/internal/benthos/utils.go
@@ -30,6 +30,7 @@ func IsCriticalError(errMsg string) bool {
 		"Duplicate entry",
 		"Cannot add or update a child row",
 		"a foreign key constraint fails",
+		"could not identify an equality operator",
 	}
 
 	for _, errStr := range criticalErrors {

diff --git a/worker/pkg/workflows/datasync/activities/gen-benthos-configs/benthos-builder.go b/worker/pkg/workflows/datasync/activities/gen-benthos-configs/benthos-builder.go
@@ -185,7 +185,7 @@ func (b *benthosBuilder) GenerateBenthosConfigs(
 		if err != nil {
 			return nil, fmt.Errorf("unable to retrieve database foreign key constraints: %w", err)
 		}
-		tableQueryMap, err := buildSelectQueryMap(db.Driver, groupedTableMapping, sourceTableOpts, fkReferenceMap, runConfigs, sqlSourceOpts.SubsetByForeignKeyConstraints)
+		tableQueryMap, err := buildSelectQueryMap(db.Driver, groupedTableMapping, sourceTableOpts, fkReferenceMap, runConfigs, sqlSourceOpts.SubsetByForeignKeyConstraints, groupedSchemas)
 		if err != nil {
 			return nil, fmt.Errorf("unable to build select queries: %w", err)
 		}

diff --git a/worker/pkg/workflows/datasync/activities/gen-benthos-configs/query-builder.go b/worker/pkg/workflows/datasync/activities/gen-benthos-configs/query-builder.go
@@ -126,6 +126,7 @@ func buildSelectJoinQuery(
 func buildSelectRecursiveQuery(
 	driver, schema, table string,
 	columns []string,
+	columnInfoMap map[string]*sql_manager.ColumnInfo,
 	dependencies []*selfReferencingCircularDependency,
 	joins []*sqlJoin,
 	whereClauses []string,
@@ -148,7 +149,12 @@ func buildSelectRecursiveQuery(
 
 	selectColumns := make([]any, len(columns))
 	for i, col := range columns {
-		selectColumns[i] = buildSqlIdentifier(schema, table, col)
+		colInfo := columnInfoMap[col]
+		if driver == sql_manager.PostgresDriver && colInfo != nil && colInfo.DataType == "json" {
+			selectColumns[i] = goqu.L("to_jsonb(?)", goqu.I(buildSqlIdentifier(schema, table, col))).As(col)
+		} else {
+			selectColumns[i] = buildSqlIdentifier(schema, table, col)
+		}
 	}
 	selectQuery := builder.From(sqltable).Select(selectColumns...)
 
@@ -220,6 +226,7 @@ func buildSelectQueryMap(
 	tableDependencies map[string][]*sql_manager.ColumnConstraint,
 	dependencyConfigs []*tabledependency.RunConfig,
 	subsetByForeignKeyConstraints bool,
+	groupedColumnInfo map[string]map[string]*sql_manager.ColumnInfo,
 ) (map[string]string, error) {
 	// map of table -> where clause
 	tableWhereMap := map[string]string{}
@@ -262,6 +269,7 @@ func buildSelectQueryMap(
 	for table := range dependencyMap {
 		tableMapping := groupedMappings[table]
 		config := subsetConfigs[table]
+		columnInfoMap := groupedColumnInfo[table]
 
 		selectCols := []string{}
 		for _, m := range tableMapping.Mappings {
@@ -273,6 +281,7 @@ func buildSelectQueryMap(
 				tableMapping.Schema,
 				tableMapping.Table,
 				buildPlainColumns(tableMapping.Mappings),
+				columnInfoMap,
 				config.SelfReferencingCircularDependency,
 				config.Joins,
 				config.WhereClauses,