From 8bf418d5c62d096c59e2f9759657a59b5c07eb21 Mon Sep 17 00:00:00 2001
From: Andrew Lamb <andrew@nerdnetworks.org>
Date: Wed, 4 Dec 2024 11:21:55 -0500
Subject: [PATCH] Fix transaction kind check

---
 datafusion/sql/src/statement.rs | 24 ++++++++++++++++++------
 1 file changed, 18 insertions(+), 6 deletions(-)

diff --git a/datafusion/sql/src/statement.rs b/datafusion/sql/src/statement.rs
index 5aec7a6b1766b..e153569e5b5f4 100644
--- a/datafusion/sql/src/statement.rs
+++ b/datafusion/sql/src/statement.rs
@@ -55,7 +55,8 @@ use datafusion_expr::{
     Volatility, WriteOp,
 };
 use sqlparser::ast::{
-    self, NullsDistinctOption, ShowStatementIn, ShowStatementOptions, SqliteOnConflict,
+    self, BeginTransactionKind, NullsDistinctOption, ShowStatementIn,
+    ShowStatementOptions, SqliteOnConflict,
 };
 use sqlparser::ast::{
     Assignment, AssignmentTarget, ColumnDef, CreateIndex, CreateTable,
@@ -904,11 +905,7 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
                         "Transaction modifier not supported: {modifier}"
                     );
                 }
-                if let Some(transaction) = transaction {
-                    return not_impl_err!(
-                        "Transaction kind not supported: {transaction}"
-                    );
-                }
+                self.validate_transaction_kind(transaction)?;
                 let isolation_level: ast::TransactionIsolationLevel = modes
                     .iter()
                     .filter_map(|m: &TransactionMode| match m {
@@ -1994,4 +1991,19 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
             .get_table_source(tables_reference)
             .is_ok()
     }
+
+    fn validate_transaction_kind(
+        &self,
+        kind: Option<BeginTransactionKind>,
+    ) -> Result<()> {
+        match kind {
+            // BEGIN
+            None => Ok(()),
+            // BEGIN TRANSACTION
+            Some(BeginTransactionKind::Transaction) => Ok(()),
+            Some(BeginTransactionKind::Work) => {
+                not_impl_err!("Transaction kind not supported: {kind:?}")
+            }
+        }
+    }
 }