Merge pull request #14 from qaspen-python/transaction_fetch_val

Implement fetch_val method for Transaction

Author: chandr-andr

python/psqlpy/_internal/__init__.pyi

@@ -412,11 +412,13 @@ class Transaction:
         querystring: str,
         parameters: list[Any] | None = None,
     ) -> SingleQueryResult:
-        """Execute the query and return first row.
+        """Fetch exaclty single row from query.
 
+        Query must return exactly one row, otherwise error will be raised.
         Querystring can contain `$<number>` parameters
         for converting them in the driver side.
 
+
         ### Parameters:
         - `querystring`: querystring to execute.
         - `parameters`: list of parameters to pass in the query.
@@ -444,6 +446,54 @@ class Transaction:
 
         # Or you can transaction as a async context manager
 
+        async def main() -> None:
+            db_pool = PSQLPool()
+            await psqlpy.startup()
+
+            transaction = await db_pool.transaction()
+            async with transaction:
+                query_result: SingleQueryResult = await transaction.execute(
+                    "SELECT username FROM users WHERE id = $1 LIMIT 1",
+                    [100],
+                )
+                dict_result: Dict[Any, Any] = query_result.result()
+            # This way transaction begins and commits by itself.
+        ```
+        """
+    async def fetch_val(
+        self: Self,
+        querystring: str,
+        parameters: list[Any] | None = None,
+    ) -> Any | None:
+        """Execute the query and return first value of the first row.
+
+        Querystring can contain `$<number>` parameters
+        for converting them in the driver side.
+
+        ### Parameters:
+        - `querystring`: querystring to execute.
+        - `parameters`: list of parameters to pass in the query.
+
+        ### Example:
+        ```python
+        import asyncio
+
+        from psqlpy import PSQLPool, QueryResult
+
+
+        async def main() -> None:
+            db_pool = PSQLPool()
+            await db_pool.startup()
+
+            transaction = await db_pool.transaction()
+            await transaction.begin()
+            value: Any | None = await transaction.execute(
+                "SELECT username FROM users WHERE id = $1",
+                [100],
+            )
+
+        # Or you can transaction as a async context manager
+
         async def main() -> None:
             db_pool = PSQLPool()
             await psqlpy.startup()

python/tests/test_transaction.py

@@ -244,9 +244,48 @@ async def test_transaction_fetch_row(
     async with connection.transaction() as transaction:
         database_single_query_result: typing.Final = (
             await transaction.fetch_row(
-                f"SELECT * FROM  {table_name}",
+                f"SELECT * FROM  {table_name} LIMIT 1",
                 [],
             )
         )
         result = database_single_query_result.result()
         assert isinstance(result, dict)
+
+
+async def test_transaction_fetch_row_more_than_one_row(
+    psql_pool: PSQLPool,
+    table_name: str,
+) -> None:
+    connection = await psql_pool.connection()
+    async with connection.transaction() as transaction:
+        with pytest.raises(RustPSQLDriverPyBaseError):
+            await transaction.fetch_row(
+                f"SELECT * FROM  {table_name}",
+                [],
+            )
+
+
+async def test_transaction_fetch_val(
+    psql_pool: PSQLPool,
+    table_name: str,
+) -> None:
+    connection = await psql_pool.connection()
+    async with connection.transaction() as transaction:
+        value: typing.Final = await transaction.fetch_val(
+            f"SELECT COUNT(*) FROM {table_name}",
+            [],
+        )
+        assert isinstance(value, int)
+
+
+async def test_transaction_fetch_val_more_than_one_row(
+    psql_pool: PSQLPool,
+    table_name: str,
+) -> None:
+    connection = await psql_pool.connection()
+    async with connection.transaction() as transaction:
+        with pytest.raises(RustPSQLDriverPyBaseError):
+            await transaction.fetch_row(
+                f"SELECT * FROM  {table_name}",
+                [],
+            )

src/driver/transaction.rs

@@ -6,7 +6,7 @@ use crate::{
     common::rustengine_future,
     exceptions::rust_errors::{RustPSQLDriverError, RustPSQLDriverPyResult},
     query_result::{PSQLDriverPyQueryResult, PSQLDriverSinglePyQueryResult},
-    value_converter::{convert_parameters, PythonDTO},
+    value_converter::{convert_parameters, postgres_to_py, PythonDTO},
 };
 use deadpool_postgres::Object;
 use futures_util::future;
@@ -17,6 +17,7 @@ use pyo3::{
 };
 use std::{collections::HashSet, sync::Arc, vec};
 use tokio_postgres::types::ToSql;
+
 /// Transaction for internal use only.
 ///
 /// It is not exposed to python.
@@ -165,7 +166,7 @@ impl RustTransaction {
         Ok(())
     }
 
-    /// Fetch single row from query.
+    /// Fetch exaclty single row from query.
     ///
     /// Method doesn't acquire lock on any structure fields.
     /// It prepares and caches querystring in the inner Object object.
@@ -178,6 +179,7 @@ impl RustTransaction {
     /// 2) Transaction is done already
     /// 3) Can not create/retrieve prepared statement
     /// 4) Can not execute statement
+    /// 5) Query returns more than one row
     pub async fn inner_fetch_row(
         &self,
         querystring: String,
@@ -210,7 +212,7 @@ impl RustTransaction {
         let statement = db_client_guard.prepare_cached(&querystring).await?;
 
         let result = db_client_guard
-            .query(&statement, &vec_parameters.into_boxed_slice())
+            .query_one(&statement, &vec_parameters.into_boxed_slice())
             .await?;
 
         Ok(PSQLDriverSinglePyQueryResult::new(result))
@@ -755,6 +757,7 @@ impl Transaction {
     /// May return Err Result if:
     /// 1) Cannot convert python parameters
     /// 2) Cannot execute querystring.
+    /// 3) Query returns more than one row.
     pub fn fetch_row<'a>(
         &'a self,
         py: Python<'a>,
@@ -773,6 +776,42 @@ impl Transaction {
         })
     }
 
+    /// Execute querystring with parameters and return first value in the first row.
+    ///
+    /// It converts incoming parameters to rust readable,
+    /// executes query with them and returns first row of response.
+    ///
+    /// # Errors
+    ///
+    /// May return Err Result if:
+    /// 1) Cannot convert python parameters
+    /// 2) Cannot execute querystring.
+    /// 3) Query returns more than one row
+    pub fn fetch_val<'a>(
+        &'a self,
+        py: Python<'a>,
+        querystring: String,
+        parameters: Option<&'a PyList>,
+    ) -> RustPSQLDriverPyResult<&PyAny> {
+        let transaction_arc = self.transaction.clone();
+        let mut params: Vec<PythonDTO> = vec![];
+        if let Some(parameters) = parameters {
+            params = convert_parameters(parameters)?;
+        }
+
+        rustengine_future(py, async move {
+            let transaction_guard = transaction_arc.read().await;
+            let first_row = transaction_guard
+                .inner_fetch_row(querystring, params)
+                .await?
+                .get_inner();
+            Python::with_gil(|py| match first_row.columns().first() {
+                Some(first_column) => postgres_to_py(py, &first_row, first_column, 0),
+                None => Ok(py.None()),
+            })
+        })
+    }
+
     /// Execute querystrings with parameters and return all results.
     ///
     /// Create pipeline of queries.

src/query_result.rs

@@ -1,10 +1,7 @@
 use pyo3::{pyclass, pymethods, types::PyDict, Py, PyAny, Python, ToPyObject};
 use tokio_postgres::Row;
 
-use crate::{
-    exceptions::rust_errors::{RustPSQLDriverError, RustPSQLDriverPyResult},
-    value_converter::postgres_to_py,
-};
+use crate::{exceptions::rust_errors::RustPSQLDriverPyResult, value_converter::postgres_to_py};
 
 /// Convert postgres `Row` into Python Dict.
 ///
@@ -82,16 +79,20 @@ impl PSQLDriverPyQueryResult {
 #[pyclass(name = "SingleQueryResult")]
 #[allow(clippy::module_name_repetitions)]
 pub struct PSQLDriverSinglePyQueryResult {
-    inner: Vec<Row>,
+    inner: Row,
 }
 
 impl PSQLDriverSinglePyQueryResult {
     #[must_use]
-    pub fn new(database_row: Vec<Row>) -> Self {
+    pub fn new(database_row: Row) -> Self {
         PSQLDriverSinglePyQueryResult {
             inner: database_row,
         }
     }
+
+    pub fn get_inner(self) -> Row {
+        self.inner
+    }
 }
 
 #[pymethods]
@@ -106,12 +107,7 @@ impl PSQLDriverSinglePyQueryResult {
     /// postgres type to python, can not set new key-value pair
     /// in python dict or there are no result.
     pub fn result(&self, py: Python<'_>) -> RustPSQLDriverPyResult<Py<PyAny>> {
-        if let Some(row) = self.inner.first() {
-            return Ok(row_to_dict(py, row)?.to_object(py));
-        }
-        Err(RustPSQLDriverError::RustToPyValueConversionError(
-            "There are not results from the query, can't return first row.".into(),
-        ))
+        Ok(row_to_dict(py, &self.inner)?.to_object(py))
     }
 
     /// Convert result from database to any class passed from Python.
@@ -126,12 +122,7 @@ impl PSQLDriverSinglePyQueryResult {
         py: Python<'a>,
         as_class: &'a PyAny,
     ) -> RustPSQLDriverPyResult<&'a PyAny> {
-        if let Some(row) = self.inner.first() {
-            let pydict: &PyDict = row_to_dict(py, row)?;
-            return Ok(as_class.call((), Some(pydict))?);
-        }
-        Err(RustPSQLDriverError::RustToPyValueConversionError(
-            "There are not results from the query, can't convert first row.".into(),
-        ))
+        let pydict: &PyDict = row_to_dict(py, &self.inner)?;
+        Ok(as_class.call((), Some(pydict))?)
     }
 }