Skip to content

Commit bdd627e

Browse files
DrewKimballDrewKimball
committed
plpgsql/parser: add parser support for RETURN QUERY
This commit adds support for `RETURN QUERY` statements to the PL/pgSQL parser. A following commit will add execution support. Informs #105240 Release note: None
1 parent cc4b900 commit bdd627e

File tree

6 files changed

+231
-69
lines changed

6 files changed

+231
-69
lines changed

‎pkg/sql/plpgsql/parser/lexer.go‎

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -384,6 +384,31 @@ func (l *lexer) ParseReturnExpr() (plpgsqltree.Expr, error) {
384384
return l.ParseExpr(exprStr)
385385
}
386386

387+
// ParseReturnQuery handles reading and parsing the query for a RETURN QUERY
388+
// statement, which can be nonexistent.
389+
func (l *lexer) ParseReturnQuery() (plpgsqltree.Statement, error) {
390+
startPos, endPos, _, err := l.readSQLConstruct(false /* isExpr */, false /* allowEmpty */, ';')
391+
if err != nil || startPos == endPos {
392+
return nil, err
393+
}
394+
queryStr := l.getStr(startPos, endPos)
395+
stmt, err := parser.ParseOne(queryStr)
396+
if err != nil {
397+
return nil, err
398+
}
399+
return &plpgsqltree.ReturnQuery{SqlStmt: stmt.AST}, nil
400+
}
401+
402+
// peekForExecute checks whether the next token is EXECUTE, used to identify
403+
// dynamic SQL statements.
404+
func (l *lexer) peekForExecute() bool {
405+
if l.parser.Lookahead() != -1 {
406+
// Push back the lookahead token so that it can be included.
407+
l.PushBack(1)
408+
}
409+
return l.Peek().id == EXECUTE
410+
}
411+
387412
func (l *lexer) ReadSqlExpr(
388413
terminator1 int, terminators ...int,
389414
) (sqlStr string, terminatorMet int, err error) {
@@ -517,6 +542,21 @@ func (l *lexer) PushBack(n int) {
517542
}
518543
}
519544

545+
// Advance advances the lexer by n tokens.
546+
func (l *lexer) Advance(n int) {
547+
if n < 0 {
548+
panic(errors.AssertionFailedf("negative n provided to Advance"))
549+
}
550+
l.lastPos += n
551+
if l.lastPos > len(l.tokens) {
552+
l.lastPos = len(l.tokens)
553+
}
554+
if n >= 1 {
555+
// Invalidate the parser lookahead token.
556+
l.parser.(*plpgsqlParserImpl).char = -1
557+
}
558+
}
559+
520560
func (l *lexer) lastToken() plpgsqlSymType {
521561
if l.lastPos < 0 {
522562
return plpgsqlSymType{}

‎pkg/sql/plpgsql/parser/plpgsql.y‎

Lines changed: 13 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -351,7 +351,7 @@ func (u *plpgsqlSymUnion) doBlockOption() tree.DoBlockOption {
351351
%type <*tree.NumVal> foreach_slice
352352
%type <plpgsqltree.ForLoopControl> for_control
353353

354-
%type <str> any_identifier opt_block_label opt_loop_label opt_label query_options
354+
%type <str> any_identifier opt_block_label opt_loop_label opt_label
355355
%type <str> opt_error_level option_type
356356

357357
%type <tree.DoBlockOptions> do_stmt_opt_list
@@ -368,6 +368,7 @@ func (u *plpgsqlSymUnion) doBlockOption() tree.DoBlockOption {
368368
%type <plpgsqltree.Statement> stmt_open stmt_fetch stmt_move stmt_close stmt_null
369369
%type <plpgsqltree.Statement> stmt_commit stmt_rollback
370370
%type <plpgsqltree.Statement> stmt_case stmt_foreach_a
371+
%type <plpgsqltree.Statement> return_query
371372

372373
%type <plpgsqltree.Statement> decl_statement
373374
%type <[]plpgsqltree.Statement> decl_sect opt_decl_stmts decl_stmts
@@ -1163,10 +1164,10 @@ stmt_return: RETURN return_expr ';'
11631164
{
11641165
$$.val = &plpgsqltree.ReturnNext{Expr: $3.expr()}
11651166
}
1166-
| RETURN_QUERY QUERY
1167-
{
1168-
return unimplemented (plpgsqllex, "return query")
1169-
}
1167+
| RETURN_QUERY QUERY return_query ';'
1168+
{
1169+
$$.val = $3.statement()
1170+
}
11701171
;
11711172

11721173
return_expr:
@@ -1179,19 +1180,19 @@ return_expr:
11791180
}
11801181
;
11811182

1182-
query_options:
1183+
return_query:
11831184
{
1184-
_, terminator, err := plpgsqllex.(*lexer).ReadSqlExpr(EXECUTE, ';')
1185-
if err != nil {
1186-
return setErr(plpgsqllex, err)
1187-
}
1188-
if terminator == EXECUTE {
1185+
if plpgsqllex.(*lexer).peekForExecute() {
1186+
// Advance the lexer by one token so that the error correctly points to
1187+
// the EXECUTE keyword.
1188+
plpgsqllex.(*lexer).Advance(1)
11891189
return unimplemented (plpgsqllex, "return dynamic sql query")
11901190
}
1191-
_, _, err = plpgsqllex.(*lexer).ReadSqlExpr(';')
1191+
retQuery, err := plpgsqllex.(*lexer).ParseReturnQuery()
11921192
if err != nil {
11931193
return setErr(plpgsqllex, err)
11941194
}
1195+
$$.val = retQuery
11951196
}
11961197
;
11971198

‎pkg/sql/plpgsql/parser/testdata/stmt_return‎

Lines changed: 0 additions & 52 deletions
Original file line numberDiff line numberDiff line change
@@ -138,58 +138,6 @@ RETURN;
138138
END;
139139
-- identifiers removed
140140

141-
error
142-
DECLARE
143-
BEGIN
144-
RETURN QUERY SELECT 1 + 1;
145-
END
146-
----
147-
----
148-
at or near "query": syntax error: unimplemented: this syntax
149-
DETAIL: source SQL:
150-
DECLARE
151-
BEGIN
152-
RETURN QUERY SELECT 1 + 1;
153-
^
154-
HINT: You have attempted to use a feature that is not yet implemented.
155-
156-
Please check the public issue tracker to check whether this problem is
157-
already tracked. If you cannot find it there, please report the error
158-
with details by creating a new issue.
159-
160-
If you would rather not post publicly, please contact us directly
161-
using the support form.
162-
163-
We appreciate your feedback.
164-
----
165-
----
166-
167-
error
168-
DECLARE
169-
BEGIN
170-
RETURN QUERY EXECUTE a dynamic command;
171-
END
172-
----
173-
----
174-
at or near "query": syntax error: unimplemented: this syntax
175-
DETAIL: source SQL:
176-
DECLARE
177-
BEGIN
178-
RETURN QUERY EXECUTE a dynamic command;
179-
^
180-
HINT: You have attempted to use a feature that is not yet implemented.
181-
182-
Please check the public issue tracker to check whether this problem is
183-
already tracked. If you cannot find it there, please report the error
184-
with details by creating a new issue.
185-
186-
If you would rather not post publicly, please contact us directly
187-
using the support form.
188-
189-
We appreciate your feedback.
190-
----
191-
----
192-
193141
error
194142
DECLARE
195143
BEGIN

‎pkg/sql/plpgsql/parser/testdata/stmt_return_query‎

Lines changed: 153 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,153 @@
1+
parse
2+
DECLARE
3+
BEGIN
4+
RETURN QUERY SELECT 1 + 1;
5+
END
6+
----
7+
DECLARE
8+
BEGIN
9+
RETURN QUERY SELECT 1 + 1;
10+
END;
11+
-- normalized!
12+
DECLARE
13+
BEGIN
14+
RETURN QUERY SELECT ((1) + (1));
15+
END;
16+
-- fully parenthesized
17+
DECLARE
18+
BEGIN
19+
RETURN QUERY SELECT _ + _;
20+
END;
21+
-- literals removed
22+
DECLARE
23+
BEGIN
24+
RETURN QUERY SELECT 1 + 1;
25+
END;
26+
-- identifiers removed
27+
28+
parse
29+
DECLARE
30+
BEGIN
31+
RETURN QUERY SELECT * FROM xy INNER JOIN ab ON x = a;
32+
END
33+
----
34+
DECLARE
35+
BEGIN
36+
RETURN QUERY SELECT * FROM xy INNER JOIN ab ON x = a;
37+
END;
38+
-- normalized!
39+
DECLARE
40+
BEGIN
41+
RETURN QUERY SELECT (*) FROM xy INNER JOIN ab ON ((x) = (a));
42+
END;
43+
-- fully parenthesized
44+
DECLARE
45+
BEGIN
46+
RETURN QUERY SELECT * FROM xy INNER JOIN ab ON x = a;
47+
END;
48+
-- literals removed
49+
DECLARE
50+
BEGIN
51+
RETURN QUERY SELECT * FROM _ INNER JOIN _ ON _ = _;
52+
END;
53+
-- identifiers removed
54+
55+
parse
56+
DECLARE
57+
BEGIN
58+
RETURN QUERY (SELECT * FROM xy INNER JOIN ab ON x = a);
59+
END
60+
----
61+
DECLARE
62+
BEGIN
63+
RETURN QUERY (SELECT * FROM xy INNER JOIN ab ON x = a);
64+
END;
65+
-- normalized!
66+
DECLARE
67+
BEGIN
68+
RETURN QUERY (SELECT (*) FROM xy INNER JOIN ab ON ((x) = (a)));
69+
END;
70+
-- fully parenthesized
71+
DECLARE
72+
BEGIN
73+
RETURN QUERY (SELECT * FROM xy INNER JOIN ab ON x = a);
74+
END;
75+
-- literals removed
76+
DECLARE
77+
BEGIN
78+
RETURN QUERY (SELECT * FROM _ INNER JOIN _ ON _ = _);
79+
END;
80+
-- identifiers removed
81+
82+
error
83+
DECLARE
84+
BEGIN
85+
RETURN QUERY;
86+
END
87+
----
88+
at or near "query": syntax error: missing SQL statement
89+
DETAIL: source SQL:
90+
DECLARE
91+
BEGIN
92+
RETURN QUERY;
93+
^
94+
95+
error
96+
DECLARE
97+
BEGIN
98+
RETURN QUERY 100;
99+
END
100+
----
101+
at or near "100": at or near "100": syntax error
102+
DETAIL: source SQL:
103+
100
104+
^
105+
--
106+
source SQL:
107+
DECLARE
108+
BEGIN
109+
RETURN QUERY 100;
110+
^
111+
112+
error
113+
DECLARE
114+
BEGIN
115+
RETURN QUERY * FROM xy INNER JOIN ab ON x = a;
116+
END
117+
----
118+
at or near "a": at or near "*": syntax error
119+
DETAIL: source SQL:
120+
* FROM xy INNER JOIN ab ON x = a
121+
^
122+
--
123+
source SQL:
124+
DECLARE
125+
BEGIN
126+
RETURN QUERY * FROM xy INNER JOIN ab ON x = a;
127+
^
128+
129+
error
130+
DECLARE
131+
BEGIN
132+
RETURN QUERY EXECUTE a dynamic command;
133+
END
134+
----
135+
----
136+
at or near "execute": syntax error: unimplemented: this syntax
137+
DETAIL: source SQL:
138+
DECLARE
139+
BEGIN
140+
RETURN QUERY EXECUTE a dynamic command;
141+
^
142+
HINT: You have attempted to use a feature that is not yet implemented.
143+
144+
Please check the public issue tracker to check whether this problem is
145+
already tracked. If you cannot find it there, please report the error
146+
with details by creating a new issue.
147+
148+
If you would rather not post publicly, please contact us directly
149+
using the support form.
150+
151+
We appreciate your feedback.
152+
----
153+
----

‎pkg/sql/sem/plpgsqltree/statements.go‎

Lines changed: 14 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -852,20 +852,30 @@ func (s *ReturnNext) WalkStmt(visitor StatementVisitor) Statement {
852852

853853
type ReturnQuery struct {
854854
StatementImpl
855-
Query Expr
856-
DynamicQuery Expr
857-
Params []Expr
855+
SqlStmt tree.Statement
856+
}
857+
858+
func (s *ReturnQuery) CopyNode() *ReturnQuery {
859+
copyNode := *s
860+
return &copyNode
858861
}
859862

860863
func (s *ReturnQuery) Format(ctx *tree.FmtCtx) {
864+
ctx.WriteString("RETURN QUERY")
865+
if s.SqlStmt != nil {
866+
ctx.WriteByte(' ')
867+
ctx.FormatNode(s.SqlStmt)
868+
}
869+
ctx.WriteString(";\n")
861870
}
862871

863872
func (s *ReturnQuery) PlpgSQLStatementTag() string {
864873
return "stmt_return_query"
865874
}
866875

867876
func (s *ReturnQuery) WalkStmt(visitor StatementVisitor) Statement {
868-
panic(unimplemented.New("plpgsql visitor", "Unimplemented PLpgSQL visitor pattern"))
877+
newStmt, _ := visitor.Visit(s)
878+
return newStmt
869879
}
870880

871881
// stmt_raise

‎pkg/sql/sem/plpgsqltree/visitor.go‎

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -200,6 +200,16 @@ func (v *SQLStmtVisitor) Visit(stmt Statement) (newStmt Statement, recurse bool)
200200
cpy.Expr = e
201201
newStmt = cpy
202202
}
203+
case *ReturnQuery:
204+
s, v.Err = v.visitStmt(t.SqlStmt)
205+
if v.Err != nil {
206+
return stmt, false
207+
}
208+
if t.SqlStmt != s {
209+
cpy := t.CopyNode()
210+
cpy.SqlStmt = s
211+
newStmt = cpy
212+
}
203213
case *Raise:
204214
for i, p := range t.Params {
205215
e, v.Err = v.visitExpr(p)
@@ -297,7 +307,7 @@ func (v *SQLStmtVisitor) Visit(stmt Statement) (newStmt Statement, recurse bool)
297307
}
298308
}
299309

300-
case *ForEachArray, *ReturnQuery, *Perform:
310+
case *ForEachArray, *Perform:
301311
panic(unimp.New("plpgsql visitor", "Unimplemented PLpgSQL visitor"))
302312
}
303313
if v.Err != nil {

0 commit comments

Comments
 (0)