From 8bfcad6662dc9515784d744de99934c59c858c42 Mon Sep 17 00:00:00 2001
From: Duncan Fairbanks <duncanfairbanks6@gmail.com>
Date: Wed, 27 Nov 2024 18:50:02 -0800
Subject: [PATCH] chore: test begin_with works for all SQLite "BEGIN"
 statements

---
 sqlx-sqlite/src/connection/mod.rs | 25 ++++++++++++++++++-
 tests/sqlite/sqlite.rs            | 40 +++++++++++++++++++++++++++++++
 2 files changed, 64 insertions(+), 1 deletion(-)

diff --git a/sqlx-sqlite/src/connection/mod.rs b/sqlx-sqlite/src/connection/mod.rs
index 53c3156e9d..9c14a316ac 100644
--- a/sqlx-sqlite/src/connection/mod.rs
+++ b/sqlx-sqlite/src/connection/mod.rs
@@ -13,7 +13,8 @@ use futures_intrusive::sync::MutexGuard;
 use futures_util::future;
 use libsqlite3_sys::{
     sqlite3, sqlite3_commit_hook, sqlite3_get_autocommit, sqlite3_progress_handler,
-    sqlite3_rollback_hook, sqlite3_update_hook, SQLITE_DELETE, SQLITE_INSERT, SQLITE_UPDATE,
+    sqlite3_rollback_hook, sqlite3_txn_state, sqlite3_update_hook, SQLITE_DELETE, SQLITE_INSERT,
+    SQLITE_TXN_NONE, SQLITE_TXN_READ, SQLITE_TXN_WRITE, SQLITE_UPDATE,
 };
 
 pub(crate) use handle::ConnectionHandle;
@@ -508,6 +509,28 @@ impl LockedSqliteHandle<'_> {
         let ret = unsafe { sqlite3_get_autocommit(self.as_raw_handle().as_ptr()) };
         ret == 0
     }
+
+    /// Calls `sqlite3_txn_state` on this handle.
+    #[cfg(test)]
+    pub fn transaction_state(&mut self) -> Result<SqliteTransactionState, Error> {
+        let state =
+            match unsafe { sqlite3_txn_state(self.as_raw_handle().as_ptr(), std::ptr::null()) } {
+                SQLITE_TXN_NONE => SqliteTransactionState::None,
+                SQLITE_TXN_READ => SqliteTransactionState::Read,
+                SQLITE_TXN_WRITE => SqliteTransactionState::Write,
+                _ => return Err(Error::Protocol("Invalid transaction state".into())),
+            };
+        Ok(state)
+    }
+}
+
+/// Transaction state enum for a SQLite connection.
+#[cfg(test)]
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub enum SqliteTransactionState {
+    None,
+    Read,
+    Write,
 }
 
 impl Drop for ConnectionState {
diff --git a/tests/sqlite/sqlite.rs b/tests/sqlite/sqlite.rs
index b733ccbb4c..acdefeab18 100644
--- a/tests/sqlite/sqlite.rs
+++ b/tests/sqlite/sqlite.rs
@@ -960,3 +960,43 @@ async fn test_multiple_set_rollback_hook_calls_drop_old_handler() -> anyhow::Res
     assert_eq!(1, Arc::strong_count(&ref_counted_object));
     Ok(())
 }
+
+#[sqlx_macros::test]
+async fn it_can_use_transaction_options() -> anyhow::Result<()> {
+    async fn check_txn_state(
+        conn: &mut SqliteConnection,
+    ) -> Result<SqliteTransactionState, sqlx::Error> {
+        conn.lock_handle().await?.transaction_state()
+    }
+
+    let mut conn = new::<Sqlite>().await?;
+
+    assert_eq!(
+        check_txn_state(&mut conn).await?,
+        SqliteTransactionState::None
+    );
+
+    let mut tx = conn.begin_with("BEGIN DEFERRED").await?;
+    assert_eq!(
+        check_txn_state(&mut *tx).await?,
+        SqliteTransactionState::None
+    );
+    drop(tx);
+
+    let mut tx = conn.begin_with("BEGIN IMMEDIATE").await?;
+    assert_eq!(
+        check_txn_state(&mut *tx).await?,
+        SqliteTransactionState::Write
+    );
+    drop(tx);
+
+    // Note: may result in database locked errors if tests are run in parallel
+    let mut tx = conn.begin_with("BEGIN EXCLUSIVE").await?;
+    assert_eq!(
+        check_txn_state(&mut *tx).await?,
+        SqliteTransactionState::Write
+    );
+    drop(tx);
+
+    Ok(())
+}