From af120cde4c2f06440cef0b20ee35d8939551c7a2 Mon Sep 17 00:00:00 2001 From: Bill Guowei Yang Date: Thu, 2 Jul 2026 10:05:00 -0400 Subject: [PATCH 1/6] Guard DuckLake ON CONFLICT merge cardinality --- transpiler/classify_wiring_test.go | 3 + transpiler/transform/onconflict.go | 195 ++++++++++++++++++++++++++- transpiler/transpiler.go | 6 +- transpiler/transpiler_test.go | 207 ++++++++++++++++++++++++++++- 4 files changed, 402 insertions(+), 9 deletions(-) diff --git a/transpiler/classify_wiring_test.go b/transpiler/classify_wiring_test.go index 3da41e97..5665096f 100644 --- a/transpiler/classify_wiring_test.go +++ b/transpiler/classify_wiring_test.go @@ -39,6 +39,9 @@ func transpileSQL(t *testing.T, tr *Transpiler, input string) string { if res.FallbackToNative { t.Fatalf("Transpile(%q) fell back to native, want transpiled SQL", input) } + if len(res.Statements) > 0 { + return res.Statements[len(res.Statements)-1] + } return res.SQL } diff --git a/transpiler/transform/onconflict.go b/transpiler/transform/onconflict.go index 5735bbb3..34f04a2a 100644 --- a/transpiler/transform/onconflict.go +++ b/transpiler/transform/onconflict.go @@ -1,6 +1,8 @@ package transform import ( + "fmt" + "math/rand/v2" "strings" pg_query "github.com/pganalyze/pg_query_go/v6" @@ -24,6 +26,13 @@ type OnConflictTransform struct { DuckLakeMode bool } +type onConflictMergeRewrite struct { + merge *pg_query.MergeStmt + sourceSelect *pg_query.SelectStmt + targetRelation *pg_query.RangeVar + conflictColumns []string +} + func NewOnConflictTransform() *OnConflictTransform { return &OnConflictTransform{DuckLakeMode: false} } @@ -57,12 +66,19 @@ func (t *OnConflictTransform) Transform(tree *pg_query.ParseResult, result *Resu return false, nil } - if mergeStmt := t.transformInsertToMerge(insert); mergeStmt != nil { + if rewrite := t.transformInsertToMerge(insert); rewrite != nil { // Replace the INSERT statement with MERGE stmt.Stmt = &pg_query.Node{ - Node: &pg_query.Node_MergeStmt{MergeStmt: mergeStmt}, + Node: &pg_query.Node_MergeStmt{MergeStmt: rewrite.merge}, } changed = true + + if statements, cleanup, err := t.buildGuardedMergeStatements(rewrite); err != nil { + return false, err + } else if len(statements) > 0 { + result.Statements = statements + result.CleanupStatements = cleanup + } } else if t.transformInsert(insert) { changed = true } @@ -73,7 +89,7 @@ func (t *OnConflictTransform) Transform(tree *pg_query.ParseResult, result *Resu } // transformInsertToMerge converts INSERT ... ON CONFLICT to MERGE for DuckLake mode -func (t *OnConflictTransform) transformInsertToMerge(insert *pg_query.InsertStmt) *pg_query.MergeStmt { +func (t *OnConflictTransform) transformInsertToMerge(insert *pg_query.InsertStmt) *onConflictMergeRewrite { if !t.DuckLakeMode { return nil } @@ -88,6 +104,7 @@ func (t *OnConflictTransform) transformInsertToMerge(insert *pg_query.InsertStmt if occ.Infer == nil || len(occ.Infer.IndexElems) == 0 { return nil } + conflictColumns := t.conflictColumnNames(occ.Infer.IndexElems) // Get column names from INSERT colNames := make([]string, len(insert.Cols)) @@ -151,12 +168,182 @@ func (t *OnConflictTransform) transformInsertToMerge(insert *pg_query.InsertStmt Node: &pg_query.Node_MergeWhenClause{MergeWhenClause: insertClause}, }) - return &pg_query.MergeStmt{ + merge := &pg_query.MergeStmt{ Relation: insert.Relation, SourceRelation: sourceRelation, JoinCondition: joinCondition, MergeWhenClauses: whenClauses, } + + return &onConflictMergeRewrite{ + merge: merge, + sourceSelect: sourceSelect, + targetRelation: insert.Relation, + conflictColumns: conflictColumns, + } +} + +func (t *OnConflictTransform) conflictColumnNames(indexElems []*pg_query.Node) []string { + cols := make([]string, 0, len(indexElems)) + for _, elem := range indexElems { + indexElem := elem.GetIndexElem() + if indexElem == nil || strings.TrimSpace(indexElem.Name) == "" { + return nil + } + cols = append(cols, indexElem.Name) + } + return cols +} + +func (t *OnConflictTransform) buildGuardedMergeStatements(rewrite *onConflictMergeRewrite) ([]string, []string, error) { + if rewrite == nil || rewrite.merge == nil || rewrite.sourceSelect == nil || rewrite.targetRelation == nil { + return nil, nil, nil + } + if len(rewrite.conflictColumns) == 0 { + return nil, nil, nil + } + + sourceSQL, err := deparseStatement(&pg_query.Node{Node: &pg_query.Node_SelectStmt{SelectStmt: rewrite.sourceSelect}}) + if err != nil { + return nil, nil, err + } + sourceTableName := t.generateOnConflictSourceTableName(rewrite.targetRelation.Relname) + sourceRelation := quoteIdent(sourceTableName) + sourceFrom := sourceRelation + " AS duckgres_on_conflict_source" + + mergeSQL, err := deparseMergeWithTempSource(rewrite.merge, sourceTableName) + if err != nil { + return nil, nil, err + } + + statements := []string{ + "CREATE TEMPORARY TABLE " + sourceRelation + " AS " + sourceSQL, + t.buildSourceDuplicateGuard(sourceFrom, rewrite.conflictColumns), + t.buildTargetDuplicateGuard(sourceFrom, rewrite.targetRelation, rewrite.conflictColumns), + mergeSQL, + } + cleanup := []string{ + "DROP TABLE IF EXISTS " + sourceRelation, + } + return statements, cleanup, nil +} + +func deparseStatement(stmt *pg_query.Node) (string, error) { + sql, err := pg_query.Deparse(&pg_query.ParseResult{ + Stmts: []*pg_query.RawStmt{{Stmt: stmt}}, + }) + if err != nil { + return "", err + } + return strings.TrimSuffix(strings.TrimSpace(sql), ";"), nil +} + +func deparseMergeWithTempSource(merge *pg_query.MergeStmt, sourceTableName string) (string, error) { + originalSourceRelation := merge.SourceRelation + merge.SourceRelation = &pg_query.Node{ + Node: &pg_query.Node_RangeVar{ + RangeVar: &pg_query.RangeVar{ + Relname: sourceTableName, + Alias: &pg_query.Alias{Aliasname: "excluded"}, + }, + }, + } + sql, err := deparseStatement(&pg_query.Node{Node: &pg_query.Node_MergeStmt{MergeStmt: merge}}) + merge.SourceRelation = originalSourceRelation + return sql, err +} + +func (t *OnConflictTransform) buildSourceDuplicateGuard(sourceFrom string, conflictColumns []string) string { + sourceRefs := qualifiedColumnRefs("duckgres_on_conflict_source", conflictColumns) + notNullPredicates := notNullPredicates(sourceRefs) + + return "SELECT CASE WHEN EXISTS (" + + "SELECT 1 FROM " + sourceFrom + " " + + "WHERE " + notNullPredicates + " " + + "GROUP BY " + strings.Join(sourceRefs, ", ") + " " + + "HAVING count(*) > 1 LIMIT 1" + + ") THEN error(" + quoteStringLiteral("ON CONFLICT source rows contain duplicate conflict keys") + ") ELSE NULL END" +} + +func (t *OnConflictTransform) buildTargetDuplicateGuard(sourceFrom string, target *pg_query.RangeVar, conflictColumns []string) string { + sourceRefs := qualifiedColumnRefs("duckgres_on_conflict_source", conflictColumns) + keyRefs := qualifiedColumnRefs("duckgres_on_conflict_keys", conflictColumns) + targetRefs := qualifiedColumnRefs("duckgres_on_conflict_target", conflictColumns) + joinPredicates := make([]string, 0, len(conflictColumns)) + for i := range conflictColumns { + joinPredicates = append(joinPredicates, targetRefs[i]+" = "+keyRefs[i]) + } + + return "SELECT CASE WHEN EXISTS (" + + "SELECT 1 FROM " + formatRangeVar(target) + " AS duckgres_on_conflict_target " + + "JOIN (" + + "SELECT DISTINCT " + strings.Join(sourceRefs, ", ") + " " + + "FROM " + sourceFrom + " " + + "WHERE " + notNullPredicates(sourceRefs) + + ") AS duckgres_on_conflict_keys " + + "ON " + strings.Join(joinPredicates, " AND ") + " " + + "GROUP BY " + strings.Join(targetRefs, ", ") + " " + + "HAVING count(*) > 1 LIMIT 1" + + ") THEN error(" + quoteStringLiteral("ON CONFLICT target contains duplicate conflict keys") + ") ELSE NULL END" +} + +func (t *OnConflictTransform) generateOnConflictSourceTableName(targetName string) string { + safe := strings.Map(func(r rune) rune { + if (r >= 'a' && r <= 'z') || (r >= '0' && r <= '9') || r == '_' { + return r + } + if r >= 'A' && r <= 'Z' { + return r + 32 + } + return '_' + }, targetName) + if safe == "" { + safe = "source" + } + + prefix := "_duckgres_on_conflict_source_" + suffix := fmt.Sprintf("%08x", rand.Uint32()) + maxNameLen := maxIdentifierLength - len(prefix) - len(suffix) - 1 + if len(safe) > maxNameLen { + safe = safe[:maxNameLen] + } + return fmt.Sprintf("%s%s_%s", prefix, safe, suffix) +} + +func qualifiedColumnRefs(alias string, columns []string) []string { + refs := make([]string, len(columns)) + for i, col := range columns { + refs[i] = alias + "." + quoteIdent(col) + } + return refs +} + +func notNullPredicates(refs []string) string { + predicates := make([]string, len(refs)) + for i, ref := range refs { + predicates[i] = ref + " IS NOT NULL" + } + return strings.Join(predicates, " AND ") +} + +func formatRangeVar(rv *pg_query.RangeVar) string { + parts := make([]string, 0, 3) + if rv.Catalogname != "" { + parts = append(parts, quoteIdent(rv.Catalogname)) + } + if rv.Schemaname != "" { + parts = append(parts, quoteIdent(rv.Schemaname)) + } + parts = append(parts, quoteIdent(rv.Relname)) + return strings.Join(parts, ".") +} + +func quoteIdent(ident string) string { + return `"` + strings.ReplaceAll(ident, `"`, `""`) + `"` +} + +func quoteStringLiteral(s string) string { + return `'` + strings.ReplaceAll(s, `'`, `''`) + `'` } // buildSourceSelect creates a SELECT statement from VALUES for use as MERGE source diff --git a/transpiler/transpiler.go b/transpiler/transpiler.go index 45aaafcd..ef78dca1 100644 --- a/transpiler/transpiler.go +++ b/transpiler/transpiler.go @@ -271,11 +271,15 @@ func (t *Transpiler) transpileWithFlags(sql string, flags TransformFlags) (*Resu // When a transform produces multiple statements, we skip remaining transforms // and return the statements directly. if len(transformResult.Statements) > 0 { + paramCount := transformResult.ParamCount + if t.config.ConvertPlaceholders { + paramCount = countParametersRegex(sql) + } return &Result{ SQL: sql, // Keep original for logging Statements: restoreLongIdentifiersAll(transformResult.Statements, longIdents), CleanupStatements: restoreLongIdentifiersAll(transformResult.CleanupStatements, longIdents), - ParamCount: transformResult.ParamCount, + ParamCount: paramCount, Warnings: transformResult.Warnings, }, nil } diff --git a/transpiler/transpiler_test.go b/transpiler/transpiler_test.go index b2c652b2..71c49f20 100644 --- a/transpiler/transpiler_test.go +++ b/transpiler/transpiler_test.go @@ -1,11 +1,13 @@ package transpiler import ( + "database/sql" "errors" "fmt" "strings" "testing" + _ "github.com/duckdb/duckdb-go/v2" "github.com/posthog/duckgres/transpiler/transform" ) @@ -2178,8 +2180,10 @@ func TestTranspile_OnConflict_DuckLakeMode(t *testing.T) { input: "INSERT INTO users (id, name) VALUES (1, 'test') ON CONFLICT (id) DO NOTHING", wantMerge: true, contains: []string{ + "CREATE TEMPORARY TABLE", + "AS SELECT 1 AS id, 'test' AS name", "MERGE INTO users", - "USING (SELECT 1 AS id, 'test' AS name) excluded", + "excluded", "ON excluded.id = users.id", "WHEN NOT MATCHED THEN INSERT", }, @@ -2193,8 +2197,10 @@ func TestTranspile_OnConflict_DuckLakeMode(t *testing.T) { input: "INSERT INTO users (id, name) VALUES (1, 'test') ON CONFLICT (id) DO UPDATE SET name = EXCLUDED.name", wantMerge: true, contains: []string{ + "CREATE TEMPORARY TABLE", + "AS SELECT 1 AS id, 'test' AS name", "MERGE INTO users", - "USING (SELECT 1 AS id, 'test' AS name) excluded", + "excluded", "ON excluded.id = users.id", "WHEN MATCHED THEN UPDATE SET name = excluded.name", "WHEN NOT MATCHED THEN INSERT", @@ -2243,6 +2249,11 @@ func TestTranspile_OnConflict_DuckLakeMode(t *testing.T) { } sql := result.SQL + finalSQL := result.SQL + if len(result.Statements) > 0 { + sql = strings.Join(result.Statements, "\n") + finalSQL = result.Statements[len(result.Statements)-1] + } // Check MERGE statement is generated if tt.wantMerge { @@ -2260,14 +2271,202 @@ func TestTranspile_OnConflict_DuckLakeMode(t *testing.T) { // Check contents that should not be present for _, notWant := range tt.notContains { - if strings.Contains(strings.ToUpper(sql), strings.ToUpper(notWant)) { - t.Errorf("Transpile(%q) = %q, should NOT contain %q", tt.input, sql, notWant) + if strings.Contains(strings.ToUpper(finalSQL), strings.ToUpper(notWant)) { + t.Errorf("Transpile(%q) final SQL = %q, should NOT contain %q", tt.input, finalSQL, notWant) } } }) } } +func TestTranspile_OnConflict_DuckLakeMode_EmitsCardinalityGuards(t *testing.T) { + tr := New(Config{DuckLakeMode: true}) + + result, err := tr.Transpile(`INSERT INTO "stripe"."price" ("id", "nickname") + SELECT "id", "nickname" FROM "stripe_price-staging-test" + ON CONFLICT ("id") DO UPDATE SET "nickname" = "excluded"."nickname"`) + if err != nil { + t.Fatalf("Transpile error: %v", err) + } + + if len(result.Statements) != 4 { + t.Fatalf("got %d statements, want source materialization, source guard, target guard, merge: %v", len(result.Statements), result.Statements) + } + if len(result.CleanupStatements) != 1 || !strings.Contains(result.CleanupStatements[0], "DROP TABLE IF EXISTS") { + t.Fatalf("cleanup = %v, want DROP TABLE for materialized source", result.CleanupStatements) + } + + checks := []struct { + stmt int + want []string + }{ + { + stmt: 0, + want: []string{ + "CREATE TEMPORARY TABLE", + "AS SELECT id, nickname FROM", + "stripe_price-staging-test", + }, + }, + { + stmt: 1, + want: []string{ + "duckgres_on_conflict_source", + "GROUP BY duckgres_on_conflict_source.\"id\"", + "HAVING count(*) > 1", + "ON CONFLICT source rows contain duplicate conflict keys", + }, + }, + { + stmt: 2, + want: []string{ + "FROM \"stripe\".\"price\" AS duckgres_on_conflict_target", + "SELECT DISTINCT duckgres_on_conflict_source.\"id\"", + "ON duckgres_on_conflict_target.\"id\" = duckgres_on_conflict_keys.\"id\"", + "ON CONFLICT target contains duplicate conflict keys", + }, + }, + { + stmt: 3, + want: []string{ + "MERGE INTO stripe.price", + "USING", + "excluded", + "WHEN MATCHED THEN UPDATE", + "WHEN NOT MATCHED THEN INSERT", + }, + }, + } + + for _, check := range checks { + for _, want := range check.want { + if !strings.Contains(result.Statements[check.stmt], want) { + t.Errorf("statement %d should contain %q\nGot: %s", check.stmt, want, result.Statements[check.stmt]) + } + } + } +} + +func TestTranspile_OnConflict_DuckLakeMode_GuardsPreserveParamCount(t *testing.T) { + tr := New(Config{DuckLakeMode: true, ConvertPlaceholders: true}) + + result, err := tr.Transpile(`INSERT INTO users (id, nickname) + VALUES ($1, $2) + ON CONFLICT (id) DO UPDATE SET nickname = excluded.nickname`) + if err != nil { + t.Fatalf("Transpile error: %v", err) + } + if len(result.Statements) != 4 { + t.Fatalf("expected guarded multi-statement rewrite, got statements=%v sql=%q", result.Statements, result.SQL) + } + if result.ParamCount != 2 { + t.Fatalf("ParamCount = %d, want 2", result.ParamCount) + } +} + +func TestTranspile_OnConflict_DuckLakeMode_GuardsPreventDuplicateMerge(t *testing.T) { + tr := New(Config{DuckLakeMode: true}) + db, err := sql.Open("duckdb", ":memory:") + if err != nil { + t.Fatalf("open duckdb: %v", err) + } + defer func() { _ = db.Close() }() + + if _, err := db.Exec(`CREATE TABLE target (id VARCHAR, nickname VARCHAR)`); err != nil { + t.Fatalf("create target: %v", err) + } + if _, err := db.Exec(`CREATE TABLE staging (id VARCHAR, nickname VARCHAR)`); err != nil { + t.Fatalf("create staging: %v", err) + } + + runStatements := func(t *testing.T, statements []string, cleanup []string) error { + t.Helper() + for _, stmt := range statements { + if _, err := db.Exec(stmt); err != nil { + for _, cleanupStmt := range cleanup { + _, _ = db.Exec(cleanupStmt) + } + return err + } + } + for _, cleanupStmt := range cleanup { + if _, err := db.Exec(cleanupStmt); err != nil { + return err + } + } + return nil + } + + result, err := tr.Transpile(`INSERT INTO target (id, nickname) + SELECT id, nickname FROM staging + ON CONFLICT (id) DO UPDATE SET nickname = excluded.nickname`) + if err != nil { + t.Fatalf("Transpile source duplicate query: %v", err) + } + if len(result.Statements) == 0 { + t.Fatalf("expected guarded multi-statement rewrite, got single SQL: %s", result.SQL) + } + + if _, err := db.Exec(`INSERT INTO staging VALUES ('price_1', 'first'), ('price_1', 'second')`); err != nil { + t.Fatalf("seed duplicate staging rows: %v", err) + } + err = runStatements(t, result.Statements, result.CleanupStatements) + if err == nil || !strings.Contains(err.Error(), "ON CONFLICT source rows contain duplicate conflict keys") { + t.Fatalf("expected source duplicate guard error, got %v", err) + } + + var rows int + if err := db.QueryRow(`SELECT count(*) FROM target`).Scan(&rows); err != nil { + t.Fatalf("count target after source guard: %v", err) + } + if rows != 0 { + t.Fatalf("source guard should stop before MERGE inserts rows, got %d target rows", rows) + } + + if _, err := db.Exec(`DELETE FROM staging`); err != nil { + t.Fatalf("clear staging: %v", err) + } + if _, err := db.Exec(`INSERT INTO target VALUES ('price_1', 'old a'), ('price_1', 'old b')`); err != nil { + t.Fatalf("seed duplicate target rows: %v", err) + } + if _, err := db.Exec(`INSERT INTO staging VALUES ('price_1', 'new')`); err != nil { + t.Fatalf("seed staging row: %v", err) + } + + err = runStatements(t, result.Statements, result.CleanupStatements) + if err == nil || !strings.Contains(err.Error(), "ON CONFLICT target contains duplicate conflict keys") { + t.Fatalf("expected target duplicate guard error, got %v", err) + } + + var updatedRows int + if err := db.QueryRow(`SELECT count(*) FROM target WHERE nickname = 'new'`).Scan(&updatedRows); err != nil { + t.Fatalf("count updated target rows: %v", err) + } + if updatedRows != 0 { + t.Fatalf("target guard should stop before fanout update, got %d updated rows", updatedRows) + } + + if _, err := db.Exec(`DELETE FROM target`); err != nil { + t.Fatalf("clear target before success case: %v", err) + } + if _, err := db.Exec(`DELETE FROM staging`); err != nil { + t.Fatalf("clear staging before success case: %v", err) + } + if _, err := db.Exec(`INSERT INTO staging VALUES ('price_2', 'valid')`); err != nil { + t.Fatalf("seed valid staging row: %v", err) + } + + if err := runStatements(t, result.Statements, result.CleanupStatements); err != nil { + t.Fatalf("guarded merge should allow unique source/target keys: %v", err) + } + if err := db.QueryRow(`SELECT count(*) FROM target WHERE id = 'price_2' AND nickname = 'valid'`).Scan(&rows); err != nil { + t.Fatalf("count inserted target rows: %v", err) + } + if rows != 1 { + t.Fatalf("guarded merge should insert the unique staging row, got %d rows", rows) + } +} + func TestTranspile_DropCascade_DuckLakeMode(t *testing.T) { // In DuckLake mode, CASCADE should be stripped from DROP TABLE/VIEW statements // See: https://github.com/duckdb/dbt-duckdb/pull/557 From ffbed9ea08e0b5024785bc545ea66fff12968660 Mon Sep 17 00:00:00 2001 From: Bill Guowei Yang Date: Thu, 2 Jul 2026 10:08:45 -0400 Subject: [PATCH 2/6] Slice bind args for multi-statement rewrites --- server/conn.go | 11 ++++++++++ server/conn_query_exec.go | 9 +++++--- server/conn_test.go | 46 +++++++++++++++++++++++++++++++++++++++ 3 files changed, 63 insertions(+), 3 deletions(-) diff --git a/server/conn.go b/server/conn.go index 7a7578f9..e61a6687 100644 --- a/server/conn.go +++ b/server/conn.go @@ -1590,6 +1590,17 @@ func countDollarParams(query string) int { return max } +func argsForStatement(query string, args []interface{}) []interface{} { + paramCount := countDollarParams(query) + if paramCount == 0 { + return nil + } + if paramCount >= len(args) { + return args + } + return args[:paramCount] +} + // isEmptyQuery and stripLeadingComments moved to server/sqlcore so the // Flight client can call them without importing server. Local thin wrappers // preserve the unexported call-site spellings used throughout this file. diff --git a/server/conn_query_exec.go b/server/conn_query_exec.go index 213964a4..718b4c7e 100644 --- a/server/conn_query_exec.go +++ b/server/conn_query_exec.go @@ -1002,7 +1002,8 @@ func (c *clientConn) executeMultiStatementExtended(statements []string, cleanup c.logger().Debug("Multi-stmt-ext setup.", "step", i+1, "total", len(statements)-1, "stmt", stmt) setupStart := time.Now() c.logQueryStarted(stmt) - result, err := c.executor.Exec(stmt, args...) + stmtArgs := argsForStatement(stmt, args) + result, err := c.executor.Exec(stmt, stmtArgs...) var setupRows int64 if result != nil { setupRows, _ = result.RowsAffected() @@ -1034,7 +1035,8 @@ func (c *clientConn) executeMultiStatementExtended(statements []string, cleanup if queryReturnsResults(finalStmt) { // Result-returning query: obtain cursor FIRST, cleanup SECOND, stream THIRD - rows, err := c.executor.Query(finalStmt, args...) + finalArgs := argsForStatement(finalStmt, args) + rows, err := c.executor.Query(finalStmt, finalArgs...) if err != nil { finalErr = err c.logger().Error("Multi-stmt-ext final query error.", "query", finalStmt, "error", err) @@ -1056,7 +1058,8 @@ func (c *clientConn) executeMultiStatementExtended(statements []string, cleanup } else { // Non-result query (DML without RETURNING, DDL, etc.): execute then cleanup - result, err := c.executor.Exec(finalStmt, args...) + finalArgs := argsForStatement(finalStmt, args) + result, err := c.executor.Exec(finalStmt, finalArgs...) if err != nil { finalErr = err c.logger().Error("Multi-stmt-ext final exec error.", "query", finalStmt, "error", err) diff --git a/server/conn_test.go b/server/conn_test.go index e2b16943..26113e89 100644 --- a/server/conn_test.go +++ b/server/conn_test.go @@ -12,6 +12,7 @@ import ( "io" "math" "net" + "reflect" "slices" "strings" "testing" @@ -3860,6 +3861,51 @@ func TestCountDollarParams(t *testing.T) { } } +func TestArgsForStatement(t *testing.T) { + args := []interface{}{"id_1", "nickname", "unused"} + + tests := []struct { + name string + query string + want []interface{} + }{ + { + name: "no placeholders", + query: "SELECT 1", + want: nil, + }, + { + name: "uses first placeholder", + query: "SELECT $1", + want: args[:1], + }, + { + name: "uses first two placeholders", + query: "INSERT INTO t VALUES ($1, $2)", + want: args[:2], + }, + { + name: "highest referenced placeholder controls argument count", + query: "SELECT $3", + want: args, + }, + { + name: "placeholder count beyond provided args does not panic", + query: "SELECT $4", + want: args, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := argsForStatement(tt.query, args) + if !reflect.DeepEqual(got, tt.want) { + t.Fatalf("argsForStatement(%q) = %#v, want %#v", tt.query, got, tt.want) + } + }) + } +} + func TestMatchPgStatActivityQuery(t *testing.T) { tests := []struct { name string From 91adac824060a2e6e09cebd39434adb541e2cc3c Mon Sep 17 00:00:00 2001 From: Bill Guowei Yang Date: Thu, 2 Jul 2026 10:19:26 -0400 Subject: [PATCH 3/6] Tighten guarded ON CONFLICT rewrites --- server/conn.go | 105 ++++++++++++++-- server/conn_query_exec.go | 31 +++-- server/conn_test.go | 55 +++++++-- transpiler/transform/onconflict.go | 91 ++++++++++++-- transpiler/transpiler_test.go | 190 +++++++++++++++++++++++++++-- 5 files changed, 421 insertions(+), 51 deletions(-) diff --git a/server/conn.go b/server/conn.go index e61a6687..fbd4e79d 100644 --- a/server/conn.go +++ b/server/conn.go @@ -1590,15 +1590,104 @@ func countDollarParams(query string) int { return max } -func argsForStatement(query string, args []interface{}) []interface{} { - paramCount := countDollarParams(query) - if paramCount == 0 { - return nil - } - if paramCount >= len(args) { - return args +func queryAndArgsForStatement(query string, args []interface{}) (string, []interface{}) { + var out strings.Builder + var stmtArgs []interface{} + inString := false + inIdent := false + inLineComment := false + inBlockComment := false + + for i := 0; i < len(query); i++ { + c := query[i] + + if inBlockComment { + out.WriteByte(c) + if c == '*' && i+1 < len(query) && query[i+1] == '/' { + out.WriteByte(query[i+1]) + inBlockComment = false + i++ + } + continue + } + + if inLineComment { + out.WriteByte(c) + if c == '\n' { + inLineComment = false + } + continue + } + + if inString { + out.WriteByte(c) + if c == '\'' { + if i+1 < len(query) && query[i+1] == '\'' { + out.WriteByte(query[i+1]) + i++ + } else { + inString = false + } + } + continue + } + + if inIdent { + out.WriteByte(c) + if c == '"' { + if i+1 < len(query) && query[i+1] == '"' { + out.WriteByte(query[i+1]) + i++ + } else { + inIdent = false + } + } + continue + } + + if c == '-' && i+1 < len(query) && query[i+1] == '-' { + out.WriteByte(c) + out.WriteByte(query[i+1]) + inLineComment = true + i++ + continue + } + if c == '/' && i+1 < len(query) && query[i+1] == '*' { + out.WriteByte(c) + out.WriteByte(query[i+1]) + inBlockComment = true + i++ + continue + } + if c == '\'' { + out.WriteByte(c) + inString = true + continue + } + if c == '"' { + out.WriteByte(c) + inIdent = true + continue + } + + if c == '$' && i+1 < len(query) && query[i+1] >= '1' && query[i+1] <= '9' { + j := i + 1 + for j < len(query) && query[j] >= '0' && query[j] <= '9' { + j++ + } + n, err := strconv.Atoi(query[i+1 : j]) + if err == nil && n >= 1 && n <= len(args) { + out.WriteByte('?') + stmtArgs = append(stmtArgs, args[n-1]) + i = j - 1 + continue + } + } + + out.WriteByte(c) } - return args[:paramCount] + + return out.String(), stmtArgs } // isEmptyQuery and stripLeadingComments moved to server/sqlcore so the diff --git a/server/conn_query_exec.go b/server/conn_query_exec.go index 718b4c7e..c23d2881 100644 --- a/server/conn_query_exec.go +++ b/server/conn_query_exec.go @@ -999,18 +999,18 @@ func (c *clientConn) executeMultiStatementExtended(statements []string, cleanup // setup work is observable per statement, not just per outer query. for i := 0; i < len(statements)-1; i++ { stmt := statements[i] - c.logger().Debug("Multi-stmt-ext setup.", "step", i+1, "total", len(statements)-1, "stmt", stmt) + stmtQuery, stmtArgs := queryAndArgsForStatement(stmt, args) + c.logger().Debug("Multi-stmt-ext setup.", "step", i+1, "total", len(statements)-1, "stmt", stmtQuery) setupStart := time.Now() - c.logQueryStarted(stmt) - stmtArgs := argsForStatement(stmt, args) - result, err := c.executor.Exec(stmt, stmtArgs...) + c.logQueryStarted(stmtQuery) + result, err := c.executor.Exec(stmtQuery, stmtArgs...) var setupRows int64 if result != nil { setupRows, _ = result.RowsAffected() } - c.logQueryFinished(stmt, setupStart, setupRows, err) + c.logQueryFinished(stmtQuery, setupStart, setupRows, err) if err != nil { - c.logger().Error("Multi-stmt-ext setup error.", "query", stmt, "error", err) + c.logger().Error("Multi-stmt-ext setup error.", "query", stmtQuery, "error", err) c.setTxError() // On error, still try to cleanup (best effort) c.executeCleanup(cleanup) @@ -1021,25 +1021,25 @@ func (c *clientConn) executeMultiStatementExtended(statements []string, cleanup // Handle final statement finalStmt := statements[len(statements)-1] + finalQuery, finalArgs := queryAndArgsForStatement(finalStmt, args) upperFinal := strings.ToUpper(strings.TrimSpace(finalStmt)) cmdType := c.getCommandType(upperFinal) - c.logger().Debug("Multi-stmt-ext final.", "stmt", finalStmt, "cmd_type", cmdType) + c.logger().Debug("Multi-stmt-ext final.", "stmt", finalQuery, "cmd_type", cmdType) finalStart := time.Now() var finalRowsAff int64 var finalErr error - c.logQueryStarted(finalStmt) + c.logQueryStarted(finalQuery) defer func() { - c.logQueryFinished(finalStmt, finalStart, finalRowsAff, finalErr) + c.logQueryFinished(finalQuery, finalStart, finalRowsAff, finalErr) }() if queryReturnsResults(finalStmt) { // Result-returning query: obtain cursor FIRST, cleanup SECOND, stream THIRD - finalArgs := argsForStatement(finalStmt, args) - rows, err := c.executor.Query(finalStmt, finalArgs...) + rows, err := c.executor.Query(finalQuery, finalArgs...) if err != nil { finalErr = err - c.logger().Error("Multi-stmt-ext final query error.", "query", finalStmt, "error", err) + c.logger().Error("Multi-stmt-ext final query error.", "query", finalQuery, "error", err) c.setTxError() c.executeCleanup(cleanup) c.sendError("ERROR", "42000", err.Error()) @@ -1054,15 +1054,14 @@ func (c *clientConn) executeMultiStatementExtended(statements []string, cleanup // is tracked by streamRowsToClientExtended; the deferred Finished // log uses 0 as an approximation (logQuery still records the // precise count via the structured channel). - c.streamRowsToClientExtended(rows, cmdType, resultFormats, described, finalStmt) + c.streamRowsToClientExtended(rows, cmdType, resultFormats, described, finalQuery) } else { // Non-result query (DML without RETURNING, DDL, etc.): execute then cleanup - finalArgs := argsForStatement(finalStmt, args) - result, err := c.executor.Exec(finalStmt, finalArgs...) + result, err := c.executor.Exec(finalQuery, finalArgs...) if err != nil { finalErr = err - c.logger().Error("Multi-stmt-ext final exec error.", "query", finalStmt, "error", err) + c.logger().Error("Multi-stmt-ext final exec error.", "query", finalQuery, "error", err) c.setTxError() c.executeCleanup(cleanup) c.sendError("ERROR", "42000", err.Error()) diff --git a/server/conn_test.go b/server/conn_test.go index 26113e89..6a38146a 100644 --- a/server/conn_test.go +++ b/server/conn_test.go @@ -3861,51 +3861,90 @@ func TestCountDollarParams(t *testing.T) { } } -func TestArgsForStatement(t *testing.T) { +func TestQueryAndArgsForStatement(t *testing.T) { args := []interface{}{"id_1", "nickname", "unused"} tests := []struct { name string query string + wantQ string want []interface{} }{ { name: "no placeholders", query: "SELECT 1", + wantQ: "SELECT 1", want: nil, }, { name: "uses first placeholder", query: "SELECT $1", + wantQ: "SELECT ?", want: args[:1], }, { name: "uses first two placeholders", query: "INSERT INTO t VALUES ($1, $2)", + wantQ: "INSERT INTO t VALUES (?, ?)", want: args[:2], }, { - name: "highest referenced placeholder controls argument count", - query: "SELECT $3", - want: args, + name: "sparse placeholders preserve referenced argument numbers", + query: "SELECT $3, $1", + wantQ: "SELECT ?, ?", + want: []interface{}{args[2], args[0]}, }, { - name: "placeholder count beyond provided args does not panic", + name: "repeated placeholders duplicate the referenced argument", + query: "SELECT $2, $2", + wantQ: "SELECT ?, ?", + want: []interface{}{args[1], args[1]}, + }, + { + name: "ignores placeholders in literals comments and identifiers", + query: "SELECT '$1', \"$2\", $3 -- $1\n/* $2 */", + wantQ: "SELECT '$1', \"$2\", ? -- $1\n/* $2 */", + want: []interface{}{args[2]}, + }, + { + name: "placeholder beyond provided args is left for downstream validation", query: "SELECT $4", - want: args, + wantQ: "SELECT $4", + want: nil, }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - got := argsForStatement(tt.query, args) + gotQ, got := queryAndArgsForStatement(tt.query, args) + if gotQ != tt.wantQ { + t.Fatalf("queryAndArgsForStatement(%q) query = %q, want %q", tt.query, gotQ, tt.wantQ) + } if !reflect.DeepEqual(got, tt.want) { - t.Fatalf("argsForStatement(%q) = %#v, want %#v", tt.query, got, tt.want) + t.Fatalf("queryAndArgsForStatement(%q) args = %#v, want %#v", tt.query, got, tt.want) } }) } } +func TestQueryAndArgsForStatementRunsSparseDollarParams(t *testing.T) { + db, err := sql.Open("duckdb", ":memory:") + if err != nil { + t.Fatalf("open duckdb: %v", err) + } + defer func() { _ = db.Close() }() + + query, args := queryAndArgsForStatement("SELECT $3 || '-' || $1", []interface{}{"first", "second", "third"}) + + var got string + if err := db.QueryRow(query, args...).Scan(&got); err != nil { + t.Fatalf("run rewritten sparse-placeholder query %q args=%v: %v", query, args, err) + } + if got != "third-first" { + t.Fatalf("rewritten sparse-placeholder query returned %q, want third-first", got) + } +} + func TestMatchPgStatActivityQuery(t *testing.T) { tests := []struct { name string diff --git a/transpiler/transform/onconflict.go b/transpiler/transform/onconflict.go index 34f04a2a..8f67ee39 100644 --- a/transpiler/transform/onconflict.go +++ b/transpiler/transform/onconflict.go @@ -30,7 +30,9 @@ type onConflictMergeRewrite struct { merge *pg_query.MergeStmt sourceSelect *pg_query.SelectStmt targetRelation *pg_query.RangeVar + insertColumns []string conflictColumns []string + action pg_query.OnConflictAction } func NewOnConflictTransform() *OnConflictTransform { @@ -58,12 +60,11 @@ func (t *OnConflictTransform) Transform(tree *pg_query.ParseResult, result *Resu // CONSTRAINT cannot be honored: there is no named constraint to // infer the conflict target from. Reject it with a clean PostgreSQL // error rather than letting it fail opaquely at DuckDB. - if t.DuckLakeMode && insert.OnConflictClause != nil && - insert.OnConflictClause.Infer != nil && - strings.TrimSpace(insert.OnConflictClause.Infer.Conname) != "" { - result.Error = NewFeatureNotSupported( - "ON CONFLICT ON CONSTRAINT is not supported: this catalog does not enforce named constraints") - return false, nil + if t.DuckLakeMode && insert.OnConflictClause != nil { + if unsupported := t.validateDuckLakeConflictTarget(insert.OnConflictClause); unsupported != nil { + result.Error = unsupported + return false, nil + } } if rewrite := t.transformInsertToMerge(insert); rewrite != nil { @@ -88,6 +89,28 @@ func (t *OnConflictTransform) Transform(tree *pg_query.ParseResult, result *Resu return changed, nil } +func (t *OnConflictTransform) validateDuckLakeConflictTarget(occ *pg_query.OnConflictClause) error { + if occ == nil || occ.Infer == nil { + return nil + } + if strings.TrimSpace(occ.Infer.Conname) != "" { + return NewFeatureNotSupported( + "ON CONFLICT ON CONSTRAINT is not supported: this catalog does not enforce named constraints") + } + if occ.Infer.WhereClause != nil { + return NewFeatureNotSupported( + "ON CONFLICT partial index predicates are not supported: this catalog does not enforce partial unique indexes") + } + for _, elem := range occ.Infer.IndexElems { + indexElem := elem.GetIndexElem() + if indexElem == nil || strings.TrimSpace(indexElem.Name) == "" { + return NewFeatureNotSupported( + "ON CONFLICT expression targets are not supported: use column names") + } + } + return nil +} + // transformInsertToMerge converts INSERT ... ON CONFLICT to MERGE for DuckLake mode func (t *OnConflictTransform) transformInsertToMerge(insert *pg_query.InsertStmt) *onConflictMergeRewrite { if !t.DuckLakeMode { @@ -179,7 +202,9 @@ func (t *OnConflictTransform) transformInsertToMerge(insert *pg_query.InsertStmt merge: merge, sourceSelect: sourceSelect, targetRelation: insert.Relation, + insertColumns: colNames, conflictColumns: conflictColumns, + action: occ.Action, } } @@ -217,17 +242,53 @@ func (t *OnConflictTransform) buildGuardedMergeStatements(rewrite *onConflictMer } statements := []string{ - "CREATE TEMPORARY TABLE " + sourceRelation + " AS " + sourceSQL, - t.buildSourceDuplicateGuard(sourceFrom, rewrite.conflictColumns), + "BEGIN", + t.buildCreateSourceTableStatement(sourceRelation, sourceSQL, rewrite), + } + if rewrite.action == pg_query.OnConflictAction_ONCONFLICT_UPDATE { + statements = append(statements, t.buildSourceDuplicateGuard(sourceFrom, rewrite.conflictColumns)) + } + statements = append(statements, t.buildTargetDuplicateGuard(sourceFrom, rewrite.targetRelation, rewrite.conflictColumns), mergeSQL, - } + ) cleanup := []string{ "DROP TABLE IF EXISTS " + sourceRelation, + "COMMIT", } return statements, cleanup, nil } +func (t *OnConflictTransform) buildCreateSourceTableStatement(sourceRelation string, sourceSQL string, rewrite *onConflictMergeRewrite) string { + columnList := quotedColumnList(rewrite.insertColumns) + if rewrite.action != pg_query.OnConflictAction_ONCONFLICT_NOTHING { + return "CREATE TEMPORARY TABLE " + sourceRelation + " " + columnList + " AS " + sourceSQL + } + + return "CREATE TEMPORARY TABLE " + sourceRelation + " " + columnList + " AS " + + t.buildDoNothingSourceSQL(sourceSQL, rewrite.insertColumns, rewrite.conflictColumns) +} + +func (t *OnConflictTransform) buildDoNothingSourceSQL(sourceSQL string, insertColumns []string, conflictColumns []string) string { + rawAlias := "duckgres_on_conflict_raw" + rankedAlias := "duckgres_on_conflict_ranked" + rankColumn := "_duckgres_on_conflict_rank" + + insertRefs := qualifiedColumnRefs(rankedAlias, insertColumns) + conflictRefs := qualifiedColumnRefs(rawAlias, conflictColumns) + preservePredicates := make([]string, 0, len(conflictColumns)+1) + for _, ref := range qualifiedColumnRefs(rankedAlias, conflictColumns) { + preservePredicates = append(preservePredicates, ref+" IS NULL") + } + preservePredicates = append(preservePredicates, rankedAlias+"."+quoteIdent(rankColumn)+" = 1") + + return "SELECT " + strings.Join(insertRefs, ", ") + " FROM (" + + "SELECT " + rawAlias + ".*, row_number() OVER (PARTITION BY " + strings.Join(conflictRefs, ", ") + ") AS " + quoteIdent(rankColumn) + " " + + "FROM (" + sourceSQL + ") AS " + rawAlias + " " + quotedColumnList(insertColumns) + + ") AS " + rankedAlias + " " + + "WHERE " + strings.Join(preservePredicates, " OR ") +} + func deparseStatement(stmt *pg_query.Node) (string, error) { sql, err := pg_query.Deparse(&pg_query.ParseResult{ Stmts: []*pg_query.RawStmt{{Stmt: stmt}}, @@ -240,6 +301,9 @@ func deparseStatement(stmt *pg_query.Node) (string, error) { func deparseMergeWithTempSource(merge *pg_query.MergeStmt, sourceTableName string) (string, error) { originalSourceRelation := merge.SourceRelation + defer func() { + merge.SourceRelation = originalSourceRelation + }() merge.SourceRelation = &pg_query.Node{ Node: &pg_query.Node_RangeVar{ RangeVar: &pg_query.RangeVar{ @@ -249,7 +313,6 @@ func deparseMergeWithTempSource(merge *pg_query.MergeStmt, sourceTableName strin }, } sql, err := deparseStatement(&pg_query.Node{Node: &pg_query.Node_MergeStmt{MergeStmt: merge}}) - merge.SourceRelation = originalSourceRelation return sql, err } @@ -318,6 +381,14 @@ func qualifiedColumnRefs(alias string, columns []string) []string { return refs } +func quotedColumnList(columns []string) string { + quoted := make([]string, len(columns)) + for i, col := range columns { + quoted[i] = quoteIdent(col) + } + return "(" + strings.Join(quoted, ", ") + ")" +} + func notNullPredicates(refs []string) string { predicates := make([]string, len(refs)) for i, ref := range refs { diff --git a/transpiler/transpiler_test.go b/transpiler/transpiler_test.go index 71c49f20..e64b7dac 100644 --- a/transpiler/transpiler_test.go +++ b/transpiler/transpiler_test.go @@ -440,6 +440,24 @@ func TestTranspile_FeatureNotSupported_Rejected(t *testing.T) { assertFeatureNotSupported(t, result) }) + t.Run(string(backend)+"/ON CONFLICT expression target", func(t *testing.T) { + result, err := tr.Transpile( + "INSERT INTO t (email, v) VALUES ('a@example.com', 2) ON CONFLICT ((lower(email))) DO UPDATE SET v = EXCLUDED.v") + if err != nil { + t.Fatalf("Transpile error: %v", err) + } + assertFeatureNotSupported(t, result) + }) + + t.Run(string(backend)+"/ON CONFLICT partial target", func(t *testing.T) { + result, err := tr.Transpile( + "INSERT INTO t (id, active, v) VALUES (1, true, 2) ON CONFLICT (id) WHERE active DO UPDATE SET v = EXCLUDED.v") + if err != nil { + t.Fatalf("Transpile error: %v", err) + } + assertFeatureNotSupported(t, result) + }) + t.Run(string(backend)+"/ALTER COLUMN TYPE USING", func(t *testing.T) { result, err := tr.Transpile( "ALTER TABLE t ALTER COLUMN c TYPE integer USING c::integer") @@ -2181,7 +2199,8 @@ func TestTranspile_OnConflict_DuckLakeMode(t *testing.T) { wantMerge: true, contains: []string{ "CREATE TEMPORARY TABLE", - "AS SELECT 1 AS id, 'test' AS name", + "SELECT 1 AS id, 'test' AS name", + "row_number() OVER", "MERGE INTO users", "excluded", "ON excluded.id = users.id", @@ -2289,11 +2308,13 @@ func TestTranspile_OnConflict_DuckLakeMode_EmitsCardinalityGuards(t *testing.T) t.Fatalf("Transpile error: %v", err) } - if len(result.Statements) != 4 { - t.Fatalf("got %d statements, want source materialization, source guard, target guard, merge: %v", len(result.Statements), result.Statements) + if len(result.Statements) != 5 { + t.Fatalf("got %d statements, want transaction, source materialization, source guard, target guard, merge: %v", len(result.Statements), result.Statements) } - if len(result.CleanupStatements) != 1 || !strings.Contains(result.CleanupStatements[0], "DROP TABLE IF EXISTS") { - t.Fatalf("cleanup = %v, want DROP TABLE for materialized source", result.CleanupStatements) + if len(result.CleanupStatements) != 2 || + !strings.Contains(result.CleanupStatements[0], "DROP TABLE IF EXISTS") || + result.CleanupStatements[1] != "COMMIT" { + t.Fatalf("cleanup = %v, want DROP TABLE then COMMIT for materialized source", result.CleanupStatements) } checks := []struct { @@ -2302,6 +2323,10 @@ func TestTranspile_OnConflict_DuckLakeMode_EmitsCardinalityGuards(t *testing.T) }{ { stmt: 0, + want: []string{"BEGIN"}, + }, + { + stmt: 1, want: []string{ "CREATE TEMPORARY TABLE", "AS SELECT id, nickname FROM", @@ -2309,7 +2334,7 @@ func TestTranspile_OnConflict_DuckLakeMode_EmitsCardinalityGuards(t *testing.T) }, }, { - stmt: 1, + stmt: 2, want: []string{ "duckgres_on_conflict_source", "GROUP BY duckgres_on_conflict_source.\"id\"", @@ -2318,7 +2343,7 @@ func TestTranspile_OnConflict_DuckLakeMode_EmitsCardinalityGuards(t *testing.T) }, }, { - stmt: 2, + stmt: 3, want: []string{ "FROM \"stripe\".\"price\" AS duckgres_on_conflict_target", "SELECT DISTINCT duckgres_on_conflict_source.\"id\"", @@ -2327,7 +2352,7 @@ func TestTranspile_OnConflict_DuckLakeMode_EmitsCardinalityGuards(t *testing.T) }, }, { - stmt: 3, + stmt: 4, want: []string{ "MERGE INTO stripe.price", "USING", @@ -2356,7 +2381,7 @@ func TestTranspile_OnConflict_DuckLakeMode_GuardsPreserveParamCount(t *testing.T if err != nil { t.Fatalf("Transpile error: %v", err) } - if len(result.Statements) != 4 { + if len(result.Statements) != 5 { t.Fatalf("expected guarded multi-statement rewrite, got statements=%v sql=%q", result.Statements, result.SQL) } if result.ParamCount != 2 { @@ -2364,6 +2389,23 @@ func TestTranspile_OnConflict_DuckLakeMode_GuardsPreserveParamCount(t *testing.T } } +func runDuckDBStatements(db *sql.DB, statements []string, cleanup []string) error { + for _, stmt := range statements { + if _, err := db.Exec(stmt); err != nil { + for _, cleanupStmt := range cleanup { + _, _ = db.Exec(cleanupStmt) + } + return err + } + } + for _, cleanupStmt := range cleanup { + if _, err := db.Exec(cleanupStmt); err != nil { + return err + } + } + return nil +} + func TestTranspile_OnConflict_DuckLakeMode_GuardsPreventDuplicateMerge(t *testing.T) { tr := New(Config{DuckLakeMode: true}) db, err := sql.Open("duckdb", ":memory:") @@ -2467,6 +2509,136 @@ func TestTranspile_OnConflict_DuckLakeMode_GuardsPreventDuplicateMerge(t *testin } } +func TestTranspile_OnConflict_DuckLakeMode_InsertSelectAliasesSourceColumns(t *testing.T) { + tr := New(Config{DuckLakeMode: true}) + db, err := sql.Open("duckdb", ":memory:") + if err != nil { + t.Fatalf("open duckdb: %v", err) + } + defer func() { _ = db.Close() }() + + if _, err := db.Exec(`CREATE TABLE target (id VARCHAR, nickname VARCHAR)`); err != nil { + t.Fatalf("create target: %v", err) + } + if _, err := db.Exec(`CREATE TABLE staging (src_id VARCHAR, src_nickname VARCHAR)`); err != nil { + t.Fatalf("create staging: %v", err) + } + if _, err := db.Exec(`INSERT INTO staging VALUES ('price_1', 'valid')`); err != nil { + t.Fatalf("seed staging: %v", err) + } + + result, err := tr.Transpile(`INSERT INTO target (id, nickname) + SELECT src_id, src_nickname FROM staging + ON CONFLICT (id) DO UPDATE SET nickname = excluded.nickname`) + if err != nil { + t.Fatalf("Transpile insert-select query: %v", err) + } + if err := runDuckDBStatements(db, result.Statements, result.CleanupStatements); err != nil { + t.Fatalf("guarded merge should map SELECT columns positionally to insert columns: %v", err) + } + + var rows int + if err := db.QueryRow(`SELECT count(*) FROM target WHERE id = 'price_1' AND nickname = 'valid'`).Scan(&rows); err != nil { + t.Fatalf("count inserted target rows: %v", err) + } + if rows != 1 { + t.Fatalf("guarded merge should insert row using target column names, got %d rows", rows) + } +} + +func TestTranspile_OnConflict_DuckLakeMode_DoNothingDeduplicatesSourceKeys(t *testing.T) { + tr := New(Config{DuckLakeMode: true}) + db, err := sql.Open("duckdb", ":memory:") + if err != nil { + t.Fatalf("open duckdb: %v", err) + } + defer func() { _ = db.Close() }() + + if _, err := db.Exec(`CREATE TABLE target (id VARCHAR, nickname VARCHAR)`); err != nil { + t.Fatalf("create target: %v", err) + } + if _, err := db.Exec(`CREATE TABLE staging (id VARCHAR, nickname VARCHAR)`); err != nil { + t.Fatalf("create staging: %v", err) + } + if _, err := db.Exec(`INSERT INTO staging VALUES ('price_1', 'first'), ('price_1', 'second'), (NULL, 'null a'), (NULL, 'null b')`); err != nil { + t.Fatalf("seed staging: %v", err) + } + + result, err := tr.Transpile(`INSERT INTO target (id, nickname) + SELECT id, nickname FROM staging + ON CONFLICT (id) DO NOTHING`) + if err != nil { + t.Fatalf("Transpile do-nothing query: %v", err) + } + if err := runDuckDBStatements(db, result.Statements, result.CleanupStatements); err != nil { + t.Fatalf("DO NOTHING should deduplicate non-null source conflict keys instead of erroring: %v", err) + } + + var duplicateKeyRows int + if err := db.QueryRow(`SELECT count(*) FROM target WHERE id = 'price_1'`).Scan(&duplicateKeyRows); err != nil { + t.Fatalf("count deduplicated target rows: %v", err) + } + if duplicateKeyRows != 1 { + t.Fatalf("DO NOTHING should insert one non-null source row per conflict key, got %d", duplicateKeyRows) + } + + var nullRows int + if err := db.QueryRow(`SELECT count(*) FROM target WHERE id IS NULL`).Scan(&nullRows); err != nil { + t.Fatalf("count null-key target rows: %v", err) + } + if nullRows != 2 { + t.Fatalf("DO NOTHING should preserve NULL conflict-key rows, got %d", nullRows) + } +} + +func TestTranspile_OnConflict_DuckLakeMode_GuardsMultiColumnConflicts(t *testing.T) { + tr := New(Config{DuckLakeMode: true}) + db, err := sql.Open("duckdb", ":memory:") + if err != nil { + t.Fatalf("open duckdb: %v", err) + } + defer func() { _ = db.Close() }() + + if _, err := db.Exec(`CREATE TABLE target (id VARCHAR, org_id VARCHAR, nickname VARCHAR)`); err != nil { + t.Fatalf("create target: %v", err) + } + if _, err := db.Exec(`CREATE TABLE staging (id VARCHAR, org_id VARCHAR, nickname VARCHAR)`); err != nil { + t.Fatalf("create staging: %v", err) + } + + result, err := tr.Transpile(`INSERT INTO target (id, org_id, nickname) + SELECT id, org_id, nickname FROM staging + ON CONFLICT (id, org_id) DO UPDATE SET nickname = excluded.nickname`) + if err != nil { + t.Fatalf("Transpile multi-column query: %v", err) + } + if len(result.Statements) == 0 { + t.Fatalf("expected guarded multi-statement rewrite, got single SQL: %s", result.SQL) + } + + if _, err := db.Exec(`INSERT INTO staging VALUES ('price_1', 'org_1', 'first'), ('price_1', 'org_1', 'second')`); err != nil { + t.Fatalf("seed duplicate staging rows: %v", err) + } + err = runDuckDBStatements(db, result.Statements, result.CleanupStatements) + if err == nil || !strings.Contains(err.Error(), "ON CONFLICT source rows contain duplicate conflict keys") { + t.Fatalf("expected multi-column source duplicate guard error, got %v", err) + } + + if _, err := db.Exec(`DELETE FROM staging`); err != nil { + t.Fatalf("clear staging: %v", err) + } + if _, err := db.Exec(`INSERT INTO target VALUES ('price_1', 'org_1', 'old a'), ('price_1', 'org_1', 'old b')`); err != nil { + t.Fatalf("seed duplicate target rows: %v", err) + } + if _, err := db.Exec(`INSERT INTO staging VALUES ('price_1', 'org_1', 'new')`); err != nil { + t.Fatalf("seed staging row: %v", err) + } + err = runDuckDBStatements(db, result.Statements, result.CleanupStatements) + if err == nil || !strings.Contains(err.Error(), "ON CONFLICT target contains duplicate conflict keys") { + t.Fatalf("expected multi-column target duplicate guard error, got %v", err) + } +} + func TestTranspile_DropCascade_DuckLakeMode(t *testing.T) { // In DuckLake mode, CASCADE should be stripped from DROP TABLE/VIEW statements // See: https://github.com/duckdb/dbt-duckdb/pull/557 From c83f3c116d670e6be1aaa45f0e4496059fa46f05 Mon Sep 17 00:00:00 2001 From: Bill Guowei Yang Date: Thu, 2 Jul 2026 10:30:34 -0400 Subject: [PATCH 4/6] Harden guarded ON CONFLICT finalization --- server/conn.go | 60 ++++++++++++++++++++ server/conn_query_exec.go | 89 +++++++++++++++++++++++++----- server/conn_test.go | 70 +++++++++++++++++++++++ transpiler/transform/onconflict.go | 40 +++++++++++++- transpiler/transpiler_test.go | 67 ++++++++++++++++++++++ 5 files changed, 310 insertions(+), 16 deletions(-) diff --git a/server/conn.go b/server/conn.go index fbd4e79d..5cd11846 100644 --- a/server/conn.go +++ b/server/conn.go @@ -1594,6 +1594,7 @@ func queryAndArgsForStatement(query string, args []interface{}) (string, []inter var out strings.Builder var stmtArgs []interface{} inString := false + inEscapeString := false inIdent := false inLineComment := false inBlockComment := false @@ -1621,12 +1622,18 @@ func queryAndArgsForStatement(query string, args []interface{}) (string, []inter if inString { out.WriteByte(c) + if inEscapeString && c == '\\' && i+1 < len(query) { + out.WriteByte(query[i+1]) + i++ + continue + } if c == '\'' { if i+1 < len(query) && query[i+1] == '\'' { out.WriteByte(query[i+1]) i++ } else { inString = false + inEscapeString = false } } continue @@ -1659,9 +1666,18 @@ func queryAndArgsForStatement(query string, args []interface{}) (string, []inter i++ continue } + if isEscapeStringLiteralStart(query, i) { + out.WriteByte(c) + out.WriteByte(query[i+1]) + inString = true + inEscapeString = true + i++ + continue + } if c == '\'' { out.WriteByte(c) inString = true + inEscapeString = false continue } if c == '"' { @@ -1683,6 +1699,13 @@ func queryAndArgsForStatement(query string, args []interface{}) (string, []inter continue } } + if c == '$' { + if end, ok := dollarQuotedLiteralEnd(query, i); ok { + out.WriteString(query[i:end]) + i = end - 1 + continue + } + } out.WriteByte(c) } @@ -1690,6 +1713,43 @@ func queryAndArgsForStatement(query string, args []interface{}) (string, []inter return out.String(), stmtArgs } +func isEscapeStringLiteralStart(query string, i int) bool { + if i+1 >= len(query) || query[i+1] != '\'' || (query[i] != 'E' && query[i] != 'e') { + return false + } + return i == 0 || !isSQLIdentPart(query[i-1]) +} + +func dollarQuotedLiteralEnd(query string, start int) (int, bool) { + if start >= len(query) || query[start] != '$' { + return 0, false + } + j := start + 1 + if j < len(query) && query[j] >= '0' && query[j] <= '9' { + return 0, false + } + for j < len(query) && isSQLIdentPart(query[j]) { + j++ + } + if j >= len(query) || query[j] != '$' { + return 0, false + } + tag := query[start : j+1] + bodyStart := j + 1 + bodyEnd := strings.Index(query[bodyStart:], tag) + if bodyEnd < 0 { + return 0, false + } + return bodyStart + bodyEnd + len(tag), true +} + +func isSQLIdentPart(c byte) bool { + return (c >= 'a' && c <= 'z') || + (c >= 'A' && c <= 'Z') || + (c >= '0' && c <= '9') || + c == '_' +} + // isEmptyQuery and stripLeadingComments moved to server/sqlcore so the // Flight client can call them without importing server. Local thin wrappers // preserve the unexported call-site spellings used throughout this file. diff --git a/server/conn_query_exec.go b/server/conn_query_exec.go index c23d2881..be39d255 100644 --- a/server/conn_query_exec.go +++ b/server/conn_query_exec.go @@ -855,6 +855,7 @@ func (c *clientConn) executeMultiStatement(statements []string, cleanup []string strings.ToUpper(strings.TrimSpace(statements[0])) == "BEGIN" && len(cleanup) > 0 && strings.ToUpper(strings.TrimSpace(cleanup[len(cleanup)-1])) == "COMMIT" + ownsTransaction := hasOurTransaction && c.txStatus != txStatusTransaction // If already in a transaction, skip our BEGIN/COMMIT wrapper if hasOurTransaction && c.txStatus == txStatusTransaction { @@ -879,8 +880,7 @@ func (c *clientConn) executeMultiStatement(statements []string, cleanup []string if err != nil { c.logger().Error("Multi-stmt setup error.", "query", stmt, "error", err) c.setTxError() - // On error, still try to cleanup (best effort) - c.executeCleanup(cleanup) + c.executeCleanupAfterError(cleanup, ownsTransaction) c.sendError("ERROR", "42000", err.Error()) _ = c.writeReadyForQuery(c.txStatus) _ = c.flushWriter() @@ -909,7 +909,7 @@ func (c *clientConn) executeMultiStatement(statements []string, cleanup []string finalErr = err c.logger().Error("Multi-stmt final query error.", "query", finalStmt, "error", err) c.setTxError() - c.executeCleanup(cleanup) + c.executeCleanupAfterError(cleanup, ownsTransaction) c.sendError("ERROR", "42000", err.Error()) _ = c.writeReadyForQuery(c.txStatus) _ = c.flushWriter() @@ -919,7 +919,13 @@ func (c *clientConn) executeMultiStatement(statements []string, cleanup []string // Execute cleanup while cursor is open (data is materialized in cursor) // DuckDB cursor holds result data even after source tables are dropped - c.executeCleanup(cleanup) + if err := c.executeCleanupAndCommit(cleanup); err != nil { + finalErr = err + c.sendError("ERROR", classifyErrorCode(err), err.Error()) + _ = c.writeReadyForQuery(c.txStatus) + _ = c.flushWriter() + return nil + } // Now stream results from cursor. streamRowsToClient counts rows // internally; we approximate Finished's rowsAff with 0 here @@ -939,7 +945,7 @@ func (c *clientConn) executeMultiStatement(statements []string, cleanup []string finalErr = err c.logger().Error("Multi-stmt final exec error.", "query", finalStmt, "error", err) c.setTxError() - c.executeCleanup(cleanup) + c.executeCleanupAfterError(cleanup, ownsTransaction) c.sendError("ERROR", "42000", err.Error()) _ = c.writeReadyForQuery(c.txStatus) _ = c.flushWriter() @@ -949,8 +955,13 @@ func (c *clientConn) executeMultiStatement(statements []string, cleanup []string finalRowsAff, _ = result.RowsAffected() } - // Execute cleanup - c.executeCleanup(cleanup) + if err := c.executeCleanupAndCommit(cleanup); err != nil { + finalErr = err + c.sendError("ERROR", classifyErrorCode(err), err.Error()) + _ = c.writeReadyForQuery(c.txStatus) + _ = c.flushWriter() + return nil + } // Send completion tag := c.buildCommandTag(cmdType, result) @@ -973,6 +984,49 @@ func (c *clientConn) executeCleanup(cleanup []string) { } } +func (c *clientConn) executeCleanupAndCommit(cleanup []string) error { + cleanupBeforeCommit, commitStmt, hasCommit := splitCommitCleanup(cleanup) + c.executeCleanup(cleanupBeforeCommit) + if !hasCommit { + return nil + } + if err := c.executeTransactionFinalizer(commitStmt); err != nil { + c.executeCleanup([]string{"ROLLBACK"}) + return err + } + return nil +} + +func (c *clientConn) executeCleanupAfterError(cleanup []string, rollbackOwnedTransaction bool) { + if rollbackOwnedTransaction { + cleanupBeforeCommit, _, _ := splitCommitCleanup(cleanup) + c.executeCleanup(cleanupBeforeCommit) + c.executeCleanup([]string{"ROLLBACK"}) + return + } + c.executeCleanup(cleanup) +} + +func (c *clientConn) executeTransactionFinalizer(stmt string) error { + c.logger().Debug("Multi-stmt transaction finalizer.", "stmt", stmt) + _, err := c.executor.Exec(stmt) + if err != nil { + c.logger().Error("Multi-stmt transaction finalizer error.", "query", stmt, "error", err) + } + return err +} + +func splitCommitCleanup(cleanup []string) ([]string, string, bool) { + if len(cleanup) == 0 { + return nil, "", false + } + last := cleanup[len(cleanup)-1] + if strings.ToUpper(strings.TrimSpace(last)) != "COMMIT" { + return cleanup, "", false + } + return cleanup[:len(cleanup)-1], last, true +} + // executeMultiStatementExtended handles execution of multi-statement query rewrites // for the extended query protocol (Parse/Bind/Execute). // Unlike executeMultiStatement, this does NOT send ReadyForQuery (that's done by Sync). @@ -987,6 +1041,7 @@ func (c *clientConn) executeMultiStatementExtended(statements []string, cleanup strings.ToUpper(strings.TrimSpace(statements[0])) == "BEGIN" && len(cleanup) > 0 && strings.ToUpper(strings.TrimSpace(cleanup[len(cleanup)-1])) == "COMMIT" + ownsTransaction := hasOurTransaction && c.txStatus != txStatusTransaction // If already in a transaction, skip our BEGIN/COMMIT wrapper if hasOurTransaction && c.txStatus == txStatusTransaction { @@ -1012,8 +1067,7 @@ func (c *clientConn) executeMultiStatementExtended(statements []string, cleanup if err != nil { c.logger().Error("Multi-stmt-ext setup error.", "query", stmtQuery, "error", err) c.setTxError() - // On error, still try to cleanup (best effort) - c.executeCleanup(cleanup) + c.executeCleanupAfterError(cleanup, ownsTransaction) c.sendError("ERROR", "42000", err.Error()) return } @@ -1041,14 +1095,18 @@ func (c *clientConn) executeMultiStatementExtended(statements []string, cleanup finalErr = err c.logger().Error("Multi-stmt-ext final query error.", "query", finalQuery, "error", err) c.setTxError() - c.executeCleanup(cleanup) + c.executeCleanupAfterError(cleanup, ownsTransaction) c.sendError("ERROR", "42000", err.Error()) return } defer func() { _ = rows.Close() }() // Execute cleanup while cursor is open (data is materialized in cursor) - c.executeCleanup(cleanup) + if err := c.executeCleanupAndCommit(cleanup); err != nil { + finalErr = err + c.sendError("ERROR", classifyErrorCode(err), err.Error()) + return + } // Stream results from cursor (extended protocol version). Row count // is tracked by streamRowsToClientExtended; the deferred Finished @@ -1063,7 +1121,7 @@ func (c *clientConn) executeMultiStatementExtended(statements []string, cleanup finalErr = err c.logger().Error("Multi-stmt-ext final exec error.", "query", finalQuery, "error", err) c.setTxError() - c.executeCleanup(cleanup) + c.executeCleanupAfterError(cleanup, ownsTransaction) c.sendError("ERROR", "42000", err.Error()) return } @@ -1071,8 +1129,11 @@ func (c *clientConn) executeMultiStatementExtended(statements []string, cleanup finalRowsAff, _ = result.RowsAffected() } - // Execute cleanup - c.executeCleanup(cleanup) + if err := c.executeCleanupAndCommit(cleanup); err != nil { + finalErr = err + c.sendError("ERROR", classifyErrorCode(err), err.Error()) + return + } // Send completion (no ReadyForQuery - that's done by Sync) tag := c.buildCommandTag(cmdType, result) diff --git a/server/conn_test.go b/server/conn_test.go index 6a38146a..428ca57a 100644 --- a/server/conn_test.go +++ b/server/conn_test.go @@ -3906,6 +3906,24 @@ func TestQueryAndArgsForStatement(t *testing.T) { wantQ: "SELECT '$1', \"$2\", ? -- $1\n/* $2 */", want: []interface{}{args[2]}, }, + { + name: "ignores placeholders in escape string literals", + query: "SELECT E'foo\\' $1', $2", + wantQ: "SELECT E'foo\\' $1', ?", + want: []interface{}{args[1]}, + }, + { + name: "ignores placeholders in dollar quoted strings", + query: "SELECT $$ $1 $$, $2", + wantQ: "SELECT $$ $1 $$, ?", + want: []interface{}{args[1]}, + }, + { + name: "ignores placeholders in tagged dollar quoted strings", + query: "SELECT $tag$ $1 $tag$, $3", + wantQ: "SELECT $tag$ $1 $tag$, ?", + want: []interface{}{args[2]}, + }, { name: "placeholder beyond provided args is left for downstream validation", query: "SELECT $4", @@ -3945,6 +3963,58 @@ func TestQueryAndArgsForStatementRunsSparseDollarParams(t *testing.T) { } } +type cleanupErrorExecutor struct { + noopProfiling + execQueries []string +} + +func (e *cleanupErrorExecutor) QueryContext(context.Context, string, ...any) (RowSet, error) { + return nil, errors.New("not implemented") +} + +func (e *cleanupErrorExecutor) ExecContext(context.Context, string, ...any) (ExecResult, error) { + return nil, errors.New("not implemented") +} + +func (e *cleanupErrorExecutor) Query(string, ...any) (RowSet, error) { + return nil, errors.New("not implemented") +} + +func (e *cleanupErrorExecutor) Exec(query string, _ ...any) (ExecResult, error) { + e.execQueries = append(e.execQueries, query) + if strings.TrimSpace(strings.ToUpper(query)) == "COMMIT" { + return nil, errors.New("commit failed") + } + return &fakeExecResult{}, nil +} + +func (e *cleanupErrorExecutor) ConnContext(context.Context) (RawConn, error) { + return nil, errors.New("not implemented") +} + +func (e *cleanupErrorExecutor) PingContext(context.Context) error { + return errors.New("not implemented") +} + +func (e *cleanupErrorExecutor) Close() error { + return nil +} + +func TestExecuteCleanupAndCommitPropagatesCommitError(t *testing.T) { + exec := &cleanupErrorExecutor{} + c := &clientConn{executor: exec} + + err := c.executeCleanupAndCommit([]string{"DROP TABLE IF EXISTS _duckgres_tmp", "COMMIT"}) + if err == nil || !strings.Contains(err.Error(), "commit failed") { + t.Fatalf("executeCleanupAndCommit error = %v, want commit failed", err) + } + + want := []string{"DROP TABLE IF EXISTS _duckgres_tmp", "COMMIT", "ROLLBACK"} + if !reflect.DeepEqual(exec.execQueries, want) { + t.Fatalf("exec queries = %#v, want %#v", exec.execQueries, want) + } +} + func TestMatchPgStatActivityQuery(t *testing.T) { tests := []struct { name string diff --git a/transpiler/transform/onconflict.go b/transpiler/transform/onconflict.go index 8f67ee39..f9f8698d 100644 --- a/transpiler/transform/onconflict.go +++ b/transpiler/transform/onconflict.go @@ -65,6 +65,10 @@ func (t *OnConflictTransform) Transform(tree *pg_query.ParseResult, result *Resu result.Error = unsupported return false, nil } + if unsupported := t.validateDuckLakeInsertColumns(insert); unsupported != nil { + result.Error = unsupported + return false, nil + } } if rewrite := t.transformInsertToMerge(insert); rewrite != nil { @@ -90,9 +94,13 @@ func (t *OnConflictTransform) Transform(tree *pg_query.ParseResult, result *Resu } func (t *OnConflictTransform) validateDuckLakeConflictTarget(occ *pg_query.OnConflictClause) error { - if occ == nil || occ.Infer == nil { + if occ == nil { return nil } + if occ.Infer == nil { + return NewFeatureNotSupported( + "ON CONFLICT without a conflict target is not supported: this catalog does not enforce unique constraints") + } if strings.TrimSpace(occ.Infer.Conname) != "" { return NewFeatureNotSupported( "ON CONFLICT ON CONSTRAINT is not supported: this catalog does not enforce named constraints") @@ -111,6 +119,34 @@ func (t *OnConflictTransform) validateDuckLakeConflictTarget(occ *pg_query.OnCon return nil } +func (t *OnConflictTransform) validateDuckLakeInsertColumns(insert *pg_query.InsertStmt) error { + if insert == nil || insert.OnConflictClause == nil { + return nil + } + if len(insert.Cols) == 0 { + return NewFeatureNotSupported( + "ON CONFLICT without an explicit insert column list is not supported") + } + insertColumns := make(map[string]struct{}, len(insert.Cols)) + for _, col := range insert.Cols { + rt := col.GetResTarget() + if rt == nil || strings.TrimSpace(rt.Name) == "" { + return NewFeatureNotSupported( + "ON CONFLICT without an explicit insert column list is not supported") + } + insertColumns[rt.Name] = struct{}{} + } + + conflictColumns := t.conflictColumnNames(insert.OnConflictClause.Infer.IndexElems) + for _, conflictColumn := range conflictColumns { + if _, ok := insertColumns[conflictColumn]; !ok { + return NewFeatureNotSupported( + "ON CONFLICT conflict columns must be present in the insert column list") + } + } + return nil +} + // transformInsertToMerge converts INSERT ... ON CONFLICT to MERGE for DuckLake mode func (t *OnConflictTransform) transformInsertToMerge(insert *pg_query.InsertStmt) *onConflictMergeRewrite { if !t.DuckLakeMode { @@ -247,9 +283,9 @@ func (t *OnConflictTransform) buildGuardedMergeStatements(rewrite *onConflictMer } if rewrite.action == pg_query.OnConflictAction_ONCONFLICT_UPDATE { statements = append(statements, t.buildSourceDuplicateGuard(sourceFrom, rewrite.conflictColumns)) + statements = append(statements, t.buildTargetDuplicateGuard(sourceFrom, rewrite.targetRelation, rewrite.conflictColumns)) } statements = append(statements, - t.buildTargetDuplicateGuard(sourceFrom, rewrite.targetRelation, rewrite.conflictColumns), mergeSQL, ) cleanup := []string{ diff --git a/transpiler/transpiler_test.go b/transpiler/transpiler_test.go index e64b7dac..a2dac332 100644 --- a/transpiler/transpiler_test.go +++ b/transpiler/transpiler_test.go @@ -458,6 +458,33 @@ func TestTranspile_FeatureNotSupported_Rejected(t *testing.T) { assertFeatureNotSupported(t, result) }) + t.Run(string(backend)+"/ON CONFLICT targetless DO NOTHING", func(t *testing.T) { + result, err := tr.Transpile( + "INSERT INTO t (id, v) VALUES (1, 2) ON CONFLICT DO NOTHING") + if err != nil { + t.Fatalf("Transpile error: %v", err) + } + assertFeatureNotSupported(t, result) + }) + + t.Run(string(backend)+"/ON CONFLICT without explicit insert column list", func(t *testing.T) { + result, err := tr.Transpile( + "INSERT INTO t VALUES (1, 2) ON CONFLICT (id) DO UPDATE SET v = EXCLUDED.v") + if err != nil { + t.Fatalf("Transpile error: %v", err) + } + assertFeatureNotSupported(t, result) + }) + + t.Run(string(backend)+"/ON CONFLICT conflict column omitted from insert column list", func(t *testing.T) { + result, err := tr.Transpile( + "INSERT INTO t (v) VALUES (2) ON CONFLICT (id) DO NOTHING") + if err != nil { + t.Fatalf("Transpile error: %v", err) + } + assertFeatureNotSupported(t, result) + }) + t.Run(string(backend)+"/ALTER COLUMN TYPE USING", func(t *testing.T) { result, err := tr.Transpile( "ALTER TABLE t ALTER COLUMN c TYPE integer USING c::integer") @@ -2591,6 +2618,46 @@ func TestTranspile_OnConflict_DuckLakeMode_DoNothingDeduplicatesSourceKeys(t *te } } +func TestTranspile_OnConflict_DuckLakeMode_DoNothingAllowsDuplicateTargetMatches(t *testing.T) { + tr := New(Config{DuckLakeMode: true}) + db, err := sql.Open("duckdb", ":memory:") + if err != nil { + t.Fatalf("open duckdb: %v", err) + } + defer func() { _ = db.Close() }() + + if _, err := db.Exec(`CREATE TABLE target (id VARCHAR, nickname VARCHAR)`); err != nil { + t.Fatalf("create target: %v", err) + } + if _, err := db.Exec(`CREATE TABLE staging (id VARCHAR, nickname VARCHAR)`); err != nil { + t.Fatalf("create staging: %v", err) + } + if _, err := db.Exec(`INSERT INTO target VALUES ('price_1', 'old a'), ('price_1', 'old b')`); err != nil { + t.Fatalf("seed duplicate target rows: %v", err) + } + if _, err := db.Exec(`INSERT INTO staging VALUES ('price_1', 'new')`); err != nil { + t.Fatalf("seed staging: %v", err) + } + + result, err := tr.Transpile(`INSERT INTO target (id, nickname) + SELECT id, nickname FROM staging + ON CONFLICT (id) DO NOTHING`) + if err != nil { + t.Fatalf("Transpile do-nothing query: %v", err) + } + if err := runDuckDBStatements(db, result.Statements, result.CleanupStatements); err != nil { + t.Fatalf("DO NOTHING should tolerate duplicate target matches: %v", err) + } + + var rows int + if err := db.QueryRow(`SELECT count(*) FROM target WHERE id = 'price_1'`).Scan(&rows); err != nil { + t.Fatalf("count target rows: %v", err) + } + if rows != 2 { + t.Fatalf("DO NOTHING should leave duplicate matched target rows unchanged, got %d rows", rows) + } +} + func TestTranspile_OnConflict_DuckLakeMode_GuardsMultiColumnConflicts(t *testing.T) { tr := New(Config{DuckLakeMode: true}) db, err := sql.Open("duckdb", ":memory:") From 7ec0959794f2de0e7ac05a9e1368933c142f25d1 Mon Sep 17 00:00:00 2001 From: Bill Guowei Yang Date: Thu, 2 Jul 2026 11:07:09 -0400 Subject: [PATCH 5/6] Reject ON CONFLICT on lake catalogs --- docs/postgres-compatibility.md | 4 +- server/conn.go | 160 ---- server/conn_query_exec.go | 117 +-- server/conn_test.go | 155 ---- .../integration/ducklake_concurrency_test.go | 71 +- transpiler/backend/profile.go | 4 +- transpiler/backend/profile_test.go | 4 +- transpiler/classify_wiring_test.go | 14 +- transpiler/transform/onconflict.go | 723 +----------------- transpiler/transform/onconflict_test.go | 702 ++--------------- transpiler/transpiler.go | 12 +- transpiler/transpiler_test.go | 582 ++------------ 12 files changed, 180 insertions(+), 2368 deletions(-) diff --git a/docs/postgres-compatibility.md b/docs/postgres-compatibility.md index fce6b90c..5f0c1b82 100644 --- a/docs/postgres-compatibility.md +++ b/docs/postgres-compatibility.md @@ -74,9 +74,9 @@ why. | DELETE (incl. `DELETE…USING`, subquery) | ✅ | `dml_test.go::TestDMLDelete`, `::TestDMLDeleteUsing`; transform `wiring_ops_test.go::TestOperatorTransform_DeleteUsing` | | | RETURNING (simple query protocol) | 🟡 | transform `wiring_ops_test.go::TestOperatorTransform_{Insert,Update,Delete}Returning`; `conn_test.go::TestContainsReturning`/`::TestIsDMLReturning` | Differential tests `TestDML{Insert,Update,Delete}Returning` are `skipIfKnown` (stale — passes at unit/client level) | | RETURNING (extended query protocol) | ⛔ | `conn_test.go::TestIsDMLReturning` | Rejected at Describe time with `0A000` by design — Describe would execute the mutation. See CLAUDE.md "DML RETURNING Detection" | -| ON CONFLICT / UPSERT (DO UPDATE) | ✅ | `dml_test.go::TestDMLInsertOnConflict`; transpiler `transform/onconflict_test.go` | Also exercised via DuckLake ON CONFLICT→MERGE rewrite: `ducklake_concurrency_test.go::TestDuckLakeConcurrentTransactions` | +| ON CONFLICT / UPSERT (DO UPDATE) | ✅ | `dml_test.go::TestDMLInsertOnConflict`; transpiler `transform/onconflict_test.go` | Constraint-less lake catalogs reject ON CONFLICT with `0A000` because they cannot enforce unique constraints | | ON CONFLICT DO NOTHING | 🟡 | `dml_test.go::TestDMLInsertOnConflict` (subtest skipped) | | -| MERGE (user-facing) | ⛔ | — | DuckDB has no MERGE; only internal ON CONFLICT→MERGE rewrite exists | +| MERGE (user-facing) | ⛔ | — | Not a PostgreSQL compatibility target; Duckgres does not use MERGE to emulate ON CONFLICT | | TRUNCATE | ✅ | `ddl_test.go::TestDDLTruncate` | | | COPY … FROM STDIN (text/CSV) | ✅ | `copy_test.go::TestCopyFromStdin`, `::TestCopyFromStdinWithSpecialChars`, `::TestCopyFromStdinMultilineJSON` | Escape sequences stored literally (documented DuckDB CSV-parser limitation) | | COPY … TO STDOUT | 🟡 | `copy_test.go::TestCopyToStdout` (skipped under lib/pq); `conn_test.go::TestCopyToStdoutRegex`; client-compat `psycopg` COPY suite | Integration skip is a lib/pq driver limitation, not a Duckgres gap | diff --git a/server/conn.go b/server/conn.go index 5cd11846..7a7578f9 100644 --- a/server/conn.go +++ b/server/conn.go @@ -1590,166 +1590,6 @@ func countDollarParams(query string) int { return max } -func queryAndArgsForStatement(query string, args []interface{}) (string, []interface{}) { - var out strings.Builder - var stmtArgs []interface{} - inString := false - inEscapeString := false - inIdent := false - inLineComment := false - inBlockComment := false - - for i := 0; i < len(query); i++ { - c := query[i] - - if inBlockComment { - out.WriteByte(c) - if c == '*' && i+1 < len(query) && query[i+1] == '/' { - out.WriteByte(query[i+1]) - inBlockComment = false - i++ - } - continue - } - - if inLineComment { - out.WriteByte(c) - if c == '\n' { - inLineComment = false - } - continue - } - - if inString { - out.WriteByte(c) - if inEscapeString && c == '\\' && i+1 < len(query) { - out.WriteByte(query[i+1]) - i++ - continue - } - if c == '\'' { - if i+1 < len(query) && query[i+1] == '\'' { - out.WriteByte(query[i+1]) - i++ - } else { - inString = false - inEscapeString = false - } - } - continue - } - - if inIdent { - out.WriteByte(c) - if c == '"' { - if i+1 < len(query) && query[i+1] == '"' { - out.WriteByte(query[i+1]) - i++ - } else { - inIdent = false - } - } - continue - } - - if c == '-' && i+1 < len(query) && query[i+1] == '-' { - out.WriteByte(c) - out.WriteByte(query[i+1]) - inLineComment = true - i++ - continue - } - if c == '/' && i+1 < len(query) && query[i+1] == '*' { - out.WriteByte(c) - out.WriteByte(query[i+1]) - inBlockComment = true - i++ - continue - } - if isEscapeStringLiteralStart(query, i) { - out.WriteByte(c) - out.WriteByte(query[i+1]) - inString = true - inEscapeString = true - i++ - continue - } - if c == '\'' { - out.WriteByte(c) - inString = true - inEscapeString = false - continue - } - if c == '"' { - out.WriteByte(c) - inIdent = true - continue - } - - if c == '$' && i+1 < len(query) && query[i+1] >= '1' && query[i+1] <= '9' { - j := i + 1 - for j < len(query) && query[j] >= '0' && query[j] <= '9' { - j++ - } - n, err := strconv.Atoi(query[i+1 : j]) - if err == nil && n >= 1 && n <= len(args) { - out.WriteByte('?') - stmtArgs = append(stmtArgs, args[n-1]) - i = j - 1 - continue - } - } - if c == '$' { - if end, ok := dollarQuotedLiteralEnd(query, i); ok { - out.WriteString(query[i:end]) - i = end - 1 - continue - } - } - - out.WriteByte(c) - } - - return out.String(), stmtArgs -} - -func isEscapeStringLiteralStart(query string, i int) bool { - if i+1 >= len(query) || query[i+1] != '\'' || (query[i] != 'E' && query[i] != 'e') { - return false - } - return i == 0 || !isSQLIdentPart(query[i-1]) -} - -func dollarQuotedLiteralEnd(query string, start int) (int, bool) { - if start >= len(query) || query[start] != '$' { - return 0, false - } - j := start + 1 - if j < len(query) && query[j] >= '0' && query[j] <= '9' { - return 0, false - } - for j < len(query) && isSQLIdentPart(query[j]) { - j++ - } - if j >= len(query) || query[j] != '$' { - return 0, false - } - tag := query[start : j+1] - bodyStart := j + 1 - bodyEnd := strings.Index(query[bodyStart:], tag) - if bodyEnd < 0 { - return 0, false - } - return bodyStart + bodyEnd + len(tag), true -} - -func isSQLIdentPart(c byte) bool { - return (c >= 'a' && c <= 'z') || - (c >= 'A' && c <= 'Z') || - (c >= '0' && c <= '9') || - c == '_' -} - // isEmptyQuery and stripLeadingComments moved to server/sqlcore so the // Flight client can call them without importing server. Local thin wrappers // preserve the unexported call-site spellings used throughout this file. diff --git a/server/conn_query_exec.go b/server/conn_query_exec.go index be39d255..213964a4 100644 --- a/server/conn_query_exec.go +++ b/server/conn_query_exec.go @@ -855,7 +855,6 @@ func (c *clientConn) executeMultiStatement(statements []string, cleanup []string strings.ToUpper(strings.TrimSpace(statements[0])) == "BEGIN" && len(cleanup) > 0 && strings.ToUpper(strings.TrimSpace(cleanup[len(cleanup)-1])) == "COMMIT" - ownsTransaction := hasOurTransaction && c.txStatus != txStatusTransaction // If already in a transaction, skip our BEGIN/COMMIT wrapper if hasOurTransaction && c.txStatus == txStatusTransaction { @@ -880,7 +879,8 @@ func (c *clientConn) executeMultiStatement(statements []string, cleanup []string if err != nil { c.logger().Error("Multi-stmt setup error.", "query", stmt, "error", err) c.setTxError() - c.executeCleanupAfterError(cleanup, ownsTransaction) + // On error, still try to cleanup (best effort) + c.executeCleanup(cleanup) c.sendError("ERROR", "42000", err.Error()) _ = c.writeReadyForQuery(c.txStatus) _ = c.flushWriter() @@ -909,7 +909,7 @@ func (c *clientConn) executeMultiStatement(statements []string, cleanup []string finalErr = err c.logger().Error("Multi-stmt final query error.", "query", finalStmt, "error", err) c.setTxError() - c.executeCleanupAfterError(cleanup, ownsTransaction) + c.executeCleanup(cleanup) c.sendError("ERROR", "42000", err.Error()) _ = c.writeReadyForQuery(c.txStatus) _ = c.flushWriter() @@ -919,13 +919,7 @@ func (c *clientConn) executeMultiStatement(statements []string, cleanup []string // Execute cleanup while cursor is open (data is materialized in cursor) // DuckDB cursor holds result data even after source tables are dropped - if err := c.executeCleanupAndCommit(cleanup); err != nil { - finalErr = err - c.sendError("ERROR", classifyErrorCode(err), err.Error()) - _ = c.writeReadyForQuery(c.txStatus) - _ = c.flushWriter() - return nil - } + c.executeCleanup(cleanup) // Now stream results from cursor. streamRowsToClient counts rows // internally; we approximate Finished's rowsAff with 0 here @@ -945,7 +939,7 @@ func (c *clientConn) executeMultiStatement(statements []string, cleanup []string finalErr = err c.logger().Error("Multi-stmt final exec error.", "query", finalStmt, "error", err) c.setTxError() - c.executeCleanupAfterError(cleanup, ownsTransaction) + c.executeCleanup(cleanup) c.sendError("ERROR", "42000", err.Error()) _ = c.writeReadyForQuery(c.txStatus) _ = c.flushWriter() @@ -955,13 +949,8 @@ func (c *clientConn) executeMultiStatement(statements []string, cleanup []string finalRowsAff, _ = result.RowsAffected() } - if err := c.executeCleanupAndCommit(cleanup); err != nil { - finalErr = err - c.sendError("ERROR", classifyErrorCode(err), err.Error()) - _ = c.writeReadyForQuery(c.txStatus) - _ = c.flushWriter() - return nil - } + // Execute cleanup + c.executeCleanup(cleanup) // Send completion tag := c.buildCommandTag(cmdType, result) @@ -984,49 +973,6 @@ func (c *clientConn) executeCleanup(cleanup []string) { } } -func (c *clientConn) executeCleanupAndCommit(cleanup []string) error { - cleanupBeforeCommit, commitStmt, hasCommit := splitCommitCleanup(cleanup) - c.executeCleanup(cleanupBeforeCommit) - if !hasCommit { - return nil - } - if err := c.executeTransactionFinalizer(commitStmt); err != nil { - c.executeCleanup([]string{"ROLLBACK"}) - return err - } - return nil -} - -func (c *clientConn) executeCleanupAfterError(cleanup []string, rollbackOwnedTransaction bool) { - if rollbackOwnedTransaction { - cleanupBeforeCommit, _, _ := splitCommitCleanup(cleanup) - c.executeCleanup(cleanupBeforeCommit) - c.executeCleanup([]string{"ROLLBACK"}) - return - } - c.executeCleanup(cleanup) -} - -func (c *clientConn) executeTransactionFinalizer(stmt string) error { - c.logger().Debug("Multi-stmt transaction finalizer.", "stmt", stmt) - _, err := c.executor.Exec(stmt) - if err != nil { - c.logger().Error("Multi-stmt transaction finalizer error.", "query", stmt, "error", err) - } - return err -} - -func splitCommitCleanup(cleanup []string) ([]string, string, bool) { - if len(cleanup) == 0 { - return nil, "", false - } - last := cleanup[len(cleanup)-1] - if strings.ToUpper(strings.TrimSpace(last)) != "COMMIT" { - return cleanup, "", false - } - return cleanup[:len(cleanup)-1], last, true -} - // executeMultiStatementExtended handles execution of multi-statement query rewrites // for the extended query protocol (Parse/Bind/Execute). // Unlike executeMultiStatement, this does NOT send ReadyForQuery (that's done by Sync). @@ -1041,7 +987,6 @@ func (c *clientConn) executeMultiStatementExtended(statements []string, cleanup strings.ToUpper(strings.TrimSpace(statements[0])) == "BEGIN" && len(cleanup) > 0 && strings.ToUpper(strings.TrimSpace(cleanup[len(cleanup)-1])) == "COMMIT" - ownsTransaction := hasOurTransaction && c.txStatus != txStatusTransaction // If already in a transaction, skip our BEGIN/COMMIT wrapper if hasOurTransaction && c.txStatus == txStatusTransaction { @@ -1054,20 +999,20 @@ func (c *clientConn) executeMultiStatementExtended(statements []string, cleanup // setup work is observable per statement, not just per outer query. for i := 0; i < len(statements)-1; i++ { stmt := statements[i] - stmtQuery, stmtArgs := queryAndArgsForStatement(stmt, args) - c.logger().Debug("Multi-stmt-ext setup.", "step", i+1, "total", len(statements)-1, "stmt", stmtQuery) + c.logger().Debug("Multi-stmt-ext setup.", "step", i+1, "total", len(statements)-1, "stmt", stmt) setupStart := time.Now() - c.logQueryStarted(stmtQuery) - result, err := c.executor.Exec(stmtQuery, stmtArgs...) + c.logQueryStarted(stmt) + result, err := c.executor.Exec(stmt, args...) var setupRows int64 if result != nil { setupRows, _ = result.RowsAffected() } - c.logQueryFinished(stmtQuery, setupStart, setupRows, err) + c.logQueryFinished(stmt, setupStart, setupRows, err) if err != nil { - c.logger().Error("Multi-stmt-ext setup error.", "query", stmtQuery, "error", err) + c.logger().Error("Multi-stmt-ext setup error.", "query", stmt, "error", err) c.setTxError() - c.executeCleanupAfterError(cleanup, ownsTransaction) + // On error, still try to cleanup (best effort) + c.executeCleanup(cleanup) c.sendError("ERROR", "42000", err.Error()) return } @@ -1075,53 +1020,48 @@ func (c *clientConn) executeMultiStatementExtended(statements []string, cleanup // Handle final statement finalStmt := statements[len(statements)-1] - finalQuery, finalArgs := queryAndArgsForStatement(finalStmt, args) upperFinal := strings.ToUpper(strings.TrimSpace(finalStmt)) cmdType := c.getCommandType(upperFinal) - c.logger().Debug("Multi-stmt-ext final.", "stmt", finalQuery, "cmd_type", cmdType) + c.logger().Debug("Multi-stmt-ext final.", "stmt", finalStmt, "cmd_type", cmdType) finalStart := time.Now() var finalRowsAff int64 var finalErr error - c.logQueryStarted(finalQuery) + c.logQueryStarted(finalStmt) defer func() { - c.logQueryFinished(finalQuery, finalStart, finalRowsAff, finalErr) + c.logQueryFinished(finalStmt, finalStart, finalRowsAff, finalErr) }() if queryReturnsResults(finalStmt) { // Result-returning query: obtain cursor FIRST, cleanup SECOND, stream THIRD - rows, err := c.executor.Query(finalQuery, finalArgs...) + rows, err := c.executor.Query(finalStmt, args...) if err != nil { finalErr = err - c.logger().Error("Multi-stmt-ext final query error.", "query", finalQuery, "error", err) + c.logger().Error("Multi-stmt-ext final query error.", "query", finalStmt, "error", err) c.setTxError() - c.executeCleanupAfterError(cleanup, ownsTransaction) + c.executeCleanup(cleanup) c.sendError("ERROR", "42000", err.Error()) return } defer func() { _ = rows.Close() }() // Execute cleanup while cursor is open (data is materialized in cursor) - if err := c.executeCleanupAndCommit(cleanup); err != nil { - finalErr = err - c.sendError("ERROR", classifyErrorCode(err), err.Error()) - return - } + c.executeCleanup(cleanup) // Stream results from cursor (extended protocol version). Row count // is tracked by streamRowsToClientExtended; the deferred Finished // log uses 0 as an approximation (logQuery still records the // precise count via the structured channel). - c.streamRowsToClientExtended(rows, cmdType, resultFormats, described, finalQuery) + c.streamRowsToClientExtended(rows, cmdType, resultFormats, described, finalStmt) } else { // Non-result query (DML without RETURNING, DDL, etc.): execute then cleanup - result, err := c.executor.Exec(finalQuery, finalArgs...) + result, err := c.executor.Exec(finalStmt, args...) if err != nil { finalErr = err - c.logger().Error("Multi-stmt-ext final exec error.", "query", finalQuery, "error", err) + c.logger().Error("Multi-stmt-ext final exec error.", "query", finalStmt, "error", err) c.setTxError() - c.executeCleanupAfterError(cleanup, ownsTransaction) + c.executeCleanup(cleanup) c.sendError("ERROR", "42000", err.Error()) return } @@ -1129,11 +1069,8 @@ func (c *clientConn) executeMultiStatementExtended(statements []string, cleanup finalRowsAff, _ = result.RowsAffected() } - if err := c.executeCleanupAndCommit(cleanup); err != nil { - finalErr = err - c.sendError("ERROR", classifyErrorCode(err), err.Error()) - return - } + // Execute cleanup + c.executeCleanup(cleanup) // Send completion (no ReadyForQuery - that's done by Sync) tag := c.buildCommandTag(cmdType, result) diff --git a/server/conn_test.go b/server/conn_test.go index 428ca57a..e2b16943 100644 --- a/server/conn_test.go +++ b/server/conn_test.go @@ -12,7 +12,6 @@ import ( "io" "math" "net" - "reflect" "slices" "strings" "testing" @@ -3861,160 +3860,6 @@ func TestCountDollarParams(t *testing.T) { } } -func TestQueryAndArgsForStatement(t *testing.T) { - args := []interface{}{"id_1", "nickname", "unused"} - - tests := []struct { - name string - query string - wantQ string - want []interface{} - }{ - { - name: "no placeholders", - query: "SELECT 1", - wantQ: "SELECT 1", - want: nil, - }, - { - name: "uses first placeholder", - query: "SELECT $1", - wantQ: "SELECT ?", - want: args[:1], - }, - { - name: "uses first two placeholders", - query: "INSERT INTO t VALUES ($1, $2)", - wantQ: "INSERT INTO t VALUES (?, ?)", - want: args[:2], - }, - { - name: "sparse placeholders preserve referenced argument numbers", - query: "SELECT $3, $1", - wantQ: "SELECT ?, ?", - want: []interface{}{args[2], args[0]}, - }, - { - name: "repeated placeholders duplicate the referenced argument", - query: "SELECT $2, $2", - wantQ: "SELECT ?, ?", - want: []interface{}{args[1], args[1]}, - }, - { - name: "ignores placeholders in literals comments and identifiers", - query: "SELECT '$1', \"$2\", $3 -- $1\n/* $2 */", - wantQ: "SELECT '$1', \"$2\", ? -- $1\n/* $2 */", - want: []interface{}{args[2]}, - }, - { - name: "ignores placeholders in escape string literals", - query: "SELECT E'foo\\' $1', $2", - wantQ: "SELECT E'foo\\' $1', ?", - want: []interface{}{args[1]}, - }, - { - name: "ignores placeholders in dollar quoted strings", - query: "SELECT $$ $1 $$, $2", - wantQ: "SELECT $$ $1 $$, ?", - want: []interface{}{args[1]}, - }, - { - name: "ignores placeholders in tagged dollar quoted strings", - query: "SELECT $tag$ $1 $tag$, $3", - wantQ: "SELECT $tag$ $1 $tag$, ?", - want: []interface{}{args[2]}, - }, - { - name: "placeholder beyond provided args is left for downstream validation", - query: "SELECT $4", - wantQ: "SELECT $4", - want: nil, - }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - gotQ, got := queryAndArgsForStatement(tt.query, args) - if gotQ != tt.wantQ { - t.Fatalf("queryAndArgsForStatement(%q) query = %q, want %q", tt.query, gotQ, tt.wantQ) - } - if !reflect.DeepEqual(got, tt.want) { - t.Fatalf("queryAndArgsForStatement(%q) args = %#v, want %#v", tt.query, got, tt.want) - } - }) - } -} - -func TestQueryAndArgsForStatementRunsSparseDollarParams(t *testing.T) { - db, err := sql.Open("duckdb", ":memory:") - if err != nil { - t.Fatalf("open duckdb: %v", err) - } - defer func() { _ = db.Close() }() - - query, args := queryAndArgsForStatement("SELECT $3 || '-' || $1", []interface{}{"first", "second", "third"}) - - var got string - if err := db.QueryRow(query, args...).Scan(&got); err != nil { - t.Fatalf("run rewritten sparse-placeholder query %q args=%v: %v", query, args, err) - } - if got != "third-first" { - t.Fatalf("rewritten sparse-placeholder query returned %q, want third-first", got) - } -} - -type cleanupErrorExecutor struct { - noopProfiling - execQueries []string -} - -func (e *cleanupErrorExecutor) QueryContext(context.Context, string, ...any) (RowSet, error) { - return nil, errors.New("not implemented") -} - -func (e *cleanupErrorExecutor) ExecContext(context.Context, string, ...any) (ExecResult, error) { - return nil, errors.New("not implemented") -} - -func (e *cleanupErrorExecutor) Query(string, ...any) (RowSet, error) { - return nil, errors.New("not implemented") -} - -func (e *cleanupErrorExecutor) Exec(query string, _ ...any) (ExecResult, error) { - e.execQueries = append(e.execQueries, query) - if strings.TrimSpace(strings.ToUpper(query)) == "COMMIT" { - return nil, errors.New("commit failed") - } - return &fakeExecResult{}, nil -} - -func (e *cleanupErrorExecutor) ConnContext(context.Context) (RawConn, error) { - return nil, errors.New("not implemented") -} - -func (e *cleanupErrorExecutor) PingContext(context.Context) error { - return errors.New("not implemented") -} - -func (e *cleanupErrorExecutor) Close() error { - return nil -} - -func TestExecuteCleanupAndCommitPropagatesCommitError(t *testing.T) { - exec := &cleanupErrorExecutor{} - c := &clientConn{executor: exec} - - err := c.executeCleanupAndCommit([]string{"DROP TABLE IF EXISTS _duckgres_tmp", "COMMIT"}) - if err == nil || !strings.Contains(err.Error(), "commit failed") { - t.Fatalf("executeCleanupAndCommit error = %v, want commit failed", err) - } - - want := []string{"DROP TABLE IF EXISTS _duckgres_tmp", "COMMIT", "ROLLBACK"} - if !reflect.DeepEqual(exec.execQueries, want) { - t.Fatalf("exec queries = %#v, want %#v", exec.execQueries, want) - } -} - func TestMatchPgStatActivityQuery(t *testing.T) { tests := []struct { name string diff --git a/tests/integration/ducklake_concurrency_test.go b/tests/integration/ducklake_concurrency_test.go index d9b41232..6fb48914 100644 --- a/tests/integration/ducklake_concurrency_test.go +++ b/tests/integration/ducklake_concurrency_test.go @@ -721,70 +721,31 @@ func runConcurrencyBenchmarks(t *testing.T, latencyMs int, openConn func(*testin t.Logf("read/write isolation: %d write conflicts, %d read errors", writeConflicts.Load(), readErrors.Load()) }) - benchSub(t, "concurrent_upsert_storm", latencyMs, func(t *testing.T, m *metric) { + benchSub(t, "on_conflict_rejected", latencyMs, func(t *testing.T, m *metric) { conn := openConn(t) defer func() { _ = conn.Close() }() - // DuckLake rewrites ON CONFLICT to MERGE — test it under concurrency mustExec(t, conn, "CREATE TABLE dl_conc_upsert (id INTEGER, counter INTEGER, last_ts BIGINT)") - // Seed some rows so upserts hit the update path - for i := range 20 { - mustExec(t, conn, fmt.Sprintf("INSERT INTO dl_conc_upsert VALUES (%d, 0, 0)", i)) - } defer func() { _, _ = conn.Exec("DROP TABLE IF EXISTS dl_conc_upsert") }() - const numWorkers = 8 - const opsPerWorker = 30 - var wg sync.WaitGroup - var conflicts atomic.Int64 - var successes atomic.Int64 - errs := make(chan error, numWorkers*opsPerWorker) - - for w := range numWorkers { - wg.Add(1) - go func(workerID int) { - defer wg.Done() - wconn := openConn(t) - defer func() { _ = wconn.Close() }() - - for i := range opsPerWorker { - // Target overlapping rows across workers - targetID := (workerID + i*3) % 30 // some IDs don't exist yet -> insert path - ts := time.Now().UnixMicro() - // Use literal values: ON CONFLICT → MERGE rewriting breaks $N placeholders - _, err := wconn.Exec(fmt.Sprintf( - `INSERT INTO dl_conc_upsert (id, counter, last_ts) VALUES (%d, 1, %d) - ON CONFLICT (id) DO UPDATE SET counter = dl_conc_upsert.counter + 1, last_ts = %d`, - targetID, ts, ts, - )) - if err != nil { - if isTransactionConflict(err) || isAbortedTransaction(err) { - conflicts.Add(1) - recoverConnection(wconn) - continue - } - // ON CONFLICT may not be supported in DuckLake — log and skip - if strings.Contains(err.Error(), "Not implemented") || - strings.Contains(err.Error(), "not supported") { - t.Logf("upsert not supported: %v", err) - return - } - errs <- fmt.Errorf("worker %d op %d: %w", workerID, i, err) - return - } - successes.Add(1) - } - }(w) + queries := []string{ + `INSERT INTO dl_conc_upsert (id, counter, last_ts) VALUES (1, 1, 1) + ON CONFLICT (id) DO UPDATE SET counter = dl_conc_upsert.counter + 1, last_ts = 1`, + `INSERT INTO dl_conc_upsert (id, counter, last_ts) VALUES (1, 1, 1) + ON CONFLICT (id) DO NOTHING`, } - - wg.Wait() - close(errs) - for err := range errs { - t.Error(err) + for _, query := range queries { + _, err := conn.Exec(query) + if err == nil { + t.Fatalf("expected ON CONFLICT to be rejected") + } + if !strings.Contains(err.Error(), "ON CONFLICT is not supported") { + t.Fatalf("expected ON CONFLICT unsupported error, got %v", err) + } } - m.Successes, m.Conflicts = successes.Load(), conflicts.Load() - t.Logf("upsert storm: %d succeeded, %d conflicts", m.Successes, m.Conflicts) + m.Successes = int64(len(queries)) + t.Logf("ON CONFLICT rejection: %d forms rejected", m.Successes) }) benchSub(t, "concurrent_ddl_while_writing", latencyMs, func(t *testing.T, m *metric) { diff --git a/transpiler/backend/profile.go b/transpiler/backend/profile.go index 662cdfe5..8caeea5b 100644 --- a/transpiler/backend/profile.go +++ b/transpiler/backend/profile.go @@ -20,7 +20,7 @@ const ( NoOpUnsupportedDDL UnsupportedDDLHandling = "noop" UseInsertOnConflict ConflictHandling = "insert_on_conflict" - RewriteToMerge ConflictHandling = "rewrite_to_merge" + RejectOnConflict ConflictHandling = "reject_on_conflict" ) type CatalogPolicy struct { @@ -138,7 +138,7 @@ func lakeProfile(name Name, physical string, mapPublicToMain, hybridDDLGuards bo WarnOnStrippedConstraints: hybridDDLGuards, ErrorOnSilentNullDefaults: hybridDDLGuards, }, - dml: DMLPolicy{ConflictHandling: RewriteToMerge}, + dml: DMLPolicy{ConflictHandling: RejectOnConflict}, metadata: MetadataPolicy{InterceptShowCreate: true}, } } diff --git a/transpiler/backend/profile_test.go b/transpiler/backend/profile_test.go index efa9208b..6bcb8076 100644 --- a/transpiler/backend/profile_test.go +++ b/transpiler/backend/profile_test.go @@ -13,8 +13,8 @@ func TestIcebergProfilePolicies(t *testing.T) { if profile.DDL().ConstraintHandling != StripConstraints { t.Fatalf("ConstraintHandling = %v, want StripConstraints", profile.DDL().ConstraintHandling) } - if profile.DML().ConflictHandling != RewriteToMerge { - t.Fatalf("ConflictHandling = %v, want RewriteToMerge", profile.DML().ConflictHandling) + if profile.DML().ConflictHandling != RejectOnConflict { + t.Fatalf("ConflictHandling = %v, want RejectOnConflict", profile.DML().ConflictHandling) } } diff --git a/transpiler/classify_wiring_test.go b/transpiler/classify_wiring_test.go index 5665096f..3e9805ee 100644 --- a/transpiler/classify_wiring_test.go +++ b/transpiler/classify_wiring_test.go @@ -39,9 +39,6 @@ func transpileSQL(t *testing.T, tr *Transpiler, input string) string { if res.FallbackToNative { t.Fatalf("Transpile(%q) fell back to native, want transpiled SQL", input) } - if len(res.Statements) > 0 { - return res.Statements[len(res.Statements)-1] - } return res.SQL } @@ -274,14 +271,17 @@ func TestClassifyWiring_OnConflictWhitespace(t *testing.T) { wantFlags(t, cfg, "INSERT INTO t (a, b) VALUES (1, 2) ON\tCONFLICT (a) DO NOTHING", FlagOnConflict) dlTr := New(Config{DuckLakeMode: true}) - out := transpileSQL(t, dlTr, "INSERT INTO t (a, b) VALUES (1, 2) ON\nCONFLICT (a) DO NOTHING") - if !strings.Contains(strings.ToUpper(out), "MERGE INTO") { - t.Errorf("DuckLake Transpile(ON\\nCONFLICT) = %q, want MERGE INTO rewrite", out) + res, err := dlTr.Transpile("INSERT INTO t (a, b) VALUES (1, 2) ON\nCONFLICT (a) DO NOTHING") + if err != nil { + t.Fatalf("DuckLake Transpile(ON\\nCONFLICT) error: %v", err) + } + if res.Error == nil { + t.Fatalf("DuckLake Transpile(ON\\nCONFLICT) should reject ON CONFLICT, got SQL=%q", res.SQL) } // False positive (CONFLICT in a string literal) must be a harmless no-op. tr := New(cfg) - out = transpileSQL(t, tr, "SELECT note FROM audit WHERE note = 'ON CONFLICT'") + out := transpileSQL(t, tr, "SELECT note FROM audit WHERE note = 'ON CONFLICT'") if strings.Contains(strings.ToUpper(out), "MERGE") { t.Errorf("Transpile(literal 'ON CONFLICT') = %q, must not rewrite", out) } diff --git a/transpiler/transform/onconflict.go b/transpiler/transform/onconflict.go index f9f8698d..2d051d82 100644 --- a/transpiler/transform/onconflict.go +++ b/transpiler/transform/onconflict.go @@ -1,46 +1,22 @@ package transform -import ( - "fmt" - "math/rand/v2" - "strings" - - pg_query "github.com/pganalyze/pg_query_go/v6" -) +import pg_query "github.com/pganalyze/pg_query_go/v6" // OnConflictTransform handles PostgreSQL ON CONFLICT (upsert) syntax. -// PostgreSQL: INSERT ... ON CONFLICT (columns) DO UPDATE SET ... -// PostgreSQL: INSERT ... ON CONFLICT (columns) DO NOTHING -// DuckDB: INSERT OR REPLACE INTO ... (for simple cases) -// DuckDB: INSERT OR IGNORE INTO ... (for DO NOTHING) -// -// Note: DuckDB's ON CONFLICT support has evolved. As of DuckDB 0.9+, -// it supports ON CONFLICT DO NOTHING and ON CONFLICT DO UPDATE. -// This transform handles cases where the syntax might differ. // -// In DuckLake mode, PRIMARY KEY and UNIQUE constraints are stripped, -// so ON CONFLICT clauses will fail with "columns not referenced by constraint". -// Instead of stripping ON CONFLICT, we convert the INSERT to a MERGE statement -// which DuckLake supports for upserts. +// DuckDB supports ON CONFLICT for catalogs that enforce unique constraints. +// Constraint-less lake catalogs do not, so they reject ON CONFLICT instead of +// trying to emulate PostgreSQL uniqueness semantics. type OnConflictTransform struct { - DuckLakeMode bool -} - -type onConflictMergeRewrite struct { - merge *pg_query.MergeStmt - sourceSelect *pg_query.SelectStmt - targetRelation *pg_query.RangeVar - insertColumns []string - conflictColumns []string - action pg_query.OnConflictAction + RejectOnConflict bool } func NewOnConflictTransform() *OnConflictTransform { - return &OnConflictTransform{DuckLakeMode: false} + return &OnConflictTransform{RejectOnConflict: false} } -func NewOnConflictTransformWithConfig(duckLakeMode bool) *OnConflictTransform { - return &OnConflictTransform{DuckLakeMode: duckLakeMode} +func NewOnConflictTransformWithConfig(rejectOnConflict bool) *OnConflictTransform { + return &OnConflictTransform{RejectOnConflict: rejectOnConflict} } func (t *OnConflictTransform) Name() string { @@ -48,691 +24,22 @@ func (t *OnConflictTransform) Name() string { } func (t *OnConflictTransform) Transform(tree *pg_query.ParseResult, result *Result) (bool, error) { - changed := false - for _, stmt := range tree.Stmts { if stmt.Stmt == nil { continue } - if insert := stmt.Stmt.GetInsertStmt(); insert != nil { - // On a constraint-less backend (MERGE rewrite enabled), ON CONFLICT ON - // CONSTRAINT cannot be honored: there is no named constraint to - // infer the conflict target from. Reject it with a clean PostgreSQL - // error rather than letting it fail opaquely at DuckDB. - if t.DuckLakeMode && insert.OnConflictClause != nil { - if unsupported := t.validateDuckLakeConflictTarget(insert.OnConflictClause); unsupported != nil { - result.Error = unsupported - return false, nil - } - if unsupported := t.validateDuckLakeInsertColumns(insert); unsupported != nil { - result.Error = unsupported - return false, nil - } - } - - if rewrite := t.transformInsertToMerge(insert); rewrite != nil { - // Replace the INSERT statement with MERGE - stmt.Stmt = &pg_query.Node{ - Node: &pg_query.Node_MergeStmt{MergeStmt: rewrite.merge}, - } - changed = true - - if statements, cleanup, err := t.buildGuardedMergeStatements(rewrite); err != nil { - return false, err - } else if len(statements) > 0 { - result.Statements = statements - result.CleanupStatements = cleanup - } - } else if t.transformInsert(insert) { - changed = true - } - } - } - - return changed, nil -} - -func (t *OnConflictTransform) validateDuckLakeConflictTarget(occ *pg_query.OnConflictClause) error { - if occ == nil { - return nil - } - if occ.Infer == nil { - return NewFeatureNotSupported( - "ON CONFLICT without a conflict target is not supported: this catalog does not enforce unique constraints") - } - if strings.TrimSpace(occ.Infer.Conname) != "" { - return NewFeatureNotSupported( - "ON CONFLICT ON CONSTRAINT is not supported: this catalog does not enforce named constraints") - } - if occ.Infer.WhereClause != nil { - return NewFeatureNotSupported( - "ON CONFLICT partial index predicates are not supported: this catalog does not enforce partial unique indexes") - } - for _, elem := range occ.Infer.IndexElems { - indexElem := elem.GetIndexElem() - if indexElem == nil || strings.TrimSpace(indexElem.Name) == "" { - return NewFeatureNotSupported( - "ON CONFLICT expression targets are not supported: use column names") - } - } - return nil -} - -func (t *OnConflictTransform) validateDuckLakeInsertColumns(insert *pg_query.InsertStmt) error { - if insert == nil || insert.OnConflictClause == nil { - return nil - } - if len(insert.Cols) == 0 { - return NewFeatureNotSupported( - "ON CONFLICT without an explicit insert column list is not supported") - } - insertColumns := make(map[string]struct{}, len(insert.Cols)) - for _, col := range insert.Cols { - rt := col.GetResTarget() - if rt == nil || strings.TrimSpace(rt.Name) == "" { - return NewFeatureNotSupported( - "ON CONFLICT without an explicit insert column list is not supported") - } - insertColumns[rt.Name] = struct{}{} - } - - conflictColumns := t.conflictColumnNames(insert.OnConflictClause.Infer.IndexElems) - for _, conflictColumn := range conflictColumns { - if _, ok := insertColumns[conflictColumn]; !ok { - return NewFeatureNotSupported( - "ON CONFLICT conflict columns must be present in the insert column list") - } - } - return nil -} - -// transformInsertToMerge converts INSERT ... ON CONFLICT to MERGE for DuckLake mode -func (t *OnConflictTransform) transformInsertToMerge(insert *pg_query.InsertStmt) *onConflictMergeRewrite { - if !t.DuckLakeMode { - return nil - } - - if insert == nil || insert.OnConflictClause == nil { - return nil - } - - occ := insert.OnConflictClause - - // Need conflict columns to build the join condition - if occ.Infer == nil || len(occ.Infer.IndexElems) == 0 { - return nil - } - conflictColumns := t.conflictColumnNames(occ.Infer.IndexElems) - - // Get column names from INSERT - colNames := make([]string, len(insert.Cols)) - for i, col := range insert.Cols { - if rt := col.GetResTarget(); rt != nil { - colNames[i] = rt.Name - } - } - - // Get the SelectStmt from INSERT - selectStmt := insert.SelectStmt.GetSelectStmt() - if selectStmt == nil { - return nil - } - - // Build source SELECT - handle both VALUES and SELECT ... FROM cases - var sourceSelect *pg_query.SelectStmt - if len(selectStmt.ValuesLists) > 0 { - // INSERT ... VALUES (...) ON CONFLICT - build SELECT from values - sourceSelect = t.buildSourceSelect(colNames, selectStmt.ValuesLists) - } else if len(selectStmt.FromClause) > 0 || len(selectStmt.TargetList) > 0 { - // INSERT ... SELECT ... FROM ... ON CONFLICT - use SELECT directly - // The SELECT already has the right columns, just use it as the source - sourceSelect = selectStmt - } - - if sourceSelect == nil { - return nil - } - - // Build source relation as a subquery with alias "excluded" - sourceRelation := &pg_query.Node{ - Node: &pg_query.Node_RangeSubselect{ - RangeSubselect: &pg_query.RangeSubselect{ - Subquery: &pg_query.Node{ - Node: &pg_query.Node_SelectStmt{SelectStmt: sourceSelect}, - }, - Alias: &pg_query.Alias{Aliasname: "excluded"}, - }, - }, - } - - // Build join condition from conflict columns - joinCondition := t.buildJoinCondition(occ.Infer.IndexElems, insert.Relation.Relname) - - // Build MERGE WHEN clauses - var whenClauses []*pg_query.Node - - // If DO UPDATE, add WHEN MATCHED THEN UPDATE - if occ.Action == pg_query.OnConflictAction_ONCONFLICT_UPDATE { - updateClause := t.buildUpdateClause(occ.TargetList, occ.WhereClause) - whenClauses = append(whenClauses, &pg_query.Node{ - Node: &pg_query.Node_MergeWhenClause{MergeWhenClause: updateClause}, - }) - } - // For DO NOTHING, we skip WHEN MATCHED (no action on match) - - // Always add WHEN NOT MATCHED THEN INSERT - insertClause := t.buildInsertClause(colNames) - whenClauses = append(whenClauses, &pg_query.Node{ - Node: &pg_query.Node_MergeWhenClause{MergeWhenClause: insertClause}, - }) - - merge := &pg_query.MergeStmt{ - Relation: insert.Relation, - SourceRelation: sourceRelation, - JoinCondition: joinCondition, - MergeWhenClauses: whenClauses, - } - - return &onConflictMergeRewrite{ - merge: merge, - sourceSelect: sourceSelect, - targetRelation: insert.Relation, - insertColumns: colNames, - conflictColumns: conflictColumns, - action: occ.Action, - } -} - -func (t *OnConflictTransform) conflictColumnNames(indexElems []*pg_query.Node) []string { - cols := make([]string, 0, len(indexElems)) - for _, elem := range indexElems { - indexElem := elem.GetIndexElem() - if indexElem == nil || strings.TrimSpace(indexElem.Name) == "" { - return nil - } - cols = append(cols, indexElem.Name) - } - return cols -} - -func (t *OnConflictTransform) buildGuardedMergeStatements(rewrite *onConflictMergeRewrite) ([]string, []string, error) { - if rewrite == nil || rewrite.merge == nil || rewrite.sourceSelect == nil || rewrite.targetRelation == nil { - return nil, nil, nil - } - if len(rewrite.conflictColumns) == 0 { - return nil, nil, nil - } - - sourceSQL, err := deparseStatement(&pg_query.Node{Node: &pg_query.Node_SelectStmt{SelectStmt: rewrite.sourceSelect}}) - if err != nil { - return nil, nil, err - } - sourceTableName := t.generateOnConflictSourceTableName(rewrite.targetRelation.Relname) - sourceRelation := quoteIdent(sourceTableName) - sourceFrom := sourceRelation + " AS duckgres_on_conflict_source" - - mergeSQL, err := deparseMergeWithTempSource(rewrite.merge, sourceTableName) - if err != nil { - return nil, nil, err - } - - statements := []string{ - "BEGIN", - t.buildCreateSourceTableStatement(sourceRelation, sourceSQL, rewrite), - } - if rewrite.action == pg_query.OnConflictAction_ONCONFLICT_UPDATE { - statements = append(statements, t.buildSourceDuplicateGuard(sourceFrom, rewrite.conflictColumns)) - statements = append(statements, t.buildTargetDuplicateGuard(sourceFrom, rewrite.targetRelation, rewrite.conflictColumns)) - } - statements = append(statements, - mergeSQL, - ) - cleanup := []string{ - "DROP TABLE IF EXISTS " + sourceRelation, - "COMMIT", - } - return statements, cleanup, nil -} - -func (t *OnConflictTransform) buildCreateSourceTableStatement(sourceRelation string, sourceSQL string, rewrite *onConflictMergeRewrite) string { - columnList := quotedColumnList(rewrite.insertColumns) - if rewrite.action != pg_query.OnConflictAction_ONCONFLICT_NOTHING { - return "CREATE TEMPORARY TABLE " + sourceRelation + " " + columnList + " AS " + sourceSQL - } - - return "CREATE TEMPORARY TABLE " + sourceRelation + " " + columnList + " AS " + - t.buildDoNothingSourceSQL(sourceSQL, rewrite.insertColumns, rewrite.conflictColumns) -} - -func (t *OnConflictTransform) buildDoNothingSourceSQL(sourceSQL string, insertColumns []string, conflictColumns []string) string { - rawAlias := "duckgres_on_conflict_raw" - rankedAlias := "duckgres_on_conflict_ranked" - rankColumn := "_duckgres_on_conflict_rank" - - insertRefs := qualifiedColumnRefs(rankedAlias, insertColumns) - conflictRefs := qualifiedColumnRefs(rawAlias, conflictColumns) - preservePredicates := make([]string, 0, len(conflictColumns)+1) - for _, ref := range qualifiedColumnRefs(rankedAlias, conflictColumns) { - preservePredicates = append(preservePredicates, ref+" IS NULL") - } - preservePredicates = append(preservePredicates, rankedAlias+"."+quoteIdent(rankColumn)+" = 1") - - return "SELECT " + strings.Join(insertRefs, ", ") + " FROM (" + - "SELECT " + rawAlias + ".*, row_number() OVER (PARTITION BY " + strings.Join(conflictRefs, ", ") + ") AS " + quoteIdent(rankColumn) + " " + - "FROM (" + sourceSQL + ") AS " + rawAlias + " " + quotedColumnList(insertColumns) + - ") AS " + rankedAlias + " " + - "WHERE " + strings.Join(preservePredicates, " OR ") -} - -func deparseStatement(stmt *pg_query.Node) (string, error) { - sql, err := pg_query.Deparse(&pg_query.ParseResult{ - Stmts: []*pg_query.RawStmt{{Stmt: stmt}}, - }) - if err != nil { - return "", err - } - return strings.TrimSuffix(strings.TrimSpace(sql), ";"), nil -} - -func deparseMergeWithTempSource(merge *pg_query.MergeStmt, sourceTableName string) (string, error) { - originalSourceRelation := merge.SourceRelation - defer func() { - merge.SourceRelation = originalSourceRelation - }() - merge.SourceRelation = &pg_query.Node{ - Node: &pg_query.Node_RangeVar{ - RangeVar: &pg_query.RangeVar{ - Relname: sourceTableName, - Alias: &pg_query.Alias{Aliasname: "excluded"}, - }, - }, - } - sql, err := deparseStatement(&pg_query.Node{Node: &pg_query.Node_MergeStmt{MergeStmt: merge}}) - return sql, err -} - -func (t *OnConflictTransform) buildSourceDuplicateGuard(sourceFrom string, conflictColumns []string) string { - sourceRefs := qualifiedColumnRefs("duckgres_on_conflict_source", conflictColumns) - notNullPredicates := notNullPredicates(sourceRefs) - - return "SELECT CASE WHEN EXISTS (" + - "SELECT 1 FROM " + sourceFrom + " " + - "WHERE " + notNullPredicates + " " + - "GROUP BY " + strings.Join(sourceRefs, ", ") + " " + - "HAVING count(*) > 1 LIMIT 1" + - ") THEN error(" + quoteStringLiteral("ON CONFLICT source rows contain duplicate conflict keys") + ") ELSE NULL END" -} - -func (t *OnConflictTransform) buildTargetDuplicateGuard(sourceFrom string, target *pg_query.RangeVar, conflictColumns []string) string { - sourceRefs := qualifiedColumnRefs("duckgres_on_conflict_source", conflictColumns) - keyRefs := qualifiedColumnRefs("duckgres_on_conflict_keys", conflictColumns) - targetRefs := qualifiedColumnRefs("duckgres_on_conflict_target", conflictColumns) - joinPredicates := make([]string, 0, len(conflictColumns)) - for i := range conflictColumns { - joinPredicates = append(joinPredicates, targetRefs[i]+" = "+keyRefs[i]) - } - - return "SELECT CASE WHEN EXISTS (" + - "SELECT 1 FROM " + formatRangeVar(target) + " AS duckgres_on_conflict_target " + - "JOIN (" + - "SELECT DISTINCT " + strings.Join(sourceRefs, ", ") + " " + - "FROM " + sourceFrom + " " + - "WHERE " + notNullPredicates(sourceRefs) + - ") AS duckgres_on_conflict_keys " + - "ON " + strings.Join(joinPredicates, " AND ") + " " + - "GROUP BY " + strings.Join(targetRefs, ", ") + " " + - "HAVING count(*) > 1 LIMIT 1" + - ") THEN error(" + quoteStringLiteral("ON CONFLICT target contains duplicate conflict keys") + ") ELSE NULL END" -} - -func (t *OnConflictTransform) generateOnConflictSourceTableName(targetName string) string { - safe := strings.Map(func(r rune) rune { - if (r >= 'a' && r <= 'z') || (r >= '0' && r <= '9') || r == '_' { - return r - } - if r >= 'A' && r <= 'Z' { - return r + 32 - } - return '_' - }, targetName) - if safe == "" { - safe = "source" - } - - prefix := "_duckgres_on_conflict_source_" - suffix := fmt.Sprintf("%08x", rand.Uint32()) - maxNameLen := maxIdentifierLength - len(prefix) - len(suffix) - 1 - if len(safe) > maxNameLen { - safe = safe[:maxNameLen] - } - return fmt.Sprintf("%s%s_%s", prefix, safe, suffix) -} - -func qualifiedColumnRefs(alias string, columns []string) []string { - refs := make([]string, len(columns)) - for i, col := range columns { - refs[i] = alias + "." + quoteIdent(col) - } - return refs -} - -func quotedColumnList(columns []string) string { - quoted := make([]string, len(columns)) - for i, col := range columns { - quoted[i] = quoteIdent(col) - } - return "(" + strings.Join(quoted, ", ") + ")" -} - -func notNullPredicates(refs []string) string { - predicates := make([]string, len(refs)) - for i, ref := range refs { - predicates[i] = ref + " IS NOT NULL" - } - return strings.Join(predicates, " AND ") -} - -func formatRangeVar(rv *pg_query.RangeVar) string { - parts := make([]string, 0, 3) - if rv.Catalogname != "" { - parts = append(parts, quoteIdent(rv.Catalogname)) - } - if rv.Schemaname != "" { - parts = append(parts, quoteIdent(rv.Schemaname)) - } - parts = append(parts, quoteIdent(rv.Relname)) - return strings.Join(parts, ".") -} - -func quoteIdent(ident string) string { - return `"` + strings.ReplaceAll(ident, `"`, `""`) + `"` -} - -func quoteStringLiteral(s string) string { - return `'` + strings.ReplaceAll(s, `'`, `''`) + `'` -} - -// buildSourceSelect creates a SELECT statement from VALUES for use as MERGE source -func (t *OnConflictTransform) buildSourceSelect(colNames []string, valuesLists []*pg_query.Node) *pg_query.SelectStmt { - if len(valuesLists) == 0 { - return nil - } - - if len(valuesLists) == 1 { - // Single row: SELECT val1 AS col1, val2 AS col2, ... - valueList := valuesLists[0].GetList() - if valueList == nil || len(valueList.Items) != len(colNames) { - return nil - } - - targetList := make([]*pg_query.Node, len(colNames)) - for i, colName := range colNames { - targetList[i] = &pg_query.Node{ - Node: &pg_query.Node_ResTarget{ - ResTarget: &pg_query.ResTarget{ - Name: colName, - Val: valueList.Items[i], - }, - }, - } - } - - return &pg_query.SelectStmt{ - TargetList: targetList, - LimitOption: pg_query.LimitOption_LIMIT_OPTION_DEFAULT, - Op: pg_query.SetOperation_SETOP_NONE, - } - } - - // Multiple rows: Use UNION ALL of SELECT statements - // First row - firstList := valuesLists[0].GetList() - if firstList == nil || len(firstList.Items) != len(colNames) { - return nil - } - - targetList := make([]*pg_query.Node, len(colNames)) - for i, colName := range colNames { - targetList[i] = &pg_query.Node{ - Node: &pg_query.Node_ResTarget{ - ResTarget: &pg_query.ResTarget{ - Name: colName, - Val: firstList.Items[i], - }, - }, - } - } - - result := &pg_query.SelectStmt{ - TargetList: targetList, - LimitOption: pg_query.LimitOption_LIMIT_OPTION_DEFAULT, - Op: pg_query.SetOperation_SETOP_NONE, - } - - // Add remaining rows as UNION ALL - for i := 1; i < len(valuesLists); i++ { - valueList := valuesLists[i].GetList() - if valueList == nil || len(valueList.Items) != len(colNames) { - continue - } - - rightTargetList := make([]*pg_query.Node, len(colNames)) - for j, colName := range colNames { - rightTargetList[j] = &pg_query.Node{ - Node: &pg_query.Node_ResTarget{ - ResTarget: &pg_query.ResTarget{ - Name: colName, - Val: valueList.Items[j], - }, - }, - } - } - - rightSelect := &pg_query.SelectStmt{ - TargetList: rightTargetList, - LimitOption: pg_query.LimitOption_LIMIT_OPTION_DEFAULT, - Op: pg_query.SetOperation_SETOP_NONE, - } - - result = &pg_query.SelectStmt{ - Op: pg_query.SetOperation_SETOP_UNION, - All: true, - Larg: result, - Rarg: rightSelect, - LimitOption: pg_query.LimitOption_LIMIT_OPTION_DEFAULT, - } - } - - return result -} - -// buildJoinCondition creates the ON condition for MERGE -func (t *OnConflictTransform) buildJoinCondition(indexElems []*pg_query.Node, tableName string) *pg_query.Node { - if len(indexElems) == 0 { - return nil - } - - // Build equality conditions for each conflict column - var conditions []*pg_query.Node - for _, elem := range indexElems { - indexElem := elem.GetIndexElem() - if indexElem == nil { + insert := stmt.Stmt.GetInsertStmt() + if insert == nil || insert.OnConflictClause == nil { continue } - colName := indexElem.Name - - // excluded.col = table.col - condition := &pg_query.Node{ - Node: &pg_query.Node_AExpr{ - AExpr: &pg_query.A_Expr{ - Kind: pg_query.A_Expr_Kind_AEXPR_OP, - Name: []*pg_query.Node{ - {Node: &pg_query.Node_String_{String_: &pg_query.String{Sval: "="}}}, - }, - Lexpr: &pg_query.Node{ - Node: &pg_query.Node_ColumnRef{ - ColumnRef: &pg_query.ColumnRef{ - Fields: []*pg_query.Node{ - {Node: &pg_query.Node_String_{String_: &pg_query.String{Sval: "excluded"}}}, - {Node: &pg_query.Node_String_{String_: &pg_query.String{Sval: colName}}}, - }, - }, - }, - }, - Rexpr: &pg_query.Node{ - Node: &pg_query.Node_ColumnRef{ - ColumnRef: &pg_query.ColumnRef{ - Fields: []*pg_query.Node{ - {Node: &pg_query.Node_String_{String_: &pg_query.String{Sval: tableName}}}, - {Node: &pg_query.Node_String_{String_: &pg_query.String{Sval: colName}}}, - }, - }, - }, - }, - }, - }, - } - conditions = append(conditions, condition) - } - - if len(conditions) == 1 { - return conditions[0] - } - // Multiple columns: combine with AND - result := conditions[0] - for i := 1; i < len(conditions); i++ { - result = &pg_query.Node{ - Node: &pg_query.Node_BoolExpr{ - BoolExpr: &pg_query.BoolExpr{ - Boolop: pg_query.BoolExprType_AND_EXPR, - Args: []*pg_query.Node{result, conditions[i]}, - }, - }, + if t.RejectOnConflict { + result.Error = NewFeatureNotSupported( + "ON CONFLICT is not supported: this catalog does not enforce unique constraints") + return false, nil } } - return result -} - -// buildUpdateClause creates WHEN MATCHED THEN UPDATE clause -func (t *OnConflictTransform) buildUpdateClause(targetList []*pg_query.Node, whereClause *pg_query.Node) *pg_query.MergeWhenClause { - clause := &pg_query.MergeWhenClause{ - MatchKind: pg_query.MergeMatchKind_MERGE_WHEN_MATCHED, - CommandType: pg_query.CmdType_CMD_UPDATE, - } - - if len(targetList) > 0 { - // Specific columns to update - SET col = val, ... - clause.TargetList = targetList - } - // If targetList is empty, it's a full row update (UPDATE without SET) - - if whereClause != nil { - clause.Condition = whereClause - } - - return clause + return false, nil } - -// buildInsertClause creates WHEN NOT MATCHED THEN INSERT clause -func (t *OnConflictTransform) buildInsertClause(colNames []string) *pg_query.MergeWhenClause { - // Build target list (column names) - targetList := make([]*pg_query.Node, len(colNames)) - for i, colName := range colNames { - targetList[i] = &pg_query.Node{ - Node: &pg_query.Node_ResTarget{ - ResTarget: &pg_query.ResTarget{ - Name: colName, - }, - }, - } - } - - // Build values list (references to excluded.col) - values := make([]*pg_query.Node, len(colNames)) - for i, colName := range colNames { - values[i] = &pg_query.Node{ - Node: &pg_query.Node_ColumnRef{ - ColumnRef: &pg_query.ColumnRef{ - Fields: []*pg_query.Node{ - {Node: &pg_query.Node_String_{String_: &pg_query.String{Sval: "excluded"}}}, - {Node: &pg_query.Node_String_{String_: &pg_query.String{Sval: colName}}}, - }, - }, - }, - } - } - - return &pg_query.MergeWhenClause{ - MatchKind: pg_query.MergeMatchKind_MERGE_WHEN_NOT_MATCHED_BY_TARGET, - CommandType: pg_query.CmdType_CMD_INSERT, - TargetList: targetList, - Values: values, - } -} - -func (t *OnConflictTransform) transformInsert(insert *pg_query.InsertStmt) bool { - if insert == nil || insert.OnConflictClause == nil { - return false - } - - // DuckDB now supports ON CONFLICT syntax similar to PostgreSQL - // However, there are some differences: - // - // 1. PostgreSQL allows ON CONFLICT ON CONSTRAINT constraint_name - // DuckDB only supports ON CONFLICT (columns) - // - // 2. PostgreSQL has EXCLUDED pseudo-table - // DuckDB also supports EXCLUDED - // - // 3. PostgreSQL allows WHERE clause in ON CONFLICT - // DuckDB supports this too (as of recent versions) - // - // For most common cases, the syntax is compatible. - // We mainly need to handle edge cases. - - occ := insert.OnConflictClause - - // Check for ON CONFLICT ON CONSTRAINT (not supported in DuckDB) - // The constraint name is in the Infer clause's Conname field - if occ.Infer != nil && occ.Infer.Conname != "" { - // This is ON CONFLICT ON CONSTRAINT constraint_name - // DuckDB doesn't support this syntax - // We could try to look up the constraint columns, but that's complex - // For now, we leave it as-is and let it error with a clear message - return false - } - - // ON CONFLICT (columns) DO NOTHING - supported in DuckDB - // ON CONFLICT (columns) DO UPDATE - supported in DuckDB - // These should work as-is - - return false -} - -// OnConflictNote documents the ON CONFLICT syntax differences: -// -// PostgreSQL: -// INSERT INTO t (a, b) VALUES (1, 2) -// ON CONFLICT (a) DO UPDATE SET b = EXCLUDED.b -// -// INSERT INTO t (a, b) VALUES (1, 2) -// ON CONFLICT (a) DO NOTHING -// -// INSERT INTO t (a, b) VALUES (1, 2) -// ON CONFLICT ON CONSTRAINT pk_name DO UPDATE SET b = EXCLUDED.b -// -// DuckDB (supported): -// INSERT INTO t (a, b) VALUES (1, 2) -// ON CONFLICT (a) DO UPDATE SET b = EXCLUDED.b -// -// INSERT INTO t (a, b) VALUES (1, 2) -// ON CONFLICT (a) DO NOTHING -// -// INSERT OR REPLACE INTO t (a, b) VALUES (1, 2) -- Alternative syntax -// INSERT OR IGNORE INTO t (a, b) VALUES (1, 2) -- Alternative syntax -// -// DuckDB (NOT supported): -// ON CONFLICT ON CONSTRAINT constraint_name -- Use column list instead diff --git a/transpiler/transform/onconflict_test.go b/transpiler/transform/onconflict_test.go index 164b1305..6c1c01c2 100644 --- a/transpiler/transform/onconflict_test.go +++ b/transpiler/transform/onconflict_test.go @@ -14,8 +14,7 @@ func TestOnConflictTransform_Name(t *testing.T) { } } -func TestOnConflictTransform_NonDuckLakeMode(t *testing.T) { - // In non-DuckLake mode, ON CONFLICT should pass through unchanged +func TestOnConflictTransform_NonRejectingMode(t *testing.T) { tr := NewOnConflictTransformWithConfig(false) tests := []struct { @@ -30,6 +29,14 @@ func TestOnConflictTransform_NonDuckLakeMode(t *testing.T) { name: "DO NOTHING passes through", input: "INSERT INTO users (id, name) VALUES (1, 'test') ON CONFLICT (id) DO NOTHING", }, + { + name: "ON CONSTRAINT passes through", + input: "INSERT INTO users (id, name) VALUES (1, 'test') ON CONFLICT ON CONSTRAINT users_pkey DO NOTHING", + }, + { + name: "targetless DO NOTHING passes through", + input: "INSERT INTO users (id, name) VALUES (1, 'test') ON CONFLICT DO NOTHING", + }, } for _, tt := range tests { @@ -45,12 +52,12 @@ func TestOnConflictTransform_NonDuckLakeMode(t *testing.T) { t.Fatalf("Transform error: %v", err) } - // Should not change in non-DuckLake mode if changed { - t.Error("Transform should not change SQL in non-DuckLake mode") + t.Error("Transform should not change SQL when ON CONFLICT rejection is disabled") + } + if result.Error != nil { + t.Fatalf("Transform returned unexpected error: %v", result.Error) } - - // Should still be an INSERT statement if tree.Stmts[0].Stmt.GetInsertStmt() == nil { t.Error("Statement should remain an INSERT") } @@ -58,168 +65,48 @@ func TestOnConflictTransform_NonDuckLakeMode(t *testing.T) { } } -func TestOnConflictTransform_DuckLakeMode_DoUpdate(t *testing.T) { +func TestOnConflictTransform_RejectingModeRejectsAllOnConflict(t *testing.T) { tr := NewOnConflictTransformWithConfig(true) tests := []struct { - name string - input string - wantContains []string - wantNotContains []string + name string + input string }{ { - name: "single row DO UPDATE", + name: "DO UPDATE", input: "INSERT INTO users (id, name) VALUES (1, 'test') ON CONFLICT (id) DO UPDATE SET name = EXCLUDED.name", - wantContains: []string{ - "MERGE INTO users", - "USING (SELECT 1 AS id, 'test' AS name)", - "excluded", - "ON excluded.id = users.id", - "WHEN MATCHED THEN UPDATE SET name = excluded.name", - "WHEN NOT MATCHED THEN INSERT", - }, - wantNotContains: []string{ - "ON CONFLICT", - "EXCLUDED", // should be lowercase "excluded" - }, }, { - name: "multiple columns in SET", - input: "INSERT INTO users (id, name, email) VALUES (1, 'test', 'test@example.com') ON CONFLICT (id) DO UPDATE SET name = EXCLUDED.name, email = EXCLUDED.email", - wantContains: []string{ - "MERGE INTO users", - "SELECT 1 AS id, 'test' AS name, 'test@example.com' AS email", - "WHEN MATCHED THEN UPDATE SET name = excluded.name, email = excluded.email", - }, + name: "DO NOTHING", + input: "INSERT INTO users (id, name) VALUES (1, 'test') ON CONFLICT (id) DO NOTHING", }, { - name: "multiple values become UNION ALL", - input: "INSERT INTO users (id, name) VALUES (1, 'first'), (2, 'second'), (3, 'third') ON CONFLICT (id) DO UPDATE SET name = EXCLUDED.name", - wantContains: []string{ - "MERGE INTO users", - "UNION ALL", - "WHEN MATCHED THEN UPDATE", - "WHEN NOT MATCHED THEN INSERT", - }, + name: "multiple values", + input: "INSERT INTO users (id, name) VALUES (1, 'first'), (2, 'second') ON CONFLICT (id) DO UPDATE SET name = EXCLUDED.name", }, { - name: "composite key conflict", - input: "INSERT INTO events (org_id, event_id, data) VALUES (1, 100, 'payload') ON CONFLICT (org_id, event_id) DO UPDATE SET data = EXCLUDED.data", - wantContains: []string{ - "MERGE INTO events", - "excluded.org_id = events.org_id", - "excluded.event_id = events.event_id", - "AND", - }, + name: "INSERT SELECT", + input: "INSERT INTO users (id, name) SELECT id, name FROM staging_users ON CONFLICT (id) DO UPDATE SET name = EXCLUDED.name", }, { - name: "INSERT SELECT FROM staging table (Fivetran pattern)", - input: `INSERT INTO users (id, name, email) SELECT id, name, email FROM staging_table ON CONFLICT (id) DO UPDATE SET name = EXCLUDED.name, email = EXCLUDED.email`, - wantContains: []string{ - "MERGE INTO users", - "USING (SELECT id, name, email FROM staging_table)", - "excluded", - "ON excluded.id = users.id", - "WHEN MATCHED THEN UPDATE SET name = excluded.name, email = excluded.email", - "WHEN NOT MATCHED THEN INSERT", - }, - wantNotContains: []string{ - "ON CONFLICT", - }, + name: "ON CONSTRAINT", + input: "INSERT INTO users (id, name) VALUES (1, 'test') ON CONFLICT ON CONSTRAINT users_pkey DO NOTHING", }, { - name: "INSERT SELECT with schema-qualified staging table", - input: `INSERT INTO "myschema"."users" (id, name) SELECT "id", "name" FROM "myschema_staging"."temp_users" ON CONFLICT (id) DO UPDATE SET name = EXCLUDED.name`, - wantContains: []string{ - "MERGE INTO myschema.users", - "FROM myschema_staging.temp_users", - "WHEN MATCHED THEN UPDATE", - "WHEN NOT MATCHED THEN INSERT", - }, - wantNotContains: []string{ - "ON CONFLICT", - }, + name: "expression target", + input: "INSERT INTO users (id, email) VALUES (1, 'a@example.com') ON CONFLICT ((lower(email))) DO NOTHING", }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - tree, err := pg_query.Parse(tt.input) - if err != nil { - t.Fatalf("Parse error: %v", err) - } - - result := &Result{} - changed, err := tr.Transform(tree, result) - if err != nil { - t.Fatalf("Transform error: %v", err) - } - - if !changed { - t.Error("Transform should change SQL in DuckLake mode") - } - - // Should now be a MERGE statement - if tree.Stmts[0].Stmt.GetMergeStmt() == nil { - t.Error("Statement should be converted to MERGE") - } - - // Deparse and check contents - sql, err := pg_query.Deparse(tree) - if err != nil { - t.Fatalf("Deparse error: %v", err) - } - - for _, want := range tt.wantContains { - if !strings.Contains(sql, want) { - t.Errorf("SQL should contain %q\nGot: %s", want, sql) - } - } - - for _, notWant := range tt.wantNotContains { - if strings.Contains(sql, notWant) { - t.Errorf("SQL should NOT contain %q\nGot: %s", notWant, sql) - } - } - }) - } -} - -func TestOnConflictTransform_DuckLakeMode_DoNothing(t *testing.T) { - tr := NewOnConflictTransformWithConfig(true) - - tests := []struct { - name string - input string - wantContains []string - wantNotContains []string - }{ { - name: "single row DO NOTHING", - input: "INSERT INTO users (id, name) VALUES (1, 'test') ON CONFLICT (id) DO NOTHING", - wantContains: []string{ - "MERGE INTO users", - "USING (SELECT 1 AS id, 'test' AS name)", - "ON excluded.id = users.id", - "WHEN NOT MATCHED THEN INSERT", - }, - wantNotContains: []string{ - "ON CONFLICT", - "WHEN MATCHED THEN UPDATE", // DO NOTHING should not have UPDATE clause - }, + name: "partial target", + input: "INSERT INTO users (id, email, active) VALUES (1, 'a@example.com', true) ON CONFLICT (email) WHERE active DO NOTHING", + }, + { + name: "targetless DO NOTHING", + input: "INSERT INTO users (id, name) VALUES (1, 'test') ON CONFLICT DO NOTHING", }, { - name: "multiple values DO NOTHING", - input: "INSERT INTO users (id, name) VALUES (1, 'first'), (2, 'second') ON CONFLICT (id) DO NOTHING", - wantContains: []string{ - "MERGE INTO users", - "UNION ALL", - "WHEN NOT MATCHED THEN INSERT", - }, - wantNotContains: []string{ - "ON CONFLICT", - "WHEN MATCHED", - }, + name: "DEFAULT VALUES", + input: "INSERT INTO users DEFAULT VALUES ON CONFLICT DO NOTHING", }, } @@ -236,25 +123,17 @@ func TestOnConflictTransform_DuckLakeMode_DoNothing(t *testing.T) { t.Fatalf("Transform error: %v", err) } - if !changed { - t.Error("Transform should change SQL in DuckLake mode") + if changed { + t.Error("Transform should reject without rewriting SQL") } - - sql, err := pg_query.Deparse(tree) - if err != nil { - t.Fatalf("Deparse error: %v", err) + if result.Error == nil { + t.Fatal("expected ON CONFLICT to be rejected") } - - for _, want := range tt.wantContains { - if !strings.Contains(sql, want) { - t.Errorf("SQL should contain %q\nGot: %s", want, sql) - } + if got := transformErrSQLState(result.Error); got != "0A000" { + t.Fatalf("SQLSTATE = %q, want 0A000", got) } - - for _, notWant := range tt.wantNotContains { - if strings.Contains(sql, notWant) { - t.Errorf("SQL should NOT contain %q\nGot: %s", notWant, sql) - } + if !strings.Contains(result.Error.Error(), "ON CONFLICT is not supported") { + t.Fatalf("error = %q, want ON CONFLICT unsupported message", result.Error.Error()) } }) } @@ -301,500 +180,9 @@ func TestOnConflictTransform_NoTransformCases(t *testing.T) { if changed { t.Error("Transform should not change SQL without ON CONFLICT") } - }) - } -} - -func TestOnConflictTransform_PreservesTableName(t *testing.T) { - tr := NewOnConflictTransformWithConfig(true) - - tests := []struct { - name string - input string - tableName string - }{ - { - name: "simple table name", - input: "INSERT INTO users (id, name) VALUES (1, 'test') ON CONFLICT (id) DO UPDATE SET name = EXCLUDED.name", - tableName: "users", - }, - { - name: "different table name", - input: "INSERT INTO events (id, data) VALUES (1, 'payload') ON CONFLICT (id) DO UPDATE SET data = EXCLUDED.data", - tableName: "events", - }, - { - name: "table name with underscore", - input: "INSERT INTO user_events (id, event_type) VALUES (1, 'click') ON CONFLICT (id) DO UPDATE SET event_type = EXCLUDED.event_type", - tableName: "user_events", - }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - tree, err := pg_query.Parse(tt.input) - if err != nil { - t.Fatalf("Parse error: %v", err) - } - - result := &Result{} - _, err = tr.Transform(tree, result) - if err != nil { - t.Fatalf("Transform error: %v", err) - } - - mergeStmt := tree.Stmts[0].Stmt.GetMergeStmt() - if mergeStmt == nil { - t.Fatal("Expected MERGE statement") - return - } - - if mergeStmt.Relation.Relname != tt.tableName { - t.Errorf("Table name = %q, want %q", mergeStmt.Relation.Relname, tt.tableName) - } - }) - } -} - -func TestOnConflictTransform_MergeWhenClauses(t *testing.T) { - tr := NewOnConflictTransformWithConfig(true) - - t.Run("DO UPDATE has both WHEN clauses", func(t *testing.T) { - input := "INSERT INTO users (id, name) VALUES (1, 'test') ON CONFLICT (id) DO UPDATE SET name = EXCLUDED.name" - tree, err := pg_query.Parse(input) - if err != nil { - t.Fatalf("Parse error: %v", err) - } - - result := &Result{} - _, err = tr.Transform(tree, result) - if err != nil { - t.Fatalf("Transform error: %v", err) - } - - mergeStmt := tree.Stmts[0].Stmt.GetMergeStmt() - if mergeStmt == nil { - t.Fatal("Expected MERGE statement") - return - } - - // Should have 2 WHEN clauses: MATCHED (UPDATE) and NOT MATCHED (INSERT) - if len(mergeStmt.MergeWhenClauses) != 2 { - t.Errorf("Expected 2 WHEN clauses, got %d", len(mergeStmt.MergeWhenClauses)) - } - - // First should be WHEN MATCHED - firstClause := mergeStmt.MergeWhenClauses[0].GetMergeWhenClause() - if firstClause.MatchKind != pg_query.MergeMatchKind_MERGE_WHEN_MATCHED { - t.Errorf("First clause should be WHEN MATCHED, got %v", firstClause.MatchKind) - } - if firstClause.CommandType != pg_query.CmdType_CMD_UPDATE { - t.Errorf("First clause command should be UPDATE, got %v", firstClause.CommandType) - } - - // Second should be WHEN NOT MATCHED - secondClause := mergeStmt.MergeWhenClauses[1].GetMergeWhenClause() - if secondClause.MatchKind != pg_query.MergeMatchKind_MERGE_WHEN_NOT_MATCHED_BY_TARGET { - t.Errorf("Second clause should be WHEN NOT MATCHED, got %v", secondClause.MatchKind) - } - if secondClause.CommandType != pg_query.CmdType_CMD_INSERT { - t.Errorf("Second clause command should be INSERT, got %v", secondClause.CommandType) - } - }) - - t.Run("DO NOTHING has only NOT MATCHED clause", func(t *testing.T) { - input := "INSERT INTO users (id, name) VALUES (1, 'test') ON CONFLICT (id) DO NOTHING" - tree, err := pg_query.Parse(input) - if err != nil { - t.Fatalf("Parse error: %v", err) - } - - result := &Result{} - _, err = tr.Transform(tree, result) - if err != nil { - t.Fatalf("Transform error: %v", err) - } - - mergeStmt := tree.Stmts[0].Stmt.GetMergeStmt() - if mergeStmt == nil { - t.Fatal("Expected MERGE statement") - return - } - - // Should have only 1 WHEN clause: NOT MATCHED (INSERT) - if len(mergeStmt.MergeWhenClauses) != 1 { - t.Errorf("Expected 1 WHEN clause, got %d", len(mergeStmt.MergeWhenClauses)) - } - - // Should be WHEN NOT MATCHED - clause := mergeStmt.MergeWhenClauses[0].GetMergeWhenClause() - if clause.MatchKind != pg_query.MergeMatchKind_MERGE_WHEN_NOT_MATCHED_BY_TARGET { - t.Errorf("Clause should be WHEN NOT MATCHED, got %v", clause.MatchKind) - } - if clause.CommandType != pg_query.CmdType_CMD_INSERT { - t.Errorf("Clause command should be INSERT, got %v", clause.CommandType) - } - }) -} - -func TestOnConflictTransform_SourceSubquery(t *testing.T) { - tr := NewOnConflictTransformWithConfig(true) - - t.Run("single row creates simple SELECT", func(t *testing.T) { - input := "INSERT INTO users (id, name) VALUES (1, 'test') ON CONFLICT (id) DO UPDATE SET name = EXCLUDED.name" - tree, err := pg_query.Parse(input) - if err != nil { - t.Fatalf("Parse error: %v", err) - } - - result := &Result{} - _, err = tr.Transform(tree, result) - if err != nil { - t.Fatalf("Transform error: %v", err) - } - - sql, _ := pg_query.Deparse(tree) - - // Should have a simple SELECT, not UNION ALL - if strings.Contains(sql, "UNION ALL") { - t.Errorf("Single row should not use UNION ALL: %s", sql) - } - if !strings.Contains(sql, "SELECT 1 AS id, 'test' AS name") { - t.Errorf("Should have SELECT with aliased columns: %s", sql) - } - }) - - t.Run("multiple rows creates UNION ALL", func(t *testing.T) { - input := "INSERT INTO users (id, name) VALUES (1, 'first'), (2, 'second') ON CONFLICT (id) DO UPDATE SET name = EXCLUDED.name" - tree, err := pg_query.Parse(input) - if err != nil { - t.Fatalf("Parse error: %v", err) - } - - result := &Result{} - _, err = tr.Transform(tree, result) - if err != nil { - t.Fatalf("Transform error: %v", err) - } - - sql, _ := pg_query.Deparse(tree) - - if !strings.Contains(sql, "UNION ALL") { - t.Errorf("Multiple rows should use UNION ALL: %s", sql) - } - }) - - t.Run("three rows creates nested UNION ALL", func(t *testing.T) { - input := "INSERT INTO users (id, name) VALUES (1, 'a'), (2, 'b'), (3, 'c') ON CONFLICT (id) DO UPDATE SET name = EXCLUDED.name" - tree, err := pg_query.Parse(input) - if err != nil { - t.Fatalf("Parse error: %v", err) - } - - result := &Result{} - _, err = tr.Transform(tree, result) - if err != nil { - t.Fatalf("Transform error: %v", err) - } - - sql, _ := pg_query.Deparse(tree) - - // Should have two UNION ALLs for three rows - count := strings.Count(sql, "UNION ALL") - if count != 2 { - t.Errorf("Three rows should have 2 UNION ALLs, got %d: %s", count, sql) - } - }) -} - -func TestOnConflictTransform_JoinCondition(t *testing.T) { - tr := NewOnConflictTransformWithConfig(true) - - t.Run("single column join", func(t *testing.T) { - input := "INSERT INTO users (id, name) VALUES (1, 'test') ON CONFLICT (id) DO UPDATE SET name = EXCLUDED.name" - tree, err := pg_query.Parse(input) - if err != nil { - t.Fatalf("Parse error: %v", err) - } - - result := &Result{} - _, err = tr.Transform(tree, result) - if err != nil { - t.Fatalf("Transform error: %v", err) - } - - sql, _ := pg_query.Deparse(tree) - - if !strings.Contains(sql, "ON excluded.id = users.id") { - t.Errorf("Should have single column join condition: %s", sql) - } - // Should NOT have AND for single column - if strings.Contains(sql, " AND ") { - t.Errorf("Single column should not have AND: %s", sql) - } - }) - - t.Run("two column join", func(t *testing.T) { - input := "INSERT INTO events (org_id, event_id, data) VALUES (1, 100, 'x') ON CONFLICT (org_id, event_id) DO UPDATE SET data = EXCLUDED.data" - tree, err := pg_query.Parse(input) - if err != nil { - t.Fatalf("Parse error: %v", err) - } - - result := &Result{} - _, err = tr.Transform(tree, result) - if err != nil { - t.Fatalf("Transform error: %v", err) - } - - sql, _ := pg_query.Deparse(tree) - - if !strings.Contains(sql, "excluded.org_id = events.org_id") { - t.Errorf("Should have org_id condition: %s", sql) - } - if !strings.Contains(sql, "excluded.event_id = events.event_id") { - t.Errorf("Should have event_id condition: %s", sql) - } - if !strings.Contains(sql, " AND ") { - t.Errorf("Two columns should be joined with AND: %s", sql) - } - }) - - t.Run("three column join", func(t *testing.T) { - input := "INSERT INTO data (a, b, c, val) VALUES (1, 2, 3, 'x') ON CONFLICT (a, b, c) DO UPDATE SET val = EXCLUDED.val" - tree, err := pg_query.Parse(input) - if err != nil { - t.Fatalf("Parse error: %v", err) - } - - result := &Result{} - _, err = tr.Transform(tree, result) - if err != nil { - t.Fatalf("Transform error: %v", err) - } - - sql, _ := pg_query.Deparse(tree) - - // Should have 2 ANDs for 3 columns - count := strings.Count(sql, " AND ") - if count != 2 { - t.Errorf("Three columns should have 2 ANDs, got %d: %s", count, sql) - } - }) -} - -func TestOnConflictTransform_ExcludedAlias(t *testing.T) { - tr := NewOnConflictTransformWithConfig(true) - - input := "INSERT INTO users (id, name, email) VALUES (1, 'test', 'test@example.com') ON CONFLICT (id) DO UPDATE SET name = EXCLUDED.name, email = EXCLUDED.email" - tree, err := pg_query.Parse(input) - if err != nil { - t.Fatalf("Parse error: %v", err) - } - - result := &Result{} - _, err = tr.Transform(tree, result) - if err != nil { - t.Fatalf("Transform error: %v", err) - } - - sql, _ := pg_query.Deparse(tree) - - // EXCLUDED should be converted to lowercase "excluded" (the alias) - if strings.Contains(sql, "EXCLUDED.") { - t.Errorf("EXCLUDED should be converted to excluded: %s", sql) - } - - // Should use excluded alias in source - if !strings.Contains(sql, ") excluded ON") { - t.Errorf("Source should have 'excluded' alias: %s", sql) - } - - // Should use excluded in UPDATE SET - if !strings.Contains(sql, "name = excluded.name") { - t.Errorf("UPDATE should reference excluded.name: %s", sql) - } - if !strings.Contains(sql, "email = excluded.email") { - t.Errorf("UPDATE should reference excluded.email: %s", sql) - } -} - -func TestOnConflictTransform_InsertClauseValues(t *testing.T) { - tr := NewOnConflictTransformWithConfig(true) - - input := "INSERT INTO users (id, name, email) VALUES (1, 'test', 'test@example.com') ON CONFLICT (id) DO UPDATE SET name = EXCLUDED.name" - tree, err := pg_query.Parse(input) - if err != nil { - t.Fatalf("Parse error: %v", err) - } - - result := &Result{} - _, err = tr.Transform(tree, result) - if err != nil { - t.Fatalf("Transform error: %v", err) - } - - sql, _ := pg_query.Deparse(tree) - - // WHEN NOT MATCHED THEN INSERT should reference all columns from excluded - if !strings.Contains(sql, "INSERT (id, name, email) VALUES (excluded.id, excluded.name, excluded.email)") { - t.Errorf("INSERT clause should reference all columns from excluded: %s", sql) - } -} - -func TestOnConflictTransform_DataTypes(t *testing.T) { - tr := NewOnConflictTransformWithConfig(true) - - tests := []struct { - name string - input string - check string - }{ - { - name: "integer values", - input: "INSERT INTO t (id, count) VALUES (1, 100) ON CONFLICT (id) DO UPDATE SET count = EXCLUDED.count", - check: "SELECT 1 AS id, 100 AS count", - }, - { - name: "string values", - input: "INSERT INTO t (id, name) VALUES (1, 'hello world') ON CONFLICT (id) DO UPDATE SET name = EXCLUDED.name", - check: "'hello world' AS name", - }, - { - name: "null values", - input: "INSERT INTO t (id, val) VALUES (1, NULL) ON CONFLICT (id) DO UPDATE SET val = EXCLUDED.val", - check: "NULL AS val", - }, - { - name: "negative numbers", - input: "INSERT INTO t (id, val) VALUES (1, -42) ON CONFLICT (id) DO UPDATE SET val = EXCLUDED.val", - check: "-42 AS val", - }, - { - name: "decimal numbers", - input: "INSERT INTO t (id, price) VALUES (1, 19.99) ON CONFLICT (id) DO UPDATE SET price = EXCLUDED.price", - check: "19.99 AS price", - }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - tree, err := pg_query.Parse(tt.input) - if err != nil { - t.Fatalf("Parse error: %v", err) - } - - result := &Result{} - _, err = tr.Transform(tree, result) - if err != nil { - t.Fatalf("Transform error: %v", err) - } - - sql, _ := pg_query.Deparse(tree) - - if !strings.Contains(sql, tt.check) { - t.Errorf("SQL should contain %q\nGot: %s", tt.check, sql) + if result.Error != nil { + t.Fatalf("Transform returned unexpected error: %v", result.Error) } }) } } - -func TestOnConflictTransform_FivetranPattern(t *testing.T) { - // Test the exact pattern Fivetran uses: INSERT INTO target SELECT FROM staging ON CONFLICT DO UPDATE - tr := NewOnConflictTransformWithConfig(true) - - // Simplified version of the Fivetran query - input := `INSERT INTO "stripe_test"."invoice" ("id", "status", "amount_due", "currency", "_fivetran_synced") - SELECT "id", "status", "amount_due", "currency", "_fivetran_synced" - FROM "stripe_test_staging"."temp_invoice" - ON CONFLICT ("id") DO UPDATE SET - "status" = "excluded"."status", - "amount_due" = "excluded"."amount_due", - "currency" = "excluded"."currency", - "_fivetran_synced" = "excluded"."_fivetran_synced"` - - tree, err := pg_query.Parse(input) - if err != nil { - t.Fatalf("Parse error: %v", err) - } - - result := &Result{} - changed, err := tr.Transform(tree, result) - if err != nil { - t.Fatalf("Transform error: %v", err) - } - - if !changed { - t.Error("Transform should change SQL in DuckLake mode") - } - - // Should be converted to MERGE - if tree.Stmts[0].Stmt.GetMergeStmt() == nil { - t.Fatal("Statement should be converted to MERGE") - } - - sql, err := pg_query.Deparse(tree) - if err != nil { - t.Fatalf("Deparse error: %v", err) - } - - // Verify key parts of the MERGE statement - checks := []string{ - "MERGE INTO stripe_test.invoice", - "USING (SELECT", - "FROM stripe_test_staging.temp_invoice", - "excluded", - "WHEN MATCHED THEN UPDATE SET", - "WHEN NOT MATCHED THEN INSERT", - } - - for _, check := range checks { - if !strings.Contains(sql, check) { - t.Errorf("SQL should contain %q\nGot: %s", check, sql) - } - } - - // Should NOT contain ON CONFLICT - if strings.Contains(sql, "ON CONFLICT") { - t.Errorf("SQL should NOT contain ON CONFLICT\nGot: %s", sql) - } - - t.Logf("Transformed SQL:\n%s", sql) -} - -func TestOnConflictTransform_InsertSelectDoNothing(t *testing.T) { - tr := NewOnConflictTransformWithConfig(true) - - input := `INSERT INTO target (id, data) SELECT id, data FROM source ON CONFLICT (id) DO NOTHING` - - tree, err := pg_query.Parse(input) - if err != nil { - t.Fatalf("Parse error: %v", err) - } - - result := &Result{} - changed, err := tr.Transform(tree, result) - if err != nil { - t.Fatalf("Transform error: %v", err) - } - - if !changed { - t.Error("Transform should change SQL in DuckLake mode") - } - - sql, err := pg_query.Deparse(tree) - if err != nil { - t.Fatalf("Deparse error: %v", err) - } - - // Should have MERGE with only WHEN NOT MATCHED (no UPDATE for DO NOTHING) - if !strings.Contains(sql, "MERGE INTO target") { - t.Errorf("Should have MERGE INTO target: %s", sql) - } - if !strings.Contains(sql, "WHEN NOT MATCHED THEN INSERT") { - t.Errorf("Should have WHEN NOT MATCHED: %s", sql) - } - if strings.Contains(sql, "WHEN MATCHED") { - t.Errorf("Should NOT have WHEN MATCHED for DO NOTHING: %s", sql) - } -} diff --git a/transpiler/transpiler.go b/transpiler/transpiler.go index ef78dca1..f57f2333 100644 --- a/transpiler/transpiler.go +++ b/transpiler/transpiler.go @@ -132,9 +132,9 @@ func New(cfg Config) *Transpiler { // 11. _pg_expandarray handling (PostgreSQL array expansion function used by JDBC) t.transforms = append(t.transforms, taggedTransform{FlagExpandArray, transform.NewExpandArrayTransform()}) - // 12. ON CONFLICT handling (rewrites ON CONFLICT to MERGE when the backend has - // no enforced unique constraints to infer against) - t.transforms = append(t.transforms, taggedTransform{FlagOnConflict, transform.NewOnConflictTransformWithConfig(dmlPolicy.ConflictHandling == backend.RewriteToMerge)}) + // 12. ON CONFLICT handling (rejects ON CONFLICT when the backend has no + // enforced unique constraints to infer against) + t.transforms = append(t.transforms, taggedTransform{FlagOnConflict, transform.NewOnConflictTransformWithConfig(dmlPolicy.ConflictHandling == backend.RejectOnConflict)}) // 13. Locking clause removal (FOR UPDATE, FOR SHARE, etc.) - DuckDB doesn't support these t.transforms = append(t.transforms, taggedTransform{FlagLocking, transform.NewLockingTransform()}) @@ -271,15 +271,11 @@ func (t *Transpiler) transpileWithFlags(sql string, flags TransformFlags) (*Resu // When a transform produces multiple statements, we skip remaining transforms // and return the statements directly. if len(transformResult.Statements) > 0 { - paramCount := transformResult.ParamCount - if t.config.ConvertPlaceholders { - paramCount = countParametersRegex(sql) - } return &Result{ SQL: sql, // Keep original for logging Statements: restoreLongIdentifiersAll(transformResult.Statements, longIdents), CleanupStatements: restoreLongIdentifiersAll(transformResult.CleanupStatements, longIdents), - ParamCount: paramCount, + ParamCount: transformResult.ParamCount, Warnings: transformResult.Warnings, }, nil } diff --git a/transpiler/transpiler_test.go b/transpiler/transpiler_test.go index a2dac332..b7f452f8 100644 --- a/transpiler/transpiler_test.go +++ b/transpiler/transpiler_test.go @@ -1,13 +1,11 @@ package transpiler import ( - "database/sql" "errors" "fmt" "strings" "testing" - _ "github.com/duckdb/duckdb-go/v2" "github.com/posthog/duckgres/transpiler/transform" ) @@ -428,62 +426,52 @@ func TestTranspile_FeatureNotSupported_Rejected(t *testing.T) { } } + onConflictQueries := []struct { + name string + input string + }{ + { + name: "ON CONFLICT DO UPDATE", + input: "INSERT INTO t (id, v) VALUES (1, 2) ON CONFLICT (id) DO UPDATE SET v = EXCLUDED.v", + }, + { + name: "ON CONFLICT DO NOTHING", + input: "INSERT INTO t (id, v) VALUES (1, 2) ON CONFLICT (id) DO NOTHING", + }, + { + name: "ON CONFLICT INSERT SELECT", + input: "INSERT INTO t (id, v) SELECT id, v FROM staging_t ON CONFLICT (id) DO UPDATE SET v = EXCLUDED.v", + }, + { + name: "ON CONFLICT ON CONSTRAINT", + input: "INSERT INTO t (id, v) VALUES (1, 2) ON CONFLICT ON CONSTRAINT t_pkey DO NOTHING", + }, + { + name: "ON CONFLICT expression target", + input: "INSERT INTO t (id, email) VALUES (1, 'a@example.com') ON CONFLICT ((lower(email))) DO NOTHING", + }, + { + name: "ON CONFLICT partial target", + input: "INSERT INTO t (id, email, active) VALUES (1, 'a@example.com', true) ON CONFLICT (email) WHERE active DO NOTHING", + }, + { + name: "ON CONFLICT targetless DO NOTHING", + input: "INSERT INTO t (id, v) VALUES (1, 2) ON CONFLICT DO NOTHING", + }, + } + for _, backend := range []StorageBackend{BackendDuckLake, BackendIceberg} { tr := New(Config{Backend: backend}) - t.Run(string(backend)+"/ON CONFLICT ON CONSTRAINT", func(t *testing.T) { - result, err := tr.Transpile( - "INSERT INTO t (id, v) VALUES (1, 2) ON CONFLICT ON CONSTRAINT t_pkey DO NOTHING") - if err != nil { - t.Fatalf("Transpile error: %v", err) - } - assertFeatureNotSupported(t, result) - }) - - t.Run(string(backend)+"/ON CONFLICT expression target", func(t *testing.T) { - result, err := tr.Transpile( - "INSERT INTO t (email, v) VALUES ('a@example.com', 2) ON CONFLICT ((lower(email))) DO UPDATE SET v = EXCLUDED.v") - if err != nil { - t.Fatalf("Transpile error: %v", err) - } - assertFeatureNotSupported(t, result) - }) - - t.Run(string(backend)+"/ON CONFLICT partial target", func(t *testing.T) { - result, err := tr.Transpile( - "INSERT INTO t (id, active, v) VALUES (1, true, 2) ON CONFLICT (id) WHERE active DO UPDATE SET v = EXCLUDED.v") - if err != nil { - t.Fatalf("Transpile error: %v", err) - } - assertFeatureNotSupported(t, result) - }) - - t.Run(string(backend)+"/ON CONFLICT targetless DO NOTHING", func(t *testing.T) { - result, err := tr.Transpile( - "INSERT INTO t (id, v) VALUES (1, 2) ON CONFLICT DO NOTHING") - if err != nil { - t.Fatalf("Transpile error: %v", err) - } - assertFeatureNotSupported(t, result) - }) - - t.Run(string(backend)+"/ON CONFLICT without explicit insert column list", func(t *testing.T) { - result, err := tr.Transpile( - "INSERT INTO t VALUES (1, 2) ON CONFLICT (id) DO UPDATE SET v = EXCLUDED.v") - if err != nil { - t.Fatalf("Transpile error: %v", err) - } - assertFeatureNotSupported(t, result) - }) - - t.Run(string(backend)+"/ON CONFLICT conflict column omitted from insert column list", func(t *testing.T) { - result, err := tr.Transpile( - "INSERT INTO t (v) VALUES (2) ON CONFLICT (id) DO NOTHING") - if err != nil { - t.Fatalf("Transpile error: %v", err) - } - assertFeatureNotSupported(t, result) - }) + for _, tt := range onConflictQueries { + t.Run(string(backend)+"/"+tt.name, func(t *testing.T) { + result, err := tr.Transpile(tt.input) + if err != nil { + t.Fatalf("Transpile error: %v", err) + } + assertFeatureNotSupported(t, result) + }) + } t.Run(string(backend)+"/ALTER COLUMN TYPE USING", func(t *testing.T) { result, err := tr.Transpile( @@ -2211,77 +2199,25 @@ func TestTranspile_OnConflict(t *testing.T) { } func TestTranspile_OnConflict_DuckLakeMode(t *testing.T) { - // In DuckLake mode, ON CONFLICT is converted to MERGE statement - // because PRIMARY KEY/UNIQUE constraints don't exist in DuckLake tests := []struct { - name string - input string - wantMerge bool - contains []string - notContains []string + name string + input string }{ { - name: "ON CONFLICT DO NOTHING becomes MERGE with only INSERT", - input: "INSERT INTO users (id, name) VALUES (1, 'test') ON CONFLICT (id) DO NOTHING", - wantMerge: true, - contains: []string{ - "CREATE TEMPORARY TABLE", - "SELECT 1 AS id, 'test' AS name", - "row_number() OVER", - "MERGE INTO users", - "excluded", - "ON excluded.id = users.id", - "WHEN NOT MATCHED THEN INSERT", - }, - notContains: []string{ - "ON CONFLICT", - "WHEN MATCHED THEN UPDATE", // DO NOTHING means no update action - }, + name: "ON CONFLICT DO NOTHING", + input: "INSERT INTO users (id, name) VALUES (1, 'test') ON CONFLICT (id) DO NOTHING", }, { - name: "ON CONFLICT DO UPDATE becomes MERGE with UPDATE and INSERT", - input: "INSERT INTO users (id, name) VALUES (1, 'test') ON CONFLICT (id) DO UPDATE SET name = EXCLUDED.name", - wantMerge: true, - contains: []string{ - "CREATE TEMPORARY TABLE", - "AS SELECT 1 AS id, 'test' AS name", - "MERGE INTO users", - "excluded", - "ON excluded.id = users.id", - "WHEN MATCHED THEN UPDATE SET name = excluded.name", - "WHEN NOT MATCHED THEN INSERT", - }, - notContains: []string{ - "ON CONFLICT", - }, + name: "ON CONFLICT DO UPDATE", + input: "INSERT INTO users (id, name) VALUES (1, 'test') ON CONFLICT (id) DO UPDATE SET name = EXCLUDED.name", }, { - name: "Multiple values become UNION ALL", - input: "INSERT INTO users (id, name) VALUES (1, 'test'), (2, 'test2') ON CONFLICT (id) DO UPDATE SET name = EXCLUDED.name", - wantMerge: true, - contains: []string{ - "MERGE INTO users", - "UNION ALL", - "WHEN MATCHED THEN UPDATE", - "WHEN NOT MATCHED THEN INSERT", - }, - notContains: []string{ - "ON CONFLICT", - }, + name: "multiple values", + input: "INSERT INTO users (id, name) VALUES (1, 'test'), (2, 'test2') ON CONFLICT (id) DO UPDATE SET name = EXCLUDED.name", }, { - name: "Multiple conflict columns", - input: "INSERT INTO users (id, org_id, name) VALUES (1, 100, 'test') ON CONFLICT (id, org_id) DO UPDATE SET name = EXCLUDED.name", - wantMerge: true, - contains: []string{ - "MERGE INTO users", - "excluded.id = users.id", - "excluded.org_id = users.org_id", - "AND", // Multiple conditions joined with AND - }, - notContains: []string{ - "ON CONFLICT", - }, + name: "multiple conflict columns", + input: "INSERT INTO users (id, org_id, name) VALUES (1, 100, 'test') ON CONFLICT (id, org_id) DO UPDATE SET name = EXCLUDED.name", }, } @@ -2294,418 +2230,20 @@ func TestTranspile_OnConflict_DuckLakeMode(t *testing.T) { t.Fatalf("Transpile(%q) error: %v", tt.input, err) } - sql := result.SQL - finalSQL := result.SQL - if len(result.Statements) > 0 { - sql = strings.Join(result.Statements, "\n") - finalSQL = result.Statements[len(result.Statements)-1] - } - - // Check MERGE statement is generated - if tt.wantMerge { - if !strings.Contains(strings.ToUpper(sql), "MERGE INTO") { - t.Errorf("Transpile(%q) = %q, should contain MERGE INTO", tt.input, sql) - } + if result.Error == nil { + t.Fatalf("expected ON CONFLICT to be rejected, got SQL=%q", result.SQL) } - - // Check expected contents - for _, want := range tt.contains { - if !strings.Contains(sql, want) { - t.Errorf("Transpile(%q) = %q, should contain %q", tt.input, sql, want) - } + coded, ok := result.Error.(interface{ SQLState() string }) + if !ok { + t.Fatalf("error does not expose SQLState: %T", result.Error) } - - // Check contents that should not be present - for _, notWant := range tt.notContains { - if strings.Contains(strings.ToUpper(finalSQL), strings.ToUpper(notWant)) { - t.Errorf("Transpile(%q) final SQL = %q, should NOT contain %q", tt.input, finalSQL, notWant) - } + if coded.SQLState() != "0A000" { + t.Fatalf("SQLState = %q, want 0A000", coded.SQLState()) } }) } } -func TestTranspile_OnConflict_DuckLakeMode_EmitsCardinalityGuards(t *testing.T) { - tr := New(Config{DuckLakeMode: true}) - - result, err := tr.Transpile(`INSERT INTO "stripe"."price" ("id", "nickname") - SELECT "id", "nickname" FROM "stripe_price-staging-test" - ON CONFLICT ("id") DO UPDATE SET "nickname" = "excluded"."nickname"`) - if err != nil { - t.Fatalf("Transpile error: %v", err) - } - - if len(result.Statements) != 5 { - t.Fatalf("got %d statements, want transaction, source materialization, source guard, target guard, merge: %v", len(result.Statements), result.Statements) - } - if len(result.CleanupStatements) != 2 || - !strings.Contains(result.CleanupStatements[0], "DROP TABLE IF EXISTS") || - result.CleanupStatements[1] != "COMMIT" { - t.Fatalf("cleanup = %v, want DROP TABLE then COMMIT for materialized source", result.CleanupStatements) - } - - checks := []struct { - stmt int - want []string - }{ - { - stmt: 0, - want: []string{"BEGIN"}, - }, - { - stmt: 1, - want: []string{ - "CREATE TEMPORARY TABLE", - "AS SELECT id, nickname FROM", - "stripe_price-staging-test", - }, - }, - { - stmt: 2, - want: []string{ - "duckgres_on_conflict_source", - "GROUP BY duckgres_on_conflict_source.\"id\"", - "HAVING count(*) > 1", - "ON CONFLICT source rows contain duplicate conflict keys", - }, - }, - { - stmt: 3, - want: []string{ - "FROM \"stripe\".\"price\" AS duckgres_on_conflict_target", - "SELECT DISTINCT duckgres_on_conflict_source.\"id\"", - "ON duckgres_on_conflict_target.\"id\" = duckgres_on_conflict_keys.\"id\"", - "ON CONFLICT target contains duplicate conflict keys", - }, - }, - { - stmt: 4, - want: []string{ - "MERGE INTO stripe.price", - "USING", - "excluded", - "WHEN MATCHED THEN UPDATE", - "WHEN NOT MATCHED THEN INSERT", - }, - }, - } - - for _, check := range checks { - for _, want := range check.want { - if !strings.Contains(result.Statements[check.stmt], want) { - t.Errorf("statement %d should contain %q\nGot: %s", check.stmt, want, result.Statements[check.stmt]) - } - } - } -} - -func TestTranspile_OnConflict_DuckLakeMode_GuardsPreserveParamCount(t *testing.T) { - tr := New(Config{DuckLakeMode: true, ConvertPlaceholders: true}) - - result, err := tr.Transpile(`INSERT INTO users (id, nickname) - VALUES ($1, $2) - ON CONFLICT (id) DO UPDATE SET nickname = excluded.nickname`) - if err != nil { - t.Fatalf("Transpile error: %v", err) - } - if len(result.Statements) != 5 { - t.Fatalf("expected guarded multi-statement rewrite, got statements=%v sql=%q", result.Statements, result.SQL) - } - if result.ParamCount != 2 { - t.Fatalf("ParamCount = %d, want 2", result.ParamCount) - } -} - -func runDuckDBStatements(db *sql.DB, statements []string, cleanup []string) error { - for _, stmt := range statements { - if _, err := db.Exec(stmt); err != nil { - for _, cleanupStmt := range cleanup { - _, _ = db.Exec(cleanupStmt) - } - return err - } - } - for _, cleanupStmt := range cleanup { - if _, err := db.Exec(cleanupStmt); err != nil { - return err - } - } - return nil -} - -func TestTranspile_OnConflict_DuckLakeMode_GuardsPreventDuplicateMerge(t *testing.T) { - tr := New(Config{DuckLakeMode: true}) - db, err := sql.Open("duckdb", ":memory:") - if err != nil { - t.Fatalf("open duckdb: %v", err) - } - defer func() { _ = db.Close() }() - - if _, err := db.Exec(`CREATE TABLE target (id VARCHAR, nickname VARCHAR)`); err != nil { - t.Fatalf("create target: %v", err) - } - if _, err := db.Exec(`CREATE TABLE staging (id VARCHAR, nickname VARCHAR)`); err != nil { - t.Fatalf("create staging: %v", err) - } - - runStatements := func(t *testing.T, statements []string, cleanup []string) error { - t.Helper() - for _, stmt := range statements { - if _, err := db.Exec(stmt); err != nil { - for _, cleanupStmt := range cleanup { - _, _ = db.Exec(cleanupStmt) - } - return err - } - } - for _, cleanupStmt := range cleanup { - if _, err := db.Exec(cleanupStmt); err != nil { - return err - } - } - return nil - } - - result, err := tr.Transpile(`INSERT INTO target (id, nickname) - SELECT id, nickname FROM staging - ON CONFLICT (id) DO UPDATE SET nickname = excluded.nickname`) - if err != nil { - t.Fatalf("Transpile source duplicate query: %v", err) - } - if len(result.Statements) == 0 { - t.Fatalf("expected guarded multi-statement rewrite, got single SQL: %s", result.SQL) - } - - if _, err := db.Exec(`INSERT INTO staging VALUES ('price_1', 'first'), ('price_1', 'second')`); err != nil { - t.Fatalf("seed duplicate staging rows: %v", err) - } - err = runStatements(t, result.Statements, result.CleanupStatements) - if err == nil || !strings.Contains(err.Error(), "ON CONFLICT source rows contain duplicate conflict keys") { - t.Fatalf("expected source duplicate guard error, got %v", err) - } - - var rows int - if err := db.QueryRow(`SELECT count(*) FROM target`).Scan(&rows); err != nil { - t.Fatalf("count target after source guard: %v", err) - } - if rows != 0 { - t.Fatalf("source guard should stop before MERGE inserts rows, got %d target rows", rows) - } - - if _, err := db.Exec(`DELETE FROM staging`); err != nil { - t.Fatalf("clear staging: %v", err) - } - if _, err := db.Exec(`INSERT INTO target VALUES ('price_1', 'old a'), ('price_1', 'old b')`); err != nil { - t.Fatalf("seed duplicate target rows: %v", err) - } - if _, err := db.Exec(`INSERT INTO staging VALUES ('price_1', 'new')`); err != nil { - t.Fatalf("seed staging row: %v", err) - } - - err = runStatements(t, result.Statements, result.CleanupStatements) - if err == nil || !strings.Contains(err.Error(), "ON CONFLICT target contains duplicate conflict keys") { - t.Fatalf("expected target duplicate guard error, got %v", err) - } - - var updatedRows int - if err := db.QueryRow(`SELECT count(*) FROM target WHERE nickname = 'new'`).Scan(&updatedRows); err != nil { - t.Fatalf("count updated target rows: %v", err) - } - if updatedRows != 0 { - t.Fatalf("target guard should stop before fanout update, got %d updated rows", updatedRows) - } - - if _, err := db.Exec(`DELETE FROM target`); err != nil { - t.Fatalf("clear target before success case: %v", err) - } - if _, err := db.Exec(`DELETE FROM staging`); err != nil { - t.Fatalf("clear staging before success case: %v", err) - } - if _, err := db.Exec(`INSERT INTO staging VALUES ('price_2', 'valid')`); err != nil { - t.Fatalf("seed valid staging row: %v", err) - } - - if err := runStatements(t, result.Statements, result.CleanupStatements); err != nil { - t.Fatalf("guarded merge should allow unique source/target keys: %v", err) - } - if err := db.QueryRow(`SELECT count(*) FROM target WHERE id = 'price_2' AND nickname = 'valid'`).Scan(&rows); err != nil { - t.Fatalf("count inserted target rows: %v", err) - } - if rows != 1 { - t.Fatalf("guarded merge should insert the unique staging row, got %d rows", rows) - } -} - -func TestTranspile_OnConflict_DuckLakeMode_InsertSelectAliasesSourceColumns(t *testing.T) { - tr := New(Config{DuckLakeMode: true}) - db, err := sql.Open("duckdb", ":memory:") - if err != nil { - t.Fatalf("open duckdb: %v", err) - } - defer func() { _ = db.Close() }() - - if _, err := db.Exec(`CREATE TABLE target (id VARCHAR, nickname VARCHAR)`); err != nil { - t.Fatalf("create target: %v", err) - } - if _, err := db.Exec(`CREATE TABLE staging (src_id VARCHAR, src_nickname VARCHAR)`); err != nil { - t.Fatalf("create staging: %v", err) - } - if _, err := db.Exec(`INSERT INTO staging VALUES ('price_1', 'valid')`); err != nil { - t.Fatalf("seed staging: %v", err) - } - - result, err := tr.Transpile(`INSERT INTO target (id, nickname) - SELECT src_id, src_nickname FROM staging - ON CONFLICT (id) DO UPDATE SET nickname = excluded.nickname`) - if err != nil { - t.Fatalf("Transpile insert-select query: %v", err) - } - if err := runDuckDBStatements(db, result.Statements, result.CleanupStatements); err != nil { - t.Fatalf("guarded merge should map SELECT columns positionally to insert columns: %v", err) - } - - var rows int - if err := db.QueryRow(`SELECT count(*) FROM target WHERE id = 'price_1' AND nickname = 'valid'`).Scan(&rows); err != nil { - t.Fatalf("count inserted target rows: %v", err) - } - if rows != 1 { - t.Fatalf("guarded merge should insert row using target column names, got %d rows", rows) - } -} - -func TestTranspile_OnConflict_DuckLakeMode_DoNothingDeduplicatesSourceKeys(t *testing.T) { - tr := New(Config{DuckLakeMode: true}) - db, err := sql.Open("duckdb", ":memory:") - if err != nil { - t.Fatalf("open duckdb: %v", err) - } - defer func() { _ = db.Close() }() - - if _, err := db.Exec(`CREATE TABLE target (id VARCHAR, nickname VARCHAR)`); err != nil { - t.Fatalf("create target: %v", err) - } - if _, err := db.Exec(`CREATE TABLE staging (id VARCHAR, nickname VARCHAR)`); err != nil { - t.Fatalf("create staging: %v", err) - } - if _, err := db.Exec(`INSERT INTO staging VALUES ('price_1', 'first'), ('price_1', 'second'), (NULL, 'null a'), (NULL, 'null b')`); err != nil { - t.Fatalf("seed staging: %v", err) - } - - result, err := tr.Transpile(`INSERT INTO target (id, nickname) - SELECT id, nickname FROM staging - ON CONFLICT (id) DO NOTHING`) - if err != nil { - t.Fatalf("Transpile do-nothing query: %v", err) - } - if err := runDuckDBStatements(db, result.Statements, result.CleanupStatements); err != nil { - t.Fatalf("DO NOTHING should deduplicate non-null source conflict keys instead of erroring: %v", err) - } - - var duplicateKeyRows int - if err := db.QueryRow(`SELECT count(*) FROM target WHERE id = 'price_1'`).Scan(&duplicateKeyRows); err != nil { - t.Fatalf("count deduplicated target rows: %v", err) - } - if duplicateKeyRows != 1 { - t.Fatalf("DO NOTHING should insert one non-null source row per conflict key, got %d", duplicateKeyRows) - } - - var nullRows int - if err := db.QueryRow(`SELECT count(*) FROM target WHERE id IS NULL`).Scan(&nullRows); err != nil { - t.Fatalf("count null-key target rows: %v", err) - } - if nullRows != 2 { - t.Fatalf("DO NOTHING should preserve NULL conflict-key rows, got %d", nullRows) - } -} - -func TestTranspile_OnConflict_DuckLakeMode_DoNothingAllowsDuplicateTargetMatches(t *testing.T) { - tr := New(Config{DuckLakeMode: true}) - db, err := sql.Open("duckdb", ":memory:") - if err != nil { - t.Fatalf("open duckdb: %v", err) - } - defer func() { _ = db.Close() }() - - if _, err := db.Exec(`CREATE TABLE target (id VARCHAR, nickname VARCHAR)`); err != nil { - t.Fatalf("create target: %v", err) - } - if _, err := db.Exec(`CREATE TABLE staging (id VARCHAR, nickname VARCHAR)`); err != nil { - t.Fatalf("create staging: %v", err) - } - if _, err := db.Exec(`INSERT INTO target VALUES ('price_1', 'old a'), ('price_1', 'old b')`); err != nil { - t.Fatalf("seed duplicate target rows: %v", err) - } - if _, err := db.Exec(`INSERT INTO staging VALUES ('price_1', 'new')`); err != nil { - t.Fatalf("seed staging: %v", err) - } - - result, err := tr.Transpile(`INSERT INTO target (id, nickname) - SELECT id, nickname FROM staging - ON CONFLICT (id) DO NOTHING`) - if err != nil { - t.Fatalf("Transpile do-nothing query: %v", err) - } - if err := runDuckDBStatements(db, result.Statements, result.CleanupStatements); err != nil { - t.Fatalf("DO NOTHING should tolerate duplicate target matches: %v", err) - } - - var rows int - if err := db.QueryRow(`SELECT count(*) FROM target WHERE id = 'price_1'`).Scan(&rows); err != nil { - t.Fatalf("count target rows: %v", err) - } - if rows != 2 { - t.Fatalf("DO NOTHING should leave duplicate matched target rows unchanged, got %d rows", rows) - } -} - -func TestTranspile_OnConflict_DuckLakeMode_GuardsMultiColumnConflicts(t *testing.T) { - tr := New(Config{DuckLakeMode: true}) - db, err := sql.Open("duckdb", ":memory:") - if err != nil { - t.Fatalf("open duckdb: %v", err) - } - defer func() { _ = db.Close() }() - - if _, err := db.Exec(`CREATE TABLE target (id VARCHAR, org_id VARCHAR, nickname VARCHAR)`); err != nil { - t.Fatalf("create target: %v", err) - } - if _, err := db.Exec(`CREATE TABLE staging (id VARCHAR, org_id VARCHAR, nickname VARCHAR)`); err != nil { - t.Fatalf("create staging: %v", err) - } - - result, err := tr.Transpile(`INSERT INTO target (id, org_id, nickname) - SELECT id, org_id, nickname FROM staging - ON CONFLICT (id, org_id) DO UPDATE SET nickname = excluded.nickname`) - if err != nil { - t.Fatalf("Transpile multi-column query: %v", err) - } - if len(result.Statements) == 0 { - t.Fatalf("expected guarded multi-statement rewrite, got single SQL: %s", result.SQL) - } - - if _, err := db.Exec(`INSERT INTO staging VALUES ('price_1', 'org_1', 'first'), ('price_1', 'org_1', 'second')`); err != nil { - t.Fatalf("seed duplicate staging rows: %v", err) - } - err = runDuckDBStatements(db, result.Statements, result.CleanupStatements) - if err == nil || !strings.Contains(err.Error(), "ON CONFLICT source rows contain duplicate conflict keys") { - t.Fatalf("expected multi-column source duplicate guard error, got %v", err) - } - - if _, err := db.Exec(`DELETE FROM staging`); err != nil { - t.Fatalf("clear staging: %v", err) - } - if _, err := db.Exec(`INSERT INTO target VALUES ('price_1', 'org_1', 'old a'), ('price_1', 'org_1', 'old b')`); err != nil { - t.Fatalf("seed duplicate target rows: %v", err) - } - if _, err := db.Exec(`INSERT INTO staging VALUES ('price_1', 'org_1', 'new')`); err != nil { - t.Fatalf("seed staging row: %v", err) - } - err = runDuckDBStatements(db, result.Statements, result.CleanupStatements) - if err == nil || !strings.Contains(err.Error(), "ON CONFLICT target contains duplicate conflict keys") { - t.Fatalf("expected multi-column target duplicate guard error, got %v", err) - } -} - func TestTranspile_DropCascade_DuckLakeMode(t *testing.T) { // In DuckLake mode, CASCADE should be stripped from DROP TABLE/VIEW statements // See: https://github.com/duckdb/dbt-duckdb/pull/557 From fd744ddf4f4c49abecb1f42a32cbd8d5b836f520 Mon Sep 17 00:00:00 2001 From: Bill Guowei Yang Date: Thu, 2 Jul 2026 11:19:29 -0400 Subject: [PATCH 6/6] Reject nested lake ON CONFLICT --- tests/integration/dml_test.go | 30 ++++++++++- tests/integration/setup_test.go | 2 +- transpiler/transform/onconflict.go | 20 +++----- transpiler/transform/onconflict_test.go | 67 +++++++++++++++++++++++++ transpiler/transform/transform.go | 9 ++++ transpiler/transpiler.go | 38 +++++++------- transpiler/transpiler_test.go | 8 +++ 7 files changed, 140 insertions(+), 34 deletions(-) diff --git a/tests/integration/dml_test.go b/tests/integration/dml_test.go index 38e44b2a..5c22bc1b 100644 --- a/tests/integration/dml_test.go +++ b/tests/integration/dml_test.go @@ -1,6 +1,7 @@ package integration import ( + "strings" "testing" ) @@ -132,6 +133,30 @@ func TestDMLInsertReturning(t *testing.T) { // TestDMLInsertOnConflict tests INSERT ... ON CONFLICT (UPSERT) func TestDMLInsertOnConflict(t *testing.T) { + assertOnConflictRejected := func(t *testing.T, query string) { + t.Helper() + _, err := testHarness.DuckgresDB.Exec(query) + if err == nil { + t.Fatal("expected ON CONFLICT to be rejected") + } + if !strings.Contains(err.Error(), "ON CONFLICT is not supported") { + t.Fatalf("expected ON CONFLICT unsupported error, got %v", err) + } + } + + if testHarness.useDuckLake { + t.Run("on_conflict_do_nothing", func(t *testing.T) { + assertOnConflictRejected(t, "INSERT INTO dml_upsert_test (id, name, counter) VALUES (1, 'Duplicate', 99) ON CONFLICT DO NOTHING") + }) + t.Run("on_conflict_do_update", func(t *testing.T) { + assertOnConflictRejected(t, ` + INSERT INTO dml_upsert_test (id, name, counter) VALUES (1, 'Updated', 20) + ON CONFLICT (id) DO UPDATE SET name = EXCLUDED.name, counter = dml_upsert_test.counter + EXCLUDED.counter + `) + }) + return + } + mustExec(t, testHarness.DuckgresDB, "CREATE TABLE dml_upsert_test (id INTEGER PRIMARY KEY, name TEXT, counter INTEGER)") mustExec(t, testHarness.DuckgresDB, "INSERT INTO dml_upsert_test VALUES (1, 'Alice', 10)") @@ -149,10 +174,11 @@ func TestDMLInsertOnConflict(t *testing.T) { }) t.Run("on_conflict_do_update", func(t *testing.T) { - _, err := testHarness.DuckgresDB.Exec(` + query := ` INSERT INTO dml_upsert_test (id, name, counter) VALUES (1, 'Updated', 20) ON CONFLICT (id) DO UPDATE SET name = EXCLUDED.name, counter = dml_upsert_test.counter + EXCLUDED.counter - `) + ` + _, err := testHarness.DuckgresDB.Exec(query) if err != nil { t.Fatalf("ON CONFLICT DO UPDATE failed: %v", err) } diff --git a/tests/integration/setup_test.go b/tests/integration/setup_test.go index 6182a4fa..284037a2 100644 --- a/tests/integration/setup_test.go +++ b/tests/integration/setup_test.go @@ -25,7 +25,7 @@ var knownFailures = map[string]string{ "insert_returning_columns": "DuckLake: RETURNING clause not yet supported", "update_returning_star": "DuckLake: RETURNING clause not yet supported", "delete_returning_star": "DuckLake: RETURNING clause not yet supported", - "on_conflict_do_nothing": "DuckLake: constraint detection requires explicit columns", + "on_conflict_do_nothing": "DuckDB: targetless DO NOTHING behavior differs", // DuckDB missing functions "where_regex_match": "DuckDB: ~* operator not available", diff --git a/transpiler/transform/onconflict.go b/transpiler/transform/onconflict.go index 2d051d82..f8d72627 100644 --- a/transpiler/transform/onconflict.go +++ b/transpiler/transform/onconflict.go @@ -24,22 +24,18 @@ func (t *OnConflictTransform) Name() string { } func (t *OnConflictTransform) Transform(tree *pg_query.ParseResult, result *Result) (bool, error) { - for _, stmt := range tree.Stmts { - if stmt.Stmt == nil { - continue - } - - insert := stmt.Stmt.GetInsertStmt() - if insert == nil || insert.OnConflictClause == nil { - continue - } + if !t.RejectOnConflict { + return false, nil + } - if t.RejectOnConflict { + WalkFunc(tree, func(node *pg_query.Node) bool { + if insert := node.GetInsertStmt(); insert != nil && insert.OnConflictClause != nil { result.Error = NewFeatureNotSupported( "ON CONFLICT is not supported: this catalog does not enforce unique constraints") - return false, nil + return false } - } + return true + }) return false, nil } diff --git a/transpiler/transform/onconflict_test.go b/transpiler/transform/onconflict_test.go index 6c1c01c2..26f7f078 100644 --- a/transpiler/transform/onconflict_test.go +++ b/transpiler/transform/onconflict_test.go @@ -139,6 +139,73 @@ func TestOnConflictTransform_RejectingModeRejectsAllOnConflict(t *testing.T) { } } +func TestOnConflictTransform_RejectingModeRejectsNestedOnConflict(t *testing.T) { + tr := NewOnConflictTransformWithConfig(true) + + tests := []struct { + name string + input string + }{ + { + name: "writable CTE under SELECT", + input: `WITH up AS ( + INSERT INTO users (id, name) VALUES (1, 'test') + ON CONFLICT (id) DO UPDATE SET name = EXCLUDED.name + RETURNING id + ) SELECT * FROM up`, + }, + { + name: "writable CTE under INSERT", + input: `WITH up AS ( + INSERT INTO users (id, name) VALUES (1, 'test') + ON CONFLICT (id) DO UPDATE SET name = EXCLUDED.name + RETURNING id + ) INSERT INTO audit_users (id) SELECT id FROM up`, + }, + { + name: "writable CTE under UPDATE", + input: `WITH up AS ( + INSERT INTO users (id, name) VALUES (1, 'test') + ON CONFLICT (id) DO UPDATE SET name = EXCLUDED.name + RETURNING id + ) UPDATE audit_users SET touched = true FROM up WHERE audit_users.id = up.id`, + }, + { + name: "writable CTE under DELETE", + input: `WITH up AS ( + INSERT INTO users (id, name) VALUES (1, 'test') + ON CONFLICT (id) DO UPDATE SET name = EXCLUDED.name + RETURNING id + ) DELETE FROM audit_users USING up WHERE audit_users.id = up.id`, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + tree, err := pg_query.Parse(tt.input) + if err != nil { + t.Fatalf("Parse error: %v", err) + } + + result := &Result{} + changed, err := tr.Transform(tree, result) + if err != nil { + t.Fatalf("Transform error: %v", err) + } + + if changed { + t.Error("Transform should reject without rewriting SQL") + } + if result.Error == nil { + t.Fatal("expected nested ON CONFLICT to be rejected") + } + if got := transformErrSQLState(result.Error); got != "0A000" { + t.Fatalf("SQLSTATE = %q, want 0A000", got) + } + }) + } +} + func TestOnConflictTransform_NoTransformCases(t *testing.T) { tr := NewOnConflictTransformWithConfig(true) diff --git a/transpiler/transform/transform.go b/transpiler/transform/transform.go index e8ddfbad..f4c77779 100644 --- a/transpiler/transform/transform.go +++ b/transpiler/transform/transform.go @@ -357,6 +357,9 @@ func walkInsertStmt(stmt *pg_query.InsertStmt, fn func(*pg_query.Node) bool) { walkNode(col, fn) } walkNode(stmt.SelectStmt, fn) + if stmt.WithClause != nil { + walkNode(&pg_query.Node{Node: &pg_query.Node_WithClause{WithClause: stmt.WithClause}}, fn) + } for _, target := range stmt.ReturningList { walkNode(target, fn) } @@ -377,6 +380,9 @@ func walkUpdateStmt(stmt *pg_query.UpdateStmt, fn func(*pg_query.Node) bool) { for _, from := range stmt.FromClause { walkNode(from, fn) } + if stmt.WithClause != nil { + walkNode(&pg_query.Node{Node: &pg_query.Node_WithClause{WithClause: stmt.WithClause}}, fn) + } for _, target := range stmt.ReturningList { walkNode(target, fn) } @@ -394,6 +400,9 @@ func walkDeleteStmt(stmt *pg_query.DeleteStmt, fn func(*pg_query.Node) bool) { for _, target := range stmt.ReturningList { walkNode(target, fn) } + if stmt.WithClause != nil { + walkNode(&pg_query.Node{Node: &pg_query.Node_WithClause{WithClause: stmt.WithClause}}, fn) + } } func walkCreateStmt(stmt *pg_query.CreateStmt, fn func(*pg_query.Node) bool) { diff --git a/transpiler/transpiler.go b/transpiler/transpiler.go index f57f2333..496f8410 100644 --- a/transpiler/transpiler.go +++ b/transpiler/transpiler.go @@ -85,57 +85,57 @@ func New(cfg Config) *Transpiler { // Core transforms - always registered // Order matters: more specific transforms should come first - // 0. Writable CTE transform - MUST BE FIRST + // 0. ON CONFLICT rejection must run before rewrites that can emit multiple + // statements and skip the rest of the transform pipeline. + t.transforms = append(t.transforms, taggedTransform{FlagOnConflict, transform.NewOnConflictTransformWithConfig(dmlPolicy.ConflictHandling == backend.RejectOnConflict)}) + + // 1. Writable CTE transform - first mutating rewrite. t.transforms = append(t.transforms, taggedTransform{FlagWritableCTE, transform.NewWritableCTETransform()}) - // 1. version() replacement - MUST run before PgCatalogTransform + // 2. version() replacement - MUST run before PgCatalogTransform t.transforms = append(t.transforms, taggedTransform{FlagVersion, transform.NewVersionTransform()}) - // 2. pg_catalog schema and view mappings + // 3. pg_catalog schema and view mappings t.transforms = append(t.transforms, taggedTransform{FlagPgCatalog, transform.NewPgCatalogTransformWithConfig(catalogPolicy.QualifyMacros)}) - // 3. information_schema mappings to compat views + // 4. information_schema mappings to compat views t.transforms = append(t.transforms, taggedTransform{FlagInfoSchema, transform.NewInformationSchemaTransform()}) - // 3.1 Map PostgreSQL "public" schema to DuckDB "main" (backends whose default + // 4.1 Map PostgreSQL "public" schema to DuckDB "main" (backends whose default // schema is "main"; disabled for Iceberg whose schema is literally "public") t.transforms = append(t.transforms, taggedTransform{FlagPublicSchema, transform.NewPublicSchemaTransform(catalogPolicy.MapPublicToMain)}) - // 3.2 Map logical database catalog references to the physical backend catalog + // 4.2 Map logical database catalog references to the physical backend catalog t.transforms = append(t.transforms, taggedTransform{FlagLogicalCatalog, transform.NewLogicalCatalogTransform(cfg.LogicalDatabaseName, catalogPolicy.PhysicalName, catalogPolicy.MapPublicToMain)}) - // 3.3 Literal rewrites (bytea \x hex, bit-string B'..') — MUST run before + // 4.3 Literal rewrites (bytea \x hex, bit-string B'..') — MUST run before // TypeMapping so it still sees the `bytea` type name. t.transforms = append(t.transforms, taggedTransform{FlagLiterals, transform.NewLiteralTransform()}) - // 4. Type mappings (JSONB->JSON, CHAR->TEXT, etc.) + // 5. Type mappings (JSONB->JSON, CHAR->TEXT, etc.) t.transforms = append(t.transforms, taggedTransform{FlagTypeMapping, transform.NewTypeMappingTransform()}) - // 5. Type casts (::regtype -> ::varchar) + // 6. Type casts (::regtype -> ::varchar) t.transforms = append(t.transforms, taggedTransform{FlagTypeCast, transform.NewTypeCastTransform()}) - // 6. Function mappings (array_agg->list, string_to_array->string_split, etc.) + // 7. Function mappings (array_agg->list, string_to_array->string_split, etc.) t.transforms = append(t.transforms, taggedTransform{FlagFunctions, transform.NewFunctionTransform()}) - // 7. Function alias normalization (current_database() -> AS current_database) + // 8. Function alias normalization (current_database() -> AS current_database) t.transforms = append(t.transforms, taggedTransform{FlagFuncAlias, transform.NewFuncAliasTransform()}) - // 8. Boolean predicate normalization for DuckLake/worker execution. + // 9. Boolean predicate normalization for DuckLake/worker execution. t.transforms = append(t.transforms, taggedTransform{FlagBooleanPredicates, transform.NewBooleanPredicateTransform()}) - // 9. Operator mappings (regex operators, etc.) + // 10. Operator mappings (regex operators, etc.) t.transforms = append(t.transforms, taggedTransform{FlagOperators, transform.NewOperatorTransform()}) - // 10. SET/SHOW command handling + // 11. SET/SHOW command handling t.transforms = append(t.transforms, taggedTransform{FlagSetShow, transform.NewSetShowTransform()}) - // 11. _pg_expandarray handling (PostgreSQL array expansion function used by JDBC) + // 12. _pg_expandarray handling (PostgreSQL array expansion function used by JDBC) t.transforms = append(t.transforms, taggedTransform{FlagExpandArray, transform.NewExpandArrayTransform()}) - // 12. ON CONFLICT handling (rejects ON CONFLICT when the backend has no - // enforced unique constraints to infer against) - t.transforms = append(t.transforms, taggedTransform{FlagOnConflict, transform.NewOnConflictTransformWithConfig(dmlPolicy.ConflictHandling == backend.RejectOnConflict)}) - // 13. Locking clause removal (FOR UPDATE, FOR SHARE, etc.) - DuckDB doesn't support these t.transforms = append(t.transforms, taggedTransform{FlagLocking, transform.NewLockingTransform()}) diff --git a/transpiler/transpiler_test.go b/transpiler/transpiler_test.go index b7f452f8..4ecc70db 100644 --- a/transpiler/transpiler_test.go +++ b/transpiler/transpiler_test.go @@ -458,6 +458,14 @@ func TestTranspile_FeatureNotSupported_Rejected(t *testing.T) { name: "ON CONFLICT targetless DO NOTHING", input: "INSERT INTO t (id, v) VALUES (1, 2) ON CONFLICT DO NOTHING", }, + { + name: "writable CTE containing ON CONFLICT", + input: `WITH up AS ( + INSERT INTO t (id, v) VALUES (1, 2) + ON CONFLICT (id) DO UPDATE SET v = EXCLUDED.v + RETURNING id + ) SELECT * FROM up`, + }, } for _, backend := range []StorageBackend{BackendDuckLake, BackendIceberg} {