Skip to content

Commit 510702c

Browse files
committed
First variant for decoders
Signed-off-by: chandr-andr (Kiselev Aleksandr) <chandr@chandr.net>
1 parent 15dcaf9 commit 510702c

File tree

5 files changed

+33
-14
lines changed

5 files changed

+33
-14
lines changed

python/psqlpy/_internal/__init__.pyi

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,10 @@ _CustomClass = TypeVar(
1111
class QueryResult:
1212
"""Result."""
1313

14-
def result(self: Self) -> list[dict[Any, Any]]:
14+
def result(
15+
self: Self,
16+
custom_decoders: dict[str, Callable[[list[int]], Any]] | None = None,
17+
) -> list[dict[Any, Any]]:
1518
"""Return result from database as a list of dicts."""
1619
def as_class(
1720
self: Self,

src/driver/connection.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -257,7 +257,7 @@ impl Connection {
257257
};
258258

259259
Python::with_gil(|gil| match result.columns().first() {
260-
Some(first_column) => postgres_to_py(gil, &result, first_column, 0),
260+
Some(first_column) => postgres_to_py(gil, &result, first_column, 0, &None),
261261
None => Ok(gil.None()),
262262
})
263263
}

src/driver/transaction.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -331,7 +331,7 @@ impl Transaction {
331331
};
332332

333333
Python::with_gil(|gil| match result.columns().first() {
334-
Some(first_column) => postgres_to_py(gil, &result, first_column, 0),
334+
Some(first_column) => postgres_to_py(gil, &result, first_column, 0, &None),
335335
None => Ok(gil.None()),
336336
})
337337
}

src/query_result.rs

Lines changed: 11 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -13,10 +13,11 @@ use crate::{exceptions::rust_errors::RustPSQLDriverPyResult, value_converter::po
1313
fn row_to_dict<'a>(
1414
py: Python<'a>,
1515
postgres_row: &'a Row,
16+
custom_decoders: &Option<Py<PyDict>>,
1617
) -> RustPSQLDriverPyResult<pyo3::Bound<'a, PyDict>> {
1718
let python_dict = PyDict::new_bound(py);
1819
for (column_idx, column) in postgres_row.columns().iter().enumerate() {
19-
let python_type = postgres_to_py(py, postgres_row, column, column_idx)?;
20+
let python_type = postgres_to_py(py, postgres_row, column, column_idx, custom_decoders)?;
2021
python_dict.set_item(column.name().to_object(py), python_type)?;
2122
}
2223
Ok(python_dict)
@@ -55,10 +56,14 @@ impl PSQLDriverPyQueryResult {
5556
/// postgres type to python or set new key-value pair
5657
/// in python dict.
5758
#[allow(clippy::needless_pass_by_value)]
58-
pub fn result(&self, py: Python<'_>) -> RustPSQLDriverPyResult<Py<PyAny>> {
59+
pub fn result(
60+
&self,
61+
py: Python<'_>,
62+
custom_decoders: Option<Py<PyDict>>,
63+
) -> RustPSQLDriverPyResult<Py<PyAny>> {
5964
let mut result: Vec<pyo3::Bound<'_, PyDict>> = vec![];
6065
for row in &self.inner {
61-
result.push(row_to_dict(py, row)?);
66+
result.push(row_to_dict(py, row, &custom_decoders)?);
6267
}
6368
Ok(result.to_object(py))
6469
}
@@ -77,7 +82,7 @@ impl PSQLDriverPyQueryResult {
7782
) -> RustPSQLDriverPyResult<Py<PyAny>> {
7883
let mut res: Vec<Py<PyAny>> = vec![];
7984
for row in &self.inner {
80-
let pydict: pyo3::Bound<'_, PyDict> = row_to_dict(py, row)?;
85+
let pydict: pyo3::Bound<'_, PyDict> = row_to_dict(py, row, &None)?;
8186
let convert_class_inst = as_class.call_bound(py, (), Some(&pydict))?;
8287
res.push(convert_class_inst);
8388
}
@@ -117,7 +122,7 @@ impl PSQLDriverSinglePyQueryResult {
117122
/// postgres type to python, can not set new key-value pair
118123
/// in python dict or there are no result.
119124
pub fn result(&self, py: Python<'_>) -> RustPSQLDriverPyResult<Py<PyAny>> {
120-
Ok(row_to_dict(py, &self.inner)?.to_object(py))
125+
Ok(row_to_dict(py, &self.inner, &None)?.to_object(py))
121126
}
122127

123128
/// Convert result from database to any class passed from Python.
@@ -133,7 +138,7 @@ impl PSQLDriverSinglePyQueryResult {
133138
py: Python<'a>,
134139
as_class: Py<PyAny>,
135140
) -> RustPSQLDriverPyResult<Py<PyAny>> {
136-
let pydict: pyo3::Bound<'_, PyDict> = row_to_dict(py, &self.inner)?;
141+
let pydict: pyo3::Bound<'_, PyDict> = row_to_dict(py, &self.inner, &None)?;
137142
Ok(as_class.call_bound(py, (), Some(&pydict))?)
138143
}
139144
}

src/value_converter.rs

Lines changed: 16 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -652,10 +652,9 @@ fn postgres_bytes_to_py(
652652
None => Ok(py.None().to_object(py)),
653653
}
654654
}
655-
_ => Ok(
656-
_composite_field_postgres_to_py::<Option<Vec<u8>>>(type_, buf, is_simple)?
657-
.to_object(py),
658-
),
655+
_ => Err(RustPSQLDriverError::RustToPyValueConversionError(
656+
format!("Cannot convert {type_} into Python type, please look at the custom_decoders functionality.")
657+
)),
659658
}
660659
}
661660

@@ -720,9 +719,21 @@ pub fn postgres_to_py(
720719
row: &Row,
721720
column: &Column,
722721
column_i: usize,
722+
custom_decoders: &Option<Py<PyDict>>,
723723
) -> RustPSQLDriverPyResult<Py<PyAny>> {
724-
let column_type = column.type_();
725724
let raw_bytes_data = row.col_buffer(column_i);
725+
726+
if let Some(custom_decoders) = custom_decoders {
727+
let py_encoder_func = custom_decoders
728+
.bind(py)
729+
.get_item(column.name().to_lowercase());
730+
731+
if let Ok(Some(py_encoder_func)) = py_encoder_func {
732+
return Ok(py_encoder_func.call((raw_bytes_data,), None)?.unbind());
733+
}
734+
}
735+
736+
let column_type = column.type_();
726737
match raw_bytes_data {
727738
Some(mut raw_bytes_data) => match column_type.kind() {
728739
Kind::Simple | Kind::Array(_) => {

0 commit comments

Comments
 (0)