postgresql: support OLD and NEW in RETURNING clauses

PostgreSQL 18 allows the RETURNING clause of INSERT, UPDATE and DELETE
statements to reference the row before and after modification through
the OLD and NEW aliases, optionally renamed with
RETURNING WITH (OLD AS ..., NEW AS ...). sqlc rejected such queries
with 'column does not exist' because neither alias resolved to a table.

Resolve both aliases to virtual copies of the statement's target table
when computing output columns and expanding star references. Columns
reached through OLD in an INSERT and through NEW in a DELETE become
nullable, since no old or new row exists for those statements. A source
table already known under the alias name shadows the virtual table,
matching PostgreSQL. The WITH (...) option list is carried through the
AST so renamed aliases resolve and format correctly.

Fixes #4556

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Q4yURXfrJTANespBca2wDE

Author: claude

internal/compiler/expand.go

@@ -84,6 +84,13 @@ func (c *Compiler) expandStmt(qc *QueryCatalog, raw *ast.RawStmt, node ast.Node)
 		return nil, err
 	}
 
+	// Virtual tables for the OLD and NEW aliases available in a RETURNING
+	// clause (PostgreSQL 18)
+	rtables, err := c.returningTables(qc, node)
+	if err != nil {
+		return nil, err
+	}
+
 	var targets *ast.List
 	switch n := node.(type) {
 	case *ast.DeleteStmt:
@@ -131,7 +138,11 @@ func (c *Compiler) expandStmt(qc *QueryCatalog, raw *ast.RawStmt, node ast.Node)
 				}
 			}
 		}
-		for _, t := range tables {
+		starTables := tables
+		if vt := returningTableForScope(tables, rtables, scope); vt != nil {
+			starTables = []*Table{vt}
+		}
+		for _, t := range starTables {
 			if scope != "" && scope != t.Rel.Name {
 				continue
 			}

internal/compiler/output_columns.go

@@ -58,6 +58,13 @@ func (c *Compiler) outputColumns(qc *QueryCatalog, node ast.Node) ([]*Column, er
 		return nil, err
 	}
 
+	// Virtual tables for the OLD and NEW aliases available in a RETURNING
+	// clause (PostgreSQL 18)
+	rtables, err := c.returningTables(qc, node)
+	if err != nil {
+		return nil, err
+	}
+
 	targets := &ast.List{}
 	switch n := node.(type) {
 	case *ast.DeleteStmt:
@@ -235,7 +242,7 @@ func (c *Compiler) outputColumns(qc *QueryCatalog, node ast.Node) ([]*Column, er
 					continue
 				}
 				if ref, ok := arg.(*ast.ColumnRef); ok {
-					columns, err := outputColumnRefs(res, tables, ref)
+					columns, err := outputColumnRefs(res, tablesForRef(ref, tables, rtables), ref)
 					if err != nil {
 						return nil, err
 					}
@@ -268,8 +275,12 @@ func (c *Compiler) outputColumns(qc *QueryCatalog, node ast.Node) ([]*Column, er
 				}
 
 				// TODO: This code is copied in func expand()
-				for _, t := range tables {
-					scope := astutils.Join(n.Fields, ".")
+				scope := astutils.Join(n.Fields, ".")
+				starTables := tables
+				if vt := returningTableForScope(tables, rtables, scope); vt != nil {
+					starTables = []*Table{vt}
+				}
+				for _, t := range starTables {
 					if scope != "" && scope != t.Rel.Name {
 						continue
 					}
@@ -297,7 +308,7 @@ func (c *Compiler) outputColumns(qc *QueryCatalog, node ast.Node) ([]*Column, er
 				continue
 			}
 
-			columns, err := outputColumnRefs(res, tables, n)
+			columns, err := outputColumnRefs(res, tablesForRef(n, tables, rtables), n)
 			if err != nil {
 				return nil, err
 			}

internal/compiler/returning.go

@@ -0,0 +1,138 @@
+package compiler
+
+import (
+	"github.com/sqlc-dev/sqlc/internal/config"
+	"github.com/sqlc-dev/sqlc/internal/sql/ast"
+)
+
+// returningTables builds virtual tables for the OLD and NEW aliases that
+// PostgreSQL 18 makes available in the RETURNING clause of INSERT, UPDATE and
+// DELETE statements. Each alias exposes the columns of the statement's target
+// table. For INSERT there is usually no old row and for DELETE there is no
+// new row, so every column reached through those aliases becomes nullable.
+func (c *Compiler) returningTables(qc *QueryCatalog, node ast.Node) ([]*Table, error) {
+	if c.conf.Engine != config.EnginePostgreSQL {
+		return nil, nil
+	}
+
+	var rv *ast.RangeVar
+	var returning *ast.List
+	oldAlias, newAlias := "old", "new"
+	var oldNullable, newNullable bool
+	switch n := node.(type) {
+	case *ast.DeleteStmt:
+		rv = firstRangeVar(n.Relations)
+		returning = n.ReturningList
+		if n.ReturningOldAlias != "" {
+			oldAlias = n.ReturningOldAlias
+		}
+		if n.ReturningNewAlias != "" {
+			newAlias = n.ReturningNewAlias
+		}
+		// A deleted row has no new value
+		newNullable = true
+	case *ast.InsertStmt:
+		rv = n.Relation
+		returning = n.ReturningList
+		if n.ReturningOldAlias != "" {
+			oldAlias = n.ReturningOldAlias
+		}
+		if n.ReturningNewAlias != "" {
+			newAlias = n.ReturningNewAlias
+		}
+		// An inserted row has no old value, except when an ON CONFLICT
+		// clause updates an existing row instead
+		oldNullable = true
+	case *ast.UpdateStmt:
+		rv = firstRangeVar(n.Relations)
+		returning = n.ReturningList
+		if n.ReturningOldAlias != "" {
+			oldAlias = n.ReturningOldAlias
+		}
+		if n.ReturningNewAlias != "" {
+			newAlias = n.ReturningNewAlias
+		}
+	default:
+		return nil, nil
+	}
+	if rv == nil || returning == nil || len(returning.Items) == 0 {
+		return nil, nil
+	}
+
+	fqn, err := ParseTableName(rv)
+	if err != nil {
+		return nil, err
+	}
+
+	build := func(alias string, nullable bool) *Table {
+		table, err := qc.GetTable(fqn)
+		if err != nil {
+			// An unresolvable target table is reported by the regular
+			// source table lookup, so ignore the error here
+			return nil
+		}
+		table.Rel = &ast.TableName{Name: alias}
+		if nullable {
+			for _, col := range table.Columns {
+				col.NotNull = false
+			}
+		}
+		return table
+	}
+
+	var tables []*Table
+	if t := build(oldAlias, oldNullable); t != nil {
+		tables = append(tables, t)
+	}
+	if t := build(newAlias, newNullable); t != nil {
+		tables = append(tables, t)
+	}
+	return tables, nil
+}
+
+func firstRangeVar(list *ast.List) *ast.RangeVar {
+	if list == nil {
+		return nil
+	}
+	for _, item := range list.Items {
+		if rv, ok := item.(*ast.RangeVar); ok && rv != nil {
+			return rv
+		}
+	}
+	return nil
+}
+
+// returningTableForScope returns the OLD or NEW virtual table named by scope.
+// A source table with the same name shadows the virtual table, matching
+// PostgreSQL, where the implicit aliases are only available when no relation
+// in the query is already known under that name.
+func returningTableForScope(tables, rtables []*Table, scope string) *Table {
+	if scope == "" {
+		return nil
+	}
+	for _, t := range tables {
+		if t.Rel.Name == scope {
+			return nil
+		}
+	}
+	for _, t := range rtables {
+		if t.Rel.Name == scope {
+			return t
+		}
+	}
+	return nil
+}
+
+// tablesForRef resolves a column reference against the source tables,
+// extended with the OLD or NEW virtual table when the reference is qualified
+// with one of their names.
+func tablesForRef(ref *ast.ColumnRef, tables, rtables []*Table) []*Table {
+	parts := stringSlice(ref.Fields)
+	if len(parts) != 2 {
+		return tables
+	}
+	if vt := returningTableForScope(tables, rtables, parts[0]); vt != nil {
+		return append(append([]*Table{}, tables...), vt)
+	}
+	return tables
+}

internal/endtoend/testdata/returning_old_new/postgresql/pgx/v5/exec.json

@@ -0,0 +1,3 @@
+{
+  "contexts": ["base"]
+}