Skip to content

Commit

Permalink
Browse files Browse the repository at this point in the history
138709: sql,plpgsql: add support for DO statements r=mgartner a=DrewKimball

#### sql: avoid collisions in dollar sign formatted strings

This commit adds a `FormatStringDollarQuotes` method to `FmtCtx` that
displays a string constant with dollar-sign quotes, like: `$$ foo $$`.
In order to ensure that the result is parsable, the new method checks
the input string for collisions; first for `$$`, then for `$funcbody$`,
then for `$funcbodyx$`, `$funcbodyxx$` and so on. This is similar to
Postgres' behavior. The new method is now used for formatting function
body strings in `CREATE FUNCTION` and `CREATE PROCEDURE` statements.

Informs cockroachdb#115268

Release note: None

#### sql: support full tree.Visitor in PL/pgSQL SQLStmtVisitor

This commit exports the `tree.WalkStmt` function and adds support for
using the full `tree.Visitor` interface to `SQLStmtVisitor`. Note that
it is still possible to use `tree.SimpleVisitFn` as before.

Epic: None

Release note: None

#### plpgsqltree: move visitor implementations

This is a mechanical commit that moves the `SQLStmtVisitor` and
`TypeRefVisitor` to the `plpgsqltree` package. This will allow the
PL/pgSQL statements in that package to make use of the visitors.

Epic: None

Release note: None

#### plpgsql/parser: parse PL/pgSQL DO statements

This commit adds support for DO statements to the PL/pgSQL parser.
The helper function `makeDoStmt` analyzes the DO statement options
(language and code string) and then calls into the PL/pgSQL parser
for the code string. A following commit will inject this helper into
the SQL parser as well.

Informs cockroachdb#115268

Release note: None

#### sql/parser: parse SQL DO statements

This commit adds support for DO statements to the SQL parser. This allows
PL/pgSQL logic to be executed inline as an anonymous function. SQL DO blocks
can be used either directly as a top-level SQL statement, or as body statements
within a SQL routine.

Informs cockroachdb#115268

Release note: None

#### sql,plpgsql: enable DO statements in PL/pgSQL routines

This commit adds support for DO statements embedded within PL/pgSQL
code. DO statements are valid within both other DO statements and PL/pgSQL
routines. They are handled similarly to a CALL statement with no arguments,
where the body is built as a RelExpr that is executed only for its side
effects.

Informs cockroachdb#115268

Release note (sql change): Added support for DO statements embedded within
PL/pgSQL routines and other DO statements. DO statements execute a block of
code inline as an anonymous function. Currently, only a PL/pgSQL body is
allowed.

#### sql: add support for SQL DO statements

This commit adds support for SQL DO statements, which allow inline execution
of a PL/pgSQL code block. DO statements are only valid as top-level
statements, meaning (for example) that they can't be used in a CTE.
A following commit will enable DO statements within SQL routines.

Informs cockroachdb#115268

Release note (sql change): Added support for DO statements in SQL, which
allow a PL/pgSQL code block to be executed inline.

#### sql: enable DO statements inside SQL routines

This commit allows DO statements to be used as body statements in a
SQL routine.

Release note (sql change): DO statements can now be used as body statements
in a SQL routine to execute PL/pgSQL code inline.

Co-authored-by: Drew Kimball <[email protected]>
  • Loading branch information
craig[bot] and DrewKimball committed Jan 10, 2025
2 parents 3ce8f44 + dbf2622 commit e0e7164
Show file tree
Hide file tree
Showing 67 changed files with 2,625 additions and 437 deletions.
1 change: 1 addition & 0 deletions docs/generated/sql/bnf/BUILD.bazel
Original file line number Diff line number Diff line change
Expand Up @@ -122,6 +122,7 @@ FILES = [
"default_value_column_level",
"delete_stmt",
"discard_stmt",
"do_stmt",
"drop_column",
"drop_constraint",
"drop_database",
Expand Down
2 changes: 2 additions & 0 deletions docs/generated/sql/bnf/do_stmt.bnf
Original file line number Diff line number Diff line change
@@ -0,0 +1,2 @@
do_stmt ::=
'DO' do_stmt_opt_list
1 change: 1 addition & 0 deletions docs/generated/sql/bnf/explainable_stmt.bnf
Original file line number Diff line number Diff line change
Expand Up @@ -3,3 +3,4 @@ explainable_stmt ::=
| comment_stmt
| execute_stmt
| call_stmt
| do_stmt
12 changes: 12 additions & 0 deletions docs/generated/sql/bnf/stmt_block.bnf
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@ stmt_without_legacy_transaction ::=
| execute_stmt
| deallocate_stmt
| discard_stmt
| do_stmt
| grant_stmt
| prepare_stmt
| revoke_stmt
Expand Down Expand Up @@ -96,6 +97,9 @@ discard_stmt ::=
| 'DISCARD' 'TEMP'
| 'DISCARD' 'TEMPORARY'

do_stmt ::=
'DO' do_stmt_opt_list

grant_stmt ::=
'GRANT' privileges 'ON' grant_targets 'TO' role_spec_list opt_with_grant_option
| 'GRANT' privilege_list 'TO' role_spec_list
Expand Down Expand Up @@ -392,6 +396,9 @@ name ::=
| unreserved_keyword
| col_name_keyword

do_stmt_opt_list ::=
( do_stmt_opt_item ) ( ( do_stmt_opt_item ) )*

privileges ::=
'ALL' opt_privileges_clause
| privilege_list
Expand Down Expand Up @@ -687,6 +694,7 @@ explainable_stmt ::=
| comment_stmt
| execute_stmt
| call_stmt
| do_stmt

explain_option_list ::=
( explain_option_name ) ( ( ',' explain_option_name ) )*
Expand Down Expand Up @@ -1600,6 +1608,10 @@ col_name_keyword ::=
| 'VIRTUAL'
| 'WORK'

do_stmt_opt_item ::=
'SCONST'
| 'LANGUAGE' non_reserved_word_or_sconst

opt_privileges_clause ::=
'PRIVILEGES'
|
Expand Down
1 change: 1 addition & 0 deletions docs/generated/sql/bnf/stmt_without_legacy_transaction.bnf
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ stmt_without_legacy_transaction ::=
| execute_stmt
| deallocate_stmt
| discard_stmt
| do_stmt
| grant_stmt
| prepare_stmt
| revoke_stmt
Expand Down
16 changes: 16 additions & 0 deletions pkg/ccl/logictestccl/testdata/logic_test/plpgsql_txn
Original file line number Diff line number Diff line change
Expand Up @@ -1009,4 +1009,20 @@ CREATE PROCEDURE p() LANGUAGE PLpgSQL AS $$ BEGIN CALL p_nested_commit(); END $$
statement error pgcode 0A000 pq: unimplemented: transaction control statements in nested routines
CREATE PROCEDURE p() LANGUAGE PLpgSQL AS $$ BEGIN CALL p_nested_rollback(); END $$;

skipif config local-mixed-24.3
statement error pgcode 0A000 pq: unimplemented: transaction control statements in nested routines
CREATE PROCEDURE p() LANGUAGE PLpgSQL AS $$ BEGIN DO $inner$ BEGIN CALL p_nested_rollback(); END $inner$; END $$;

skipif config local-mixed-24.3
statement error pgcode 0A000 pq: unimplemented: transaction control statements in nested routines
DO $$ BEGIN CALL p_nested_commit(); END $$;

skipif config local-mixed-24.3
statement error pgcode 0A000 pq: unimplemented: transaction control statements in DO blocks
CREATE PROCEDURE p() LANGUAGE PLpgSQL AS $$ BEGIN DO $inner$ BEGIN COMMIT; END $inner$; END $$;

skipif config local-mixed-24.3
statement error pgcode 0A000 pq: unimplemented: transaction control statements in DO blocks
DO $$ BEGIN COMMIT; END $$;

subtest end
117 changes: 117 additions & 0 deletions pkg/ccl/logictestccl/testdata/logic_test/udf_rewrite
Original file line number Diff line number Diff line change
Expand Up @@ -582,3 +582,120 @@ statement ok
DROP PROCEDURE p_rewrite;

subtest end

subtest do_block

skipif config local-mixed-24.3
statement ok
CREATE FUNCTION f_rewrite() RETURNS INT AS $$
BEGIN
DO $foo$
BEGIN
SELECT nextval('seq');
SELECT 'wednesday'::weekday;
RAISE NOTICE 'foo';
DO $bar$
BEGIN
RAISE NOTICE 'bar';
SELECT nextval('seq');
SELECT 'wednesday'::weekday;
END;
$bar$;
END
$foo$;
RETURN 100;
END
$$ LANGUAGE PLpgSQL;

skipif config local-mixed-24.3
query T
SELECT get_body_str('f_rewrite');
----
"BEGIN\nDO $funcbody$\nBEGIN\nSELECT nextval(106:::REGCLASS);\nSELECT b'\\x80':::@100107;\nRAISE NOTICE 'foo';\nDO $$\nBEGIN\nRAISE NOTICE 'bar';\nSELECT nextval(106:::REGCLASS);\nSELECT b'\\x80':::@100107;\nEND;\n$$;\nEND;\n$funcbody$;\nRETURN 100;\nEND;\n"

skipif config local-mixed-24.3
query TT
SHOW CREATE FUNCTION f_rewrite;
----
f_rewrite CREATE FUNCTION public.f_rewrite()
RETURNS INT8
VOLATILE
NOT LEAKPROOF
CALLED ON NULL INPUT
LANGUAGE plpgsql
SECURITY INVOKER
AS $funcbodyx$
BEGIN
DO $funcbody$
BEGIN
SELECT nextval('public.seq'::REGCLASS);
SELECT 'wednesday':::public.weekday;
RAISE NOTICE 'foo';
DO $$
BEGIN
RAISE NOTICE 'bar';
SELECT nextval('public.seq'::REGCLASS);
SELECT 'wednesday':::public.weekday;
END;
$$;
END;
$funcbody$;
RETURN 100;
END;
$funcbodyx$

statement ok
ALTER TYPE weekday RENAME VALUE 'friday' TO 'humpday';

statement ok
ALTER TYPE weekday RENAME TO workday;

skipif config local-mixed-24.3
query T
SELECT get_body_str('f_rewrite');
----
"BEGIN\nDO $funcbody$\nBEGIN\nSELECT nextval(106:::REGCLASS);\nSELECT b'\\x80':::@100107;\nRAISE NOTICE 'foo';\nDO $$\nBEGIN\nRAISE NOTICE 'bar';\nSELECT nextval(106:::REGCLASS);\nSELECT b'\\x80':::@100107;\nEND;\n$$;\nEND;\n$funcbody$;\nRETURN 100;\nEND;\n"

skipif config local-mixed-24.3
query TT
SHOW CREATE FUNCTION f_rewrite;
----
f_rewrite CREATE FUNCTION public.f_rewrite()
RETURNS INT8
VOLATILE
NOT LEAKPROOF
CALLED ON NULL INPUT
LANGUAGE plpgsql
SECURITY INVOKER
AS $funcbodyx$
BEGIN
DO $funcbody$
BEGIN
SELECT nextval('public.seq'::REGCLASS);
SELECT 'wednesday':::public.workday;
RAISE NOTICE 'foo';
DO $$
BEGIN
RAISE NOTICE 'bar';
SELECT nextval('public.seq'::REGCLASS);
SELECT 'wednesday':::public.workday;
END;
$$;
END;
$funcbody$;
RETURN 100;
END;
$funcbodyx$

# Reset types for subtest.
statement ok
ALTER TYPE workday RENAME TO weekday;

statement ok
ALTER TYPE weekday RENAME VALUE 'humpday' TO 'friday';

skipif config local-mixed-24.3
statement ok
DROP FUNCTION f_rewrite;

subtest end
7 changes: 7 additions & 0 deletions pkg/ccl/logictestccl/tests/3node-tenant/generated_test.go

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

1 change: 1 addition & 0 deletions pkg/gen/bnf.bzl
Original file line number Diff line number Diff line change
Expand Up @@ -122,6 +122,7 @@ BNF_SRCS = [
"//docs/generated/sql/bnf:default_value_column_level.bnf",
"//docs/generated/sql/bnf:delete_stmt.bnf",
"//docs/generated/sql/bnf:discard_stmt.bnf",
"//docs/generated/sql/bnf:do_stmt.bnf",
"//docs/generated/sql/bnf:drop_column.bnf",
"//docs/generated/sql/bnf:drop_constraint.bnf",
"//docs/generated/sql/bnf:drop_database.bnf",
Expand Down
1 change: 1 addition & 0 deletions pkg/gen/diagrams.bzl
Original file line number Diff line number Diff line change
Expand Up @@ -122,6 +122,7 @@ DIAGRAMS_SRCS = [
"//docs/generated/sql/bnf:default_value_column_level.html",
"//docs/generated/sql/bnf:delete.html",
"//docs/generated/sql/bnf:discard.html",
"//docs/generated/sql/bnf:do.html",
"//docs/generated/sql/bnf:drop.html",
"//docs/generated/sql/bnf:drop_column.html",
"//docs/generated/sql/bnf:drop_constraint.html",
Expand Down
1 change: 1 addition & 0 deletions pkg/gen/docs.bzl
Original file line number Diff line number Diff line change
Expand Up @@ -135,6 +135,7 @@ DOCS_SRCS = [
"//docs/generated/sql/bnf:default_value_column_level.bnf",
"//docs/generated/sql/bnf:delete_stmt.bnf",
"//docs/generated/sql/bnf:discard_stmt.bnf",
"//docs/generated/sql/bnf:do_stmt.bnf",
"//docs/generated/sql/bnf:drop_column.bnf",
"//docs/generated/sql/bnf:drop_constraint.bnf",
"//docs/generated/sql/bnf:drop_database.bnf",
Expand Down
1 change: 0 additions & 1 deletion pkg/sql/BUILD.bazel
Original file line number Diff line number Diff line change
Expand Up @@ -495,7 +495,6 @@ go_library(
"//pkg/sql/sem/catid",
"//pkg/sql/sem/eval",
"//pkg/sql/sem/plpgsqltree",
"//pkg/sql/sem/plpgsqltree/utils",
"//pkg/sql/sem/semenumpb",
"//pkg/sql/sem/transform",
"//pkg/sql/sem/tree",
Expand Down
1 change: 0 additions & 1 deletion pkg/sql/catalog/rewrite/BUILD.bazel
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,6 @@ go_library(
"//pkg/sql/schemachanger/screl",
"//pkg/sql/sem/catid",
"//pkg/sql/sem/plpgsqltree",
"//pkg/sql/sem/plpgsqltree/utils",
"//pkg/sql/sem/tree",
"//pkg/sql/types",
"//pkg/util/hlc",
Expand Down
3 changes: 1 addition & 2 deletions pkg/sql/catalog/rewrite/rewrite.go
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,6 @@ import (
"github.com/cockroachdb/cockroach/pkg/sql/schemachanger/screl"
"github.com/cockroachdb/cockroach/pkg/sql/sem/catid"
"github.com/cockroachdb/cockroach/pkg/sql/sem/plpgsqltree"
"github.com/cockroachdb/cockroach/pkg/sql/sem/plpgsqltree/utils"
"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
"github.com/cockroachdb/cockroach/pkg/sql/types"
"github.com/cockroachdb/cockroach/pkg/util/hlc"
Expand Down Expand Up @@ -622,7 +621,7 @@ func rewriteSequencesInFunction(
if err != nil {
return "", err
}
v := utils.SQLStmtVisitor{Fn: replaceSeqFunc}
v := plpgsqltree.SQLStmtVisitor{Fn: replaceSeqFunc}
newStmt := plpgsqltree.Walk(&v, stmt.AST)
fmtCtx.FormatNode(newStmt)

Expand Down
7 changes: 3 additions & 4 deletions pkg/sql/create_view.go
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,6 @@ import (
"github.com/cockroachdb/cockroach/pkg/sql/privilege"
"github.com/cockroachdb/cockroach/pkg/sql/sem/eval"
"github.com/cockroachdb/cockroach/pkg/sql/sem/plpgsqltree"
"github.com/cockroachdb/cockroach/pkg/sql/sem/plpgsqltree/utils"
"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
"github.com/cockroachdb/cockroach/pkg/sql/sqlerrors"
"github.com/cockroachdb/cockroach/pkg/sql/sqltelemetry"
Expand Down Expand Up @@ -514,7 +513,7 @@ func replaceSeqNamesWithIDsLang(
}
stmts = plstmt.AST

v := utils.SQLStmtVisitor{Fn: replaceSeqFunc}
v := plpgsqltree.SQLStmtVisitor{Fn: replaceSeqFunc}
newStmt := plpgsqltree.Walk(&v, stmts)
fmtCtx.FormatNode(newStmt)
}
Expand Down Expand Up @@ -661,12 +660,12 @@ func serializeUserDefinedTypesLang(
}
stmts = plstmt.AST

v := utils.SQLStmtVisitor{Fn: replaceFunc}
v := plpgsqltree.SQLStmtVisitor{Fn: replaceFunc}
newStmt := plpgsqltree.Walk(&v, stmts)
// Some PLpgSQL statements (i.e., declarations), may contain type
// annotations containing the UDT. We need to walk the AST to replace them,
// too.
v2 := utils.TypeRefVisitor{Fn: replaceTypeFunc}
v2 := plpgsqltree.TypeRefVisitor{Fn: replaceTypeFunc}
newStmt = plpgsqltree.Walk(&v2, newStmt)
fmtCtx.FormatNode(newStmt)
}
Expand Down
Loading

0 comments on commit e0e7164

Please sign in to comment.