sql: support * in udf bodies

This change allows `*` usage in UDF bodies. In order to facilitate
schema changes after a UDF is created but should not affect the UDF
output, we rewrite UDF ASTs in place to reference tables and columns by
ID instead of by name.

Informs: cockroachdb#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'
 
@@ -142,9 +124,9 @@ CREATE FUNCTION public.f(IN a test.public.notmyworkday)
   CALLED ON NULL INPUT
   LANGUAGE SQL
   AS $$
-  SELECT a FROM test.public.t;
-  SELECT b FROM test.public.t@t_idx_b;
-  SELECT c FROM test.public.t@t_idx_c;
+  SELECT a FROM [113(1, 2, 3) AS t];
+  SELECT b FROM [113(1, 2, 3) AS t]@t_idx_b;
+  SELECT c FROM [113(1, 2, 3) AS t]@t_idx_c;
   SELECT nextval('public.sq1'::REGCLASS);
 $$
 
@@ -2321,7 +2303,7 @@ CREATE FUNCTION public.get_l(IN i INT8)
     CALLED ON NULL INPUT
     LANGUAGE SQL
     AS $$
-    SELECT v FROM test.public.kv WHERE k = i;
+    SELECT v FROM [225(1, 2) AS kv] WHERE k = i;
 $$
 
 query T
@@ -2788,7 +2770,7 @@ SELECT oid, proname, pronamespace, proowner, prolang, proleakproof, proisstrict,
 FROM pg_catalog.pg_proc WHERE proname IN ('f_93314', 'f_93314_alias', 'f_93314_comp', 'f_93314_comp_t')
 ORDER BY oid;
 ----
-100257  f_93314         105  1546506610  14  false  false  false  v  0  100256  ·  {}  NULL  SELECT i, e FROM test.public.t_93314 ORDER BY i LIMIT 1;
-100259  f_93314_alias   105  1546506610  14  false  false  false  v  0  100258  ·  {}  NULL  SELECT i, e FROM test.public.t_93314_alias ORDER BY i LIMIT 1;
+100257  f_93314         105  1546506610  14  false  false  false  v  0  100256  ·  {}  NULL  SELECT i, e FROM [256(1, 2) AS t_93314] ORDER BY i LIMIT 1;
+100259  f_93314_alias   105  1546506610  14  false  false  false  v  0  100258  ·  {}  NULL  SELECT i, e FROM [258(1, 2) AS t_93314_alias] ORDER BY i LIMIT 1;
 100263  f_93314_comp    105  1546506610  14  false  false  false  v  0  100260  ·  {}  NULL  SELECT (1, 2);
-100264  f_93314_comp_t  105  1546506610  14  false  false  false  v  0  100262  ·  {}  NULL  SELECT a, c FROM test.public.t_93314_comp LIMIT 1;
+100264  f_93314_comp_t  105  1546506610  14  false  false  false  v  0  100262  ·  {}  NULL  SELECT a, c FROM [262(1, 2) AS t_93314_comp] LIMIT 1;

pkg/sql/logictest/testdata/logic_test/udf_star

@@ -0,0 +1,110 @@
+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;
+
+query TTT
+SELECT oid, proname, prosrc
+FROM pg_catalog.pg_proc WHERE proname LIKE 'f\_%' ORDER BY oid;
+----
+100108  f_unqualified_onecol   SELECT * FROM [106(1) AS t_onecol];
+100109  f_subquery  SELECT * FROM (SELECT a FROM (SELECT * FROM [106(1) AS t_onecol]) AS foo) AS bar;
+100110  f_unqualified_twocol   SELECT * FROM [107(1, 2) AS t_twocol];
+100111  f_allcolsel   SELECT t_twocol.* FROM [107(1, 2) AS t_twocol];
+100112  f_allcolsel_alias   SELECT t1.* FROM [107(1, 2) AS t_twocol] AS t1, [107(1, 2) AS t_twocol] AS t2 WHERE t1.a = t2.a;
+100113  f_tuplestar   SELECT (t_twocol.*).* FROM [107(1, 2) AS t_twocol];
+
+query I
+SELECT f_unqualified_onecol()
+----
+1
+
+query I
+SELECT f_subquery()
+----
+1
+
+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
+
+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): Postgres returns an error after adding a column when the table
+# is used as the return type.
+query T
+SELECT f_unqualified_twocol()
+----
+(1,2)
+
+# 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;

pkg/sql/opt/optbuilder/builder.go

@@ -143,6 +143,10 @@ type Builder struct {
 	// using AST annotations.
 	qualifyDataSourceNamesInAST bool
 
+	// If set, the data source names in the AST are rewritten to the table ID and
+	// all the visible column IDs.
+	useIDsInAST bool
+
 	// isCorrelated is set to true if we already reported to telemetry that the
 	// query contains a correlated subquery.
 	isCorrelated bool

pkg/sql/opt/optbuilder/create_function.go

@@ -27,6 +27,7 @@ import (
 
 func (b *Builder) buildCreateFunction(cf *tree.CreateFunction, inScope *scope) (outScope *scope) {
 	b.DisableMemoReuse = true
+	b.useIDsInAST = true
 	if cf.FuncName.ExplicitCatalog {
 		if string(cf.FuncName.CatalogName) != b.evalCtx.SessionData().Database {
 			panic(unimplemented.New("CREATE FUNCTION", "cross-db references not supported"))

pkg/sql/opt/optbuilder/join.go

@@ -33,7 +33,7 @@ import (
 func (b *Builder) buildJoin(
 	join *tree.JoinTableExpr, locking lockingSpec, inScope *scope,
 ) (outScope *scope) {
-	leftScope := b.buildDataSource(join.Left, nil /* indexFlags */, locking, inScope)
+	leftScope := b.buildDataSource(&join.Left, nil /* indexFlags */, locking, inScope)
 
 	inScopeRight := inScope
 	isLateral := b.exprIsLateral(join.Right)
@@ -45,7 +45,7 @@ func (b *Builder) buildJoin(
 		inScopeRight.context = exprKindLateralJoin
 	}
 
-	rightScope := b.buildDataSource(join.Right, nil /* indexFlags */, locking, inScopeRight)
+	rightScope := b.buildDataSource(&join.Right, nil /* indexFlags */, locking, inScopeRight)
 
 	// Check that the same table name is not used on both sides.
 	b.validateJoinTableNames(leftScope, rightScope)

pkg/sql/opt/optbuilder/select.go

@@ -12,7 +12,6 @@ package optbuilder
 
 import (
 	"context"
-
 	"github.com/cockroachdb/cockroach/pkg/server/telemetry"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/colinfo"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/descpb"
@@ -44,14 +43,14 @@ import (
 // See Builder.buildStmt for a description of the remaining input and
 // return values.
 func (b *Builder) buildDataSource(
-	texpr tree.TableExpr, indexFlags *tree.IndexFlags, locking lockingSpec, inScope *scope,
+	texpr *tree.TableExpr, indexFlags *tree.IndexFlags, locking lockingSpec, inScope *scope,
 ) (outScope *scope) {
 	defer func(prevAtRoot bool) {
 		inScope.atRoot = prevAtRoot
 	}(inScope.atRoot)
 	inScope.atRoot = false
 	// NB: The case statements are sorted lexicographically.
-	switch source := texpr.(type) {
+	switch source := (*texpr).(type) {
 	case *tree.AliasedTableExpr:
 		if source.IndexFlags != nil {
 			telemetry.Inc(sqltelemetry.IndexHintUseCounter)
@@ -64,7 +63,7 @@ func (b *Builder) buildDataSource(
 			locking = locking.filter(source.As.Alias)
 		}
 
-		outScope = b.buildDataSource(source.Expr, indexFlags, locking, inScope)
+		outScope = b.buildDataSource(&source.Expr, indexFlags, locking, inScope)
 
 		if source.Ordinality {
 			outScope = b.buildWithOrdinality(outScope)
@@ -110,6 +109,24 @@ func (b *Builder) buildDataSource(
 		}
 
 		ds, depName, resName := b.resolveDataSource(tn, privilege.SELECT)
+		if b.useIDsInAST {
+			switch t := ds.(type) {
+			case cat.Table:
+				if t.IsVirtualTable() {
+					break
+				}
+				var cols []tree.ColumnID
+				for i := 0; i < t.ColumnCount(); i++ {
+					if t.Column(i).Visibility() == cat.Visible {
+						cols = append(cols, tree.ColumnID(t.Column(i).ColID()))
+					}
+				}
+				expansion := &tree.TableRef{TableID: int64(ds.ID()),
+					Columns: cols,
+					As:      tree.AliasClause{Alias: ds.Name()}}
+				*texpr = expansion
+			}
+		}
 
 		locking = locking.filter(tn.ObjectName)
 		if locking.isSet() {
@@ -142,7 +159,7 @@ func (b *Builder) buildDataSource(
 		}
 
 	case *tree.ParenTableExpr:
-		return b.buildDataSource(source.Expr, indexFlags, locking, inScope)
+		return b.buildDataSource(&source.Expr, indexFlags, locking, inScope)
 
 	case *tree.RowsFromExpr:
 		return b.buildZip(source.Items, inScope)
@@ -1216,7 +1233,7 @@ func (b *Builder) buildFromTables(
 func (b *Builder) buildFromTablesRightDeep(
 	tables tree.TableExprs, locking lockingSpec, inScope *scope,
 ) (outScope *scope) {
-	outScope = b.buildDataSource(tables[0], nil /* indexFlags */, locking, inScope)
+	outScope = b.buildDataSource(&tables[0], nil /* indexFlags */, locking, inScope)
 
 	// Recursively build table join.
 	tables = tables[1:]
@@ -1266,7 +1283,7 @@ func (b *Builder) exprIsLateral(t tree.TableExpr) bool {
 func (b *Builder) buildFromWithLateral(
 	tables tree.TableExprs, locking lockingSpec, inScope *scope,
 ) (outScope *scope) {
-	outScope = b.buildDataSource(tables[0], nil /* indexFlags */, locking, inScope)
+	outScope = b.buildDataSource(&tables[0], nil /* indexFlags */, locking, inScope)
 	for i := 1; i < len(tables); i++ {
 		scope := inScope
 		// Lateral expressions need to be able to refer to the expressions that
@@ -1275,7 +1292,7 @@ func (b *Builder) buildFromWithLateral(
 			scope = outScope
 			scope.context = exprKindLateralJoin
 		}
-		tableScope := b.buildDataSource(tables[i], nil /* indexFlags */, locking, scope)
+		tableScope := b.buildDataSource(&tables[i], nil /* indexFlags */, locking, scope)
 
 		// Check that the same table name is not used multiple times.
 		b.validateJoinTableNames(outScope, tableScope)