Skip to content

Added new as_tuple parameter to QueryResult #147

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 16 additions & 2 deletions docs/components/results.md
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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
Expand Down Expand Up @@ -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

Expand All @@ -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
Expand Down
51 changes: 48 additions & 3 deletions python/psqlpy/_internal/__init__.pyi
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
import types
import typing
from enum import Enum
from io import BytesIO
from ipaddress import IPv4Address, IPv6Address
Expand All @@ -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.
Expand Down Expand Up @@ -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.
Expand Down
76 changes: 76 additions & 0 deletions python/tests/test_query_results.py
Original file line number Diff line number Diff line change
@@ -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"
29 changes: 29 additions & 0 deletions python/tests/test_value_converter.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
60 changes: 53 additions & 7 deletions src/query_result.rs
Original file line number Diff line number Diff line change
@@ -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};
Expand All @@ -15,7 +20,7 @@ fn row_to_dict<'a>(
py: Python<'a>,
postgres_row: &'a Row,
custom_decoders: &Option<Py<PyDict>>,
) -> PSQLPyResult<pyo3::Bound<'a, PyDict>> {
) -> PSQLPyResult<Bound<'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, custom_decoders)?;
Expand All @@ -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<Py<PyDict>>,
) -> PSQLPyResult<Bound<'a, PyTuple>> {
let mut rows: Vec<Bound<'_, PyTuple>> = 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 {
Expand Down Expand Up @@ -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<Py<PyDict>>,
as_tuple: Option<bool>,
) -> PSQLPyResult<Py<PyAny>> {
let mut result: Vec<pyo3::Bound<'_, PyDict>> = vec![];
let as_tuple = as_tuple.unwrap_or(false);

if as_tuple {
let mut tuple_rows: Vec<Bound<'_, PyTuple>> = 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<Bound<'_, PyDict>> = 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.
Expand Down Expand Up @@ -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<Py<PyDict>>,
as_tuple: Option<bool>,
) -> PSQLPyResult<Py<PyAny>> {
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)?)
}

Expand Down
Loading