postgresql: support OLD and NEW in RETURNING clauses (#4562)

Author: kyleconroy

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