Add context manager support for database connections (#96)

penberg · web-flow · commit 696a23d7f155 · 2025-07-16T10:05:41.000+03:00
Implements __enter__ and __exit__ methods to enable using Connection objects as context managers. On clean exit, transactions are automatically committed. On exception, transactions are automatically rolled back. This provides sqlite3-compatible behavior and safer transaction handling. Closes #95
diff --git a/docs/api.md b/docs/api.md
@@ -32,6 +32,16 @@ Rolls back the current transaction and starts a new one.
 
 Closes the database connection.
 
+### `with` statement
+
+Connection objects can be used as context managers to ensure that transactions are properly committed or rolled back. When entering the context, the connection object is returned. When exiting:
+- Without exception: automatically commits the transaction
+- With exception: automatically rolls back the transaction
+
+This behavior is compatible with Python's `sqlite3` module. Context managers work correctly in both transactional and autocommit modes.
+
+When mixing manual transaction control with context managers, the context manager's commit/rollback will apply to any active transaction at the time of exit. Manual calls to `commit()` or `rollback()` within the context are allowed and will start a new transaction as usual.
+
 ### execute(sql, parameters=())
 
 Create a new cursor object and executes the SQL statement.
diff --git a/src/lib.rs b/src/lib.rs
@@ -3,7 +3,7 @@ use pyo3::create_exception;
 use pyo3::exceptions::PyValueError;
 use pyo3::prelude::*;
 use pyo3::types::{PyList, PyTuple};
-use std::cell::{OnceCell, RefCell};
+use std::cell::RefCell;
 use std::sync::{Arc, OnceLock};
 use std::time::Duration;
 use tokio::runtime::{Handle, Runtime};
@@ -38,14 +38,14 @@ fn is_remote_path(path: &str) -> bool {
 
 #[pyfunction]
 #[cfg(not(Py_3_12))]
-#[pyo3(signature = (database, timeout=5.0, isolation_level="DEFERRED".to_string(), check_same_thread=true, uri=false, sync_url=None, sync_interval=None, auth_token="", encryption_key=None))]
+#[pyo3(signature = (database, timeout=5.0, isolation_level="DEFERRED".to_string(), _check_same_thread=true, _uri=false, sync_url=None, sync_interval=None, auth_token="", encryption_key=None))]
 fn connect(
     py: Python<'_>,
     database: String,
     timeout: f64,
     isolation_level: Option<String>,
-    check_same_thread: bool,
-    uri: bool,
+    _check_same_thread: bool,
+    _uri: bool,
     sync_url: Option<String>,
     sync_interval: Option<f64>,
     auth_token: &str,
@@ -56,8 +56,8 @@ fn connect(
         database,
         timeout,
         isolation_level,
-        check_same_thread,
-        uri,
+        _check_same_thread,
+        _uri,
         sync_url,
         sync_interval,
         auth_token,
@@ -68,14 +68,14 @@ fn connect(
 
 #[pyfunction]
 #[cfg(Py_3_12)]
-#[pyo3(signature = (database, timeout=5.0, isolation_level="DEFERRED".to_string(), check_same_thread=true, uri=false, sync_url=None, sync_interval=None, auth_token="", encryption_key=None, autocommit = LEGACY_TRANSACTION_CONTROL))]
+#[pyo3(signature = (database, timeout=5.0, isolation_level="DEFERRED".to_string(), _check_same_thread=true, _uri=false, sync_url=None, sync_interval=None, auth_token="", encryption_key=None, autocommit = LEGACY_TRANSACTION_CONTROL))]
 fn connect(
     py: Python<'_>,
     database: String,
     timeout: f64,
     isolation_level: Option<String>,
-    check_same_thread: bool,
-    uri: bool,
+    _check_same_thread: bool,
+    _uri: bool,
     sync_url: Option<String>,
     sync_interval: Option<f64>,
     auth_token: &str,
@@ -87,8 +87,8 @@ fn connect(
         database,
         timeout,
         isolation_level.clone(),
-        check_same_thread,
-        uri,
+        _check_same_thread,
+        _uri,
         sync_url,
         sync_interval,
         auth_token,
@@ -111,8 +111,8 @@ fn _connect_core(
     database: String,
     timeout: f64,
     isolation_level: Option<String>,
-    check_same_thread: bool,
-    uri: bool,
+    _check_same_thread: bool,
+    _uri: bool,
     sync_url: Option<String>,
     sync_interval: Option<f64>,
     auth_token: &str,
@@ -220,7 +220,7 @@ unsafe impl Send for Connection {}
 
 #[pymethods]
 impl Connection {
-    fn close(self_: PyRef<'_, Self>, py: Python<'_>) -> PyResult<()> {
+    fn close(self_: PyRef<'_, Self>, _py: Python<'_>) -> PyResult<()> {
         self_.conn.replace(None);
         Ok(())
     }
@@ -330,11 +330,14 @@ impl Connection {
     fn in_transaction(self_: PyRef<'_, Self>) -> PyResult<bool> {
         #[cfg(Py_3_12)]
         {
-            return Ok(
+            Ok(
                 !self_.conn.borrow().as_ref().unwrap().is_autocommit() || self_.autocommit == 0
-            );
+            )
+        }
+        #[cfg(not(Py_3_12))]
+        {
+            Ok(!self_.conn.borrow().as_ref().unwrap().is_autocommit())
         }
-        Ok(!self_.conn.borrow().as_ref().unwrap().is_autocommit())
     }
 
     #[getter]
@@ -354,6 +357,26 @@ impl Connection {
         self_.autocommit = autocommit;
         Ok(())
     }
+
+    fn __enter__(slf: PyRef<'_, Self>) -> PyResult<PyRef<'_, Self>> {
+        Ok(slf)
+    }
+
+    fn __exit__(
+        self_: PyRef<'_, Self>,
+        exc_type: Option<&PyAny>,
+        _exc_val: Option<&PyAny>,
+        _exc_tb: Option<&PyAny>,
+    ) -> PyResult<bool> {
+        if exc_type.is_none() {
+            // Commit on clean exit
+            Connection::commit(self_)?;
+        } else {
+            // Rollback on error
+            Connection::rollback(self_)?;
+        }
+        Ok(false) // Always propagate exceptions
+    }
 }
 
 #[pyclass]
diff --git a/tests/test_suite.py b/tests/test_suite.py
