feat(pg-pg): introduce new MERGE statement leveraging jsonb_to_record

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

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
 }

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")
+}