Skip to content

Commit ee26552

Browse files
committed
Added new as_tuple parameter to QueryResult
1 parent 91aa22e commit ee26552

File tree

5 files changed

+222
-12
lines changed

5 files changed

+222
-12
lines changed

docs/components/results.md

Lines changed: 16 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ Currently there are two results:
1515

1616
#### Parameters
1717
- `custom_decoders`: custom decoders for unsupported types. [Read more](/usage/types/advanced_type_usage.md)
18+
- `as_tuple`: return result as a tuple instead of dict.
1819

1920
Get the result as a list of dicts
2021

@@ -27,7 +28,13 @@ async def main() -> None:
2728
[],
2829
)
2930

30-
result: List[Dict[str, Any]] = query_result.result()
31+
# Result as dict
32+
list_dict_result: List[Dict[str, Any]] = query_result.result()
33+
34+
# Result as tuple
35+
list_tuple_result: List[Tuple[Tuple[str, typing.Any], ...]] = query_result.result(
36+
as_tuple=True,
37+
)
3138
```
3239

3340
### As class
@@ -72,6 +79,7 @@ async def main() -> None:
7279

7380
#### Parameters
7481
- `custom_decoders`: custom decoders for unsupported types. [Read more](/usage/types/advanced_type_usage.md)
82+
- `as_tuple`: return result as a tuple instead of dict.
7583

7684
Get the result as a dict
7785

@@ -84,7 +92,13 @@ async def main() -> None:
8492
[100],
8593
)
8694

87-
result: Dict[str, Any] = query_result.result()
95+
# Result as dict
96+
dict_result: Dict[str, Any] = query_result.result()
97+
98+
# Result as tuple
99+
tuple_result: Tuple[Tuple[str, typing.Any], ...] = query_result.result(
100+
as_tuple=True,
101+
)
88102
```
89103

90104
### As class

python/psqlpy/_internal/__init__.pyi

Lines changed: 48 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
import types
2+
import typing
23
from enum import Enum
34
from io import BytesIO
45
from ipaddress import IPv4Address, IPv6Address
@@ -18,11 +19,33 @@ ParamsT: TypeAlias = Sequence[Any] | Mapping[str, Any] | None
1819
class QueryResult:
1920
"""Result."""
2021

22+
@typing.overload
2123
def result(
2224
self: Self,
25+
as_tuple: typing.Literal[None] = None,
2326
custom_decoders: dict[str, Callable[[bytes], Any]] | None = None,
24-
) -> list[dict[Any, Any]]:
25-
"""Return result from database as a list of dicts.
27+
) -> list[dict[str, Any]]: ...
28+
@typing.overload
29+
def result(
30+
self: Self,
31+
as_tuple: typing.Literal[False],
32+
custom_decoders: dict[str, Callable[[bytes], Any]] | None = None,
33+
) -> list[dict[str, Any]]: ...
34+
@typing.overload
35+
def result(
36+
self: Self,
37+
as_tuple: typing.Literal[True],
38+
custom_decoders: dict[str, Callable[[bytes], Any]] | None = None,
39+
) -> list[tuple[tuple[str, typing.Any], ...]]: ...
40+
@typing.overload
41+
def result(
42+
self: Self,
43+
custom_decoders: dict[str, Callable[[bytes], Any]] | None = None,
44+
as_tuple: bool | None = None,
45+
) -> list[dict[str, Any]]:
46+
"""Return result from database.
47+
48+
By default it returns result as a list of dicts.
2649
2750
`custom_decoders` must be used when you use
2851
PostgreSQL Type which isn't supported, read more in our docs.
@@ -84,11 +107,33 @@ class QueryResult:
84107
class SingleQueryResult:
85108
"""Single result."""
86109

110+
@typing.overload
87111
def result(
88112
self: Self,
113+
as_tuple: typing.Literal[None] = None,
89114
custom_decoders: dict[str, Callable[[bytes], Any]] | None = None,
115+
) -> dict[str, Any]: ...
116+
@typing.overload
117+
def result(
118+
self: Self,
119+
as_tuple: typing.Literal[False],
120+
custom_decoders: dict[str, Callable[[bytes], Any]] | None = None,
121+
) -> dict[str, Any]: ...
122+
@typing.overload
123+
def result(
124+
self: Self,
125+
as_tuple: typing.Literal[True],
126+
custom_decoders: dict[str, Callable[[bytes], Any]] | None = None,
127+
) -> tuple[tuple[str, typing.Any]]: ...
128+
@typing.overload
129+
def result(
130+
self: Self,
131+
custom_decoders: dict[str, Callable[[bytes], Any]] | None = None,
132+
as_tuple: bool | None = None,
90133
) -> dict[Any, Any]:
91-
"""Return result from database as a dict.
134+
"""Return result from database.
135+
136+
By default it returns result as a dict.
92137
93138
`custom_decoders` must be used when you use
94139
PostgreSQL Type which isn't supported, read more in our docs.

python/tests/test_query_results.py

Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,76 @@
1+
from __future__ import annotations
2+
3+
import pytest
4+
from psqlpy import ConnectionPool, QueryResult, SingleQueryResult
5+
6+
pytestmark = pytest.mark.anyio
7+
8+
9+
async def test_result_as_dict(
10+
psql_pool: ConnectionPool,
11+
table_name: str,
12+
) -> None:
13+
"""Test that single connection can execute queries."""
14+
connection = await psql_pool.connection()
15+
16+
conn_result = await connection.execute(
17+
querystring=f"SELECT * FROM {table_name}",
18+
)
19+
result_list_dicts = conn_result.result()
20+
single_dict_row = result_list_dicts[0]
21+
22+
assert isinstance(conn_result, QueryResult)
23+
assert isinstance(single_dict_row, dict)
24+
assert single_dict_row.get("id")
25+
26+
27+
async def test_result_as_tuple(
28+
psql_pool: ConnectionPool,
29+
table_name: str,
30+
) -> None:
31+
"""Test that single connection can execute queries."""
32+
connection = await psql_pool.connection()
33+
34+
conn_result = await connection.execute(
35+
querystring=f"SELECT * FROM {table_name}",
36+
)
37+
result_tuple = conn_result.result(as_tuple=True)
38+
single_tuple_row = result_tuple[0]
39+
40+
assert isinstance(conn_result, QueryResult)
41+
assert isinstance(single_tuple_row, tuple)
42+
assert single_tuple_row[0][0] == "id"
43+
44+
45+
async def test_single_result_as_dict(
46+
psql_pool: ConnectionPool,
47+
table_name: str,
48+
) -> None:
49+
"""Test that single connection can execute queries."""
50+
connection = await psql_pool.connection()
51+
52+
conn_result = await connection.fetch_row(
53+
querystring=f"SELECT * FROM {table_name} LIMIT 1",
54+
)
55+
result_dict = conn_result.result()
56+
57+
assert isinstance(conn_result, SingleQueryResult)
58+
assert isinstance(result_dict, dict)
59+
assert result_dict.get("id")
60+
61+
62+
async def test_single_result_as_tuple(
63+
psql_pool: ConnectionPool,
64+
table_name: str,
65+
) -> None:
66+
"""Test that single connection can execute queries."""
67+
connection = await psql_pool.connection()
68+
69+
conn_result = await connection.fetch_row(
70+
querystring=f"SELECT * FROM {table_name} LIMIT 1",
71+
)
72+
result_tuple = conn_result.result(as_tuple=True)
73+
74+
assert isinstance(conn_result, SingleQueryResult)
75+
assert isinstance(result_tuple, tuple)
76+
assert result_tuple[0][0] == "id"

python/tests/test_value_converter.py

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -646,6 +646,35 @@ def point_encoder(point_bytes: bytes) -> str: # noqa: ARG001
646646
assert result[0]["geo_point"] == "Just An Example"
647647

648648

649+
async def test_custom_decoder_as_tuple_result(
650+
psql_pool: ConnectionPool,
651+
) -> None:
652+
def point_encoder(point_bytes: bytes) -> str: # noqa: ARG001
653+
return "Just An Example"
654+
655+
async with psql_pool.acquire() as conn:
656+
await conn.execute("DROP TABLE IF EXISTS for_test")
657+
await conn.execute(
658+
"CREATE TABLE for_test (geo_point POINT)",
659+
)
660+
661+
await conn.execute(
662+
"INSERT INTO for_test VALUES ('(1, 1)')",
663+
)
664+
665+
qs_result = await conn.execute(
666+
"SELECT * FROM for_test",
667+
)
668+
result = qs_result.result(
669+
custom_decoders={
670+
"geo_point": point_encoder,
671+
},
672+
as_tuple=True,
673+
)
674+
675+
assert result[0][0][1] == "Just An Example"
676+
677+
649678
async def test_row_factory_query_result(
650679
psql_pool: ConnectionPool,
651680
table_name: str,

src/query_result.rs

Lines changed: 53 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,9 @@
1-
use pyo3::{prelude::*, pyclass, pymethods, types::PyDict, IntoPyObjectExt, Py, PyAny, Python};
1+
use pyo3::{
2+
prelude::*,
3+
pyclass, pymethods,
4+
types::{PyDict, PyTuple},
5+
IntoPyObjectExt, Py, PyAny, Python,
6+
};
27
use tokio_postgres::Row;
38

49
use crate::{exceptions::rust_errors::PSQLPyResult, value_converter::to_python::postgres_to_py};
@@ -15,7 +20,7 @@ fn row_to_dict<'a>(
1520
py: Python<'a>,
1621
postgres_row: &'a Row,
1722
custom_decoders: &Option<Py<PyDict>>,
18-
) -> PSQLPyResult<pyo3::Bound<'a, PyDict>> {
23+
) -> PSQLPyResult<Bound<'a, PyDict>> {
1924
let python_dict = PyDict::new(py);
2025
for (column_idx, column) in postgres_row.columns().iter().enumerate() {
2126
let python_type = postgres_to_py(py, postgres_row, column, column_idx, custom_decoders)?;
@@ -24,6 +29,29 @@ fn row_to_dict<'a>(
2429
Ok(python_dict)
2530
}
2631

32+
/// Convert postgres `Row` into Python Tuple.
33+
///
34+
/// # Errors
35+
///
36+
/// May return Err Result if can not convert
37+
/// postgres type to python or set new key-value pair
38+
/// in python dict.
39+
#[allow(clippy::ref_option)]
40+
fn row_to_tuple<'a>(
41+
py: Python<'a>,
42+
postgres_row: &'a Row,
43+
custom_decoders: &Option<Py<PyDict>>,
44+
) -> PSQLPyResult<Bound<'a, PyTuple>> {
45+
let mut rows: Vec<Bound<'_, PyTuple>> = vec![];
46+
47+
for (column_idx, column) in postgres_row.columns().iter().enumerate() {
48+
let python_type = postgres_to_py(py, postgres_row, column, column_idx, custom_decoders)?;
49+
let timed_tuple = PyTuple::new(py, vec![column.name().into_py_any(py)?, python_type])?;
50+
rows.push(timed_tuple);
51+
}
52+
Ok(PyTuple::new(py, rows)?)
53+
}
54+
2755
#[pyclass(name = "QueryResult")]
2856
#[allow(clippy::module_name_repetitions)]
2957
pub struct PSQLDriverPyQueryResult {
@@ -56,18 +84,29 @@ impl PSQLDriverPyQueryResult {
5684
/// May return Err Result if can not convert
5785
/// postgres type to python or set new key-value pair
5886
/// in python dict.
59-
#[pyo3(signature = (custom_decoders=None))]
87+
#[pyo3(signature = (custom_decoders=None, as_tuple=None))]
6088
#[allow(clippy::needless_pass_by_value)]
6189
pub fn result(
6290
&self,
6391
py: Python<'_>,
6492
custom_decoders: Option<Py<PyDict>>,
93+
as_tuple: Option<bool>,
6594
) -> PSQLPyResult<Py<PyAny>> {
66-
let mut result: Vec<pyo3::Bound<'_, PyDict>> = vec![];
95+
let as_tuple = as_tuple.unwrap_or(false);
96+
97+
if as_tuple {
98+
let mut tuple_rows: Vec<Bound<'_, PyTuple>> = vec![];
99+
for row in &self.inner {
100+
tuple_rows.push(row_to_tuple(py, row, &custom_decoders)?);
101+
}
102+
return Ok(tuple_rows.into_py_any(py)?);
103+
}
104+
105+
let mut dict_rows: Vec<Bound<'_, PyDict>> = vec![];
67106
for row in &self.inner {
68-
result.push(row_to_dict(py, row, &custom_decoders)?);
107+
dict_rows.push(row_to_dict(py, row, &custom_decoders)?);
69108
}
70-
Ok(result.into_py_any(py)?)
109+
Ok(dict_rows.into_py_any(py)?)
71110
}
72111

73112
/// Convert result from database to any class passed from Python.
@@ -143,12 +182,19 @@ impl PSQLDriverSinglePyQueryResult {
143182
/// postgres type to python, can not set new key-value pair
144183
/// in python dict or there are no result.
145184
#[allow(clippy::needless_pass_by_value)]
146-
#[pyo3(signature = (custom_decoders=None))]
185+
#[pyo3(signature = (custom_decoders=None, as_tuple=None))]
147186
pub fn result(
148187
&self,
149188
py: Python<'_>,
150189
custom_decoders: Option<Py<PyDict>>,
190+
as_tuple: Option<bool>,
151191
) -> PSQLPyResult<Py<PyAny>> {
192+
let as_tuple = as_tuple.unwrap_or(false);
193+
194+
if as_tuple {
195+
return Ok(row_to_tuple(py, &self.inner, &custom_decoders)?.into_py_any(py)?);
196+
}
197+
152198
Ok(row_to_dict(py, &self.inner, &custom_decoders)?.into_py_any(py)?)
153199
}
154200

0 commit comments

Comments
 (0)