From aa5a432878be8e23617669529423f284206a4ca6 Mon Sep 17 00:00:00 2001 From: LucaCappelletti94 Date: Mon, 26 Oct 2026 17:13:02 +0100 Subject: [PATCH 1/3] PostgreSQL: Support BEGIN ATOMIC function bodies --- .../postgres_create_function_begin_atomic.sql | 1 + src/ast/ddl.rs | 28 ++++++--- src/ast/mod.rs | 4 ++ src/parser/mod.rs | 11 ++++ tests/sqlparser_postgres.rs | 57 +++++++++++++++++++ 5 files changed, 92 insertions(+), 9 deletions(-) create mode 100644 fuzz/fuzz_seeds/postgres_create_function_begin_atomic.sql diff --git a/fuzz/fuzz_seeds/postgres_create_function_begin_atomic.sql b/fuzz/fuzz_seeds/postgres_create_function_begin_atomic.sql new file mode 100644 index 0000000000..ee57fff6c0 --- /dev/null +++ b/fuzz/fuzz_seeds/postgres_create_function_begin_atomic.sql @@ -0,0 +1 @@ +CREATE FUNCTION f(a INT) RETURNS INT LANGUAGE sql BEGIN ATOMIC INSERT INTO t VALUES (a); RETURN a + 1; END diff --git a/src/ast/ddl.rs b/src/ast/ddl.rs index c1780e2a43..ad3f6e66c8 100644 --- a/src/ast/ddl.rs +++ b/src/ast/ddl.rs @@ -36,19 +36,19 @@ use sqlparser_derive::{Visit, VisitMut}; use crate::ast::value::escape_single_quote_string; use crate::ast::{ - display_comma_separated, display_separated, + display_comma_separated, display_separated, format_statement_list, table_constraints::{ CheckConstraint, ForeignKeyConstraint, PrimaryKeyConstraint, TableConstraint, UniqueConstraint, }, - ArgMode, AttachedToken, CommentDef, ConditionalStatements, CreateFunctionBody, - CreateFunctionUsing, CreateServerOption, CreateTableLikeKind, CreateTableOptions, - CreateViewParams, DataType, Expr, FileFormat, FunctionBehavior, FunctionCalledOnNull, - FunctionDefinitionSetParam, FunctionDesc, FunctionDeterminismSpecifier, FunctionParallel, - FunctionSecurity, HiveDistributionStyle, HiveFormat, HiveIOFormat, HiveRowFormat, - HiveSetLocation, Ident, InitializeKind, MySQLColumnPosition, ObjectName, OnCommit, - OneOrManyWithParens, OperateFunctionArg, OrderByExpr, ProjectionSelect, Query, RefreshModeKind, - ResetConfig, RowAccessPolicy, SequenceOptions, Spanned, SqlOption, Statement, + ArgMode, AttachedToken, BeginEndStatements, CommentDef, ConditionalStatements, + CreateFunctionBody, CreateFunctionUsing, CreateServerOption, CreateTableLikeKind, + CreateTableOptions, CreateViewParams, DataType, Expr, FileFormat, FunctionBehavior, + FunctionCalledOnNull, FunctionDefinitionSetParam, FunctionDesc, FunctionDeterminismSpecifier, + FunctionParallel, FunctionSecurity, HiveDistributionStyle, HiveFormat, HiveIOFormat, + HiveRowFormat, HiveSetLocation, Ident, InitializeKind, MySQLColumnPosition, ObjectName, + OnCommit, OneOrManyWithParens, OperateFunctionArg, OrderByExpr, ProjectionSelect, Query, + RefreshModeKind, ResetConfig, RowAccessPolicy, SequenceOptions, Spanned, SqlOption, Statement, StorageLifecyclePolicy, StorageSerializationPolicy, TableVersion, Tag, TriggerEvent, TriggerExecBody, TriggerObject, TriggerPeriod, TriggerReferencing, Value, ValueWithSpan, WrappedCollection, @@ -3863,6 +3863,16 @@ impl fmt::Display for CreateFunction { if let Some(CreateFunctionBody::Return(function_body)) = &self.function_body { write!(f, " RETURN {function_body}")?; } + if let Some(CreateFunctionBody::BeginAtomic(BeginEndStatements { statements, .. })) = + &self.function_body + { + write!(f, " BEGIN ATOMIC")?; + if !statements.is_empty() { + write!(f, " ")?; + format_statement_list(f, statements)?; + } + write!(f, " END")?; + } if let Some(CreateFunctionBody::AsReturnExpr(function_body)) = &self.function_body { write!(f, " AS RETURN {function_body}")?; } diff --git a/src/ast/mod.rs b/src/ast/mod.rs index 9b5441738a..b49bab1c4a 100644 --- a/src/ast/mod.rs +++ b/src/ast/mod.rs @@ -10496,6 +10496,10 @@ pub enum CreateFunctionBody { /// /// [PostgreSQL]: https://www.postgresql.org/docs/current/sql-createfunction.html Return(Expr), + /// SQL-standard function body enclosed by `BEGIN ATOMIC` and `END`. + /// + /// [PostgreSQL]: https://www.postgresql.org/docs/current/sql-createfunction.html + BeginAtomic(BeginEndStatements), /// Function body expression using the 'AS RETURN' keywords /// diff --git a/src/parser/mod.rs b/src/parser/mod.rs index c5a6f4e854..06fc9015d2 100644 --- a/src/parser/mod.rs +++ b/src/parser/mod.rs @@ -6145,6 +6145,17 @@ impl<'a> Parser<'a> { } else if self.parse_keyword(Keyword::RETURN) { ensure_not_set(&body.function_body, "RETURN")?; body.function_body = Some(CreateFunctionBody::Return(self.parse_expr()?)); + } else if self.peek_keywords(&[Keyword::BEGIN, Keyword::ATOMIC]) { + ensure_not_set(&body.function_body, "BEGIN ATOMIC")?; + let begin_token = self.expect_keyword(Keyword::BEGIN)?; + self.expect_keyword_is(Keyword::ATOMIC)?; + let statements = self.parse_statement_list(&[Keyword::END])?; + let end_token = self.expect_keyword(Keyword::END)?; + body.function_body = Some(CreateFunctionBody::BeginAtomic(BeginEndStatements { + begin_token: AttachedToken(begin_token), + statements, + end_token: AttachedToken(end_token), + })); } else { break; } diff --git a/tests/sqlparser_postgres.rs b/tests/sqlparser_postgres.rs index 11a4703f55..885f8e67f4 100644 --- a/tests/sqlparser_postgres.rs +++ b/tests/sqlparser_postgres.rs @@ -11057,3 +11057,60 @@ fn parse_role_signed_connection_limit() { "sql parser error: Expected: option, found: CONNECTION" ); } + +#[test] +fn parse_create_function_begin_atomic() { + let sql = "CREATE FUNCTION f() RETURNS INT LANGUAGE sql BEGIN ATOMIC SELECT 1; END"; + match pg_and_generic().verified_stmt(sql) { + Statement::CreateFunction(CreateFunction { + function_body: + Some(CreateFunctionBody::BeginAtomic(BeginEndStatements { statements, .. })), + .. + }) => assert_eq!(statements, vec![pg_and_generic().verified_stmt("SELECT 1")]), + stmt => panic!("unexpected statement: {stmt:?}"), + } + + for sql in [ + "CREATE FUNCTION f(a INT) RETURNS INT LANGUAGE sql IMMUTABLE BEGIN ATOMIC SELECT a; RETURN a + 1; END", + "CREATE OR REPLACE FUNCTION f(a INT) RETURNS INT LANGUAGE sql BEGIN ATOMIC SELECT CASE WHEN a > 0 THEN 1 END; END", + "CREATE FUNCTION f() RETURNS VOID LANGUAGE sql BEGIN ATOMIC END", + ] { + pg_and_generic().verified_stmt(sql); + } + + pg_and_generic().one_statement_parses_to( + r#"CREATE OR REPLACE FUNCTION public.f(a integer) + RETURNS integer + LANGUAGE sql +BEGIN ATOMIC + INSERT INTO t VALUES (a); + RETURN a; +END"#, + "CREATE OR REPLACE FUNCTION public.f(a INTEGER) RETURNS INTEGER LANGUAGE sql BEGIN ATOMIC INSERT INTO t VALUES (a); RETURN a; END", + ); + + for (sql, err) in [ + ( + "CREATE FUNCTION f() RETURNS INT LANGUAGE sql BEGIN ATOMIC SELECT 1;", + "Expected: END, found: EOF", + ), + ( + "CREATE FUNCTION f() RETURNS INT LANGUAGE sql BEGIN ATOMIC SELECT 1 END", + "Expected: ;, found: END", + ), + ( + "CREATE FUNCTION f() RETURNS INT LANGUAGE sql RETURN 1 BEGIN ATOMIC SELECT 1; END", + "BEGIN ATOMIC specified more than once", + ), + ( + "CREATE FUNCTION f() RETURNS INT LANGUAGE sql BEGIN SELECT 1; END", + "Expected: end of statement, found: BEGIN", + ), + ] { + assert_eq!( + pg_and_generic().parse_sql_statements(sql).unwrap_err(), + ParserError::ParserError(err.to_string()), + "{sql}" + ); + } +} From b1eed5e950dc007078346b7f5a8daf1cfbc6080f Mon Sep 17 00:00:00 2001 From: LucaCappelletti94 Date: Wed, 28 Oct 2026 12:33:12 +0100 Subject: [PATCH 2/3] Take the peeked BEGIN ATOMIC tokens without re-checking them --- src/parser/mod.rs | 4 ++-- tests/sqlparser_postgres.rs | 17 +++++++++++++---- 2 files changed, 15 insertions(+), 6 deletions(-) diff --git a/src/parser/mod.rs b/src/parser/mod.rs index 06fc9015d2..23bcc4d528 100644 --- a/src/parser/mod.rs +++ b/src/parser/mod.rs @@ -6147,8 +6147,8 @@ impl<'a> Parser<'a> { body.function_body = Some(CreateFunctionBody::Return(self.parse_expr()?)); } else if self.peek_keywords(&[Keyword::BEGIN, Keyword::ATOMIC]) { ensure_not_set(&body.function_body, "BEGIN ATOMIC")?; - let begin_token = self.expect_keyword(Keyword::BEGIN)?; - self.expect_keyword_is(Keyword::ATOMIC)?; + let begin_token = self.next_token(); + self.advance_token(); let statements = self.parse_statement_list(&[Keyword::END])?; let end_token = self.expect_keyword(Keyword::END)?; body.function_body = Some(CreateFunctionBody::BeginAtomic(BeginEndStatements { diff --git a/tests/sqlparser_postgres.rs b/tests/sqlparser_postgres.rs index 885f8e67f4..6d553a2a35 100644 --- a/tests/sqlparser_postgres.rs +++ b/tests/sqlparser_postgres.rs @@ -11061,12 +11061,21 @@ fn parse_role_signed_connection_limit() { #[test] fn parse_create_function_begin_atomic() { let sql = "CREATE FUNCTION f() RETURNS INT LANGUAGE sql BEGIN ATOMIC SELECT 1; END"; - match pg_and_generic().verified_stmt(sql) { + pg_and_generic().verified_stmt(sql); + match pg_and_generic().run_parser_method(sql, |parser| parser.parse_statement().unwrap()) { Statement::CreateFunction(CreateFunction { - function_body: - Some(CreateFunctionBody::BeginAtomic(BeginEndStatements { statements, .. })), + function_body: Some(CreateFunctionBody::BeginAtomic(body)), .. - }) => assert_eq!(statements, vec![pg_and_generic().verified_stmt("SELECT 1")]), + }) => { + assert_eq!( + body.statements, + vec![pg_and_generic().verified_stmt("SELECT 1")] + ); + assert_eq!( + body.span(), + Span::new(Location::new(1, 46), Location::new(1, 72)) + ); + } stmt => panic!("unexpected statement: {stmt:?}"), } From 30d237143ef417acf9742fba0bd18c4c52c4a866 Mon Sep 17 00:00:00 2001 From: LucaCappelletti94 Date: Wed, 28 Oct 2026 12:38:44 +0100 Subject: [PATCH 3/3] Test that BEGIN ATOMIC display propagates write errors --- tests/sqlparser_postgres.rs | 24 ++++++++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/tests/sqlparser_postgres.rs b/tests/sqlparser_postgres.rs index 6d553a2a35..1e32b20289 100644 --- a/tests/sqlparser_postgres.rs +++ b/tests/sqlparser_postgres.rs @@ -11123,3 +11123,27 @@ END"#, ); } } + +#[test] +fn display_create_function_begin_atomic_propagates_write_errors() { + use core::fmt::{self, Write}; + + struct RejectingWriter(&'static str); + + impl Write for RejectingWriter { + fn write_str(&mut self, value: &str) -> fmt::Result { + if value == self.0 { + Err(fmt::Error) + } else { + Ok(()) + } + } + } + + let stmt = pg_and_generic() + .verified_stmt("CREATE FUNCTION f() RETURNS INT LANGUAGE sql BEGIN ATOMIC SELECT 1; END"); + for rejected in [" BEGIN ATOMIC", " ", ";", " END"] { + let mut writer = RejectingWriter(rejected); + assert_eq!(write!(writer, "{stmt}"), Err(fmt::Error), "{rejected:?}"); + } +}