Pulled

Signed-off-by: chandr-andr (Kiselev Aleksandr) <chandr@chandr.net>

Author: chandr-andr

python/psqlpy/__init__.py

@@ -1,23 +1,7 @@
-from ._internal import (
-    # Connection,
-    # ConnRecyclingMethod,
-    # Cursor,
-    # IsolationLevel,
-    # PSQLPool,
-    # QueryResult,
-    # ReadVariant,
-    # Transaction,
-    ConnectionPool,
-)
+from psqlpy._internal import Connection, ConnectionPool, Transaction
 
 __all__ = [
     "ConnectionPool",
-    # "PSQLPool",
-    # "QueryResult",
-    # "Transaction",
-    # "IsolationLevel",
-    # "ReadVariant",
-    # "Connection",
-    # "Cursor",
-    # "ConnRecyclingMethod",
+    "Transaction",
+    "Connection",
 ]

python/psqlpy/_internal/__init__.pyi

@@ -941,7 +941,6 @@ class PSQLPool:
         It acquires new connection from the database pool.
         """
 
-
 class ConnectionPool:
     """Connection pool for executing queries.
 
@@ -1017,9 +1016,8 @@ class ConnectionPool:
             # it will be dropped on Rust side.
         ```
         """
-
     async def connection(self: Self) -> Connection:
         """Create new connection.
 
         It acquires new connection from the database pool.
-        """
+        """

src/common.rs

@@ -1,4 +1,11 @@
-use pyo3::{types::PyModule, PyResult, Python};
+use deadpool_postgres::Object;
+use pyo3::{types::PyModule, PyAny, PyResult, Python};
+
+use crate::{
+    exceptions::rust_errors::RustPSQLDriverPyResult,
+    query_result::{PSQLDriverPyQueryResult, PSQLDriverSinglePyQueryResult},
+    value_converter::{convert_parameters, PythonDTO, QueryParameter},
+};
 
 /// Add new module to the parent one.
 ///
@@ -23,3 +30,104 @@ pub fn add_module(
     )?;
     Ok(())
 }
+
+pub trait BaseDataBaseQuery {
+    fn psqlpy_query_one(
+        &self,
+        querystring: String,
+        parameters: Option<pyo3::Py<PyAny>>,
+        prepared: Option<bool>,
+    ) -> impl std::future::Future<Output = RustPSQLDriverPyResult<PSQLDriverSinglePyQueryResult>> + Send;
+
+    fn psqlpy_query(
+        &self,
+        querystring: String,
+        parameters: Option<pyo3::Py<PyAny>>,
+        prepared: Option<bool>,
+    ) -> impl std::future::Future<Output = RustPSQLDriverPyResult<PSQLDriverPyQueryResult>> + Send;
+
+    fn psqlpy_query_simple(
+        &self,
+        querystring: String,
+    ) -> impl std::future::Future<Output = RustPSQLDriverPyResult<()>> + Send;
+}
+
+impl BaseDataBaseQuery for Object {
+    async fn psqlpy_query_one(
+        &self,
+        querystring: String,
+        parameters: Option<pyo3::Py<PyAny>>,
+        prepared: Option<bool>,
+    ) -> RustPSQLDriverPyResult<PSQLDriverSinglePyQueryResult> {
+        let mut params: Vec<PythonDTO> = vec![];
+        if let Some(parameters) = parameters {
+            params = convert_parameters(parameters)?;
+        }
+        let prepared = prepared.unwrap_or(true);
+
+        let result = if prepared {
+            self.query_one(
+                &self.prepare_cached(&querystring).await?,
+                &params
+                    .iter()
+                    .map(|param| param as &QueryParameter)
+                    .collect::<Vec<&QueryParameter>>()
+                    .into_boxed_slice(),
+            )
+            .await?
+        } else {
+            self.query_one(
+                &querystring,
+                &params
+                    .iter()
+                    .map(|param| param as &QueryParameter)
+                    .collect::<Vec<&QueryParameter>>()
+                    .into_boxed_slice(),
+            )
+            .await?
+        };
+
+        Ok(PSQLDriverSinglePyQueryResult::new(result))
+    }
+
+    async fn psqlpy_query(
+        &self,
+        querystring: String,
+        parameters: Option<pyo3::Py<PyAny>>,
+        prepared: Option<bool>,
+    ) -> RustPSQLDriverPyResult<PSQLDriverPyQueryResult> {
+        let mut params: Vec<PythonDTO> = vec![];
+        if let Some(parameters) = parameters {
+            params = convert_parameters(parameters)?;
+        }
+        let prepared = prepared.unwrap_or(true);
+
+        let result = if prepared {
+            self.query(
+                &self.prepare_cached(&querystring).await?,
+                &params
+                    .iter()
+                    .map(|param| param as &QueryParameter)
+                    .collect::<Vec<&QueryParameter>>()
+                    .into_boxed_slice(),
+            )
+            .await?
+        } else {
+            self.query(
+                &querystring,
+                &params
+                    .iter()
+                    .map(|param| param as &QueryParameter)
+                    .collect::<Vec<&QueryParameter>>()
+                    .into_boxed_slice(),
+            )
+            .await?
+        };
+
+        Ok(PSQLDriverPyQueryResult::new(result))
+    }
+
+    async fn psqlpy_query_simple(&self, querystring: String) -> RustPSQLDriverPyResult<()> {
+        Ok(self.batch_execute(querystring.as_str()).await?)
+    }
+}

src/driver/connection.rs

@@ -1,13 +1,18 @@
 use deadpool_postgres::Object;
 use pyo3::{pyclass, pymethods, Py, PyAny, Python};
-use std::{sync::Arc, vec};
+use std::{collections::HashSet, sync::Arc, vec};
 
 use crate::{
     exceptions::rust_errors::{RustPSQLDriverError, RustPSQLDriverPyResult},
     query_result::{PSQLDriverPyQueryResult, PSQLDriverSinglePyQueryResult},
     value_converter::{convert_parameters, postgres_to_py, PythonDTO, QueryParameter},
 };
 
+use super::{
+    transaction::Transaction,
+    transaction_options::{IsolationLevel, ReadVariant},
+};
+
 #[pyclass]
 pub struct Connection {
     pub db_client: Arc<Object>,
@@ -34,8 +39,8 @@ impl Connection {
     pub async fn execute(
         self_: pyo3::Py<Self>,
         querystring: String,
-        prepared: Option<bool>,
         parameters: Option<pyo3::Py<PyAny>>,
+        prepared: Option<bool>,
     ) -> RustPSQLDriverPyResult<PSQLDriverPyQueryResult> {
         let db_client = pyo3::Python::with_gil(|gil| self_.borrow(gil).db_client.clone());
 
@@ -255,4 +260,22 @@ impl Connection {
             None => Ok(gil.None()),
         })
     }
+
+    #[must_use]
+    pub fn transaction(
+        &self,
+        isolation_level: Option<IsolationLevel>,
+        read_variant: Option<ReadVariant>,
+        deferrable: Option<bool>,
+    ) -> Transaction {
+        Transaction::new(
+            self.db_client.clone(),
+            false,
+            false,
+            HashSet::new(),
+            isolation_level,
+            read_variant,
+            deferrable,
+        )
+    }
 }