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..23bcc4d528 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.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 { + 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..1e32b20289 100644 --- a/tests/sqlparser_postgres.rs +++ b/tests/sqlparser_postgres.rs @@ -11057,3 +11057,93 @@ 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"; + 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(body)), + .. + }) => { + 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:?}"), + } + + 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}" + ); + } +} + +#[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:?}"); + } +}