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,194 @@
+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 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 a FROM t_onecol;
+$$ LANGUAGE SQL;
+
+statement ok
+CREATE FUNCTION f_unqualified_doublestar() RETURNS INT AS
+$$
+ SELECT *, * FROM t_onecol;
+ SELECT a FROM t_onecol;
+$$ LANGUAGE SQL;
+
+statement ok
+CREATE FUNCTION f_exprstar() RETURNS STRING AS
+$$
+ SELECT word FROM (SELECT (pg_get_keywords()).* ORDER BY word LIMIT 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 a FROM test.public.t_onecol;
+100115 f_unqualified_doublestar SELECT t_onecol.a, t_onecol.a FROM test.public.t_onecol;
+ SELECT a FROM test.public.t_onecol;
+100116 f_exprstar SELECT word FROM (SELECT (pg_get_keywords()).word, (pg_get_keywords()).catcode, (pg_get_keywords()).catdesc ORDER BY word LIMIT 1);
+
+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;
+
+# 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);

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()
 

pkg/sql/opt/optbuilder/mutation_builder.go

@@ -1019,7 +1019,7 @@ func (mb *mutationBuilder) mapToReturnColID(tabOrd int) opt.ColumnID {
 
 // buildReturning wraps the input expression with a Project operator that
 // projects the given RETURNING expressions.
-func (mb *mutationBuilder) buildReturning(returning tree.ReturningExprs) {
+func (mb *mutationBuilder) buildReturning(returning *tree.ReturningExprs) {
  // Handle case of no RETURNING clause.
  if returning == nil {
  expr := mb.outScope.expr