Merge pull request #54 from qaspen-python/feature/add_row_factory

Added row_factory method for query result

Author: chandr-andr

python/psqlpy/_internal/__init__.pyi

@@ -8,6 +8,9 @@ from typing_extensions import Self
 _CustomClass = TypeVar(
     "_CustomClass",
 )
+_RowFactoryRV = TypeVar(
+    "_RowFactoryRV",
+)
 
 class QueryResult:
     """Result."""
@@ -16,7 +19,11 @@ class QueryResult:
         self: Self,
         custom_decoders: dict[str, Callable[[bytes], Any]] | None = None,
     ) -> list[dict[Any, Any]]:
-        """Return result from database as a list of dicts."""
+        """Return result from database as a list of dicts.
+
+        `custom_decoders` must be used when you use
+        PostgreSQL Type which isn't supported, read more in our docs.
+        """
     def as_class(
         self: Self,
         as_class: Callable[..., _CustomClass],
@@ -52,12 +59,39 @@ class QueryResult:
             )
         ```
         """
+    def row_factory(
+        self,
+        row_factory: Callable[[dict[str, Any]], _RowFactoryRV],
+        custom_decoders: dict[str, Callable[[bytes], Any]] | None = None,
+    ) -> list[_RowFactoryRV]:
+        """Use custom function to convert results from database.
+
+        `custom_decoders` must be used when you use
+        PostgreSQL Type isn't supported, read more in the docs.
+
+        Argument order: firstly we apply `custom_decoders` (if specified),
+        then we apply `row_factory`.
+
+        ### Parameters:
+        - `row_factory`: function which takes `dict[str, Any]` as an argument.
+        - `custom_decoders`: functions for custom decoding.
+
+        ### Returns:
+        List of type that return passed `row_factory`.
+        """
 
 class SingleQueryResult:
     """Single result."""
 
-    def result(self: Self) -> dict[Any, Any]:
-        """Return result from database as a dict."""
+    def result(
+        self: Self,
+        custom_decoders: dict[str, Callable[[bytes], Any]] | None = None,
+    ) -> dict[Any, Any]:
+        """Return result from database as a dict.
+
+        `custom_decoders` must be used when you use
+        PostgreSQL Type which isn't supported, read more in our docs.
+        """
     def as_class(
         self: Self,
         as_class: Callable[..., _CustomClass],
@@ -96,6 +130,26 @@ class SingleQueryResult:
             )
         ```
         """
+    def row_factory(
+        self,
+        row_factory: Callable[[dict[str, Any]], _RowFactoryRV],
+        custom_decoders: dict[str, Callable[[bytes], Any]] | None = None,
+    ) -> _RowFactoryRV:
+        """Use custom function to convert results from database.
+
+        `custom_decoders` must be used when you use
+        PostgreSQL Type isn't supported, read more in our docs.
+
+        Argument order: firstly we apply `custom_decoders` (if specified),
+        then we apply `row_factory`.
+
+        ### Parameters:
+        - `row_factory`: function which takes `list[dict[str, Any]]` as an argument.
+        - `custom_decoders`: functions for custom decoding.
+
+        ### Returns:
+        Type that return passed function.
+        """
 
 class IsolationLevel(Enum):
     """Class for Isolation Level for transactions."""

python/tests/test_value_converter.py

@@ -571,3 +571,44 @@ def point_encoder(point_bytes: bytes) -> str:
     )
 
     assert result[0]["geo_point"] == "Just An Example"
+
+
+async def test_row_factory_query_result(
+    psql_pool: ConnectionPool,
+    table_name: str,
+    number_database_records: int,
+) -> None:
+    select_result = await psql_pool.execute(
+        f"SELECT * FROM {table_name}",
+    )
+
+    def row_factory(db_result: Dict[str, Any]) -> List[str]:
+        return list(db_result.keys())
+
+    as_row_factory = select_result.row_factory(
+        row_factory=row_factory,
+    )
+    assert len(as_row_factory) == number_database_records
+
+    assert isinstance(as_row_factory[0], list)
+
+
+async def test_row_factory_single_query_result(
+    psql_pool: ConnectionPool,
+    table_name: str,
+) -> None:
+    connection = await psql_pool.connection()
+    select_result = await connection.fetch_row(
+        f"SELECT * FROM {table_name} LIMIT 1",
+    )
+
+    def row_factory(db_result: Dict[str, Any]) -> List[str]:
+        return list(db_result.keys())
+
+    as_row_factory = select_result.row_factory(
+        row_factory=row_factory,
+    )
+    expected_number_of_elements_in_result = 2
+    assert len(as_row_factory) == expected_number_of_elements_in_result
+
+    assert isinstance(as_row_factory, list)

src/query_result.rs

@@ -89,6 +89,28 @@ impl PSQLDriverPyQueryResult {
 
         Ok(res.to_object(py))
     }
+
+    /// Convert result from database with function passed from Python.
+    ///
+    /// # Errors
+    ///
+    /// May return Err Result if can not convert
+    /// postgres type with custom function.
+    #[allow(clippy::needless_pass_by_value)]
+    pub fn row_factory<'a>(
+        &'a self,
+        py: Python<'a>,
+        row_factory: Py<PyAny>,
+        custom_decoders: Option<Py<PyDict>>,
+    ) -> RustPSQLDriverPyResult<Py<PyAny>> {
+        let mut res: Vec<Py<PyAny>> = vec![];
+        for row in &self.inner {
+            let pydict: pyo3::Bound<'_, PyDict> = row_to_dict(py, row, &custom_decoders)?;
+            let row_factory_class = row_factory.call_bound(py, (pydict,), None)?;
+            res.push(row_factory_class);
+        }
+        Ok(res.to_object(py))
+    }
 }
 
 #[pyclass(name = "SingleQueryResult")]
@@ -121,8 +143,13 @@ impl PSQLDriverSinglePyQueryResult {
     /// May return Err Result if can not convert
     /// 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>> {
-        Ok(row_to_dict(py, &self.inner, &None)?.to_object(py))
+    #[allow(clippy::needless_pass_by_value)]
+    pub fn result(
+        &self,
+        py: Python<'_>,
+        custom_decoders: Option<Py<PyDict>>,
+    ) -> RustPSQLDriverPyResult<Py<PyAny>> {
+        Ok(row_to_dict(py, &self.inner, &custom_decoders)?.to_object(py))
     }
 
     /// Convert result from database to any class passed from Python.
@@ -141,4 +168,21 @@ impl PSQLDriverSinglePyQueryResult {
         let pydict: pyo3::Bound<'_, PyDict> = row_to_dict(py, &self.inner, &None)?;
         Ok(as_class.call_bound(py, (), Some(&pydict))?)
     }
+
+    /// Convert result from database with function passed from Python.
+    ///
+    /// # Errors
+    ///
+    /// May return Err Result if can not convert
+    /// postgres type with custom function
+    #[allow(clippy::needless_pass_by_value)]
+    pub fn row_factory<'a>(
+        &'a self,
+        py: Python<'a>,
+        row_factory: Py<PyAny>,
+        custom_decoders: Option<Py<PyDict>>,
+    ) -> RustPSQLDriverPyResult<Py<PyAny>> {
+        let pydict = row_to_dict(py, &self.inner, &custom_decoders)?.to_object(py);
+        Ok(row_factory.call_bound(py, (pydict,), None)?)
+    }
 }