sql: implement RETURNS TABLE syntax

This commit implements `RETURNS TABLE` for UDFs. `RETURNS TABLE` is syntactic sugar for `RETURNS SETOF` with: - RECORD if there are multiple TABLE parameters, or - the type of the single TABLE parameter. The TABLE parameters are added to the list of routine parameters. Fixes #100226 Release note (sql change): Added support for `RETURNS TABLE` syntax when creating a UDF.
cockroachdb · Dec 12, 2024 · 074371d · 074371d
1 parent 3ad8c0f
commit 074371d
Show file tree

Hide file tree

Showing 7 changed files with 239 additions and 26 deletions.
diff --git a/docs/generated/sql/bnf/create_func.bnf b/docs/generated/sql/bnf/create_func.bnf
@@ -1,3 +1,4 @@
 create_func_stmt ::=
 	'CREATE' ( 'OR' 'REPLACE' |  ) 'FUNCTION' routine_create_name '(' ( ( ( ( routine_param | routine_param   | routine_param   ) ) ( ( ',' ( routine_param | routine_param   | routine_param   ) ) )* ) |  ) ')' 'RETURNS' ( 'SETOF' |  ) routine_return_type ( ( ( ( 'AS' routine_body_str  | 'LANGUAGE' ('SQL' | 'PLPGSQL') | ( 'CALLED' 'ON' 'NULL' 'INPUT' | 'RETURNS' 'NULL' 'ON' 'NULL' 'INPUT' | 'STRICT' | 'IMMUTABLE' | 'STABLE' | 'VOLATILE' | 'EXTERNAL' 'SECURITY' 'DEFINER' | 'EXTERNAL' 'SECURITY' 'INVOKER' | 'SECURITY' 'DEFINER' | 'SECURITY' 'INVOKER' | 'LEAKPROOF' | 'NOT' 'LEAKPROOF' ) ) ) ( ( ( 'AS' routine_body_str  | 'LANGUAGE' ('SQL' | 'PLPGSQL') | ( 'CALLED' 'ON' 'NULL' 'INPUT' | 'RETURNS' 'NULL' 'ON' 'NULL' 'INPUT' | 'STRICT' | 'IMMUTABLE' | 'STABLE' | 'VOLATILE' | 'EXTERNAL' 'SECURITY' 'DEFINER' | 'EXTERNAL' 'SECURITY' 'INVOKER' | 'SECURITY' 'DEFINER' | 'SECURITY' 'INVOKER' | 'LEAKPROOF' | 'NOT' 'LEAKPROOF' ) ) ) )* ) |  ) 
+	| 'CREATE' ( 'OR' 'REPLACE' |  ) 'FUNCTION' routine_create_name '(' ( ( ( ( routine_param | routine_param   | routine_param   ) ) ( ( ',' ( routine_param | routine_param   | routine_param   ) ) )* ) |  ) ')' 'RETURNS' 'TABLE' '(' table_func_column_list ')' ( ( ( ( 'AS' routine_body_str  | 'LANGUAGE' ('SQL' | 'PLPGSQL') | ( 'CALLED' 'ON' 'NULL' 'INPUT' | 'RETURNS' 'NULL' 'ON' 'NULL' 'INPUT' | 'STRICT' | 'IMMUTABLE' | 'STABLE' | 'VOLATILE' | 'EXTERNAL' 'SECURITY' 'DEFINER' | 'EXTERNAL' 'SECURITY' 'INVOKER' | 'SECURITY' 'DEFINER' | 'SECURITY' 'INVOKER' | 'LEAKPROOF' | 'NOT' 'LEAKPROOF' ) ) ) ( ( ( 'AS' routine_body_str  | 'LANGUAGE' ('SQL' | 'PLPGSQL') | ( 'CALLED' 'ON' 'NULL' 'INPUT' | 'RETURNS' 'NULL' 'ON' 'NULL' 'INPUT' | 'STRICT' | 'IMMUTABLE' | 'STABLE' | 'VOLATILE' | 'EXTERNAL' 'SECURITY' 'DEFINER' | 'EXTERNAL' 'SECURITY' 'INVOKER' | 'SECURITY' 'DEFINER' | 'SECURITY' 'INVOKER' | 'LEAKPROOF' | 'NOT' 'LEAKPROOF' ) ) ) )* ) |  ) 
 	| 'CREATE' ( 'OR' 'REPLACE' |  ) 'FUNCTION' routine_create_name '(' ( ( ( ( routine_param | routine_param   | routine_param   ) ) ( ( ',' ( routine_param | routine_param   | routine_param   ) ) )* ) |  ) ')' ( ( ( ( 'AS' routine_body_str  | 'LANGUAGE' ('SQL' | 'PLPGSQL') | ( 'CALLED' 'ON' 'NULL' 'INPUT' | 'RETURNS' 'NULL' 'ON' 'NULL' 'INPUT' | 'STRICT' | 'IMMUTABLE' | 'STABLE' | 'VOLATILE' | 'EXTERNAL' 'SECURITY' 'DEFINER' | 'EXTERNAL' 'SECURITY' 'INVOKER' | 'SECURITY' 'DEFINER' | 'SECURITY' 'INVOKER' | 'LEAKPROOF' | 'NOT' 'LEAKPROOF' ) ) ) ( ( ( 'AS' routine_body_str  | 'LANGUAGE' ('SQL' | 'PLPGSQL') | ( 'CALLED' 'ON' 'NULL' 'INPUT' | 'RETURNS' 'NULL' 'ON' 'NULL' 'INPUT' | 'STRICT' | 'IMMUTABLE' | 'STABLE' | 'VOLATILE' | 'EXTERNAL' 'SECURITY' 'DEFINER' | 'EXTERNAL' 'SECURITY' 'INVOKER' | 'SECURITY' 'DEFINER' | 'SECURITY' 'INVOKER' | 'LEAKPROOF' | 'NOT' 'LEAKPROOF' ) ) ) )* ) |  ) 
diff --git a/docs/generated/sql/bnf/stmt_block.bnf b/docs/generated/sql/bnf/stmt_block.bnf
@@ -1823,6 +1823,7 @@ create_sequence_stmt ::=
 
 create_func_stmt ::=
 	'CREATE' opt_or_replace 'FUNCTION' routine_create_name '(' opt_routine_param_with_default_list ')' 'RETURNS' opt_return_set routine_return_type opt_create_routine_opt_list opt_routine_body
+	| 'CREATE' opt_or_replace 'FUNCTION' routine_create_name '(' opt_routine_param_with_default_list ')' 'RETURNS' 'TABLE' '(' table_func_column_list ')' opt_create_routine_opt_list opt_routine_body
 	| 'CREATE' opt_or_replace 'FUNCTION' routine_create_name '(' opt_routine_param_with_default_list ')' opt_create_routine_opt_list opt_routine_body
 
 create_proc_stmt ::=
@@ -2630,6 +2631,9 @@ opt_routine_body ::=
 	| 'BEGIN' 'ATOMIC' routine_body_stmt_list 'END'
 	| 
 
+table_func_column_list ::=
+	( table_func_column ) ( ( ',' table_func_column ) )*
+
 trigger_action_time ::=
 	'BEFORE'
 	| 'AFTER'
@@ -3276,6 +3280,9 @@ routine_return_stmt ::=
 routine_body_stmt_list ::=
 	(  ) ( ( routine_body_stmt ';' ) )*
 
+table_func_column ::=
+	param_name routine_param_type
+
 trigger_event ::=
 	'INSERT'
 	| 'DELETE'
@@ -3698,6 +3705,9 @@ routine_body_stmt ::=
 	stmt_without_legacy_transaction
 	| routine_return_stmt
 
+param_name ::=
+	type_function_name
+
 trigger_transition ::=
 	transition_is_new transition_is_row opt_as table_alias_name
 
@@ -4658,9 +4668,6 @@ routine_param_class ::=
 	| 'INOUT'
 	| 'IN' 'OUT'
 
-param_name ::=
-	type_function_name
-
 opt_float ::=
 	'(' 'ICONST' ')'
 	| 

diff --git a/pkg/ccl/logictestccl/testdata/logic_test/plpgsql_unsupported b/pkg/ccl/logictestccl/testdata/logic_test/plpgsql_unsupported
@@ -10,6 +10,12 @@ CREATE OR REPLACE PROCEDURE foo() AS $$
   END
 $$ LANGUAGE PLpgSQL;
 
+statement error pq: unimplemented: set-returning PL/pgSQL functions
+CREATE OR REPLACE FUNCTION bar() RETURNS SETOF INT LANGUAGE PLpgSQL AS $$ BEGIN RETURN NEXT 100; END $$;
+
+statement error pq: unimplemented: set-returning PL/pgSQL functions
+CREATE OR REPLACE FUNCTION bar() RETURNS TABLE (x INT) LANGUAGE PLpgSQL AS $$ BEGIN RETURN NEXT 100; END $$;
+
 subtest error_detail
 
 # Regression test for #123672 - annotate "unsupported" errors with the

diff --git a/pkg/sql/logictest/testdata/logic_test/udf_setof b/pkg/sql/logictest/testdata/logic_test/udf_setof
@@ -237,3 +237,73 @@ SELECT * FROM f128403();
 10  10
 
 subtest end
+
+# RETURNS TABLE is syntactic sugar for RETURNS SETOF with:
+# - RECORD if there are multiple TABLE parameters, or
+# - the type of the single TABLE parameter.
+subtest returns_table
+
+statement error pgcode 42601 pq: OUT and INOUT arguments aren't allowed in TABLE functions
+CREATE FUNCTION f_table1(OUT x INT, OUT y TEXT) RETURNS TABLE(x INT, y TEXT) AS $$
+  SELECT t, t::TEXT FROM generate_series(1, 10) g(t);
+$$ LANGUAGE SQL;
+
+statement ok
+CREATE FUNCTION f_table1() RETURNS TABLE(x INT, y TEXT) AS $$
+  SELECT t, t::TEXT FROM generate_series(1, 10) g(t);
+$$ LANGUAGE SQL;
+
+query T rowsort
+select f_table1();
+----
+(1,1)
+(2,2)
+(3,3)
+(4,4)
+(5,5)
+(6,6)
+(7,7)
+(8,8)
+(9,9)
+(10,10)
+
+query IT rowsort
+SELECT * FROM f_table1();
+----
+1   1
+2   2
+3   3
+4   4
+5   5
+6   6
+7   7
+8   8
+9   9
+10  10
+
+# Case with a single TABLE parameter.
+statement ok
+CREATE FUNCTION f_table2() RETURNS TABLE(x INT) AS $$
+  SELECT t FROM generate_series(1, 10) g(t);
+$$ LANGUAGE SQL;
+
+query I rowsort
+select f_table2();
+----
+1
+2
+3
+4
+5
+6
+7
+8
+9
+10
+
+statement error pgcode 42P13 return type mismatch in function declared to return int\nDETAIL: Actual return type is record
+CREATE FUNCTION err() RETURNS TABLE (x INT) STRICT LANGUAGE SQL AS $$
+  SELECT a, b FROM ab ORDER BY a
+$$
+
+subtest end
diff --git a/pkg/sql/parser/lexer.go b/pkg/sql/parser/lexer.go
@@ -394,6 +394,13 @@ func (l *lexer) setErr(err error) {
 	l.populateErrorDetails()
 }
 
+// setErrNoDetails is similar to setErr, but is used for an error that should
+// not be further annotated with details.
+func (l *lexer) setErrNoDetails(err error) {
+	err = pgerror.WithCandidateCode(err, pgcode.Syntax)
+	l.lastError = err
+}
+
 func (l *lexer) Error(e string) {
 	e = strings.TrimPrefix(e, "syntax error: ") // we'll add it again below.
 	l.lastError = pgerror.WithCandidateCode(errors.Newf("%s", e), pgcode.Syntax)

diff --git a/pkg/sql/parser/sql.y b/pkg/sql/parser/sql.y
@@ -66,6 +66,11 @@ func setErr(sqllex sqlLexer, err error) int {
     return 1
 }
 
+func setErrNoDetails(sqllex sqlLexer, err error) int {
+    sqllex.(*lexer).setErrNoDetails(err)
+    return 1
+}
+
 func unimplementedWithIssue(sqllex sqlLexer, issue int) int {
     sqllex.(*lexer).UnimplementedWithIssue(issue)
     return 1
@@ -1725,10 +1730,11 @@ func (u *sqlSymUnion) triggerForEach() tree.TriggerForEach {
 %type <privilege.TargetObjectType> target_object_type
 
 // Routine (UDF/SP) relevant components.
-%type <bool> opt_or_replace opt_return_table opt_return_set opt_no
+%type <bool> opt_or_replace opt_return_set opt_no
 %type <str> param_name routine_as
-%type <tree.RoutineParams> opt_routine_param_with_default_list routine_param_with_default_list func_params func_params_list
-%type <tree.RoutineParam> routine_param_with_default routine_param
+%type <tree.RoutineParams> opt_routine_param_with_default_list routine_param_with_default_list
+%type <tree.RoutineParams> func_params func_params_list table_func_column_list
+%type <tree.RoutineParam> routine_param_with_default routine_param table_func_column
 %type <tree.ResolvableTypeReference> routine_return_type routine_param_type
 %type <tree.RoutineOptions> opt_create_routine_opt_list create_routine_opt_list alter_func_opt_list
 %type <tree.RoutineOption> create_routine_opt_item common_routine_opt_item
@@ -4869,8 +4875,7 @@ create_extension_stmt:
 // %SeeAlso: WEBDOCS/create-function.html
 create_func_stmt:
   CREATE opt_or_replace FUNCTION routine_create_name '(' opt_routine_param_with_default_list ')'
-  RETURNS opt_return_table opt_return_set routine_return_type
-  opt_create_routine_opt_list opt_routine_body
+  RETURNS opt_return_set routine_return_type opt_create_routine_opt_list opt_routine_body
   {
     name := $4.unresolvedObjectName().ToRoutineName()
     $$.val = &tree.CreateRoutine{
@@ -4879,11 +4884,43 @@ create_func_stmt:
       Name: name,
       Params: $6.routineParams(),
       ReturnType: &tree.RoutineReturnType{
-        Type: $11.typeReference(),
-        SetOf: $10.bool(),
+        Type: $10.typeReference(),
+        SetOf: $9.bool(),
       },
-      Options: $12.routineOptions(),
-      RoutineBody: $13.routineBody(),
+      Options: $11.routineOptions(),
+      RoutineBody: $12.routineBody(),
+    }
+  }
+| CREATE opt_or_replace FUNCTION routine_create_name '(' opt_routine_param_with_default_list ')'
+  RETURNS TABLE '(' table_func_column_list ')' opt_create_routine_opt_list opt_routine_body
+  {
+    // RETURNS TABLE is syntactic sugar for RETURNS SETOF with:
+    // - RECORD if there are multiple TABLE parameters, or
+    // - the type of the single TABLE parameter.
+    // The TABLE parameters are added to the list of routine parameters.
+    tableParams := $11.routineParams()
+    returnType := tree.ResolvableTypeReference(types.AnyTuple)
+    if len(tableParams) == 1 {
+      returnType = tableParams[0].Type
+    }
+    routineParams := $6.routineParams()
+    for i := range routineParams {
+      // OUT parameters are not allowed in table functions.
+      if tree.IsOutParamClass(routineParams[i].Class) {
+        return setErrNoDetails(sqllex, errors.New("OUT and INOUT arguments aren't allowed in TABLE functions"))
+      }
+    }
+    $$.val = &tree.CreateRoutine{
+      IsProcedure: false,
+      Replace: $2.bool(),
+      Name: $4.unresolvedObjectName().ToRoutineName(),
+      Params: append(routineParams, tableParams...),
+      ReturnType: &tree.RoutineReturnType{
+        Type: returnType,
+        SetOf: true,
+      },
+      Options: $13.routineOptions(),
+      RoutineBody: $14.routineBody(),
     }
   }
 | CREATE opt_or_replace FUNCTION routine_create_name '(' opt_routine_param_with_default_list ')'
@@ -4932,10 +4969,6 @@ opt_or_replace:
   OR REPLACE { $$.val = true }
 | /* EMPTY */ { $$.val = false }
 
-opt_return_table:
-  TABLE { return unimplementedWithIssueDetail(sqllex, 100226, "UDF returning TABLE") }
-| /* EMPTY */ { $$.val = false }
-
 opt_return_set:
   SETOF { $$.val = true}
 | /* EMPTY */ { $$.val = false }
@@ -5022,6 +5055,25 @@ routine_param_type:
 routine_return_type:
   routine_param_type
 
+table_func_column: param_name routine_param_type
+  {
+    $$.val = tree.RoutineParam{
+      Name: tree.Name($1),
+      Type: $2.typeReference(),
+      Class: tree.RoutineParamOut,
+    }
+  }
+
+table_func_column_list:
+  table_func_column
+  {
+    $$.val = tree.RoutineParams{$1.routineParam()}
+  }
+| table_func_column_list ',' table_func_column
+  {
+    $$.val = append($1.routineParams(), $3.routineParam())
+  }
+
 opt_create_routine_opt_list:
   create_routine_opt_list { $$.val = $1.routineOptions() }
 | /* EMPTY */ { $$.val = tree.RoutineOptions{} }

diff --git a/pkg/sql/parser/testdata/create_function b/pkg/sql/parser/testdata/create_function
@@ -469,16 +469,6 @@ CREATE FUNCTION _()
 	LANGUAGE plpgsql
 	AS $$_$$ -- identifiers removed
 
-error
-CREATE FUNCTION f() RETURNS TABLE 'SELECT 1' LANGUAGE SQL
-----
-at or near "table": syntax error: unimplemented: this syntax
-DETAIL: source SQL:
-CREATE FUNCTION f() RETURNS TABLE 'SELECT 1' LANGUAGE SQL
-                            ^
-HINT: You have attempted to use a feature that is not yet implemented.
-See: https://go.crdb.dev/issue-v/100226/
-
 parse
 CREATE OR REPLACE FUNCTION f(a int = 7) RETURNS INT EXTERNAL SECURITY DEFINER AS 'SELECT 1' LANGUAGE SQL
 ----
@@ -574,3 +564,83 @@ CREATE OR REPLACE FUNCTION _(_ INT8 DEFAULT 7)
 	SECURITY INVOKER
 	LANGUAGE SQL
 	AS $$_$$ -- identifiers removed
+
+parse
+CREATE FUNCTION f(OUT x INT) RETURNS SETOF INT LANGUAGE SQL AS $$ SELECT 1 $$;
+----
+CREATE FUNCTION f(OUT x INT8)
+	RETURNS SETOF INT8
+	LANGUAGE SQL
+	AS $$ SELECT 1 $$ -- normalized!
+CREATE FUNCTION f(OUT x INT8)
+	RETURNS SETOF INT8
+	LANGUAGE SQL
+	AS $$ SELECT 1 $$ -- fully parenthesized
+CREATE FUNCTION f(OUT x INT8)
+	RETURNS SETOF INT8
+	LANGUAGE SQL
+	AS $$_$$ -- literals removed
+CREATE FUNCTION _(OUT _ INT8)
+	RETURNS SETOF INT8
+	LANGUAGE SQL
+	AS $$_$$ -- identifiers removed
+
+parse
+CREATE FUNCTION f(OUT x INT, OUT y INT) RETURNS SETOF RECORD LANGUAGE SQL AS $$ SELECT 1, 2 $$;
+----
+CREATE FUNCTION f(OUT x INT8, OUT y INT8)
+	RETURNS SETOF RECORD
+	LANGUAGE SQL
+	AS $$ SELECT 1, 2 $$ -- normalized!
+CREATE FUNCTION f(OUT x INT8, OUT y INT8)
+	RETURNS SETOF RECORD
+	LANGUAGE SQL
+	AS $$ SELECT 1, 2 $$ -- fully parenthesized
+CREATE FUNCTION f(OUT x INT8, OUT y INT8)
+	RETURNS SETOF RECORD
+	LANGUAGE SQL
+	AS $$_$$ -- literals removed
+CREATE FUNCTION _(OUT _ INT8, OUT _ INT8)
+	RETURNS SETOF RECORD
+	LANGUAGE SQL
+	AS $$_$$ -- identifiers removed
+
+parse
+CREATE FUNCTION f() RETURNS TABLE(x INT) LANGUAGE SQL AS $$ SELECT 1 $$;
+----
+CREATE FUNCTION f(OUT x INT8)
+	RETURNS SETOF INT8
+	LANGUAGE SQL
+	AS $$ SELECT 1 $$ -- normalized!
+CREATE FUNCTION f(OUT x INT8)
+	RETURNS SETOF INT8
+	LANGUAGE SQL
+	AS $$ SELECT 1 $$ -- fully parenthesized
+CREATE FUNCTION f(OUT x INT8)
+	RETURNS SETOF INT8
+	LANGUAGE SQL
+	AS $$_$$ -- literals removed
+CREATE FUNCTION _(OUT _ INT8)
+	RETURNS SETOF INT8
+	LANGUAGE SQL
+	AS $$_$$ -- identifiers removed
+
+parse
+CREATE FUNCTION f() RETURNS TABLE(x INT, y INT) LANGUAGE SQL AS $$ SELECT 1, 2 $$;
+----
+CREATE FUNCTION f(OUT x INT8, OUT y INT8)
+	RETURNS SETOF RECORD
+	LANGUAGE SQL
+	AS $$ SELECT 1, 2 $$ -- normalized!
+CREATE FUNCTION f(OUT x INT8, OUT y INT8)
+	RETURNS SETOF RECORD
+	LANGUAGE SQL
+	AS $$ SELECT 1, 2 $$ -- fully parenthesized
+CREATE FUNCTION f(OUT x INT8, OUT y INT8)
+	RETURNS SETOF RECORD
+	LANGUAGE SQL
+	AS $$_$$ -- literals removed
+CREATE FUNCTION _(OUT _ INT8, OUT _ INT8)
+	RETURNS SETOF RECORD
+	LANGUAGE SQL
+	AS $$_$$ -- identifiers removed