sql: support * in udf bodies

This change allows `*` usage in UDF bodies. We rewrite UDF ASTs in place
to expand `*`s into the columns they reference.

Informs: #90080

Epic: CRDB-19496
Release note (sql change): Allow `*` expressions in UDFs.

pkg/sql/logictest/testdata/logic_test/udf

@@ -52,24 +52,6 @@ CREATE FUNCTION err(i INT) RETURNS INT LANGUAGE SQL AS 'SELECT j'
 statement error pgcode 42703 column \"j\" does not exist
 CREATE FUNCTION err(i INT) RETURNS INT LANGUAGE SQL AS 'SELECT a FROM ab WHERE a = j'
 
-statement error pgcode 0A000 functions do not currently support \* expressions
-CREATE FUNCTION err(i INT) RETURNS ab LANGUAGE SQL AS 'SELECT * FROM ab'
-
-statement error pgcode 0A000 functions do not currently support \* expressions
-CREATE FUNCTION err(i INT) RETURNS ab LANGUAGE SQL AS 'SELECT ab.* FROM ab'
-
-statement error pgcode 0A000 functions do not currently support \* expressions
-CREATE FUNCTION err(i INT) RETURNS ab LANGUAGE SQL AS $$
-  SELECT 1;
-  SELECT * FROM ab;
-$$
-
-statement error pgcode 0A000 functions do not currently support \* expressions
-CREATE FUNCTION err(i INT) RETURNS INT LANGUAGE SQL AS $$
-  SELECT * FROM ab;
-  SELECT 1;
-$$
-
 statement ok
 CREATE FUNCTION d(i INT2) RETURNS INT4 LANGUAGE SQL AS 'SELECT i'
 

pkg/sql/logictest/testdata/logic_test/udf_star

@@ -0,0 +1,222 @@
+statement ok
+CREATE TABLE t_onecol (a INT);
+INSERT INTO t_onecol VALUES (1)
+
+statement ok
+CREATE TABLE t_twocol (a INT, b INT);
+INSERT INTO t_twocol VALUES (1,2)
+
+statement ok
+CREATE FUNCTION f_unqualified_onecol() RETURNS INT AS
+$$
+  SELECT * FROM t_onecol;
+$$ LANGUAGE SQL;
+
+statement ok
+CREATE FUNCTION f_subquery() RETURNS INT AS
+$$
+  SELECT * FROM (SELECT a FROM (SELECT * FROM t_onecol) AS foo) AS bar;
+$$ LANGUAGE SQL;
+
+statement error pq: unimplemented: unaliased subquery inside a function definition
+CREATE FUNCTION f_subquery_unaliased() RETURNS INT AS
+$$
+  SELECT * FROM (SELECT a FROM (SELECT * FROM t_onecol));
+$$ LANGUAGE SQL;
+
+statement ok
+CREATE FUNCTION f_unqualified_twocol() RETURNS t_twocol AS
+$$
+  SELECT * FROM t_twocol;
+$$ LANGUAGE SQL;
+
+statement ok
+CREATE FUNCTION f_allcolsel() RETURNS t_twocol AS
+$$
+  SELECT t_twocol.* FROM t_twocol;
+$$ LANGUAGE SQL;
+
+statement ok
+CREATE FUNCTION f_allcolsel_alias() RETURNS t_twocol AS
+$$
+  SELECT t1.* FROM t_twocol AS t1, t_twocol AS t2 WHERE t1.a = t2.a;
+$$ LANGUAGE SQL;
+
+statement ok
+CREATE FUNCTION f_tuplestar() RETURNS t_twocol AS
+$$
+  SELECT (t_twocol.*).* FROM t_twocol;
+$$ LANGUAGE SQL;
+
+statement ok
+CREATE FUNCTION f_unqualified_multicol() RETURNS INT AS
+$$
+  SELECT *, a FROM t_onecol;
+  SELECT 1;
+$$ LANGUAGE SQL;
+
+statement ok
+CREATE FUNCTION f_unqualified_doublestar() RETURNS INT AS
+$$
+  SELECT *, * FROM t_onecol;
+  SELECT 1;
+$$ LANGUAGE SQL;
+
+statement ok
+CREATE FUNCTION f_exprstar() RETURNS STRING AS
+$$
+  SELECT word FROM (SELECT (pg_get_keywords()).* ORDER BY word LIMIT 1) AS foo;
+$$ LANGUAGE SQL;
+
+statement error pq: unimplemented: unaliased subquery inside a function definition
+CREATE FUNCTION f_ambiguous() RETURNS INT AS
+$$
+  SELECT * FROM (SELECT a FROM t_onecol) JOIN (SELECT a FROM t_twocol) ON true;
+  SELECT 1;
+$$ LANGUAGE SQL;
+
+statement error pq: column reference "a" is ambiguous
+CREATE FUNCTION f_ambiguous() RETURNS INT AS
+$$
+  SELECT a FROM (SELECT * FROM (SELECT a FROM t_onecol) AS foo JOIN (SELECT a FROM t_twocol) AS bar ON true) AS baz;
+  SELECT 1;
+$$ LANGUAGE SQL;
+
+query TTT
+SELECT oid, proname, prosrc
+FROM pg_catalog.pg_proc WHERE proname LIKE 'f\_%' ORDER BY oid;
+----
+100108  f_unqualified_onecol      SELECT t_onecol.a FROM test.public.t_onecol;
+100109  f_subquery                SELECT bar.a FROM (SELECT a FROM (SELECT t_onecol.a FROM test.public.t_onecol) AS foo) AS bar;
+100110  f_unqualified_twocol      SELECT t_twocol.a, t_twocol.b FROM test.public.t_twocol;
+100111  f_allcolsel               SELECT t_twocol.a, t_twocol.b FROM test.public.t_twocol;
+100112  f_allcolsel_alias         SELECT t1.a, t1.b FROM test.public.t_twocol AS t1, test.public.t_twocol AS t2 WHERE t1.a = t2.a;
+100113  f_tuplestar               SELECT t_twocol.a, t_twocol.b FROM test.public.t_twocol;
+100114  f_unqualified_multicol    SELECT t_onecol.a, a FROM test.public.t_onecol;
+                                  SELECT 1;
+100115  f_unqualified_doublestar  SELECT t_onecol.a, t_onecol.a FROM test.public.t_onecol;
+                                  SELECT 1;
+100116  f_exprstar                SELECT word FROM (SELECT (pg_get_keywords()).word, (pg_get_keywords()).catcode, (pg_get_keywords()).catdesc ORDER BY word LIMIT 1) AS foo;
+
+
+query TT
+SHOW CREATE FUNCTION f_subquery
+----
+f_subquery  CREATE FUNCTION public.f_subquery()
+              RETURNS INT8
+              VOLATILE
+              NOT LEAKPROOF
+              CALLED ON NULL INPUT
+              LANGUAGE SQL
+              AS $$
+              SELECT bar.a FROM (SELECT a FROM (SELECT t_onecol.a FROM test.public.t_onecol) AS foo) AS bar;
+            $$
+
+query TT
+SHOW CREATE FUNCTION f_allcolsel_alias
+----
+f_allcolsel_alias  CREATE FUNCTION public.f_allcolsel_alias()
+                     RETURNS T_TWOCOL
+                     VOLATILE
+                     NOT LEAKPROOF
+                     CALLED ON NULL INPUT
+                     LANGUAGE SQL
+                     AS $$
+                     SELECT t1.a, t1.b FROM test.public.t_twocol AS t1, test.public.t_twocol AS t2 WHERE t1.a = t2.a;
+                   $$
+
+query I
+SELECT f_unqualified_onecol()
+----
+1
+
+query I
+SELECT f_subquery()
+----
+1
+
+query T
+SELECT f_exprstar()
+----
+abort
+
+# Adding a column to a table should not change the UDFs that reference it.
+statement ok
+ALTER TABLE t_onecol ADD COLUMN b INT DEFAULT 5;
+
+query I
+SELECT f_unqualified_onecol()
+----
+1
+
+query I
+SELECT f_subquery()
+----
+1
+
+# It's ok to drop a column that was not used by the original UDF.
+statement ok
+ALTER TABLE t_onecol DROP COLUMN b;
+
+query T
+SELECT f_unqualified_twocol()
+----
+(1,2)
+
+query T
+SELECT f_allcolsel()
+----
+(1,2)
+
+query T
+SELECT f_allcolsel_alias()
+----
+(1,2)
+
+statement ok
+ALTER TABLE t_twocol ADD COLUMN c INT DEFAULT 5;
+
+# TODO(#95558): With early binding, postgres returns an error after adding a
+# column when the table is used as the return type. Note that this behavior is
+# ok for late binding.
+query T
+SELECT f_unqualified_twocol()
+----
+(1,2)
+
+# Altering a column type is not allowed in postgres or CRDB.
+statement error pq: cannot alter type of column "b" because function "f_unqualified_twocol" depends on it
+ALTER TABLE t_twocol ALTER b TYPE FLOAT;
+
+# TODO(harding): Postgres allows column renaming when only referenced by UDFs.
+statement error pq: cannot rename column "a" because function "f_unqualified_twocol" depends on it
+ALTER TABLE t_twocol RENAME COLUMN a TO d;
+
+# TODO(harding): Postgres allows table renaming when only referenced by UDFs.
+statement error pq: cannot rename relation "t_twocol" because function "f_unqualified_twocol" depends on it
+ALTER TABLE t_twocol RENAME TO t_twocol_prime;
+
+# Dropping a column a UDF depends on is not allowed.
+statement error pq: cannot drop column "b" because function "f_unqualified_twocol" depends on it
+ALTER TABLE t_twocol DROP COLUMN b;
+
+# Drop all but one of the functions with an implicit record return value.
+# TODO(96368): Allow these UDFs to be dropped in the CASCADE when the cross-
+# references are fixed instead.
+statement ok
+DROP FUNCTION f_tuplestar;
+DROP FUNCTION f_allcolsel_alias;
+
+# Dropping a column using CASCADE is ok.
+statement ok
+ALTER TABLE t_twocol DROP COLUMN b CASCADE;
+
+statement ok
+DROP TABLE t_onecol CASCADE;
+
+# The only remaining function should not reference the tables.
+query TTT
+SELECT oid, proname, prosrc
+FROM pg_catalog.pg_proc WHERE proname LIKE 'f\_%' ORDER BY oid;
+----
+100116  f_exprstar                SELECT word FROM (SELECT (pg_get_keywords()).word, (pg_get_keywords()).catcode, (pg_get_keywords()).catdesc ORDER BY word LIMIT 1) AS foo;

pkg/sql/opt/optbuilder/delete.go

@@ -66,7 +66,7 @@ func (b *Builder) buildDelete(del *tree.Delete, inScope *scope) (outScope *scope
 
 	// Build the final delete statement, including any returned expressions.
 	if resultsNeeded(del.Returning) {
-		mb.buildDelete(*del.Returning.(*tree.ReturningExprs))
+		mb.buildDelete(del.Returning.(*tree.ReturningExprs))
 	} else {
 		mb.buildDelete(nil /* returning */)
 	}
@@ -76,7 +76,7 @@ func (b *Builder) buildDelete(del *tree.Delete, inScope *scope) (outScope *scope
 
 // buildDelete constructs a Delete operator, possibly wrapped by a Project
 // operator that corresponds to the given RETURNING clause.
-func (mb *mutationBuilder) buildDelete(returning tree.ReturningExprs) {
+func (mb *mutationBuilder) buildDelete(returning *tree.ReturningExprs) {
 	mb.buildFKChecksAndCascadesForDelete()
 
 	// Project partial index DEL boolean columns.

pkg/sql/opt/optbuilder/insert.go

@@ -281,9 +281,9 @@ func (b *Builder) buildInsert(ins *tree.Insert, inScope *scope) (outScope *scope
 	// See mutationBuilder.buildCheckInputScan.
 	mb.insertExpr = mb.outScope.expr
 
-	var returning tree.ReturningExprs
+	var returning *tree.ReturningExprs
 	if resultsNeeded(ins.Returning) {
-		returning = *ins.Returning.(*tree.ReturningExprs)
+		returning = ins.Returning.(*tree.ReturningExprs)
 	}
 
 	switch {
@@ -669,7 +669,7 @@ func (mb *mutationBuilder) addSynthesizedColsForInsert() {
 
 // buildInsert constructs an Insert operator, possibly wrapped by a Project
 // operator that corresponds to the given RETURNING clause.
-func (mb *mutationBuilder) buildInsert(returning tree.ReturningExprs) {
+func (mb *mutationBuilder) buildInsert(returning *tree.ReturningExprs) {
 	// Disambiguate names so that references in any expressions, such as a
 	// check constraint, refer to the correct columns.
 	mb.disambiguateColumns()
@@ -874,7 +874,7 @@ func (mb *mutationBuilder) setUpsertCols(insertCols tree.NameList) {
 
 // buildUpsert constructs an Upsert operator, possibly wrapped by a Project
 // operator that corresponds to the given RETURNING clause.
-func (mb *mutationBuilder) buildUpsert(returning tree.ReturningExprs) {
+func (mb *mutationBuilder) buildUpsert(returning *tree.ReturningExprs) {
 	// Merge input insert and update columns using CASE expressions.
 	mb.projectUpsertColumns()