Extend exception handling

src/ast/mod.rs

@@ -2982,6 +2982,63 @@ impl From<Set> for Statement {
     }
 }
 
+/// An exception representing exception handling with the `EXCEPTION` keyword.
+///
+/// Snowflake: <https://docs.snowflake.com/en/sql-reference/snowflake-scripting/exception>
+/// BigQuery: <https://cloud.google.com/bigquery/docs/reference/standard-sql/procedural-language#beginexceptionend>
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub struct Exception {
+    pub when: Vec<ExceptionWhen>,
+    pub raises: Option<Box<Statement>>,
+}
+
+impl Display for Exception {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        write!(f, " EXCEPTION")?;
+        for w in &self.when {
+            write!(f, "{w}")?;
+        }
+
+        if let Some(raises) = &self.raises {
+            write!(f, " {raises};")?;
+        }
+
+        Ok(())
+    }
+}
+
+/// A representation of a `WHEN` arm with all the identifiers catched and the statements to execute
+/// for the arm.
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub struct ExceptionWhen {
+    pub idents: Vec<Ident>,
+    pub statements: Vec<Statement>,
+}
+
+impl Display for ExceptionWhen {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        let idents = self
+            .idents
+            .iter()
+            .map(ToString::to_string)
+            .collect::<Vec<_>>()
+            .join(" OR ");
+
+        write!(f, " WHEN {idents} THEN", idents = idents)?;
+
+        if !self.statements.is_empty() {
+            write!(f, " ")?;
+            format_statement_list(f, &self.statements)?;
+        }
+
+        Ok(())
+    }
+}
+
 /// A top-level statement (SELECT, INSERT, CREATE, etc.)
 #[allow(clippy::large_enum_variant)]
 #[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
@@ -3670,17 +3727,24 @@ pub enum Statement {
         /// END;
         /// ```
         statements: Vec<Statement>,
-        /// Statements of an exception clause.
+        /// Exception handling with exception clauses and raises.
         /// Example:
         /// ```sql
         /// BEGIN
         ///     SELECT 1;
-        /// EXCEPTION WHEN ERROR THEN
-        ///     SELECT 2;
-        ///     SELECT 3;
+        /// EXCEPTION
+        ///     WHEN EXCEPTION_1 THEN
+        ///         SELECT 2;
+        ///     WHEN EXCEPTION_2 OR EXCEPTION_3 THEN
+        ///         SELECT 3;
+        ///     WHEN OTHER THEN
+        ///         SELECT 4;
+        /// RAISE;
         /// END;
+        /// ```
         /// <https://cloud.google.com/bigquery/docs/reference/standard-sql/procedural-language#beginexceptionend>
-        exception_statements: Option<Vec<Statement>>,
+        /// <https://docs.snowflake.com/en/sql-reference/snowflake-scripting/exception>
+        exception: Option<Exception>,
         /// TRUE if the statement has an `END` keyword.
         has_end_keyword: bool,
     },
@@ -5525,7 +5589,7 @@ impl fmt::Display for Statement {
                 transaction,
                 modifier,
                 statements,
-                exception_statements,
+                exception,
                 has_end_keyword,
             } => {
                 if *syntax_begin {
@@ -5547,12 +5611,8 @@ impl fmt::Display for Statement {
                     write!(f, " ")?;
                     format_statement_list(f, statements)?;
                 }
-                if let Some(exception_statements) = exception_statements {
-                    write!(f, " EXCEPTION WHEN ERROR THEN")?;
-                    if !exception_statements.is_empty() {
-                        write!(f, " ")?;
-                        format_statement_list(f, exception_statements)?;
-                    }
+                if let Some(exception) = exception {
+                    write!(f, "{exception}")?;
                 }
                 if *has_end_keyword {
                     write!(f, " END")?;

src/dialect/bigquery.rs

@@ -15,10 +15,9 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use crate::ast::Statement;
 use crate::dialect::Dialect;
 use crate::keywords::Keyword;
-use crate::parser::{Parser, ParserError};
+use crate::parser::Parser;
 
 /// These keywords are disallowed as column identifiers. Such that
 /// `SELECT 5 AS <col> FROM T` is rejected by BigQuery.
@@ -45,10 +44,6 @@ const RESERVED_FOR_COLUMN_ALIAS: &[Keyword] = &[
 pub struct BigQueryDialect;
 
 impl Dialect for BigQueryDialect {
-    fn parse_statement(&self, parser: &mut Parser) -> Option<Result<Statement, ParserError>> {
-        self.maybe_parse_statement(parser)
-    }
-
     /// See <https://cloud.google.com/bigquery/docs/reference/standard-sql/lexical#identifiers>
     fn is_delimited_identifier_start(&self, ch: char) -> bool {
         ch == '`'
@@ -141,48 +136,3 @@ impl Dialect for BigQueryDialect {
         true
     }
 }
-
-impl BigQueryDialect {
-    fn maybe_parse_statement(&self, parser: &mut Parser) -> Option<Result<Statement, ParserError>> {
-        if parser.peek_keyword(Keyword::BEGIN) {
-            return Some(self.parse_begin(parser));
-        }
-        None
-    }
-
-    /// Parse a `BEGIN` statement.
-    /// <https://cloud.google.com/bigquery/docs/reference/standard-sql/procedural-language#beginexceptionend>
-    fn parse_begin(&self, parser: &mut Parser) -> Result<Statement, ParserError> {
-        parser.expect_keyword(Keyword::BEGIN)?;
-
-        let statements = parser.parse_statement_list(&[Keyword::EXCEPTION, Keyword::END])?;
-
-        let has_exception_when_clause = parser.parse_keywords(&[
-            Keyword::EXCEPTION,
-            Keyword::WHEN,
-            Keyword::ERROR,
-            Keyword::THEN,
-        ]);
-        let exception_statements = if has_exception_when_clause {
-            if !parser.peek_keyword(Keyword::END) {
-                Some(parser.parse_statement_list(&[Keyword::END])?)
-            } else {
-                Some(Default::default())
-            }
-        } else {
-            None
-        };
-
-        parser.expect_keyword(Keyword::END)?;
-
-        Ok(Statement::StartTransaction {
-            begin: true,
-            statements,
-            exception_statements,
-            has_end_keyword: true,
-            transaction: None,
-            modifier: None,
-            modes: Default::default(),
-        })
-    }
-}

src/keywords.rs

@@ -646,6 +646,7 @@ define_keywords!(
     ORDER,
     ORDINALITY,
     ORGANIZATION,
+    OTHER,
     OUT,
     OUTER,
     OUTPUT,

src/parser/mod.rs

@@ -15096,12 +15096,16 @@ impl<'a> Parser<'a> {
             transaction: Some(BeginTransactionKind::Transaction),
             modifier: None,
             statements: vec![],
-            exception_statements: None,
+            exception: None,
             has_end_keyword: false,
         })
     }
 
     pub fn parse_begin(&mut self) -> Result<Statement, ParserError> {
+        if dialect_of!(self is SnowflakeDialect | BigQueryDialect) {
+            return self.parse_begin_exception_end();
+        }
+
         let modifier = if !self.dialect.supports_start_transaction_modifier() {
             None
         } else if self.parse_keyword(Keyword::DEFERRED) {
@@ -15128,11 +15132,69 @@ impl<'a> Parser<'a> {
             transaction,
             modifier,
             statements: vec![],
-            exception_statements: None,
+            exception: None,
             has_end_keyword: false,
         })
     }
 
+    pub fn parse_begin_exception_end(&mut self) -> Result<Statement, ParserError> {
+        let statements = self.parse_statement_list(&[Keyword::EXCEPTION, Keyword::END])?;
+
+        let exception = if self.parse_keyword(Keyword::EXCEPTION) {
+            let mut when = Vec::new();
+
+            // We can have multiple `WHEN` arms so we consume all cases until `END` or an exception
+            // is `RAISE`ed.
+            while self
+                .peek_one_of_keywords(&[Keyword::END, Keyword::RAISE])
+                .is_none()
+            {
+                self.expect_keyword(Keyword::WHEN)?;
+
+                // Each `WHEN` case can have one or more conditions, e.g.
+                // WHEN EXCEPTION_1 [OR EXCEPTION_2] THEN
+                // So we parse identifiers until the `THEN` keyword.
+                let mut idents = Vec::new();
+
+                while !self.parse_keyword(Keyword::THEN) {
+                    let ident = self.parse_identifier()?;
+                    idents.push(ident);
+
+                    self.maybe_parse(|p| p.expect_keyword(Keyword::OR))?;
+                }
+
+                let statements =
+                    self.parse_statement_list(&[Keyword::WHEN, Keyword::RAISE, Keyword::END])?;
+
+                when.push(ExceptionWhen { idents, statements });
+            }
+
+            let raises = if self.peek_keyword(Keyword::RAISE) {
+                let raises = Some(Box::new(self.parse_raise_stmt()?));
+                self.expect_token(&Token::SemiColon)?;
+                raises
+            } else {
+                None
+            };
+
+            Some(Exception { when, raises })
+        } else {
+            None
+        };
+
+        self.expect_keyword(Keyword::END)?;
+
+        Ok(Statement::StartTransaction {
+            begin: true,
+            statements,
+            exception,
+            has_end_keyword: true,
+            transaction: None,
+            modifier: None,
+            modes: Default::default(),
+        })
+    }
+
     pub fn parse_end(&mut self) -> Result<Statement, ParserError> {
         let modifier = if !self.dialect.supports_end_transaction_modifier() {
             None

tests/sqlparser_bigquery.rs

@@ -264,15 +264,18 @@ fn parse_begin() {
     let sql = r#"BEGIN SELECT 1; EXCEPTION WHEN ERROR THEN SELECT 2; END"#;
     let Statement::StartTransaction {
         statements,
-        exception_statements,
+        exception,
         has_end_keyword,
         ..
     } = bigquery().verified_stmt(sql)
     else {
         unreachable!();
     };
     assert_eq!(1, statements.len());
-    assert_eq!(1, exception_statements.unwrap().len());
+    assert!(exception.is_some());
+
+    let exception = exception.unwrap();
+    assert_eq!(1, exception.when.len());
     assert!(has_end_keyword);
 
     bigquery().verified_stmt(

tests/sqlparser_common.rs

@@ -8593,7 +8593,7 @@ fn lateral_function() {
 fn parse_start_transaction() {
     let dialects = all_dialects_except(|d|
         // BigQuery does not support this syntax
-        d.is::<BigQueryDialect>());
+        d.is::<BigQueryDialect>() || d.is::<SnowflakeDialect>());
     match dialects
         .verified_stmt("START TRANSACTION READ ONLY, READ WRITE, ISOLATION LEVEL SERIALIZABLE")
     {

tests/sqlparser_snowflake.rs

@@ -4066,3 +4066,67 @@ fn parse_connect_by_root_operator() {
         "sql parser error: Expected an expression, found: FROM"
     );
 }
+
+#[test]
+fn test_begin_exception_end() {
+    for sql in [
+        "BEGIN SELECT 1; EXCEPTION WHEN OTHER THEN SELECT 2; RAISE; END",
+        "BEGIN SELECT 1; EXCEPTION WHEN OTHER THEN SELECT 2; RAISE EX_1; END",
+        "BEGIN SELECT 1; EXCEPTION WHEN FOO THEN SELECT 2; WHEN OTHER THEN SELECT 3; RAISE; END",
+        "BEGIN BEGIN SELECT 1; EXCEPTION WHEN OTHER THEN SELECT 2; RAISE; END; END",
+    ] {
+        snowflake().verified_stmt(sql);
+    }
+
+    let sql = r#"
+DECLARE
+  EXCEPTION_1 EXCEPTION (-20001, 'I caught the expected exception.');
+  EXCEPTION_2 EXCEPTION (-20002, 'Not the expected exception!');
+  EXCEPTION_3 EXCEPTION (-20003, 'The worst exception...');
+BEGIN
+    BEGIN
+        SELECT 1;
+    EXCEPTION
+        WHEN EXCEPTION_1 THEN
+            SELECT 1;
+        WHEN EXCEPTION_2 OR EXCEPTION_3 THEN
+            SELECT 2;
+            SELECT 3;
+        WHEN OTHER THEN
+            SELECT 4;
+    RAISE;
+    END;
+END
+"#;
+
+    // Outer `BEGIN` of the two nested `BEGIN` statements.
+    let Statement::StartTransaction { mut statements, .. } = snowflake()
+        .parse_sql_statements(sql)
+        .unwrap()
+        .pop()
+        .unwrap()
+    else {
+        unreachable!();
+    };
+
+    // Inner `BEGIN` of the two nested `BEGIN` statements.
+    let Statement::StartTransaction {
+        statements,
+        exception,
+        has_end_keyword,
+        ..
+    } = statements.pop().unwrap()
+    else {
+        unreachable!();
+    };
+
+    assert_eq!(1, statements.len());
+    assert!(has_end_keyword);
+
+    let exception = exception.unwrap();
+    assert_eq!(3, exception.when.len());
+    assert_eq!(1, exception.when[0].idents.len());
+    assert_eq!(1, exception.when[0].statements.len());
+    assert_eq!(2, exception.when[1].idents.len());
+    assert_eq!(2, exception.when[1].statements.len());
+}