continue implementing transaction

Author: insani7y

src/common.rs

@@ -2,6 +2,7 @@ use deadpool_postgres::Object;
 use pyo3::{types::PyModule, PyAny, PyResult, Python};
 
 use crate::{
+    driver::transaction_options::{IsolationLevel, ReadVariant},
     exceptions::rust_errors::RustPSQLDriverPyResult,
     query_result::{PSQLDriverPyQueryResult, PSQLDriverSinglePyQueryResult},
     value_converter::{convert_parameters, PythonDTO, QueryParameter},
@@ -31,6 +32,56 @@ pub fn add_module(
     Ok(())
 }
 
+pub trait BaseTransactionQuery {
+    fn start_transaction(
+        &self,
+        isolation_level: Option<IsolationLevel>,
+        read_variant: Option<ReadVariant>,
+        defferable: Option<bool>,
+    ) -> impl std::future::Future<Output = RustPSQLDriverPyResult<()>> + Send;
+    fn commit(&self) -> impl std::future::Future<Output = RustPSQLDriverPyResult<()>> + Send;
+    fn rollback(&self) -> impl std::future::Future<Output = RustPSQLDriverPyResult<()>> + Send;
+}
+
+impl BaseTransactionQuery for Object {
+    async fn start_transaction(
+        &self,
+        isolation_level: Option<IsolationLevel>,
+        read_variant: Option<ReadVariant>,
+        deferrable: Option<bool>,
+    ) -> RustPSQLDriverPyResult<()> {
+        let mut querystring = "START TRANSACTION".to_string();
+
+        if let Some(level) = isolation_level {
+            let level = &level.to_str_level();
+            querystring.push_str(format!(" ISOLATION LEVEL {level}").as_str());
+        };
+
+        querystring.push_str(match read_variant {
+            Some(ReadVariant::ReadOnly) => " READ ONLY",
+            Some(ReadVariant::ReadWrite) => " READ WRITE",
+            None => "",
+        });
+
+        querystring.push_str(match deferrable {
+            Some(true) => " DEFERRABLE",
+            Some(false) => " NOT DEFERRABLE",
+            None => "",
+        });
+        self.batch_execute(&querystring).await?;
+
+        Ok(())
+    }
+    async fn commit(&self) -> RustPSQLDriverPyResult<()> {
+        self.batch_execute("COMMIT;").await?;
+        Ok(())
+    }
+    async fn rollback(&self) -> RustPSQLDriverPyResult<()> {
+        self.batch_execute("ROLLBACK;").await?;
+        Ok(())
+    }
+}
+
 pub trait BaseDataBaseQuery {
     fn psqlpy_query_one(
         &self,

src/driver/transaction.rs

@@ -4,20 +4,11 @@ use pyo3::{prelude::*, pyclass};
 use crate::{
     exceptions::rust_errors::{RustPSQLDriverError, RustPSQLDriverPyResult},
     query_result::PSQLDriverPyQueryResult,
-    value_converter::{convert_parameters, PythonDTO},
 };
 
-use super::{
-    connection::Connection,
-    transaction_options::{IsolationLevel, ReadVariant},
-};
-use crate::common::BaseDataBaseQuery;
-use std::{
-    borrow::Borrow,
-    collections::HashSet,
-    ops::{Deref, DerefMut},
-    sync::Arc,
-};
+use super::transaction_options::{IsolationLevel, ReadVariant};
+use crate::common::{BaseDataBaseQuery, BaseTransactionQuery};
+use std::{collections::HashSet, sync::Arc};
 
 // use super::connection::RustConnection;
 
@@ -1252,7 +1243,6 @@ use std::{
 // }
 
 #[pyclass]
-#[derive(Clone)]
 pub struct Transaction {
     pub db_client: Arc<Object>,
     is_started: bool,
@@ -1276,9 +1266,28 @@ impl Transaction {
         slf
     }
 
-    async fn __aenter__<'a>(&mut self) -> RustPSQLDriverPyResult<Transaction> {
-        self.begin().await?;
-        Ok(self.clone())
+    async fn __aenter__<'a>(slf: Py<Self>) -> RustPSQLDriverPyResult<Py<Self>> {
+        let (is_transaction_ready, isolation_level, read_variant, deferrable, db_client) =
+            pyo3::Python::with_gil(|gil| {
+                let self_ = slf.borrow(gil);
+                (
+                    self_.check_is_transaction_ready(),
+                    self_.isolation_level,
+                    self_.read_variant,
+                    self_.deferrable,
+                    self_.db_client.clone(),
+                )
+            });
+        is_transaction_ready?;
+        db_client
+            .start_transaction(isolation_level, read_variant, deferrable)
+            .await?;
+
+        Python::with_gil(|gil| {
+            let mut self_ = slf.borrow_mut(gil);
+            self_.is_started = true;
+        });
+        Ok(slf)
     }
 
     #[allow(clippy::needless_pass_by_value)]
@@ -1288,21 +1297,30 @@ impl Transaction {
         exception: Py<PyAny>,
         _traceback: Py<PyAny>,
     ) -> RustPSQLDriverPyResult<()> {
-        let (is_exception_none, py_err, mut transaction) = pyo3::Python::with_gil(|gil| {
-            (
-                exception.is_none(gil),
-                PyErr::from_value_bound(exception.into_bound(gil)),
-                slf.borrow_mut(gil).clone(),
-            )
-        });
-
-        if is_exception_none {
-            transaction.commit().await?;
+        let (is_transaction_ready, is_exception_none, py_err, db_client) =
+            pyo3::Python::with_gil(|gil| {
+                let self_ = slf.borrow(gil);
+                (
+                    self_.check_is_transaction_ready(),
+                    exception.is_none(gil),
+                    PyErr::from_value_bound(exception.into_bound(gil)),
+                    self_.db_client.clone(),
+                )
+            });
+        is_transaction_ready?;
+        let exit_result = if is_exception_none {
+            db_client.commit().await?;
             Ok(())
         } else {
-            transaction.rollback().await?;
+            db_client.rollback().await?;
             Err(RustPSQLDriverError::PyError(py_err))
-        }
+        };
+
+        pyo3::Python::with_gil(|gil| {
+            let mut self_ = slf.borrow_mut(gil);
+            self_.is_done = true;
+        });
+        exit_result
     }
 
     /// Commit the transaction.
@@ -1317,9 +1335,8 @@ impl Transaction {
     /// 3) Cannot execute `COMMIT` command
     pub async fn commit(&mut self) -> RustPSQLDriverPyResult<()> {
         self.check_is_transaction_ready()?;
-        self.db_client.batch_execute("COMMIT;").await?;
+        self.db_client.commit().await?;
         self.is_done = true;
-
         Ok(())
     }
 
@@ -1350,15 +1367,17 @@ impl Transaction {
     /// 1) Cannot convert python parameters
     /// 2) Cannot execute querystring.
     pub async fn execute(
-        self_: pyo3::Py<Self>,
+        slf: Py<Self>,
         querystring: String,
         parameters: Option<pyo3::Py<PyAny>>,
         prepared: Option<bool>,
     ) -> RustPSQLDriverPyResult<PSQLDriverPyQueryResult> {
-        let transaction = pyo3::Python::with_gil(|gil| self_.borrow(gil).clone());
-        transaction.check_is_transaction_ready()?;
-        transaction
-            .db_client
+        let (is_transaction_ready, db_client) = pyo3::Python::with_gil(|gil| {
+            let self_ = slf.borrow(gil);
+            (self_.check_is_transaction_ready(), self_.db_client.clone())
+        });
+        is_transaction_ready?;
+        db_client
             .psqlpy_query(querystring, parameters, prepared)
             .await
     }
@@ -1443,7 +1462,6 @@ impl Transaction {
 
         Ok(())
     }
-
     fn check_is_transaction_ready(&self) -> RustPSQLDriverPyResult<()> {
         if !self.is_started {
             return Err(RustPSQLDriverError::DataBaseTransactionError(
@@ -1455,7 +1473,6 @@ impl Transaction {
                 "Transaction is already committed or rolled back".into(),
             ));
         }
-
         Ok(())
     }
 }