feat(postgresql): add MERGE support

Author: taua-almeida

docs/howto/merge.md

@@ -0,0 +1,77 @@
+# Merging rows
+
+sqlc supports PostgreSQL's [`MERGE`](https://www.postgresql.org/docs/current/sql-merge.html)
+statement (PostgreSQL 15+). Parameters are inferred in every clause: the source
+subquery, the join condition, `WHEN` conditions, `UPDATE SET` assignments, and
+`INSERT ... VALUES` lists.
+
+```sql
+CREATE TABLE inventory (
+  sku        text PRIMARY KEY,
+  quantity   integer NOT NULL,
+  updated_by text
+);
+
+CREATE TABLE inventory_updates (
+  sku   text PRIMARY KEY,
+  delta integer NOT NULL
+);
+```
+
+```sql
+-- name: SyncInventory :exec
+MERGE INTO inventory AS t
+USING inventory_updates AS u
+ON t.sku = u.sku
+WHEN MATCHED AND t.quantity + u.delta <= $1 THEN
+    DELETE
+WHEN MATCHED THEN
+    UPDATE SET quantity = t.quantity + u.delta, updated_by = $2
+WHEN NOT MATCHED THEN
+    INSERT (sku, quantity, updated_by) VALUES (u.sku, u.delta, $2);
+```
+
+```go
+type SyncInventoryParams struct {
+	Quantity  int32
+	UpdatedBy sql.NullString
+}
+
+func (q *Queries) SyncInventory(ctx context.Context, arg SyncInventoryParams) error {
+	_, err := q.db.ExecContext(ctx, syncInventory, arg.Quantity, arg.UpdatedBy)
+	return err
+}
+```
+
+## Returning merged rows
+
+On PostgreSQL 17 and later, `MERGE` supports a `RETURNING` clause, which can be
+used with `:one` or `:many`. The clause can reference columns from both the
+target and the source relations.
+
+```sql
+-- name: MergeInventory :many
+MERGE INTO inventory AS t
+USING inventory_updates AS u
+ON t.sku = u.sku
+WHEN MATCHED THEN
+    UPDATE SET quantity = u.delta
+WHEN NOT MATCHED THEN
+    INSERT (sku, quantity) VALUES (u.sku, u.delta)
+RETURNING t.sku, t.quantity;
+```
+
+```go
+type MergeInventoryRow struct {
+	Sku      string
+	Quantity int32
+}
+
+func (q *Queries) MergeInventory(ctx context.Context) ([]MergeInventoryRow, error) {
+	// ...
+}
+```
+
+Note that sqlc does not validate the server version: generating code for a
+`MERGE ... RETURNING` query succeeds, but running it against a server older
+than PostgreSQL 17 returns a syntax error at execution time.

docs/index.rst

@@ -58,6 +58,7 @@ code ever again.
    howto/insert.md
    howto/update.md
    howto/delete.md
+   howto/merge.md
 
    howto/prepared_query.md
    howto/transactions.md

internal/compiler/expand.go

@@ -24,6 +24,7 @@ func (c *Compiler) expand(qc *QueryCatalog, raw *ast.RawStmt) ([]source.Edit, er
 		switch node.(type) {
 		case *ast.DeleteStmt:
 		case *ast.InsertStmt:
+		case *ast.MergeStmt:
 		case *ast.SelectStmt:
 		case *ast.UpdateStmt:
 		default:
@@ -90,6 +91,8 @@ func (c *Compiler) expandStmt(qc *QueryCatalog, raw *ast.RawStmt, node ast.Node)
 		targets = n.ReturningList
 	case *ast.InsertStmt:
 		targets = n.ReturningList
+	case *ast.MergeStmt:
+		targets = n.ReturningList
 	case *ast.SelectStmt:
 		targets = n.TargetList
 	case *ast.UpdateStmt:

internal/compiler/find_params.go

@@ -7,6 +7,30 @@ import (
 	"github.com/sqlc-dev/sqlc/internal/sql/astutils"
 )
 
+// paramRefForAssignment returns the ParamRef assigned to a single UPDATE SET
+// target, handling both `col = $1` and the multi-column form
+// `(a, b) = ($1, $2)`, where each column is a separate ResTarget carrying a
+// MultiAssignRef with its own Colno pointing into a shared RowExpr source.
+func paramRefForAssignment(target *ast.ResTarget) (*ast.ParamRef, bool) {
+	switch val := target.Val.(type) {
+	case *ast.ParamRef:
+		return val, true
+	case *ast.MultiAssignRef:
+		row, ok := val.Source.(*ast.RowExpr)
+		if !ok || row.Args == nil {
+			return nil, false
+		}
+		idx := val.Colno - 1
+		if idx < 0 || idx >= len(row.Args.Items) {
+			return nil, false
+		}
+		if ref, ok := row.Args.Items[idx].(*ast.ParamRef); ok {
+			return ref, true
+		}
+	}
+	return nil, false
+}
+
 func findParameters(root ast.Node) ([]paramRef, []error) {
 	refs := make([]paramRef, 0)
 	errors := make([]error, 0)
@@ -110,6 +134,66 @@ func (p paramSearch) Visit(node ast.Node) astutils.Visitor {
 			}
 		}
 
+	case *ast.MergeStmt:
+		if n.MergeWhenClauses == nil {
+			break
+		}
+		for _, item := range n.MergeWhenClauses.Items {
+			clause, ok := item.(*ast.MergeWhenClause)
+			if !ok {
+				continue
+			}
+			switch clause.CommandType {
+			case ast.CmdTypeUpdate:
+				// WHEN MATCHED THEN UPDATE SET col = $1
+				if clause.TargetList == nil {
+					continue
+				}
+				for _, item := range clause.TargetList.Items {
+					target, ok := item.(*ast.ResTarget)
+					if !ok {
+						continue
+					}
+					ref, ok := paramRefForAssignment(target)
+					if !ok {
+						continue
+					}
+					*p.refs = append(*p.refs, paramRef{parent: target, ref: ref, rv: n.Relation})
+					p.seen[ref.Location] = struct{}{}
+				}
+			case ast.CmdTypeInsert:
+				// WHEN NOT MATCHED THEN INSERT (a, b) VALUES ($1, $2)
+				if clause.Values == nil {
+					continue
+				}
+				if clause.TargetList == nil || len(clause.TargetList.Items) == 0 {
+					// Without an explicit column list sqlc cannot map the
+					// positional VALUES parameters to target columns.
+					params := astutils.Search(clause.Values, func(node ast.Node) bool {
+						_, ok := node.(*ast.ParamRef)
+						return ok
+					})
+					if len(params.Items) > 0 {
+						*p.errs = append(*p.errs, fmt.Errorf("MERGE INSERT with parameters requires an explicit column list"))
+						return p
+					}
+					continue
+				}
+				for i, v := range clause.Values.Items {
+					ref, ok := v.(*ast.ParamRef)
+					if !ok {
+						continue
+					}
+					if len(clause.TargetList.Items) <= i {
+						*p.errs = append(*p.errs, fmt.Errorf("MERGE INSERT has more expressions than target columns"))
+						return p
+					}
+					*p.refs = append(*p.refs, paramRef{parent: clause.TargetList.Items[i], ref: ref, rv: n.Relation})
+					p.seen[ref.Location] = struct{}{}
+				}
+			}
+		}
+
 	case *ast.UpdateStmt:
 		for _, item := range n.TargetList.Items {
 			target, ok := item.(*ast.ResTarget)

internal/compiler/output_columns.go

@@ -63,6 +63,8 @@ func (c *Compiler) outputColumns(qc *QueryCatalog, node ast.Node) ([]*Column, er
 		targets = n.ReturningList
 	case *ast.InsertStmt:
 		targets = n.ReturningList
+	case *ast.MergeStmt:
+		targets = n.ReturningList
 	case *ast.SelectStmt:
 		targets = n.TargetList
 		isUnion := len(targets.Items) == 0 && n.Larg != nil
@@ -404,9 +406,73 @@ func (c *Compiler) outputColumns(qc *QueryCatalog, node ast.Node) ([]*Column, er
 		}
 	}
 
+	if m, ok := node.(*ast.MergeStmt); ok && mergeHasNotMatchedBySource(m) {
+		// A WHEN NOT MATCHED BY SOURCE action returns NULL for every column of
+		// the source relation (there is no matching source row), so any source
+		// column in the RETURNING list is nullable. The target is always
+		// m.Relation, so treat every provenance-bearing column that is not the
+		// target as a source column (this covers subquery sources too).
+		for _, col := range cols {
+			if !col.NotNull || col.skipTableRequiredCheck {
+				continue
+			}
+			// Columns without table provenance (e.g. merge_action()) are never
+			// NULL from a NOT MATCHED BY SOURCE action.
+			if col.Table == nil && col.TableAlias == "" {
+				continue
+			}
+			if !columnIsMergeTarget(col, m.Relation) {
+				col.NotNull = false
+			}
+		}
+	}
+
 	return cols, nil
 }
 
+// mergeHasNotMatchedBySource reports whether the MERGE has a WHEN NOT MATCHED
+// BY SOURCE clause, which makes source-relation columns nullable in RETURNING.
+func mergeHasNotMatchedBySource(m *ast.MergeStmt) bool {
+	if m.MergeWhenClauses == nil {
+		return false
+	}
+	for _, item := range m.MergeWhenClauses.Items {
+		if clause, ok := item.(*ast.MergeWhenClause); ok {
+			if clause.MatchKind == ast.MergeWhenNotMatchedBySource {
+				return true
+			}
+		}
+	}
+	return false
+}
+
+// columnIsMergeTarget reports whether col is qualified as (or resolves to) the
+// MERGE target relation.
+func columnIsMergeTarget(col *Column, target *ast.RangeVar) bool {
+	if target == nil {
+		return false
+	}
+	var name, alias string
+	if target.Relname != nil {
+		name = *target.Relname
+	}
+	if target.Alias != nil && target.Alias.Aliasname != nil {
+		alias = *target.Alias.Aliasname
+	}
+	if col.TableAlias != "" {
+		// Qualified: it is the target only if the qualifier is the target's
+		// alias (or, when the target is unaliased, its table name).
+		if alias != "" {
+			return col.TableAlias == alias
+		}
+		return col.TableAlias == name
+	}
+	if col.Table != nil && name != "" {
+		return col.Table.Name == name
+	}
+	return false
+}
+
 const (
 	tableNotFound = iota
 	tableRequired
@@ -494,6 +560,19 @@ func (c *Compiler) sourceTables(qc *QueryCatalog, node ast.Node) ([]*Table, erro
 		list = &ast.List{
 			Items: []ast.Node{n.Relation},
 		}
+	case *ast.MergeStmt:
+		// PostgreSQL expands MERGE ... RETURNING * to all source columns
+		// followed by all target columns, so collect the source relation
+		// before the target relation.
+		// https://www.postgresql.org/docs/17/sql-merge.html
+		var tv tableVisitor
+		if n.SourceRelation != nil {
+			astutils.Walk(&tv, n.SourceRelation)
+		}
+		if n.Relation != nil {
+			astutils.Walk(&tv, n.Relation)
+		}
+		list = &tv.list
 	case *ast.SelectStmt:
 		var tv tableVisitor
 		astutils.Walk(&tv, n.FromClause)

internal/compiler/query_catalog.go

@@ -21,6 +21,8 @@ func (comp *Compiler) buildQueryCatalog(c *catalog.Catalog, node ast.Node, embed
 		with = n.WithClause
 	case *ast.InsertStmt:
 		with = n.WithClause
+	case *ast.MergeStmt:
+		with = n.WithClause
 	case *ast.UpdateStmt:
 		with = n.WithClause
 	case *ast.SelectStmt:

internal/endtoend/testdata/merge/postgresql/parse/exec.json

@@ -0,0 +1,5 @@
+{
+  "command": "parse",
+  "args": ["--dialect", "postgresql", "query.sql"],
+  "contexts": ["base"]
+}

internal/endtoend/testdata/merge/postgresql/parse/query.sql

@@ -0,0 +1,2 @@
+-- name: SyncInventory :exec
+MERGE INTO inventory AS t USING inventory_updates AS u ON t.sku = u.sku WHEN MATCHED AND t.quantity <= $1 THEN DELETE WHEN MATCHED THEN UPDATE SET quantity = u.delta WHEN NOT MATCHED THEN INSERT (sku, quantity) VALUES (u.sku, $2);