Merge pull request #11 from qaspen-python/feature/add_convertion_to_class_support

Added support for any custom python class in QueryResult

Author: chandr-andr

README.md

@@ -119,14 +119,48 @@ As connection can be closed in different situations on various sides you can sel
   `DEALLOCATE ALL` and `DISCARD PLAN`, so that the statement cache is not
   rendered ineffective.
 
+## Results from querying
+You have some options to get results from the query.
+`execute()` method, for example, returns `QueryResult` and this class can be converted into `list` of `dict`s - `list[dict[Any, Any]]` or into any Python class (`pydantic` model, as an example).
+
+Let's see some code:
+```python
+from typing import Any
+
+from pydantic import BaseModel
+from psqlpy import PSQLPool, QueryResult
+
+
+class ExampleModel(BaseModel):
+    id: int
+    username: str
+
+
+db_pool = PSQLPool(
+    dsn="postgres://postgres:postgres@localhost:5432/postgres",
+    max_db_pool_size=2,
+)
+
+async def main() -> None:
+    await db_pool.startup()
+
+    res: QueryResult = await db_pool.execute(
+        "SELECT * FROM users",
+    )
+
+    pydantic_res: list[ExampleModel] = res.as_class(
+        as_class=ExampleModel,
+    )
+```
+
 ## Query parameters
 
 You can pass parameters into queries.
 Parameters can be passed in any `execute` method as the second parameter, it must be a list.
 Any placeholder must be marked with `$< num>`.
 
 ```python
-    res: list[dict[str, Any]] = await db_pool.execute(
+    res: QueryResult = await db_pool.execute(
         "SELECT * FROM users WHERE user_id = $1 AND first_name = $2",
         [100, "RustDriver"],
     )

python/psqlpy/_internal/__init__.pyi

@@ -1,20 +1,101 @@
 import types
 from enum import Enum
-from typing import Any, Optional
+from typing import Any, Callable, Optional, TypeVar
 
 from typing_extensions import Self
 
+_CustomClass = TypeVar(
+    "_CustomClass",
+)
+
 class QueryResult:
     """Result."""
 
     def result(self: Self) -> list[dict[Any, Any]]:
         """Return result from database as a list of dicts."""
+    def as_class(
+        self: Self,
+        as_class: Callable[..., _CustomClass],
+    ) -> list[_CustomClass]:
+        """Convert results to passed class.
+
+        The main goal of this method is pydantic,
+        msgspec and dataclasses support.
+
+        ### Parameters:
+        - `as_class`: Any callable python class for the results.
+
+        ### Example:
+        ```python
+        import asyncio
+
+        from psqlpy import PSQLPool, QueryResult
+
+
+        class ExampleOfAsClass:
+            def __init__(self, username: str) -> None:
+                self.username = username
+
+
+        async def main() -> None:
+            db_pool = PSQLPool()
+            await db_pool.startup()
+
+            query_result: QueryResult = await db_pool.execute(
+                "SELECT username FROM users WHERE id = $1",
+                [100],
+            )
+            class_results: List[ExampleOfAsClass] = query_result.as_class(
+                as_class=ExampleOfAsClass,
+            )
+        ```
+        """
 
 class SingleQueryResult:
     """Single result."""
 
     def result(self: Self) -> dict[Any, Any]:
         """Return result from database as a dict."""
+    def as_class(
+        self: Self,
+        as_class: Callable[..., _CustomClass],
+    ) -> list[_CustomClass]:
+        """Convert results to passed class.
+
+        The main goal of this method is pydantic,
+        msgspec and dataclasses support.
+
+        ### Parameters:
+        - `as_class`: Any callable python class for the results.
+
+        ### Example:
+        ```python
+        import asyncio
+
+        from psqlpy import PSQLPool, QueryResult
+
+
+        class ExampleOfAsClass:
+            def __init__(self, username: str) -> None:
+                self.username = username
+
+
+        async def main() -> None:
+            db_pool = PSQLPool()
+            await db_pool.startup()
+
+            connection = await db_pool.connection()
+            async with connection.transaction() as trans:
+                query_result: SingleQueryResult = await trans.fetch_row(
+                    "SELECT username FROM users WHERE id = $1",
+                    [100],
+                )
+
+            class_result: ExampleOfAsClass = query_result.as_class(
+                as_class=ExampleOfAsClass,
+            )
+        ```
+        """
 
 class IsolationLevel(Enum):
     """Class for Isolation Level for transactions."""

python/tests/conftest.py

@@ -3,10 +3,26 @@
 from typing import AsyncGenerator
 
 import pytest
+from pydantic import BaseModel
 
 from psqlpy import Cursor, PSQLPool
 
 
+class DefaultPydanticModel(BaseModel):
+    """Validation model for test data based on Pydantic."""
+
+    id: int
+    name: str
+
+
+class DefaultPythonModelClass:
+    """Validation model for test data based on default Python class."""
+
+    def __init__(self, id: int, name: str) -> None:
+        self.id = id
+        self.name = name
+
+
 @pytest.fixture()
 def anyio_backend() -> str:
     """

python/tests/test_value_converter.py

@@ -0,0 +1,34 @@
+import pytest
+from tests.conftest import DefaultPydanticModel, DefaultPythonModelClass
+
+from psqlpy import PSQLPool
+
+pytestmark = pytest.mark.anyio
+
+
+async def test_as_class(
+    psql_pool: PSQLPool,
+    table_name: str,
+    number_database_records: int,
+) -> None:
+    """Test `as_class()` method."""
+    select_result = await psql_pool.execute(
+        f"SELECT * FROM {table_name}",
+    )
+
+    as_pydantic = select_result.as_class(
+        as_class=DefaultPydanticModel,
+    )
+    assert len(as_pydantic) == number_database_records
+
+    for single_record in as_pydantic:
+        assert isinstance(single_record, DefaultPydanticModel)
+
+    as_py_class = select_result.as_class(
+        as_class=DefaultPythonModelClass,
+    )
+
+    assert len(as_py_class) == number_database_records
+
+    for single_py_record in as_py_class:
+        assert isinstance(single_py_record, DefaultPythonModelClass)

src/query_result.rs

@@ -1,7 +1,26 @@
 use pyo3::{pyclass, pymethods, types::PyDict, Py, PyAny, Python, ToPyObject};
 use tokio_postgres::Row;
 
-use crate::{exceptions::rust_errors::RustPSQLDriverPyResult, value_converter::postgres_to_py};
+use crate::{
+    exceptions::rust_errors::{RustPSQLDriverError, RustPSQLDriverPyResult},
+    value_converter::postgres_to_py,
+};
+
+/// Convert postgres `Row` into Python Dict.
+///
+/// # Errors
+///
+/// May return Err Result if can not convert
+/// postgres type to python or set new key-value pair
+/// in python dict.
+fn row_to_dict<'a>(py: Python<'a>, postgres_row: &'a Row) -> RustPSQLDriverPyResult<&'a PyDict> {
+    let python_dict = PyDict::new(py);
+    for (column_idx, column) in postgres_row.columns().iter().enumerate() {
+        let python_type = postgres_to_py(py, postgres_row, column, column_idx)?;
+        python_dict.set_item(column.name().to_object(py), python_type)?;
+    }
+    Ok(python_dict)
+}
 
 #[pyclass(name = "QueryResult")]
 #[allow(clippy::module_name_repetitions)]
@@ -33,15 +52,31 @@ impl PSQLDriverPyQueryResult {
     pub fn result(&self, py: Python<'_>) -> RustPSQLDriverPyResult<Py<PyAny>> {
         let mut result: Vec<&PyDict> = vec![];
         for row in &self.inner {
-            let python_dict = PyDict::new(py);
-            for (column_idx, column) in row.columns().iter().enumerate() {
-                let python_type = postgres_to_py(py, row, column, column_idx)?;
-                python_dict.set_item(column.name().to_object(py), python_type)?;
-            }
-            result.push(python_dict);
+            result.push(row_to_dict(py, row)?);
         }
         Ok(result.to_object(py))
     }
+
+    /// Convert result from database to any class passed from Python.
+    ///
+    /// # Errors
+    ///
+    /// May return Err Result if can not convert
+    /// postgres type to python or create new Python class.
+    pub fn as_class<'a>(
+        &'a self,
+        py: Python<'a>,
+        as_class: &'a PyAny,
+    ) -> RustPSQLDriverPyResult<Py<PyAny>> {
+        let mut res: Vec<&PyAny> = vec![];
+        for row in &self.inner {
+            let pydict: &PyDict = row_to_dict(py, row)?;
+            let convert_class_inst = as_class.call((), Some(pydict))?;
+            res.push(convert_class_inst);
+        }
+
+        Ok(res.to_object(py))
+    }
 }
 
 #[pyclass(name = "SingleQueryResult")]
@@ -68,16 +103,35 @@ impl PSQLDriverSinglePyQueryResult {
     /// # Errors
     ///
     /// May return Err Result if can not convert
-    /// postgres type to python or set new key-value pair
-    /// in python dict.
+    /// 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>> {
-        let python_dict = PyDict::new(py);
         if let Some(row) = self.inner.first() {
-            for (column_idx, column) in row.columns().iter().enumerate() {
-                let python_type = postgres_to_py(py, row, column, column_idx)?;
-                python_dict.set_item(column.name().to_object(py), python_type)?;
-            }
+            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(),
+        ))
+    }
+
+    /// Convert result from database to any class passed from Python.
+    ///
+    /// # Errors
+    ///
+    /// May return Err Result if can not convert
+    /// postgres type to python, can not create new Python class
+    /// or there are no results.
+    pub fn as_class<'a>(
+        &'a self,
+        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))?);
         }
-        Ok(python_dict.to_object(py))
+        Err(RustPSQLDriverError::RustToPyValueConversionError(
+            "There are not results from the query, can't convert first row.".into(),
+        ))
     }
 }

tox.ini

@@ -21,6 +21,7 @@ deps =
     pytest>=7,<8
     anyio>=3,<4
     maturin>=1,<2
+    pydantic>=2
 allowlist_externals = maturin
 commands_pre =
     maturin develop