Fix transaction kind check

apache · Dec 4, 2024 · 8bf418d · 8bf418d
1 parent 9b9df23
commit 8bf418d
Showing 1 changed file with 18 additions and 6 deletions.
diff --git 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:?}")
+            }
+        }
+    }
 }