opt: check that generator functions are not used in CASE or COALESCE

This patch adds checks during type-checking to ensure that generator functions
are not used in the arguments of `CASE`, `IF`, `COALESCE`, or `IFNULL`
expressions. This mirrors postgres behavior. This patch also corrects the error
message that is returned in these cases to say "set-returning" instead of
"generator".

Fixes cockroachdb#97119
Fixes cockroachdb#94890

Release note (bug fix): CASE, IF, COALESCE, and IFNULL expressions now return
an error when passed a generator function as an argument. This mirrors postgres
behavior.

pkg/cmd/roachtest/tests/follower_reads.go

@@ -530,9 +530,11 @@ func initFollowerReadsDB(
 			// parsing the replica_localities array using the same pattern as the
 			// one used by SHOW REGIONS.
 			const q2 = `
-			SELECT
-				count(distinct substring(unnest(replica_localities), 'region=([^,]*)'))
-			FROM [SHOW RANGES FROM TABLE test.test]`
+			SELECT count(DISTINCT substring(unnested, 'region=([^,]*)'))
+			FROM (
+				SELECT unnest(replica_localities) AS unnested
+				FROM [SHOW RANGES FROM TABLE test.test]
+			)`
 
 			var distinctRegions int
 			require.NoError(t, db.QueryRowContext(ctx, q2).Scan(&distinctRegions))

pkg/sql/delegate/show_functions.go

@@ -56,7 +56,7 @@ func (d *delegator) delegateShowFunctions(n *tree.ShowFunctions) (tree.Statement
 SELECT n.nspname as schema_name,
   p.proname as function_name,
   p.prorettype::REGTYPE::TEXT as result_data_type,
-	COALESCE((SELECT trim('{}' FROM replace(array_agg(unnest(proargtypes)::REGTYPE::TEXT)::TEXT, ',', ', '))), '') as argument_data_types,
+	COALESCE((SELECT trim('{}' FROM replace((SELECT array_agg(unnested::REGTYPE::TEXT) FROM unnest(proargtypes) AS unnested)::TEXT, ',', ', '))), '') as argument_data_types,
   CASE p.prokind
 	  WHEN 'a' THEN 'agg'
 	  WHEN 'w' THEN 'window'

pkg/sql/logictest/testdata/logic_test/aggregate

@@ -2880,10 +2880,8 @@ SELECT string_agg('foo', CAST ((SELECT NULL) AS BYTES)) OVER ();
 foo
 
 # Regression test for #30166.
-query T
+statement error pq: array_agg\(\): set-returning functions are not allowed in aggregate
 SELECT array_agg(generate_series(1, 2))
-----
-{1,2}
 
 # Regression test for #31882.
 

pkg/sql/logictest/testdata/logic_test/create_as

@@ -105,7 +105,7 @@ CREATE TABLE foo2 (x) AS (VALUES(ROW()))
 statement error pq: nested array unsupported as column type: int\[\]\[\]
 CREATE TABLE foo2 (x) AS (VALUES(ARRAY[ARRAY[1]]))
 
-statement error generator functions are not allowed in VALUES
+statement error pq: generate_series\(\): set-returning functions are not allowed in VALUES
 CREATE TABLE foo2 (x) AS (VALUES(generate_series(1,3)))
 
 statement error pq: value type unknown cannot be used for table columns

pkg/sql/logictest/testdata/logic_test/join

@@ -1069,7 +1069,7 @@ FROM
 NULL  NULL
 
 # Regression test for #28817. Do not allow special functions in ON clause.
-query error generator functions are not allowed in ON
+query error pq: generate_series\(\): set-returning functions are not allowed in ON
 SELECT * FROM foo JOIN bar ON generate_series(0, 1) < 2
 
 query error aggregate functions are not allowed in JOIN conditions

pkg/sql/logictest/testdata/logic_test/partial_index

@@ -49,7 +49,7 @@ statement error pgcode 42P20 window functions are not allowed in INDEX PREDICATE
 CREATE TABLE error (a INT, INDEX (a) WHERE row_number() OVER () > 1)
 
 # Don't allow set-returning functions.
-statement error pgcode 0A000 generator functions are not allowed in INDEX PREDICATE
+statement error pgcode 0A000 set-returning functions are not allowed in INDEX PREDICATE
 CREATE TABLE error (a INT, INDEX (a) WHERE generate_series(1, 1))
 
 # Fail on bad types.

pkg/sql/logictest/testdata/logic_test/srfs

@@ -167,7 +167,7 @@ SELECT * FROM generate_series(1, 1) WITH ORDINALITY AS c(x, y)
 x y
 1 1
 
-query error generator functions are not allowed in LIMIT
+query error set-returning functions are not allowed in LIMIT
 SELECT * FROM (VALUES (1)) LIMIT generate_series(1, 3)
 
 query I colnames,nosort
@@ -838,31 +838,31 @@ generate_subscripts
 
 subtest srf_errors
 
-query error generator functions are not allowed in ORDER BY
+query error set-returning functions are not allowed in ORDER BY
 SELECT * FROM t ORDER BY generate_series(1, 3)
 
-query error generator functions are not allowed in WHERE
+query error set-returning functions are not allowed in WHERE
 SELECT * FROM t WHERE generate_series(1, 3) < 3
 
-query error generator functions are not allowed in HAVING
+query error set-returning functions are not allowed in HAVING
 SELECT * FROM t HAVING generate_series(1, 3) < 3
 
-query error generator functions are not allowed in LIMIT
+query error set-returning functions are not allowed in LIMIT
 SELECT * FROM t LIMIT generate_series(1, 3)
 
-query error generator functions are not allowed in OFFSET
+query error set-returning functions are not allowed in OFFSET
 SELECT * FROM t OFFSET generate_series(1, 3)
 
-query error generator functions are not allowed in VALUES
+query error set-returning functions are not allowed in VALUES
 VALUES (generate_series(1,3))
 
-statement error generator functions are not allowed in DEFAULT
+statement error set-returning functions are not allowed in DEFAULT
 CREATE TABLE uu (x INT DEFAULT generate_series(1, 3))
 
-statement error generator functions are not allowed in CHECK
+statement error set-returning functions are not allowed in CHECK
 CREATE TABLE uu (x INT CHECK (generate_series(1, 3) < 3))
 
-statement error generate_series\(\): generator functions are not allowed in STORED COMPUTED COLUMN
+statement error generate_series\(\): set-returning functions are not allowed in STORED COMPUTED COLUMN
 CREATE TABLE uu (x INT AS (generate_series(1, 3)) STORED)
 
 subtest correlated_srf
@@ -1317,3 +1317,91 @@ SELECT 1,c1,c2 FROM t95315 JOIN ROWS FROM (CAST(c1 AS INT), CAST(c2 AS INT)) USI
 ----
 1  3  6
 1  4  7
+
+subtest generator-syntax
+
+# Regression test for #97119 and #94890 - return syntax error when CASE or
+# COALESCE is used with a set-generating function as argument.
+statement error pq: set-returning functions are not allowed in CASE
+SELECT CASE generate_series(1, 3) WHEN 3 THEN 0 ELSE 1 END;
+
+statement error pq: set-returning functions are not allowed in CASE
+SELECT CASE WHEN true THEN generate_series(1, 3) ELSE 1 END;
+
+statement error pq: set-returning functions are not allowed in CASE
+SELECT CASE WHEN false THEN 1 ELSE generate_series(1, 3) END;
+
+statement error pq: set-returning functions are not allowed in COALESCE
+SELECT COALESCE(generate_series(1, 10));
+
+# A subquery with a generator function is allowed within CASE and COALESCE.
+query I rowsort
+SELECT CASE WHEN true THEN (SELECT * FROM generate_series(1, 3) LIMIT 1) ELSE 1 END;
+----
+1
+
+query I rowsort
+SELECT COALESCE((SELECT * FROM generate_series(1, 3) LIMIT 1));
+----
+1
+
+# Aggregate functions within CASE and COALESCE are allowed.
+query R rowsort
+SELECT CASE WHEN true THEN sum(x) ELSE 1 END FROM xy;
+----
+15
+
+query R rowsort
+SELECT COALESCE(sum(x)) FROM xy;
+----
+15
+
+# Window functions within CASE and COALESCE are allowed.
+query R rowsort
+SELECT CASE WHEN true THEN sum(x) OVER () ELSE 1 END FROM xy;
+----
+15
+15
+15
+15
+15
+
+query R rowsort
+SELECT COALESCE(sum(x) OVER ()) FROM xy;
+----
+15
+15
+15
+15
+15
+
+# IF does not allow generator functions.
+statement error pq: set-returning functions are not allowed in IF
+SELECT IF(x > y, generate_series(1, 3), 0) FROM xy;
+
+# IFNULL does not allow generator functions.
+statement error pq: set-returning functions are not allowed in COALESCE
+SELECT IFNULL(1, generate_series(1, 2));
+
+# NULLIF allows generator functions.
+query I rowsort
+SELECT NULLIF(generate_series(1, x), generate_series(1, 3)) from xy;
+----
+NULL
+NULL
+NULL
+NULL
+NULL
+NULL
+NULL
+NULL
+NULL
+NULL
+NULL
+NULL
+4
+NULL
+NULL
+NULL
+4
+5

pkg/sql/logictest/testdata/logic_test/update

@@ -550,7 +550,7 @@ UPDATE t32477 SET x = count(x)
 statement error window functions are not allowed in UPDATE SET
 UPDATE t32477 SET x = rank() OVER ()
 
-statement error generator functions are not allowed in UPDATE SET
+statement error pq: generate_series\(\): set-returning functions are not allowed in UPDATE SET
 UPDATE t32477 SET x = generate_series(1,2)
 
 #regression test for #32054

pkg/sql/opt/optbuilder/srfs.go

@@ -45,6 +45,11 @@ func (s *srf) Walk(v tree.Visitor) tree.Expr {
 func (s *srf) TypeCheck(
 	_ context.Context, ctx *tree.SemaContext, desired *types.T,
 ) (tree.TypedExpr, error) {
+	if ctx.Properties.IsSet(tree.RejectGenerators) {
+		// srf replacement can happen before type-checking, so we need to check
+		// invalid usage here.
+		return nil, tree.NewInvalidFunctionUsageError(tree.GeneratorClass, ctx.TypeCheckContext())
+	}
 	if ctx.Properties.Derived.SeenGenerator {
 		// This error happens if this srf struct is nested inside a raw srf that
 		// has not yet been replaced. This is possible since scope.replaceSRF first

pkg/sql/opt/optbuilder/testdata/aggregate

@@ -3633,17 +3633,7 @@ scalar-group-by
 build
 SELECT array_agg(generate_series(1, 2))
 ----
-scalar-group-by
- ├── columns: array_agg:2
- ├── project-set
- │    ├── columns: generate_series:1
- │    ├── values
- │    │    └── ()
- │    └── zip
- │         └── generate_series(1, 2)
- └── aggregations
-      └── array-agg [as=array_agg:2]
-           └── generate_series:1
+error (0A000): array_agg(): set-returning functions are not allowed in aggregate
 
 # Regression test for #37317: duplicate column in GROUP BY
 build format=show-all
@@ -4024,7 +4014,7 @@ scalar-group-by
 build
 SELECT * FROM ROWS FROM (count(json_each('[]')))
 ----
-error (0A000): count(): json_each(): generator functions are not allowed in aggregate
+error (0A000): count(): json_each(): set-returning functions are not allowed in aggregate
 
 # Tests for projecting non-grouping columns when we group by a PK.
 build

pkg/sql/opt/optbuilder/testdata/insert

@@ -466,7 +466,7 @@ error (42803): sum(): aggregate functions are not allowed in RETURNING
 build
 INSERT INTO abcde VALUES (1) RETURNING generate_series(1, 10)
 ----
-error (0A000): generate_series(): generator functions are not allowed in RETURNING
+error (0A000): generate_series(): set-returning functions are not allowed in RETURNING
 
 # Try to use non-returning INSERT as expression.
 build
@@ -641,7 +641,7 @@ error (0A000): WITH clause "cte" does not return any columns
 build
 INSERT INTO abcde VALUES (1) RETURNING generate_series(1, 100)
 ----
-error (0A000): generate_series(): generator functions are not allowed in RETURNING
+error (0A000): generate_series(): set-returning functions are not allowed in RETURNING
 
 # ------------------------------------------------------------------------------
 # Tests with target column names.

pkg/sql/opt/optbuilder/testdata/join

@@ -2413,7 +2413,7 @@ error (42804): argument of ON must be type bool, not type float
 build
 SELECT * FROM foo JOIN bar ON generate_series(0, 1) < 2
 ----
-error (0A000): generate_series(): generator functions are not allowed in ON
+error (0A000): generate_series(): set-returning functions are not allowed in ON
 
 build
 SELECT * FROM foo JOIN bar ON max(foo.c) < 2

pkg/sql/opt/optbuilder/testdata/srfs

@@ -66,7 +66,7 @@ ordinality
 build
 SELECT * FROM (VALUES (1)) LIMIT generate_series(1, 3)
 ----
-error (0A000): generate_series(): generator functions are not allowed in LIMIT
+error (0A000): generate_series(): set-returning functions are not allowed in LIMIT
 
 # multiple_SRFs
 
@@ -709,7 +709,7 @@ error (0A000): generate_series(): generate_series(): set-returning functions mus
 build
 SELECT max(a) FROM t HAVING max(a::int) > generate_series(0, a::int)
 ----
-error (0A000): generate_series(): generator functions are not allowed in HAVING
+error (0A000): generate_series(): set-returning functions are not allowed in HAVING
 
 build
 SELECT max(a) FROM t HAVING max(a::int) > (SELECT generate_series(0, b::int) FROM u limit 1)
@@ -861,7 +861,7 @@ error (42803): column "unnest" must appear in the GROUP BY clause or be used in
 build
 SELECT 0, unnest(ARRAY[0]) GROUP BY 1, 2
 ----
-error (0A000): unnest(): generator functions are not allowed in GROUP BY
+error (0A000): unnest(): set-returning functions are not allowed in GROUP BY
 
 build
 SELECT 0, information_schema._pg_expandarray(ARRAY[0]) GROUP BY 1

pkg/sql/opt/optbuilder/testdata/update

@@ -805,7 +805,7 @@ error (42803): sum(): aggregate functions are not allowed in RETURNING
 build
 UPDATE abcde SET c=1 RETURNING generate_series(1, 10)
 ----
-error (0A000): generate_series(): generator functions are not allowed in RETURNING
+error (0A000): generate_series(): set-returning functions are not allowed in RETURNING
 
 # ------------------------------------------------------------------------------
 # Test DEFAULT values.

pkg/sql/sem/builtins/pg_builtins.go

@@ -566,16 +566,16 @@ func makeToRegOverload(typ *types.T, helpText string) builtinDefinition {
 
 // Format the array {type,othertype} as type, othertype.
 // If there are no args, output the empty string.
-const getFunctionArgStringQuery = `SELECT 
-										COALESCE(
-										    (SELECT trim('{}' FROM replace(
-										        array_agg(unnest(proargtypes)::REGTYPE::TEXT)::TEXT,
-										        ',', ', ')))
-										    , '')
-                    FROM pg_catalog.pg_proc
-                    WHERE oid=$1
-                    GROUP BY oid, proargtypes
-                    LIMIT 1`
+const getFunctionArgStringQuery = `
+SELECT COALESCE(
+    (SELECT trim('{}' FROM replace(
+        (
+            SELECT array_agg(unnested::REGTYPE::TEXT)
+            FROM unnest(proargtypes) AS unnested
+        )::TEXT, ',', ', '))
+    ), '')
+FROM pg_catalog.pg_proc WHERE oid=$1 GROUP BY oid, proargtypes LIMIT 1
+`
 
 var pgBuiltins = map[string]builtinDefinition{
 	// See https://www.postgresql.org/docs/9.6/static/functions-info.html.