Refactor logic to mod.rs

Signed-off-by: Guan-Ming (Wesley) Chiu <105915352+guan404ming@users.noreply.github.com>

Author: guan404ming

src/dialect/mssql.rs

@@ -146,20 +146,27 @@ impl Dialect for MsSqlDialect {
 
     fn parse_statement(&self, parser: &mut Parser) -> Option<Result<Statement, ParserError>> {
         if parser.parse_keyword(Keyword::BEGIN) {
-            if parser.peek_keyword(Keyword::TRANSACTION)
-                || parser.peek_keyword(Keyword::WORK)
-                || parser.peek_keyword(Keyword::TRY)
-                || parser.peek_keyword(Keyword::CATCH)
-                || parser.peek_keyword(Keyword::DEFERRED)
-                || parser.peek_keyword(Keyword::IMMEDIATE)
-                || parser.peek_keyword(Keyword::EXCLUSIVE)
-                || parser.peek_token_ref().token == Token::SemiColon
-                || parser.peek_token_ref().token == Token::EOF
-            {
+            // Check if this is a BEGIN...END block rather than BEGIN TRANSACTION
+            let is_block = parser
+                .maybe_parse(|p| {
+                    if p.parse_transaction_modifier().is_some()
+                        || p.parse_one_of_keywords(&[Keyword::TRANSACTION, Keyword::WORK])
+                            .is_some()
+                        || matches!(p.peek_token_ref().token, Token::SemiColon | Token::EOF)
+                    {
+                        p.expected("statement", p.peek_token())
+                    } else {
+                        Ok(())
+                    }
+                })
+                .unwrap_or(None)
+                .is_some();
+            if is_block {
+                Some(parser.parse_begin_exception_end())
+            } else {
                 parser.prev_token();
-                return None;
+                None
             }
-            Some(parser.parse_begin_exception_end())
         } else if parser.peek_keyword(Keyword::IF) {
             Some(self.parse_if_stmt(parser))
         } else if parser.parse_keywords(&[Keyword::CREATE, Keyword::TRIGGER]) {

src/parser/mod.rs

@@ -17925,9 +17925,9 @@ impl<'a> Parser<'a> {
         })
     }
 
-    /// Parse a 'BEGIN' statement
-    pub fn parse_begin(&mut self) -> Result<Statement, ParserError> {
-        let modifier = if !self.dialect.supports_start_transaction_modifier() {
+    /// Parse a transaction modifier keyword that can follow a BEGIN statement.
+    pub fn parse_transaction_modifier(&mut self) -> Option<TransactionModifier> {
+        if !self.dialect.supports_start_transaction_modifier() {
             None
         } else if self.parse_keyword(Keyword::DEFERRED) {
             Some(TransactionModifier::Deferred)
@@ -17941,7 +17941,12 @@ impl<'a> Parser<'a> {
             Some(TransactionModifier::Catch)
         } else {
             None
-        };
+        }
+    }
+
+    /// Parse a 'BEGIN' statement
+    pub fn parse_begin(&mut self) -> Result<Statement, ParserError> {
+        let modifier = self.parse_transaction_modifier();
         let transaction = match self.parse_one_of_keywords(&[Keyword::TRANSACTION, Keyword::WORK]) {
             Some(Keyword::TRANSACTION) => Some(BeginTransactionKind::Transaction),
             Some(Keyword::WORK) => Some(BeginTransactionKind::Work),