diff --git a/fuzz/fuzz_seeds/postgres_create_procedure_begin_atomic.sql b/fuzz/fuzz_seeds/postgres_create_procedure_begin_atomic.sql new file mode 100644 index 0000000000..27c2d6e28b --- /dev/null +++ b/fuzz/fuzz_seeds/postgres_create_procedure_begin_atomic.sql @@ -0,0 +1 @@ +CREATE PROCEDURE p (a INT) LANGUAGE sql BEGIN ATOMIC INSERT INTO t VALUES (a); UPDATE t SET b = a; END diff --git a/src/ast/ddl.rs b/src/ast/ddl.rs index 6d651e39a0..16c7723682 100644 --- a/src/ast/ddl.rs +++ b/src/ast/ddl.rs @@ -36,7 +36,7 @@ use sqlparser_derive::{Visit, VisitMut}; use crate::ast::value::escape_single_quote_string; use crate::ast::{ - display_comma_separated, display_separated, format_statement_list, + display_comma_separated, display_separated, format_begin_atomic, table_constraints::{ CheckConstraint, ForeignKeyConstraint, PrimaryKeyConstraint, TableConstraint, UniqueConstraint, @@ -3898,12 +3898,7 @@ impl fmt::Display for CreateFunction { 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")?; + format_begin_atomic(f, statements)?; } 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 653a3581e5..bc87dd0131 100644 --- a/src/ast/mod.rs +++ b/src/ast/mod.rs @@ -209,6 +209,16 @@ fn format_statement_list(f: &mut fmt::Formatter, statements: &[Statement]) -> fm write!(f, ";") } +/// Writes ` BEGIN ATOMIC END`. +fn format_begin_atomic(f: &mut fmt::Formatter, statements: &[Statement]) -> fmt::Result { + write!(f, " BEGIN ATOMIC")?; + if !statements.is_empty() { + write!(f, " ")?; + format_statement_list(f, statements)?; + } + write!(f, " END") +} + /// A item `T` enclosed in a pair of parentheses #[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)] #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))] @@ -4584,8 +4594,8 @@ pub enum Statement { params: Option>, /// Optional language identifier. language: Option, - /// Procedure body statements. - body: ConditionalStatements, + /// Procedure body. + body: CreateProcedureBody, }, /// ```sql /// CREATE MACRO @@ -5601,7 +5611,12 @@ impl fmt::Display for Statement { write!(f, " LANGUAGE {language}")?; } - write!(f, " AS {body}") + match body { + CreateProcedureBody::As(body) => write!(f, " AS {body}"), + CreateProcedureBody::BeginAtomic(BeginEndStatements { statements, .. }) => { + format_begin_atomic(f, statements) + } + } } Statement::CreateMacro { or_replace, @@ -10574,6 +10589,19 @@ pub enum CreateFunctionBody { AsReturnSelect(Select), } +/// Body of a `CREATE PROCEDURE` statement. +#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)] +#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))] +#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))] +pub enum CreateProcedureBody { + /// Statements after `AS`. + As(ConditionalStatements), + /// SQL-standard body enclosed by `BEGIN ATOMIC` and `END`. + /// + /// [PostgreSQL]: https://www.postgresql.org/docs/current/sql-createprocedure.html + BeginAtomic(BeginEndStatements), +} + #[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)] #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))] #[cfg_attr(feature = "visitor", derive(Visit, VisitMut))] diff --git a/src/parser/mod.rs b/src/parser/mod.rs index 1b35306d3f..ba23f339d6 100644 --- a/src/parser/mod.rs +++ b/src/parser/mod.rs @@ -6148,17 +6148,9 @@ 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]) { + } else if let Some(statements) = self.maybe_parse_begin_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), - })); + body.function_body = Some(CreateFunctionBody::BeginAtomic(statements)); } else { break; } @@ -6428,6 +6420,22 @@ impl<'a> Parser<'a> { Ok((name, args)) } + /// Parses `BEGIN ATOMIC END` when the next tokens are `BEGIN ATOMIC`. + fn maybe_parse_begin_atomic(&mut self) -> Result, ParserError> { + if !self.peek_keywords(&[Keyword::BEGIN, Keyword::ATOMIC]) { + return Ok(None); + } + 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)?; + Ok(Some(BeginEndStatements { + begin_token: AttachedToken(begin_token), + statements, + end_token: AttachedToken(end_token), + })) + } + fn parse_function_arg(&mut self) -> Result { let mode = if self.parse_keyword(Keyword::IN) { Some(ArgMode::In) @@ -21490,9 +21498,13 @@ impl<'a> Parser<'a> { None }; - self.expect_keyword_is(Keyword::AS)?; - - let body = self.parse_conditional_statements(&[Keyword::END])?; + let body = match self.maybe_parse_begin_atomic()? { + Some(statements) => CreateProcedureBody::BeginAtomic(statements), + None => { + self.expect_keyword_is(Keyword::AS)?; + CreateProcedureBody::As(self.parse_conditional_statements(&[Keyword::END])?) + } + }; Ok(Statement::CreateProcedure { name, diff --git a/tests/sqlparser_common.rs b/tests/sqlparser_common.rs index db7786d681..8558fc6aa1 100644 --- a/tests/sqlparser_common.rs +++ b/tests/sqlparser_common.rs @@ -22908,3 +22908,91 @@ fn parse_alter_table_reset_options() { ); } } + +#[test] +fn parse_create_procedure_begin_atomic() { + let sql = "CREATE PROCEDURE p () LANGUAGE sql BEGIN ATOMIC INSERT INTO t VALUES (1); END"; + verified_stmt(sql); + match all_dialects().run_parser_method(sql, |parser| parser.parse_statement().unwrap()) { + Statement::CreateProcedure { + body: CreateProcedureBody::BeginAtomic(body), + .. + } => { + assert_eq!( + body.statements, + vec![verified_stmt("INSERT INTO t VALUES (1)")] + ); + assert_eq!( + body.span(), + Span::new(Location::new(1, 36), Location::new(1, 78)) + ); + } + stmt => panic!("unexpected statement: {stmt:?}"), + } + + for sql in [ + "CREATE PROCEDURE p (a INT) LANGUAGE sql BEGIN ATOMIC INSERT INTO t VALUES (a); UPDATE t SET b = a; END", + "CREATE PROCEDURE p () BEGIN ATOMIC SELECT 1; END", + "CREATE PROCEDURE p () LANGUAGE sql BEGIN ATOMIC END", + ] { + verified_stmt(sql); + } + + one_statement_parses_to( + r#"CREATE PROCEDURE public.p(a integer) + LANGUAGE sql + BEGIN ATOMIC + INSERT INTO public.t (a) + VALUES (p.a); +END"#, + "CREATE PROCEDURE public.p (a INTEGER) LANGUAGE sql BEGIN ATOMIC INSERT INTO public.t (a) VALUES (p.a); END", + ); + + for (sql, err) in [ + ( + "CREATE PROCEDURE p () LANGUAGE sql BEGIN ATOMIC SELECT 1;", + "Expected: END, found: EOF", + ), + ( + "CREATE PROCEDURE p () LANGUAGE sql BEGIN ATOMIC SELECT 1 AS a END", + "Expected: ;, found: END", + ), + ( + "CREATE PROCEDURE p () LANGUAGE sql BEGIN SELECT 1; END", + "Expected: AS, found: BEGIN", + ), + ( + "CREATE PROCEDURE p () LANGUAGE sql AS BEGIN SELECT 1;", + "Expected: END, found: EOF", + ), + ] { + assert_eq!( + parse_sql_statements(sql).unwrap_err(), + ParserError::ParserError(err.to_string()), + "{sql}" + ); + } +} + +#[test] +fn display_create_procedure_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 = verified_stmt("CREATE PROCEDURE p () 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:?}"); + } +} diff --git a/tests/sqlparser_mssql.rs b/tests/sqlparser_mssql.rs index 71e3b6cf7b..6c6058feee 100644 --- a/tests/sqlparser_mssql.rs +++ b/tests/sqlparser_mssql.rs @@ -134,7 +134,7 @@ fn parse_create_procedure() { ms().verified_stmt(sql), Statement::CreateProcedure { or_alter: true, - body: ConditionalStatements::BeginEnd(BeginEndStatements { + body: CreateProcedureBody::As(ConditionalStatements::BeginEnd(BeginEndStatements { begin_token: AttachedToken::empty(), statements: vec![Statement::Query(Box::new(Query { with: None, @@ -176,7 +176,7 @@ fn parse_create_procedure() { }))) }))], end_token: AttachedToken::empty(), - }), + })), params: Some(vec![ ProcedureParam { name: Ident {