PostgreSQL: Support BEGIN ATOMIC procedure bodies

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

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,
@@ -3885,12 +3885,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}")?;

src/ast/mod.rs

@@ -209,6 +209,16 @@ fn format_statement_list(f: &mut fmt::Formatter, statements: &[Statement]) -> fm
     write!(f, ";")
 }
 
+/// Writes ` BEGIN ATOMIC <statements> 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<Vec<ProcedureParam>>,
         /// Optional language identifier.
         language: Option<Ident>,
-        /// 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,
@@ -10573,6 +10588,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))]

src/parser/mod.rs

@@ -6185,17 +6185,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;
             }
@@ -6465,6 +6457,22 @@ impl<'a> Parser<'a> {
         Ok((name, args))
     }
 
+    /// Parses `BEGIN ATOMIC <statements> END` when the next tokens are `BEGIN ATOMIC`.
+    fn maybe_parse_begin_atomic(&mut self) -> Result<Option<BeginEndStatements>, 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<OperateFunctionArg, ParserError> {
         let mode = if self.parse_keyword(Keyword::IN) {
             Some(ArgMode::In)
@@ -21496,9 +21504,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,

tests/sqlparser_common.rs

@@ -22848,3 +22848,91 @@ fn parse_grant_on_tables_named_after_object_types() {
         verified_stmt(&format!("GRANT SELECT ON {name}.t TO r"));
     }
 }
+
+#[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:?}");
+    }
+}

tests/sqlparser_mssql.rs

@@ -133,7 +133,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,
@@ -175,7 +175,7 @@ fn parse_create_procedure() {
                     })))
                 }))],
                 end_token: AttachedToken::empty(),
-            }),
+            })),
             params: Some(vec![
                 ProcedureParam {
                     name: Ident {