plpgsql/parser: add parser support for RETURN NEXT

This commit adds support for `RETURN NEXT` statements to the PL/pgSQL
parser. A following commit will add execution support.

Informs #105240

Release note: None

Author: DrewKimball

pkg/sql/plpgsql/parser/lexer.go

@@ -373,6 +373,17 @@ func makeDoStmt(options tree.DoBlockOptions) (*plpgsqltree.DoBlock, error) {
 	return &plpgsqltree.DoBlock{Block: parsedStmt.AST}, nil
 }
 
+// ParseReturnExpr handles reading and parsing the expression for a RETURN or
+// RETURN NEXT statement, which can be nonexistent.
+func (l *lexer) ParseReturnExpr() (plpgsqltree.Expr, error) {
+	startPos, endPos, _, err := l.readSQLConstruct(true /* isExpr */, true /* allowEmpty */, ';')
+	if err != nil || startPos == endPos {
+		return nil, err
+	}
+	exprStr := l.getStr(startPos, endPos)
+	return l.ParseExpr(exprStr)
+}
+
 func (l *lexer) ReadSqlExpr(
 	terminator1 int, terminators ...int,
 ) (sqlStr string, terminatorMet int, err error) {
@@ -383,17 +394,6 @@ func (l *lexer) ReadSqlExpr(
 	return l.getStr(startPos, endPos), terminatorMet, err
 }
 
-// ReadReturnExpr handles reading the expression for a RETURN statement, which
-// can be nonexistent.
-func (l *lexer) ReadReturnExpr() (sqlStr string, err error) {
-	var startPos, endPos int
-	startPos, endPos, _, err = l.readSQLConstruct(true /* isExpr */, true /* allowEmpty */, ';')
-	if err != nil || startPos == endPos {
-		return "", err
-	}
-	return l.getStr(startPos, endPos), err
-}
-
 func (l *lexer) ReadSqlStatement(
 	terminator1 int, terminators ...int,
 ) (sqlStr string, terminatorMet int, err error) {

pkg/sql/plpgsql/parser/plpgsql.y

@@ -338,14 +338,14 @@ func (u *plpgsqlSymUnion) doBlockOption() tree.DoBlockOption {
 }
 
 %type <str> decl_varname decl_defkey
-%type <bool>	decl_const decl_notnull
+%type <bool> decl_const decl_notnull
 %type <plpgsqltree.Expr>	decl_defval decl_cursor_query
 %type <tree.ResolvableTypeReference>	decl_datatype
-%type <str>		decl_collate
+%type <str>	decl_collate
 
-%type <str>	expr_until_semi expr_until_paren stmt_until_semi return_expr
+%type <str>	expr_until_semi expr_until_paren stmt_until_semi
 %type <str>	expr_until_then expr_until_loop opt_expr_until_when
-%type <plpgsqltree.Expr>	opt_exitcond
+%type <plpgsqltree.Expr> return_expr opt_exitcond
 
 %type <[]plpgsqltree.Variable> for_target
 %type <*tree.NumVal> foreach_slice
@@ -1157,19 +1157,11 @@ stmt_continue: CONTINUE opt_label opt_exitcond
 
 stmt_return: RETURN return_expr ';'
   {
-    var expr plpgsqltree.Expr
-    if $2 != "" {
-      var err error
-      expr, err = plpgsqllex.(*lexer).ParseExpr($2)
-      if err != nil {
-        return setErr(plpgsqllex, err)
-      }
-    }
-    $$.val = &plpgsqltree.Return{Expr: expr}
+    $$.val = &plpgsqltree.Return{Expr: $2.expr()}
   }
-| RETURN_NEXT NEXT
+| RETURN_NEXT NEXT return_expr ';'
   {
-    return unimplemented(plpgsqllex, "return next")
+    $$.val = &plpgsqltree.ReturnNext{Expr: $3.expr()}
   }
 | RETURN_QUERY QUERY
  {
@@ -1179,11 +1171,11 @@ stmt_return: RETURN return_expr ';'
 
 return_expr:
   {
-    sqlStr, err := plpgsqllex.(*lexer).ReadReturnExpr()
+    retExpr, err := plpgsqllex.(*lexer).ParseReturnExpr()
     if err != nil {
       return setErr(plpgsqllex, err)
     }
-    $$ = sqlStr
+    $$.val = retExpr
   }
 ;
 

pkg/sql/plpgsql/parser/testdata/stmt_return

@@ -190,32 +190,6 @@ We appreciate your feedback.
 ----
 ----
 
-error
-DECLARE
-BEGIN
-  RETURN NEXT 1 + 1;
-END
-----
-----
-at or near "next": syntax error: unimplemented: this syntax
-DETAIL: source SQL:
-DECLARE
-BEGIN
-  RETURN NEXT 1 + 1;
-         ^
-HINT: You have attempted to use a feature that is not yet implemented.
-
-Please check the public issue tracker to check whether this problem is
-already tracked. If you cannot find it there, please report the error
-with details by creating a new issue.
-
-If you would rather not post publicly, please contact us directly
-using the support form.
-
-We appreciate your feedback.
-----
-----
-
 error
 DECLARE
 BEGIN
@@ -263,22 +237,22 @@ BEGIN
   RETURN 1, 'string';
 END
 ----
-at or near ";": syntax error: query returned 2 columns
+at or near "string": syntax error: query returned 2 columns
 DETAIL: source SQL:
 DECLARE
 BEGIN
   RETURN 1, 'string';
-                    ^
+            ^
 
 error
 DECLARE
 BEGIN
   RETURN 1, (2, 3, 4, 5);
 END
 ----
-at or near ";": syntax error: query returned 2 columns
+at or near ")": syntax error: query returned 2 columns
 DETAIL: source SQL:
 DECLARE
 BEGIN
   RETURN 1, (2, 3, 4, 5);
-                        ^
+                       ^

pkg/sql/plpgsql/parser/testdata/stmt_return_next

@@ -0,0 +1,137 @@
+parse
+BEGIN
+  RETURN NEXT foo;
+END
+----
+BEGIN
+RETURN NEXT foo;
+END;
+ -- normalized!
+BEGIN
+RETURN NEXT (foo);
+END;
+ -- fully parenthesized
+BEGIN
+RETURN NEXT foo;
+END;
+ -- literals removed
+BEGIN
+RETURN NEXT _;
+END;
+ -- identifiers removed
+
+parse
+BEGIN
+  RETURN NEXT 1 + foo + 3;
+END
+----
+BEGIN
+RETURN NEXT (1 + foo) + 3;
+END;
+ -- normalized!
+BEGIN
+RETURN NEXT ((((1) + (foo))) + (3));
+END;
+ -- fully parenthesized
+BEGIN
+RETURN NEXT (_ + foo) + _;
+END;
+ -- literals removed
+BEGIN
+RETURN NEXT (1 + _) + 3;
+END;
+ -- identifiers removed
+
+parse
+BEGIN
+  RETURN NEXT;
+END
+----
+BEGIN
+RETURN NEXT;
+END;
+ -- normalized!
+BEGIN
+RETURN NEXT;
+END;
+ -- fully parenthesized
+BEGIN
+RETURN NEXT;
+END;
+ -- literals removed
+BEGIN
+RETURN NEXT;
+END;
+ -- identifiers removed
+
+parse
+BEGIN
+  RETURN NEXT (1 + 2) + 3;
+END
+----
+BEGIN
+RETURN NEXT (1 + 2) + 3;
+END;
+ -- normalized!
+BEGIN
+RETURN NEXT (((((1) + (2)))) + (3));
+END;
+ -- fully parenthesized
+BEGIN
+RETURN NEXT (_ + _) + _;
+END;
+ -- literals removed
+BEGIN
+RETURN NEXT (1 + 2) + 3;
+END;
+ -- identifiers removed
+
+parse
+BEGIN
+  RETURN NEXT (SELECT * FROM xy);
+END
+----
+BEGIN
+RETURN NEXT (SELECT * FROM xy);
+END;
+ -- normalized!
+BEGIN
+RETURN NEXT ((SELECT (*) FROM xy));
+END;
+ -- fully parenthesized
+BEGIN
+RETURN NEXT (SELECT * FROM xy);
+END;
+ -- literals removed
+BEGIN
+RETURN NEXT (SELECT * FROM _);
+END;
+ -- identifiers removed
+
+error
+BEGIN
+  RETURN NEXT (foo;
+END
+----
+at or near "EOF": syntax error: mismatched parentheses
+DETAIL: source SQL:
+BEGIN
+  RETURN NEXT (foo;
+END
+   ^
+
+error
+BEGIN
+  RETURN NEXT SELECT * FROM xy;
+END
+----
+at or near "xy": at or near "select": syntax error
+DETAIL: source SQL:
+SET ROW (SELECT * FROM xy)
+         ^
+--
+source SQL:
+BEGIN
+  RETURN NEXT SELECT * FROM xy;
+                            ^
+HINT: try \h SET SESSION

pkg/sql/sem/plpgsqltree/statements.go

@@ -824,19 +824,30 @@ func (s *Return) WalkStmt(visitor StatementVisitor) Statement {
 
 type ReturnNext struct {
 	StatementImpl
-	Expr   Expr
-	RetVar Variable
+	Expr Expr
+}
+
+func (s *ReturnNext) CopyNode() *ReturnNext {
+	copyNode := *s
+	return &copyNode
 }
 
 func (s *ReturnNext) Format(ctx *tree.FmtCtx) {
+	ctx.WriteString("RETURN NEXT")
+	if s.Expr != nil {
+		ctx.WriteByte(' ')
+		ctx.FormatNode(s.Expr)
+	}
+	ctx.WriteString(";\n")
 }
 
 func (s *ReturnNext) PlpgSQLStatementTag() string {
 	return "stmt_return_next"
 }
 
 func (s *ReturnNext) WalkStmt(visitor StatementVisitor) Statement {
-	panic(unimplemented.New("plpgsql visitor", "Unimplemented PLpgSQL visitor pattern"))
+	newStmt, _ := visitor.Visit(s)
+	return newStmt
 }
 
 type ReturnQuery struct {

pkg/sql/sem/plpgsqltree/visitor.go

@@ -190,6 +190,16 @@ func (v *SQLStmtVisitor) Visit(stmt Statement) (newStmt Statement, recurse bool)
 			cpy.Expr = e
 			newStmt = cpy
 		}
+	case *ReturnNext:
+		e, v.Err = v.visitExpr(t.Expr)
+		if v.Err != nil {
+			return stmt, false
+		}
+		if t.Expr != e {
+			cpy := t.CopyNode()
+			cpy.Expr = e
+			newStmt = cpy
+		}
 	case *Raise:
 		for i, p := range t.Params {
 			e, v.Err = v.visitExpr(p)
@@ -287,8 +297,7 @@ func (v *SQLStmtVisitor) Visit(stmt Statement) (newStmt Statement, recurse bool)
 			}
 		}
 
-	case *ForEachArray, *ReturnNext,
-		*ReturnQuery, *Perform:
+	case *ForEachArray, *ReturnQuery, *Perform:
 		panic(unimp.New("plpgsql visitor", "Unimplemented PLpgSQL visitor"))
 	}
 	if v.Err != nil {