MSSQL: Support standalone BEGIN...END blocks (apache#2186)

Signed-off-by: Guan-Ming (Wesley) Chiu <105915352+guan404ming@users.noreply.github.com>
Co-authored-by: Ifeanyi Ubah <ify1992@yahoo.com>

Author: guan404ming

src/dialect/mssql.rs

@@ -145,7 +145,29 @@ impl Dialect for MsSqlDialect {
     }
 
     fn parse_statement(&self, parser: &mut Parser) -> Option<Result<Statement, ParserError>> {
-        if parser.peek_keyword(Keyword::IF) {
+        if parser.parse_keyword(Keyword::BEGIN) {
+            // 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();
+                None
+            }
+        } else if parser.peek_keyword(Keyword::IF) {
             Some(self.parse_if_stmt(parser))
         } else if parser.parse_keywords(&[Keyword::CREATE, Keyword::TRIGGER]) {
             Some(self.parse_create_trigger(parser, false))

src/parser/mod.rs

@@ -17902,9 +17902,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(crate) 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)
@@ -17918,7 +17918,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),

tests/sqlparser_mssql.rs

@@ -2554,3 +2554,74 @@ fn test_sql_keywords_as_column_aliases() {
         }
     }
 }
+
+#[test]
+fn parse_mssql_begin_end_block() {
+    // Single statement
+    let sql = "BEGIN SELECT 1; END";
+    let stmt = ms().verified_stmt(sql);
+    match &stmt {
+        Statement::StartTransaction {
+            begin,
+            has_end_keyword,
+            statements,
+            transaction,
+            modifier,
+            ..
+        } => {
+            assert!(begin);
+            assert!(has_end_keyword);
+            assert!(transaction.is_none());
+            assert!(modifier.is_none());
+            assert_eq!(statements.len(), 1);
+        }
+        _ => panic!("Expected StartTransaction, got: {stmt:?}"),
+    }
+
+    // Multiple statements
+    let sql = "BEGIN SELECT 1; SELECT 2; END";
+    let stmt = ms().verified_stmt(sql);
+    match &stmt {
+        Statement::StartTransaction {
+            statements,
+            has_end_keyword,
+            ..
+        } => {
+            assert!(has_end_keyword);
+            assert_eq!(statements.len(), 2);
+        }
+        _ => panic!("Expected StartTransaction, got: {stmt:?}"),
+    }
+
+    // DML inside BEGIN/END
+    let sql = "BEGIN INSERT INTO t VALUES (1); UPDATE t SET x = 2; END";
+    let stmt = ms().verified_stmt(sql);
+    match &stmt {
+        Statement::StartTransaction {
+            statements,
+            has_end_keyword,
+            ..
+        } => {
+            assert!(has_end_keyword);
+            assert_eq!(statements.len(), 2);
+        }
+        _ => panic!("Expected StartTransaction, got: {stmt:?}"),
+    }
+
+    // BEGIN TRANSACTION still works
+    let sql = "BEGIN TRANSACTION";
+    let stmt = ms().verified_stmt(sql);
+    match &stmt {
+        Statement::StartTransaction {
+            begin,
+            has_end_keyword,
+            transaction,
+            ..
+        } => {
+            assert!(begin);
+            assert!(!has_end_keyword);
+            assert!(transaction.is_some());
+        }
+        _ => panic!("Expected StartTransaction, got: {stmt:?}"),
+    }
+}