diff --git a/flow/connectors/postgres/client.go b/flow/connectors/postgres/client.go index 2270c7fc9..5052c7666 100644 --- a/flow/connectors/postgres/client.go +++ b/flow/connectors/postgres/client.go @@ -54,13 +54,27 @@ const ( getTableNameToUnchangedToastColsSQL = `SELECT _peerdb_destination_table_name, ARRAY_AGG(DISTINCT _peerdb_unchanged_toast_columns) FROM %s.%s WHERE _peerdb_batch_id=$1 AND _peerdb_record_type!=2 GROUP BY _peerdb_destination_table_name` + mergeStatementSQLJsonbToRecord = `WITH src_rank AS ( + SELECT r.*,_peerdb_record_type,_peerdb_unchanged_toast_columns, _peerdb_timestamp, + RANK() OVER (PARTITION BY %s ORDER BY _peerdb_timestamp DESC) AS _peerdb_rank + FROM %s.%s, jsonb_to_record(_peerdb_data) AS r(%s) + WHERE _peerdb_batch_id = $1 AND _peerdb_destination_table_name = $2 + ) + MERGE INTO %s dst + USING (SELECT %s,_peerdb_record_type,_peerdb_unchanged_toast_columns + FROM src_rank WHERE _peerdb_rank=1 ORDER BY _peerdb_timestamp) src + ON %s + WHEN NOT MATCHED AND src._peerdb_record_type!=2 THEN + INSERT (%s) VALUES (%s) %s + WHEN MATCHED AND src._peerdb_record_type=2 THEN %s` mergeStatementSQL = `WITH src_rank AS ( - SELECT _peerdb_data,_peerdb_record_type,_peerdb_unchanged_toast_columns,_peerdb_timestamp, + SELECT _peerdb_data,_peerdb_record_type,_peerdb_unchanged_toast_columns, _peerdb_timestamp, RANK() OVER (PARTITION BY %s ORDER BY _peerdb_timestamp DESC) AS _peerdb_rank FROM %s.%s WHERE _peerdb_batch_id = $1 AND _peerdb_destination_table_name=$2 ) MERGE INTO %s dst - USING (SELECT %s,_peerdb_record_type,_peerdb_unchanged_toast_columns FROM src_rank WHERE _peerdb_rank=1 ORDER BY _peerdb_timestamp) src + USING (SELECT %s,_peerdb_record_type,_peerdb_unchanged_toast_columns + FROM src_rank WHERE _peerdb_rank=1 ORDER BY _peerdb_timestamp) src ON %s WHEN NOT MATCHED AND src._peerdb_record_type!=2 THEN INSERT (%s) VALUES (%s) %s diff --git a/flow/connectors/postgres/normalize_stmt_generator.go b/flow/connectors/postgres/normalize_stmt_generator.go index 0c190303a..e71332bf3 100644 --- a/flow/connectors/postgres/normalize_stmt_generator.go +++ b/flow/connectors/postgres/normalize_stmt_generator.go @@ -153,26 +153,55 @@ func (n *normalizeStmtGenerator) generateMergeStatement( columnCount := len(normalizedTableSchema.Columns) quotedColumnNames := make([]string, columnCount) - flattenedCastsSQLArray := make([]string, 0, columnCount) parsedDstTable, _ := common.ParseTableIdentifier(dstTableName) - primaryKeyColumnCasts := make(map[string]string) + primaryKeyColumnCasts := make(map[string]string, len(normalizedTableSchema.PrimaryKeyColumns)) primaryKeySelectSQLArray := make([]string, 0, len(normalizedTableSchema.PrimaryKeyColumns)) + + // For jsonb_to_record path we build: + // recordDefs – column definitions for the AS clause of jsonb_to_record (in the CTE) + // selectExprs – the SELECT list in the USING subquery, referencing columns from src_rank. + // json/jsonb columns are wrapped with _peerdb_parse_jsonb/_peerdb_parse_json + // to unwrap PeerDB's stringified representation. + // For legacy path we only build selectExprs (flattened casts via ->>). + selectExprs := make([]string, 0, columnCount) + recordDefs := make([]string, 0, columnCount) + useJsonbToRecord := normalizedTableSchema.System == protos.TypeSystem_PG for i, column := range normalizedTableSchema.Columns { - genericColumnType := column.Type quotedCol := common.QuoteIdentifier(column.Name) stringCol := utils.QuoteLiteral(column.Name) quotedColumnNames[i] = quotedCol pgType := n.columnTypeToPg(normalizedTableSchema, column) - expr := n.generateExpr(normalizedTableSchema, genericColumnType, stringCol, pgType) - flattenedCastsSQLArray = append(flattenedCastsSQLArray, fmt.Sprintf("%s AS %s", expr, quotedCol)) + if useJsonbToRecord { + // json/jsonb columns are stored as stringified text inside _peerdb_data, + // so jsonb_to_record extracts them as a jsonb string wrapper. + // Include them in the record as jsonb, then unwrap in the USING SELECT. + switch column.Type { + case "json": + recordDefs = append(recordDefs, quotedCol+" jsonb") + selectExprs = append(selectExprs, fmt.Sprintf("(%s #>> '{}')::json AS %s", quotedCol, quotedCol)) + case "jsonb": + recordDefs = append(recordDefs, quotedCol+" jsonb") + selectExprs = append(selectExprs, fmt.Sprintf("(%s #>> '{}')::jsonb AS %s", quotedCol, quotedCol)) + default: + recordDefs = append(recordDefs, fmt.Sprintf("%s %s", quotedCol, pgType)) + selectExprs = append(selectExprs, quotedCol) + } + } else { + genericColumnType := column.Type + expr := n.generateExpr(normalizedTableSchema, genericColumnType, stringCol, pgType) + selectExprs = append(selectExprs, fmt.Sprintf("%s AS %s", expr, quotedCol)) + } + if slices.Contains(normalizedTableSchema.PrimaryKeyColumns, column.Name) { - primaryKeyColumnCasts[column.Name] = fmt.Sprintf("(_peerdb_data->>%s)::%s", stringCol, pgType) + if !useJsonbToRecord { + primaryKeyColumnCasts[column.Name] = fmt.Sprintf("(_peerdb_data->>%s)::%s", stringCol, pgType) + } primaryKeySelectSQLArray = append(primaryKeySelectSQLArray, fmt.Sprintf("src.%s=dst.%s", quotedCol, quotedCol)) } } - flattenedCastsSQL := strings.Join(flattenedCastsSQLArray, ",") + selectExprsSQL := strings.Join(selectExprs, ",") insertValuesSQLArray := make([]string, 0, columnCount+2) for _, quotedCol := range quotedColumnNames { insertValuesSQLArray = append(insertValuesSQLArray, "src."+quotedCol) @@ -207,19 +236,43 @@ func (n *normalizeStmtGenerator) generateMergeStatement( } } - mergeStmt := fmt.Sprintf( - mergeStatementSQL, - strings.Join(slices.Collect(maps.Values(primaryKeyColumnCasts)), ","), - n.metadataSchema, - n.rawTableName, - parsedDstTable.String(), - flattenedCastsSQL, - strings.Join(primaryKeySelectSQLArray, " AND "), - insertColumnsSQL, - insertValuesSQL, - updateStringToastCols, - conflictPart, - ) + var mergeStmt string + if useJsonbToRecord { + // PARTITION BY uses quoted PK column names directly — jsonb_to_record + // already extracted them in the CTE via r.* + primaryKeyQuotedNames := make([]string, 0, len(normalizedTableSchema.PrimaryKeyColumns)) + for _, pkCol := range normalizedTableSchema.PrimaryKeyColumns { + primaryKeyQuotedNames = append(primaryKeyQuotedNames, common.QuoteIdentifier(pkCol)) + } + mergeStmt = fmt.Sprintf( + mergeStatementSQLJsonbToRecord, + strings.Join(primaryKeyQuotedNames, ","), + n.metadataSchema, + n.rawTableName, + strings.Join(recordDefs, ","), + parsedDstTable.String(), + selectExprsSQL, + strings.Join(primaryKeySelectSQLArray, " AND "), + insertColumnsSQL, + insertValuesSQL, + updateStringToastCols, + conflictPart, + ) + } else { + mergeStmt = fmt.Sprintf( + mergeStatementSQL, + strings.Join(slices.Collect(maps.Values(primaryKeyColumnCasts)), ","), + n.metadataSchema, + n.rawTableName, + parsedDstTable.String(), + selectExprsSQL, + strings.Join(primaryKeySelectSQLArray, " AND "), + insertColumnsSQL, + insertValuesSQL, + updateStringToastCols, + conflictPart, + ) + } return mergeStmt } diff --git a/flow/connectors/postgres/normalize_stmt_generator_test.go b/flow/connectors/postgres/normalize_stmt_generator_test.go index f1c7e9f4c..00065660a 100644 --- a/flow/connectors/postgres/normalize_stmt_generator_test.go +++ b/flow/connectors/postgres/normalize_stmt_generator_test.go @@ -4,6 +4,8 @@ import ( "reflect" "testing" + "github.com/stretchr/testify/require" + "github.com/PeerDB-io/peerdb/flow/connectors/utils" "github.com/PeerDB-io/peerdb/flow/generated/protos" ) @@ -138,3 +140,213 @@ func TestGenerateMergeUpdateStatement_WithUnchangedToastColsAndSoftDelete(t *tes t.Errorf("Unexpected result. Expected: %v, but got: %v", expected, result) } } + +// helper to build a simple PG-typed TableSchema for merge tests +func buildTableSchema(columns []*protos.FieldDescription, pks []string) *protos.TableSchema { + return &protos.TableSchema{ + TableIdentifier: "test_table", + System: protos.TypeSystem_PG, + PrimaryKeyColumns: pks, + Columns: columns, + } +} + +func normalizeSQL(s string) string { + return utils.RemoveSpacesTabsNewlines(s) +} + +func TestGenerateMergeStatement_BasicColumns(t *testing.T) { + schema := buildTableSchema([]*protos.FieldDescription{ + {Name: "id", Type: "integer"}, + {Name: "name", Type: "text"}, + {Name: "value", Type: "numeric"}, + }, []string{"id"}) + + gen := normalizeStmtGenerator{ + rawTableName: "_peerdb_raw_test", + tableSchemaMapping: map[string]*protos.TableSchema{"public.test_table": schema}, + unchangedToastColumnsMap: map[string][]string{"public.test_table": {""}}, + peerdbCols: &protos.PeerDBColumns{ + SyncedAtColName: "_peerdb_synced_at", + SoftDeleteColName: "", + }, + metadataSchema: "_peerdb_internal", + supportsMerge: true, + } + + result := gen.generateMergeStatement("public.test_table", schema, []string{""}) + + // CTE uses jsonb_to_record with record definitions + require.Contains(t, result, "jsonb_to_record(_peerdb_data)") + require.Contains(t, result, `"id" integer`) + require.Contains(t, result, `"name" text`) + require.Contains(t, result, `"value" numeric`) + // USING select just references the column names directly (no ->> casts) + require.Contains(t, normalizeSQL(result), normalizeSQL(`"id","name","value",_peerdb_record_type,_peerdb_unchanged_toast_columns`)) + // MERGE ON condition + require.Contains(t, result, `src."id"=dst."id"`) + require.NotContains(t, result, "_peerdb_data->>") +} + +func TestGenerateMergeStatement_JsonColumns(t *testing.T) { + schema := buildTableSchema([]*protos.FieldDescription{ + {Name: "id", Type: "integer"}, + {Name: "metadata", Type: "json"}, + {Name: "config", Type: "jsonb"}, + }, []string{"id"}) + + gen := normalizeStmtGenerator{ + rawTableName: "_peerdb_raw_test", + tableSchemaMapping: map[string]*protos.TableSchema{"public.test_table": schema}, + unchangedToastColumnsMap: map[string][]string{"public.test_table": {""}}, + peerdbCols: &protos.PeerDBColumns{ + SyncedAtColName: "_peerdb_synced_at", + SoftDeleteColName: "", + }, + metadataSchema: "_peerdb_internal", + supportsMerge: true, + } + + result := gen.generateMergeStatement("public.test_table", schema, []string{""}) + + // json/jsonb columns: record defs should declare them as jsonb + require.Contains(t, result, `"metadata" jsonb`) + require.Contains(t, result, `"config" jsonb`) + // USING select should unwrap them via #>> + require.Contains(t, normalizeSQL(result), normalizeSQL(`("metadata" #>> '{}')::json AS "metadata"`)) + require.Contains(t, normalizeSQL(result), normalizeSQL(`("config" #>> '{}')::jsonb AS "config"`)) +} + +func TestGenerateMergeStatement_SoftDelete(t *testing.T) { + schema := buildTableSchema([]*protos.FieldDescription{ + {Name: "id", Type: "integer"}, + {Name: "data", Type: "text"}, + }, []string{"id"}) + + gen := normalizeStmtGenerator{ + rawTableName: "_peerdb_raw_test", + tableSchemaMapping: map[string]*protos.TableSchema{"public.test_table": schema}, + unchangedToastColumnsMap: map[string][]string{"public.test_table": {""}}, + peerdbCols: &protos.PeerDBColumns{ + SyncedAtColName: "_peerdb_synced_at", + SoftDeleteColName: "_peerdb_soft_delete", + }, + metadataSchema: "_peerdb_internal", + supportsMerge: true, + } + + result := gen.generateMergeStatement("public.test_table", schema, []string{""}) + + // soft delete: WHEN NOT MATCHED with record_type=2 inserts with soft delete TRUE + require.Contains(t, normalizeSQL(result), + normalizeSQL(`WHEN NOT MATCHED AND (src._peerdb_record_type=2) THEN INSERT`)) + require.Contains(t, normalizeSQL(result), normalizeSQL(`"_peerdb_soft_delete"`)) + // conflict part should be UPDATE SET soft_delete=TRUE + require.Contains(t, normalizeSQL(result), + normalizeSQL(`UPDATE SET "_peerdb_soft_delete"=TRUE,"_peerdb_synced_at"=CURRENT_TIMESTAMP`)) +} + +func TestGenerateMergeStatement_CompositePK(t *testing.T) { + schema := buildTableSchema([]*protos.FieldDescription{ + {Name: "tenant_id", Type: "integer"}, + {Name: "user_id", Type: "integer"}, + {Name: "email", Type: "text"}, + }, []string{"tenant_id", "user_id"}) + + gen := normalizeStmtGenerator{ + rawTableName: "_peerdb_raw_test", + tableSchemaMapping: map[string]*protos.TableSchema{"public.test_table": schema}, + unchangedToastColumnsMap: map[string][]string{"public.test_table": {""}}, + peerdbCols: &protos.PeerDBColumns{ + SyncedAtColName: "_peerdb_synced_at", + SoftDeleteColName: "", + }, + metadataSchema: "_peerdb_internal", + supportsMerge: true, + } + + result := gen.generateMergeStatement("public.test_table", schema, []string{""}) + + // composite PK: PARTITION BY should use both PK columns + require.Contains(t, normalizeSQL(result), normalizeSQL(`PARTITION BY "tenant_id","user_id"`)) + // ON clause should join on both + require.Contains(t, normalizeSQL(result), normalizeSQL(`src."tenant_id"=dst."tenant_id"`)) + require.Contains(t, normalizeSQL(result), normalizeSQL(`src."user_id"=dst."user_id"`)) +} + +func TestGenerateMergeStatement_ToastColumns(t *testing.T) { + schema := buildTableSchema([]*protos.FieldDescription{ + {Name: "id", Type: "integer"}, + {Name: "small_col", Type: "text"}, + {Name: "big_col", Type: "text"}, + }, []string{"id"}) + + gen := normalizeStmtGenerator{ + rawTableName: "_peerdb_raw_test", + tableSchemaMapping: map[string]*protos.TableSchema{"public.test_table": schema}, + unchangedToastColumnsMap: map[string][]string{"public.test_table": {"", "big_col"}}, + peerdbCols: &protos.PeerDBColumns{ + SyncedAtColName: "_peerdb_synced_at", + SoftDeleteColName: "", + }, + metadataSchema: "_peerdb_internal", + supportsMerge: true, + } + + result := gen.generateMergeStatement("public.test_table", schema, []string{"", "big_col"}) + + normalized := normalizeSQL(result) + // Should have an update branch for unchanged toast col = '' (all cols updated) + require.Contains(t, normalized, normalizeSQL(`_peerdb_unchanged_toast_columns=''`)) + // Should have an update branch for unchanged toast col = 'big_col' (only id and small_col updated) + require.Contains(t, normalized, normalizeSQL(`_peerdb_unchanged_toast_columns='big_col'`)) + require.Contains(t, normalized, normalizeSQL(`"id"=src."id","small_col"=src."small_col"`)) +} + +func TestGenerateMergeStatement_UserDefinedType(t *testing.T) { + schema := buildTableSchema([]*protos.FieldDescription{ + {Name: "id", Type: "integer"}, + {Name: "status", Type: "my_enum", TypeSchemaName: "my_schema"}, + }, []string{"id"}) + + gen := normalizeStmtGenerator{ + rawTableName: "_peerdb_raw_test", + tableSchemaMapping: map[string]*protos.TableSchema{"public.test_table": schema}, + unchangedToastColumnsMap: map[string][]string{"public.test_table": {""}}, + peerdbCols: &protos.PeerDBColumns{ + SyncedAtColName: "_peerdb_synced_at", + SoftDeleteColName: "", + }, + metadataSchema: "_peerdb_internal", + supportsMerge: true, + } + + result := gen.generateMergeStatement("public.test_table", schema, []string{""}) + + // User-defined types should be schema-qualified in the record definition + require.Contains(t, result, `"status" "my_schema"."my_enum"`) +} + +func TestGenerateNormalizeStatements_Merge(t *testing.T) { + schema := buildTableSchema([]*protos.FieldDescription{ + {Name: "id", Type: "integer"}, + {Name: "name", Type: "text"}, + }, []string{"id"}) + + gen := normalizeStmtGenerator{ + rawTableName: "_peerdb_raw_test", + tableSchemaMapping: map[string]*protos.TableSchema{"public.test_table": schema}, + unchangedToastColumnsMap: map[string][]string{"public.test_table": {""}}, + peerdbCols: &protos.PeerDBColumns{ + SyncedAtColName: "_peerdb_synced_at", + SoftDeleteColName: "", + }, + metadataSchema: "_peerdb_internal", + supportsMerge: true, + } + + stmts := gen.generateNormalizeStatements("public.test_table") + require.Len(t, stmts, 1) + require.Contains(t, stmts[0], "jsonb_to_record") + require.Contains(t, stmts[0], "MERGE INTO") +} diff --git a/flow/e2e/postgres_test.go b/flow/e2e/postgres_test.go index 57af4be6e..4a6a0925c 100644 --- a/flow/e2e/postgres_test.go +++ b/flow/e2e/postgres_test.go @@ -85,11 +85,11 @@ func (s PeerFlowE2ETestSuitePG) Test_Geospatial_PG() { RequireEnvCanceled(s.t, env) } -func (s PeerFlowE2ETestSuitePG) Test_Types() { +func (s PeerFlowE2ETestSuitePG) testTypes(suffix string, system protos.TypeSystem) { tc := NewTemporalClient(s.t) - srcTableName := s.attachSchemaSuffix("test_types_pg") - dstTableName := s.attachSchemaSuffix("test_types_pg_dst") + srcTableName := s.attachSchemaSuffix("test_types_" + suffix) + dstTableName := s.attachSchemaSuffix("test_types_" + suffix + "_dst") _, err := s.Conn().Exec(s.t.Context(), fmt.Sprintf(` CREATE TABLE IF NOT EXISTS %s (id serial PRIMARY KEY,c1 BIGINT,c2 BYTEA,c4 BOOLEAN, @@ -103,7 +103,7 @@ func (s PeerFlowE2ETestSuitePG) Test_Types() { require.NoError(s.t, err) connectionGen := FlowConnectionGenerationConfig{ - FlowJobName: s.attachSuffix("test_types_pg"), + FlowJobName: s.attachSuffix("test_types_" + suffix), TableNameMapping: map[string]string{srcTableName: dstTableName}, Destination: s.Peer().Name, } @@ -112,6 +112,7 @@ func (s PeerFlowE2ETestSuitePG) Test_Types() { flowConnConfig.MaxBatchSize = 100 flowConnConfig.SoftDeleteColName = "" flowConnConfig.SyncedAtColName = "" + flowConnConfig.System = system env := ExecutePeerflow(s.t, tc, flowConnConfig) SetupCDCFlowStatusQuery(s.t, env, flowConnConfig) @@ -138,7 +139,7 @@ func (s PeerFlowE2ETestSuitePG) Test_Types() { allCols := strings.Join([]string{ "c1", "c2", "c4", "c40", "id", "c9", "c11", "c12", "c13", "c14", "c15", - "c21", "c29", "c33", "c34", "c35", "c37", + "c21", "c29", "c33", "c34", "c35", "c7", "c8", "c32", "c42", "c43", "c44", "c45", "c46", "c47", "c48", "c49", "c50", }, ",") EnvWaitFor(s.t, env, 3*time.Minute, "normalize types", func() bool { @@ -148,15 +149,27 @@ func (s PeerFlowE2ETestSuitePG) Test_Types() { } return err == nil }) - // c36 converted to UTC, losing tz info, so does not compare equal + // Q type system converts timetz to UTC, PG type system preserves original offset var c36 string require.NoError(s.t, s.Conn().QueryRow(s.t.Context(), "select c36 from "+dstTableName).Scan(&c36)) - require.Equal(s.t, "06:25:00+00", c36) + if system == protos.TypeSystem_PG { + require.Equal(s.t, "09:25:00+03", c36) + } else { + require.Equal(s.t, "06:25:00+00", c36) + } env.Cancel(s.t.Context()) RequireEnvCanceled(s.t, env) } +func (s PeerFlowE2ETestSuitePG) Test_Types_QValue() { + s.testTypes("q", protos.TypeSystem_Q) +} + +func (s PeerFlowE2ETestSuitePG) Test_Types_PG() { + s.testTypes("pg", protos.TypeSystem_PG) +} + func (s PeerFlowE2ETestSuitePG) Test_PgVector() { srcTableName := "pg_pgvector" srcFullName := s.attachSchemaSuffix(srcTableName)