From c761f0babbeefdc7b2e8fff5bf0e7bb02988ad03 Mon Sep 17 00:00:00 2001
From: Michael Victor Zink <michael.z@readyset.io>
Date: Tue, 3 Dec 2024 17:10:28 -0800
Subject: [PATCH] Fix displaying WORK or TRANSACTION after BEGIN (#1565)

---
 src/ast/mod.rs            | 29 ++++++++++++++++++++++++++---
 src/parser/mod.rs         |  8 +++++++-
 tests/sqlparser_common.rs |  6 +++---
 tests/sqlparser_mysql.rs  |  5 +++++
 tests/sqlparser_sqlite.rs |  6 +++---
 5 files changed, 44 insertions(+), 10 deletions(-)

diff --git a/src/ast/mod.rs b/src/ast/mod.rs
index d4278e4f9..326375b5f 100644
--- a/src/ast/mod.rs
+++ b/src/ast/mod.rs
@@ -2944,6 +2944,7 @@ pub enum Statement {
     StartTransaction {
         modes: Vec<TransactionMode>,
         begin: bool,
+        transaction: Option<BeginTransactionKind>,
         /// Only for SQLite
         modifier: Option<TransactionModifier>,
     },
@@ -4519,16 +4520,20 @@ impl fmt::Display for Statement {
             Statement::StartTransaction {
                 modes,
                 begin: syntax_begin,
+                transaction,
                 modifier,
             } => {
                 if *syntax_begin {
                     if let Some(modifier) = *modifier {
-                        write!(f, "BEGIN {} TRANSACTION", modifier)?;
+                        write!(f, "BEGIN {}", modifier)?;
                     } else {
-                        write!(f, "BEGIN TRANSACTION")?;
+                        write!(f, "BEGIN")?;
                     }
                 } else {
-                    write!(f, "START TRANSACTION")?;
+                    write!(f, "START")?;
+                }
+                if let Some(transaction) = transaction {
+                    write!(f, " {transaction}")?;
                 }
                 if !modes.is_empty() {
                     write!(f, " {}", display_comma_separated(modes))?;
@@ -5023,6 +5028,24 @@ pub enum TruncateCascadeOption {
     Restrict,
 }
 
+/// Transaction started with [ TRANSACTION | WORK ]
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub enum BeginTransactionKind {
+    Transaction,
+    Work,
+}
+
+impl Display for BeginTransactionKind {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        match self {
+            BeginTransactionKind::Transaction => write!(f, "TRANSACTION"),
+            BeginTransactionKind::Work => write!(f, "WORK"),
+        }
+    }
+}
+
 /// Can use to describe options in  create sequence or table column type identity
 /// [ MINVALUE minvalue | NO MINVALUE ] [ MAXVALUE maxvalue | NO MAXVALUE ]
 #[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
diff --git a/src/parser/mod.rs b/src/parser/mod.rs
index 32e7e3743..7b175f1d9 100644
--- a/src/parser/mod.rs
+++ b/src/parser/mod.rs
@@ -12123,6 +12123,7 @@ impl<'a> Parser<'a> {
         Ok(Statement::StartTransaction {
             modes: self.parse_transaction_modes()?,
             begin: false,
+            transaction: Some(BeginTransactionKind::Transaction),
             modifier: None,
         })
     }
@@ -12139,10 +12140,15 @@ impl<'a> Parser<'a> {
         } else {
             None
         };
-        let _ = self.parse_one_of_keywords(&[Keyword::TRANSACTION, Keyword::WORK]);
+        let transaction = match self.parse_one_of_keywords(&[Keyword::TRANSACTION, Keyword::WORK]) {
+            Some(Keyword::TRANSACTION) => Some(BeginTransactionKind::Transaction),
+            Some(Keyword::WORK) => Some(BeginTransactionKind::Work),
+            _ => None,
+        };
         Ok(Statement::StartTransaction {
             modes: self.parse_transaction_modes()?,
             begin: true,
+            transaction,
             modifier,
         })
     }
diff --git a/tests/sqlparser_common.rs b/tests/sqlparser_common.rs
index f146b298e..e80223807 100644
--- a/tests/sqlparser_common.rs
+++ b/tests/sqlparser_common.rs
@@ -7736,9 +7736,9 @@ fn parse_start_transaction() {
     }
 
     verified_stmt("START TRANSACTION");
-    one_statement_parses_to("BEGIN", "BEGIN TRANSACTION");
-    one_statement_parses_to("BEGIN WORK", "BEGIN TRANSACTION");
-    one_statement_parses_to("BEGIN TRANSACTION", "BEGIN TRANSACTION");
+    verified_stmt("BEGIN");
+    verified_stmt("BEGIN WORK");
+    verified_stmt("BEGIN TRANSACTION");
 
     verified_stmt("START TRANSACTION ISOLATION LEVEL READ UNCOMMITTED");
     verified_stmt("START TRANSACTION ISOLATION LEVEL READ COMMITTED");
diff --git a/tests/sqlparser_mysql.rs b/tests/sqlparser_mysql.rs
index f20a759af..f7a21f99b 100644
--- a/tests/sqlparser_mysql.rs
+++ b/tests/sqlparser_mysql.rs
@@ -3032,3 +3032,8 @@ fn parse_longblob_type() {
     mysql_and_generic().verified_stmt("CREATE TABLE foo (bar MEDIUMTEXT)");
     mysql_and_generic().verified_stmt("CREATE TABLE foo (bar LONGTEXT)");
 }
+
+#[test]
+fn parse_begin_without_transaction() {
+    mysql().verified_stmt("BEGIN");
+}
diff --git a/tests/sqlparser_sqlite.rs b/tests/sqlparser_sqlite.rs
index c3cfb7a63..4f23979c5 100644
--- a/tests/sqlparser_sqlite.rs
+++ b/tests/sqlparser_sqlite.rs
@@ -527,9 +527,9 @@ fn parse_start_transaction_with_modifier() {
     sqlite_and_generic().verified_stmt("BEGIN DEFERRED TRANSACTION");
     sqlite_and_generic().verified_stmt("BEGIN IMMEDIATE TRANSACTION");
     sqlite_and_generic().verified_stmt("BEGIN EXCLUSIVE TRANSACTION");
-    sqlite_and_generic().one_statement_parses_to("BEGIN DEFERRED", "BEGIN DEFERRED TRANSACTION");
-    sqlite_and_generic().one_statement_parses_to("BEGIN IMMEDIATE", "BEGIN IMMEDIATE TRANSACTION");
-    sqlite_and_generic().one_statement_parses_to("BEGIN EXCLUSIVE", "BEGIN EXCLUSIVE TRANSACTION");
+    sqlite_and_generic().verified_stmt("BEGIN DEFERRED");
+    sqlite_and_generic().verified_stmt("BEGIN IMMEDIATE");
+    sqlite_and_generic().verified_stmt("BEGIN EXCLUSIVE");
 
     let unsupported_dialects = TestedDialects::new(
         all_dialects()