Added Isolation Level for Transactions

Author: chandr-andr

python/psql_rust_driver/__init__.py

@@ -1,4 +1,4 @@
-from ._internal import PSQLPool, QueryResult, Transaction
+from ._internal import PSQLPool, QueryResult, Transaction, IsolationLevel
 
 __all__ = [
     "PSQLPool",
@@ -7,4 +7,5 @@
     "SmallInt",
     "Integer",
     "BigInt",
+    "IsolationLevel",
 ]

python/psql_rust_driver/_internal/__init__.pyi

@@ -1,4 +1,5 @@
 from __future__ import annotations
+from enum import Enum
 from typing_extensions import Self
 
 from typing import Dict, Optional, Any, List
@@ -11,6 +12,15 @@ class QueryResult:
         """Return result from database as a list of dicts."""
 
 
+class IsolationLevel(Enum):
+    """Class for Isolation Level for transactions."""
+
+    ReadUncommitted = 1
+    ReadCommitted = 2
+    RepeatableRead = 3
+    Serializable = 4
+
+
 class Transaction:
     """Single connection for executing queries.
 
@@ -39,7 +49,7 @@ class Transaction:
     async def execute(
         self: Self,
         querystring: str,
-        parameters: List[Any],
+        parameters: List[Any] | None = None,
     ) -> QueryResult:
         """Execute the query.
         
@@ -206,12 +216,12 @@ class PSQLPool:
 
     def __init__(
         self: Self,
-        username: Optional[str],
-        password: Optional[str],
-        host: Optional[str],
-        port: Optional[int],
-        db_name: Optional[str],
-        max_db_pool_size: Optional[str],
+        username: Optional[str] = None,
+        password: Optional[str] = None,
+        host: Optional[str] = None,
+        port: Optional[int] = None,
+        db_name: Optional[str] = None,
+        max_db_pool_size: Optional[str] = None,
     ) -> None:
         """Create new PostgreSQL connection pool.
         
@@ -243,7 +253,7 @@ class PSQLPool:
     async def execute(
         self: Self,
         querystring: str,
-        parameters: List[Any],
+        parameters: List[Any] | None = None,
     ) -> QueryResult:
         """Execute the query.
         
@@ -274,9 +284,15 @@ class PSQLPool:
         """
         ...
 
-    async def transaction(self) -> Transaction:
+    async def transaction(
+        self,
+        isolation_level: IsolationLevel | None = IsolationLevel.ReadCommitted,
+    ) -> Transaction:
         """Create new transaction.
         
         It acquires new connection from the database pool
         and make it acts as transaction.
+
+        ### Parameters:
+        - `isolation_level`: configure isolation level of the transaction.
         """

src/driver/connection_pool.rs

@@ -11,6 +11,8 @@ use crate::{
     value_converter::{convert_parameters, PythonDTO},
 };
 
+use super::transaction_options::IsolationLevel;
+
 /// PSQLPool for internal use only.
 ///
 /// It is not exposed to python.
@@ -89,7 +91,10 @@ impl RustPSQLPool {
     ///
     /// # Errors:
     /// May return Err Result if cannot retrieve connection from the pool.
-    pub async fn inner_transaction<'a>(&'a self) -> RustPSQLDriverPyResult<Transaction> {
+    pub async fn inner_transaction<'a>(
+        &'a self,
+        isolation_level: Option<IsolationLevel>,
+    ) -> RustPSQLDriverPyResult<Transaction> {
         let db_pool_arc = self.db_pool.clone();
         let db_pool_guard = db_pool_arc.read().await;
 
@@ -106,6 +111,7 @@ impl RustPSQLPool {
             is_started: Arc::new(tokio::sync::RwLock::new(false)),
             is_done: Arc::new(tokio::sync::RwLock::new(false)),
             rollback_savepoint: Arc::new(tokio::sync::RwLock::new(HashSet::new())),
+            isolation_level: isolation_level,
         };
 
         Ok(Transaction {
@@ -221,13 +227,17 @@ impl PSQLPool {
     ///
     /// # Errors:
     /// May return Err Result if `inner_transaction` returns error.
-    pub fn transaction<'a>(&'a self, py: Python<'a>) -> RustPSQLDriverPyResult<&'a PyAny> {
+    pub fn transaction<'a>(
+        &'a self,
+        py: Python<'a>,
+        isolation_level: Option<IsolationLevel>,
+    ) -> RustPSQLDriverPyResult<&'a PyAny> {
         let psql_pool_arc = self.rust_psql_pool.clone();
 
         rustengine_future(py, async move {
             let psql_pool_guard = psql_pool_arc.write().await;
 
-            let transaction = psql_pool_guard.inner_transaction().await?;
+            let transaction = psql_pool_guard.inner_transaction(isolation_level).await?;
 
             Ok(transaction)
         })

src/driver/mod.rs

@@ -1,2 +1,3 @@
 pub mod connection_pool;
 pub mod transaction;
+pub mod transaction_options;

src/driver/transaction.rs

@@ -10,6 +10,8 @@ use crate::{
     value_converter::{convert_parameters, PythonDTO},
 };
 
+use super::transaction_options::IsolationLevel;
+
 /// Transaction for internal use only.
 ///
 /// It is not exposed to python.
@@ -18,6 +20,8 @@ pub struct RustTransaction {
     pub is_started: Arc<tokio::sync::RwLock<bool>>,
     pub is_done: Arc<tokio::sync::RwLock<bool>>,
     pub rollback_savepoint: Arc<tokio::sync::RwLock<HashSet<String>>>,
+
+    pub isolation_level: Option<IsolationLevel>,
 }
 
 impl RustTransaction {
@@ -73,6 +77,27 @@ impl RustTransaction {
         Ok(PSQLDriverPyQueryResult::new(result))
     }
 
+    /// Start transaction with isolation level if specified
+    ///
+    /// # Errors:
+    /// May return Err Result if cannot execute querystring.
+    pub async fn start_transaction<'a>(&'a self) -> RustPSQLDriverPyResult<()> {
+        let mut querystring = "START TRANSACTION".to_string();
+
+        if let Some(level) = self.isolation_level {
+            querystring.push_str(" ISOLATION LEVEL ");
+            let level = &level.to_str_level();
+            querystring.push_str(level);
+        };
+
+        let db_client_arc = self.db_client.clone();
+        let db_client_guard = db_client_arc.read().await;
+
+        db_client_guard.batch_execute(&querystring).await?;
+
+        Ok(())
+    }
+
     /// Start the transaction.
     ///
     /// Execute `BEGIN` commands and mark transaction as `started`.
@@ -84,7 +109,6 @@ impl RustTransaction {
     /// 2) Transaction is done.
     /// 3) Cannot execute `BEGIN` command.
     pub async fn inner_begin<'a>(&'a self) -> RustPSQLDriverPyResult<()> {
-        let db_client_arc = self.db_client.clone();
         let is_started_arc = self.is_started.clone();
         let is_done_arc = self.is_done.clone();
 
@@ -108,8 +132,7 @@ impl RustTransaction {
             ));
         }
 
-        let db_client_guard = db_client_arc.read().await;
-        db_client_guard.batch_execute("BEGIN").await?;
+        self.start_transaction().await?;
         let mut is_started_write_guard = is_started_arc.write().await;
         *is_started_write_guard = true;
 

src/driver/transaction_options.rs

@@ -0,0 +1,22 @@
+use pyo3::pyclass;
+
+#[pyclass]
+#[derive(Clone, Copy)]
+pub enum IsolationLevel {
+    ReadUncommitted = 1,
+    ReadCommitted = 2,
+    RepeatableRead = 3,
+    Serializable = 4,
+}
+
+impl IsolationLevel {
+    /// Return isolation level as String literal.
+    pub fn to_str_level(&self) -> String {
+        match self {
+            IsolationLevel::ReadUncommitted => "READ UNCOMMITTED".into(),
+            IsolationLevel::ReadCommitted => "READ COMMITTED".into(),
+            IsolationLevel::RepeatableRead => "REPEATABLE READ".into(),
+            IsolationLevel::Serializable => "SERIALIZABLE".into(),
+        }
+    }
+}

src/lib.rs

@@ -15,6 +15,7 @@ use pyo3::{pymodule, types::PyModule, PyResult, Python};
 fn psql_rust_engine(py: Python<'_>, pymod: &PyModule) -> PyResult<()> {
     pymod.add_class::<driver::connection_pool::PSQLPool>()?;
     pymod.add_class::<driver::transaction::Transaction>()?;
+    pymod.add_class::<driver::transaction_options::IsolationLevel>()?;
     pymod.add_class::<query_result::PSQLDriverPyQueryResult>()?;
     add_module(py, pymod, "extra_types", extra_types_module)?;
     add_module(py, pymod, "exceptions", python_exceptions_module)?;