From 1318e497ac4798e9c3203f20f9393dd86ac2fa43 Mon Sep 17 00:00:00 2001
From: Takahiro Ebato <taka2207@i.softbank.jp>
Date: Fri, 15 Dec 2023 23:26:31 +0900
Subject: [PATCH] add support for BEGIN TRANSACTION modifiers in sqlite

Support the following syntaxes
- BEGIN DEFERRED
- BEGIN IMMEDIATE
- BEGIN EXCLUSIVE
---
 src/ast/mod.rs            | 32 +++++++++++++++++++++++++++++++-
 src/keywords.rs           |  3 +++
 src/parser/mod.rs         | 13 +++++++++++++
 tests/sqlparser_common.rs | 34 ++++++++++++++++++++++++++++++++++
 4 files changed, 81 insertions(+), 1 deletion(-)

diff --git a/src/ast/mod.rs b/src/ast/mod.rs
index 4769ea9bd..95491ad7b 100644
--- a/src/ast/mod.rs
+++ b/src/ast/mod.rs
@@ -1819,6 +1819,8 @@ pub enum Statement {
     StartTransaction {
         modes: Vec<TransactionMode>,
         begin: bool,
+        // Only for sqlite
+        modifier: Option<TransactionModifier>,
     },
     /// `SET TRANSACTION ...`
     SetTransaction {
@@ -3107,9 +3109,14 @@ impl fmt::Display for Statement {
             Statement::StartTransaction {
                 modes,
                 begin: syntax_begin,
+                modifier,
             } => {
                 if *syntax_begin {
-                    write!(f, "BEGIN TRANSACTION")?;
+                    if let Some(modifier) = *modifier {
+                        write!(f, "BEGIN {} TRANSACTION", modifier)?;
+                    } else {
+                        write!(f, "BEGIN TRANSACTION")?;
+                    }
                 } else {
                     write!(f, "START TRANSACTION")?;
                 }
@@ -4292,6 +4299,29 @@ impl fmt::Display for TransactionIsolationLevel {
     }
 }
 
+/// Sqlite specific syntax
+///
+/// https://sqlite.org/lang_transaction.html
+#[derive(Debug, Copy, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub enum TransactionModifier {
+    Deferred,
+    Immediate,
+    Exclusive,
+}
+
+impl fmt::Display for TransactionModifier {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        use TransactionModifier::*;
+        f.write_str(match self {
+            Deferred => "DEFERRED",
+            Immediate => "IMMEDIATE",
+            Exclusive => "EXCLUSIVE",
+        })
+    }
+}
+
 #[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/keywords.rs b/src/keywords.rs
index 2de36562f..7f13e3990 100644
--- a/src/keywords.rs
+++ b/src/keywords.rs
@@ -209,6 +209,7 @@ define_keywords!(
     DECIMAL,
     DECLARE,
     DEFAULT,
+    DEFERRED,
     DELETE,
     DELIMITED,
     DELIMITER,
@@ -254,6 +255,7 @@ define_keywords!(
     EVERY,
     EXCEPT,
     EXCLUDE,
+    EXCLUSIVE,
     EXEC,
     EXECUTE,
     EXISTS,
@@ -321,6 +323,7 @@ define_keywords!(
     IF,
     IGNORE,
     ILIKE,
+    IMMEDIATE,
     IMMUTABLE,
     IN,
     INCLUDE,
diff --git a/src/parser/mod.rs b/src/parser/mod.rs
index d9d4761c3..9842ca4ba 100644
--- a/src/parser/mod.rs
+++ b/src/parser/mod.rs
@@ -7834,14 +7834,27 @@ impl<'a> Parser<'a> {
         Ok(Statement::StartTransaction {
             modes: self.parse_transaction_modes()?,
             begin: false,
+            modifier: None,
         })
     }
 
     pub fn parse_begin(&mut self) -> Result<Statement, ParserError> {
+        let modifier = if !dialect_of!(self is SQLiteDialect) {
+            None
+        } else if self.parse_keyword(Keyword::DEFERRED) {
+            Some(TransactionModifier::Deferred)
+        } else if self.parse_keyword(Keyword::IMMEDIATE) {
+            Some(TransactionModifier::Immediate)
+        } else if self.parse_keyword(Keyword::EXCLUSIVE) {
+            Some(TransactionModifier::Exclusive)
+        } else {
+            None
+        };
         let _ = self.parse_one_of_keywords(&[Keyword::TRANSACTION, Keyword::WORK]);
         Ok(Statement::StartTransaction {
             modes: self.parse_transaction_modes()?,
             begin: true,
+            modifier,
         })
     }
 
diff --git a/tests/sqlparser_common.rs b/tests/sqlparser_common.rs
index 1d0923b4f..fc21a686e 100644
--- a/tests/sqlparser_common.rs
+++ b/tests/sqlparser_common.rs
@@ -6232,6 +6232,40 @@ fn parse_start_transaction() {
     );
 }
 
+#[test]
+fn parse_start_transaction_sqlite() {
+    let dialect = SQLiteDialect {};
+
+    let check = |sql: &str, expected_modifier: Option<TransactionModifier>| match Parser::parse_sql(
+        &dialect, &sql,
+    )
+    .unwrap()
+    .pop()
+    .unwrap()
+    {
+        Statement::StartTransaction { modifier, .. } => assert_eq!(modifier, expected_modifier),
+        _ => panic!("{}", sql),
+    };
+
+    let sql = "BEGIN DEFERRED";
+    check(sql, Some(TransactionModifier::Deferred));
+
+    let sql = "BEGIN DEFERRED TRANSACTION";
+    check(sql, Some(TransactionModifier::Deferred));
+
+    let sql = "BEGIN IMMEDIATE";
+    check(sql, Some(TransactionModifier::Immediate));
+
+    let sql = "BEGIN IMMEDIATE TRANSACTION";
+    check(sql, Some(TransactionModifier::Immediate));
+
+    let sql = "BEGIN EXCLUSIVE";
+    check(sql, Some(TransactionModifier::Exclusive));
+
+    let sql = "BEGIN EXCLUSIVE TRANSACTION";
+    check(sql, Some(TransactionModifier::Exclusive));
+}
+
 #[test]
 fn parse_set_transaction() {
     // SET TRANSACTION shares transaction mode parsing code with START