diff --git a/docs/components/results.md b/docs/components/results.md index 15bc3690..7cea19a5 100644 --- a/docs/components/results.md +++ b/docs/components/results.md @@ -15,6 +15,7 @@ Currently there are two results: #### Parameters - `custom_decoders`: custom decoders for unsupported types. [Read more](/usage/types/advanced_type_usage.md) +- `as_tuple`: return result as a tuple instead of dict. Get the result as a list of dicts @@ -27,7 +28,13 @@ async def main() -> None: [], ) - result: List[Dict[str, Any]] = query_result.result() + # Result as dict + list_dict_result: List[Dict[str, Any]] = query_result.result() + + # Result as tuple + list_tuple_result: List[Tuple[Tuple[str, typing.Any], ...]] = query_result.result( + as_tuple=True, + ) ``` ### As class @@ -72,6 +79,7 @@ async def main() -> None: #### Parameters - `custom_decoders`: custom decoders for unsupported types. [Read more](/usage/types/advanced_type_usage.md) +- `as_tuple`: return result as a tuple instead of dict. Get the result as a dict @@ -84,7 +92,13 @@ async def main() -> None: [100], ) - result: Dict[str, Any] = query_result.result() + # Result as dict + dict_result: Dict[str, Any] = query_result.result() + + # Result as tuple + tuple_result: Tuple[Tuple[str, typing.Any], ...] = query_result.result( + as_tuple=True, + ) ``` ### As class diff --git a/python/psqlpy/_internal/__init__.pyi b/python/psqlpy/_internal/__init__.pyi index ddb74de1..2665678b 100644 --- a/python/psqlpy/_internal/__init__.pyi +++ b/python/psqlpy/_internal/__init__.pyi @@ -1,4 +1,5 @@ import types +import typing from enum import Enum from io import BytesIO from ipaddress import IPv4Address, IPv6Address @@ -18,11 +19,33 @@ ParamsT: TypeAlias = Sequence[Any] | Mapping[str, Any] | None class QueryResult: """Result.""" + @typing.overload def result( self: Self, + as_tuple: typing.Literal[None] = None, custom_decoders: dict[str, Callable[[bytes], Any]] | None = None, - ) -> list[dict[Any, Any]]: - """Return result from database as a list of dicts. + ) -> list[dict[str, Any]]: ... + @typing.overload + def result( + self: Self, + as_tuple: typing.Literal[False], + custom_decoders: dict[str, Callable[[bytes], Any]] | None = None, + ) -> list[dict[str, Any]]: ... + @typing.overload + def result( + self: Self, + as_tuple: typing.Literal[True], + custom_decoders: dict[str, Callable[[bytes], Any]] | None = None, + ) -> list[tuple[tuple[str, typing.Any], ...]]: ... + @typing.overload + def result( + self: Self, + custom_decoders: dict[str, Callable[[bytes], Any]] | None = None, + as_tuple: bool | None = None, + ) -> list[dict[str, Any]]: + """Return result from database. + + By default it returns result as a list of dicts. `custom_decoders` must be used when you use PostgreSQL Type which isn't supported, read more in our docs. @@ -84,11 +107,33 @@ class QueryResult: class SingleQueryResult: """Single result.""" + @typing.overload def result( self: Self, + as_tuple: typing.Literal[None] = None, custom_decoders: dict[str, Callable[[bytes], Any]] | None = None, + ) -> dict[str, Any]: ... + @typing.overload + def result( + self: Self, + as_tuple: typing.Literal[False], + custom_decoders: dict[str, Callable[[bytes], Any]] | None = None, + ) -> dict[str, Any]: ... + @typing.overload + def result( + self: Self, + as_tuple: typing.Literal[True], + custom_decoders: dict[str, Callable[[bytes], Any]] | None = None, + ) -> tuple[tuple[str, typing.Any]]: ... + @typing.overload + def result( + self: Self, + custom_decoders: dict[str, Callable[[bytes], Any]] | None = None, + as_tuple: bool | None = None, ) -> dict[Any, Any]: - """Return result from database as a dict. + """Return result from database. + + By default it returns result as a dict. `custom_decoders` must be used when you use PostgreSQL Type which isn't supported, read more in our docs. diff --git a/python/tests/test_query_results.py b/python/tests/test_query_results.py new file mode 100644 index 00000000..95de93c7 --- /dev/null +++ b/python/tests/test_query_results.py @@ -0,0 +1,76 @@ +from __future__ import annotations + +import pytest +from psqlpy import ConnectionPool, QueryResult, SingleQueryResult + +pytestmark = pytest.mark.anyio + + +async def test_result_as_dict( + psql_pool: ConnectionPool, + table_name: str, +) -> None: + """Test that single connection can execute queries.""" + connection = await psql_pool.connection() + + conn_result = await connection.execute( + querystring=f"SELECT * FROM {table_name}", + ) + result_list_dicts = conn_result.result() + single_dict_row = result_list_dicts[0] + + assert isinstance(conn_result, QueryResult) + assert isinstance(single_dict_row, dict) + assert single_dict_row.get("id") + + +async def test_result_as_tuple( + psql_pool: ConnectionPool, + table_name: str, +) -> None: + """Test that single connection can execute queries.""" + connection = await psql_pool.connection() + + conn_result = await connection.execute( + querystring=f"SELECT * FROM {table_name}", + ) + result_tuple = conn_result.result(as_tuple=True) + single_tuple_row = result_tuple[0] + + assert isinstance(conn_result, QueryResult) + assert isinstance(single_tuple_row, tuple) + assert single_tuple_row[0][0] == "id" + + +async def test_single_result_as_dict( + psql_pool: ConnectionPool, + table_name: str, +) -> None: + """Test that single connection can execute queries.""" + connection = await psql_pool.connection() + + conn_result = await connection.fetch_row( + querystring=f"SELECT * FROM {table_name} LIMIT 1", + ) + result_dict = conn_result.result() + + assert isinstance(conn_result, SingleQueryResult) + assert isinstance(result_dict, dict) + assert result_dict.get("id") + + +async def test_single_result_as_tuple( + psql_pool: ConnectionPool, + table_name: str, +) -> None: + """Test that single connection can execute queries.""" + connection = await psql_pool.connection() + + conn_result = await connection.fetch_row( + querystring=f"SELECT * FROM {table_name} LIMIT 1", + ) + result_tuple = conn_result.result(as_tuple=True) + + assert isinstance(conn_result, SingleQueryResult) + assert isinstance(result_tuple, tuple) + assert result_tuple[0][0] == "id" diff --git a/python/tests/test_value_converter.py b/python/tests/test_value_converter.py index ce2f05ed..07833848 100644 --- a/python/tests/test_value_converter.py +++ b/python/tests/test_value_converter.py @@ -646,6 +646,35 @@ def point_encoder(point_bytes: bytes) -> str: # noqa: ARG001 assert result[0]["geo_point"] == "Just An Example" +async def test_custom_decoder_as_tuple_result( + psql_pool: ConnectionPool, +) -> None: + def point_encoder(point_bytes: bytes) -> str: # noqa: ARG001 + return "Just An Example" + + async with psql_pool.acquire() as conn: + await conn.execute("DROP TABLE IF EXISTS for_test") + await conn.execute( + "CREATE TABLE for_test (geo_point POINT)", + ) + + await conn.execute( + "INSERT INTO for_test VALUES ('(1, 1)')", + ) + + qs_result = await conn.execute( + "SELECT * FROM for_test", + ) + result = qs_result.result( + custom_decoders={ + "geo_point": point_encoder, + }, + as_tuple=True, + ) + + assert result[0][0][1] == "Just An Example" + + async def test_row_factory_query_result( psql_pool: ConnectionPool, table_name: str, diff --git a/src/query_result.rs b/src/query_result.rs index b17acad9..a5af132d 100644 --- a/src/query_result.rs +++ b/src/query_result.rs @@ -1,4 +1,9 @@ -use pyo3::{prelude::*, pyclass, pymethods, types::PyDict, IntoPyObjectExt, Py, PyAny, Python}; +use pyo3::{ + prelude::*, + pyclass, pymethods, + types::{PyDict, PyTuple}, + IntoPyObjectExt, Py, PyAny, Python, +}; use tokio_postgres::Row; use crate::{exceptions::rust_errors::PSQLPyResult, value_converter::to_python::postgres_to_py}; @@ -15,7 +20,7 @@ fn row_to_dict<'a>( py: Python<'a>, postgres_row: &'a Row, custom_decoders: &Option>, -) -> PSQLPyResult> { +) -> PSQLPyResult> { 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, custom_decoders)?; @@ -24,6 +29,29 @@ fn row_to_dict<'a>( Ok(python_dict) } +/// Convert postgres `Row` into Python Tuple. +/// +/// # Errors +/// +/// May return Err Result if can not convert +/// postgres type to python or set new key-value pair +/// in python dict. +#[allow(clippy::ref_option)] +fn row_to_tuple<'a>( + py: Python<'a>, + postgres_row: &'a Row, + custom_decoders: &Option>, +) -> PSQLPyResult> { + let mut rows: Vec> = vec![]; + + for (column_idx, column) in postgres_row.columns().iter().enumerate() { + let python_type = postgres_to_py(py, postgres_row, column, column_idx, custom_decoders)?; + let timed_tuple = PyTuple::new(py, vec![column.name().into_py_any(py)?, python_type])?; + rows.push(timed_tuple); + } + Ok(PyTuple::new(py, rows)?) +} + #[pyclass(name = "QueryResult")] #[allow(clippy::module_name_repetitions)] pub struct PSQLDriverPyQueryResult { @@ -56,18 +84,29 @@ impl PSQLDriverPyQueryResult { /// May return Err Result if can not convert /// postgres type to python or set new key-value pair /// in python dict. - #[pyo3(signature = (custom_decoders=None))] + #[pyo3(signature = (custom_decoders=None, as_tuple=None))] #[allow(clippy::needless_pass_by_value)] pub fn result( &self, py: Python<'_>, custom_decoders: Option>, + as_tuple: Option, ) -> PSQLPyResult> { - let mut result: Vec> = vec![]; + let as_tuple = as_tuple.unwrap_or(false); + + if as_tuple { + let mut tuple_rows: Vec> = vec![]; + for row in &self.inner { + tuple_rows.push(row_to_tuple(py, row, &custom_decoders)?); + } + return Ok(tuple_rows.into_py_any(py)?); + } + + let mut dict_rows: Vec> = vec![]; for row in &self.inner { - result.push(row_to_dict(py, row, &custom_decoders)?); + dict_rows.push(row_to_dict(py, row, &custom_decoders)?); } - Ok(result.into_py_any(py)?) + Ok(dict_rows.into_py_any(py)?) } /// Convert result from database to any class passed from Python. @@ -143,12 +182,19 @@ impl PSQLDriverSinglePyQueryResult { /// postgres type to python, can not set new key-value pair /// in python dict or there are no result. #[allow(clippy::needless_pass_by_value)] - #[pyo3(signature = (custom_decoders=None))] + #[pyo3(signature = (custom_decoders=None, as_tuple=None))] pub fn result( &self, py: Python<'_>, custom_decoders: Option>, + as_tuple: Option, ) -> PSQLPyResult> { + let as_tuple = as_tuple.unwrap_or(false); + + if as_tuple { + return Ok(row_to_tuple(py, &self.inner, &custom_decoders)?.into_py_any(py)?); + } + Ok(row_to_dict(py, &self.inner, &custom_decoders)?.into_py_any(py)?) }