8000 Added support for multi-dimensional arrays · psqlpy-python/psqlpy@6c993b4 · GitHub
[go: up one dir, main page]
Skip to content

Commit 6c993b4

Browse files
chandr-andrchandr-andr
committed
Added support for multi-dimensional arrays
Signed-off-by: chandr-andr (Kiselev Aleksandr) <chandr@chandr.net>
1 parent 15bfc0c commit 6c993b4

File tree

3 files changed

+245
-12
lines changed

3 files changed

+245
-12
lines changed

‎Cargo.lock

Lines changed: 12 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎Cargo.toml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -50,3 +50,4 @@ openssl = { version = "0.10.64", features = ["vendored"] }
5050
itertools = "0.12.1"
5151
openssl-src = "300.2.2"
5252
openssl-sys = "0.9.102"
53+
postgres_array = { git = "https://github.com/chandr-andr/rust-postgres-array.git", branch = "psqlpy" }

‎src/value_converter.rs

Lines changed: 232 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@ use pyo3::{
1414
sync::GILOnceCell,
1515
types::{
1616
PyAnyMethods, PyBool, PyBytes, PyDate, PyDateTime, PyDict, PyDictMethods, PyFloat, PyInt,
17-
PyList, PyListMethods, PySet, PyString, PyTime, PyTuple, PyType, PyTypeMethods,
17+
PyList, PyListMethods, PySequence, PySet, PyString, PyTime, PyTuple, PyType, PyTypeMethods,
1818
},
1919
Bound, IntoPy, Py, PyAny, PyObject, PyResult, Python, ToPyObject,
2020
};
@@ -35,6 +35,7 @@ use crate::{
3535
SmallInt,
3636
},
3737
};
38+
use postgres_array::{array::Array, Dimension};
3839

3940
static DECIMAL_CLS: GILOnceCell<Py<PyType>> = GILOnceCell::new();
4041

@@ -96,6 +97,7 @@ pub enum PythonDTO {
9697
PyDateTimeTz(DateTime<FixedOffset>),
9798
PyIpAddress(IpAddr),
9899
PyList(Vec<PythonDTO>),
100+
PyArray(Array<PythonDTO>),
99101
PyTuple(Vec<PythonDTO>),
100102
PyJsonb(Value),
101103
PyJson(Value),
@@ -111,6 +113,24 @@ pub enum PythonDTO {
111113
PyCircle(Circle),
112114
}
113115

116+
impl ToPyObject for PythonDTO {
117+
fn to_object(&self, py: Python<'_>) -> PyObject {
118+
match self {
119+
PythonDTO::PyNone => py.None(),
120+
PythonDTO::PyBool(pybool) => pybool.to_object(py),
121+
PythonDTO::PyString(py_string)
122+
| PythonDTO::PyText(py_string)
123+
| PythonDTO::PyVarChar(py_string) => py_string.to_object(py),
124+
PythonDTO::PyIntI32(pyint) => pyint.to_object(py),
125+
PythonDTO::PyIntI64(pyint) => pyint.to_object(py),
126+
PythonDTO::PyIntU64(pyint) => pyint.to_object(py),
127+
PythonDTO::PyFloat32(pyfloat) => pyfloat.to_object(py),
128+
PythonDTO::PyFloat64(pyfloat) => pyfloat.to_object(py),
129+
_ => unreachable!(),
130+
}
131+
}
132+
}
133+
114134
impl PythonDTO {
115135
/// Return type of the Array for `PostgreSQL`.
116136
///
@@ -183,6 +203,26 @@ impl PythonDTO {
183203

184204
Ok(json!(vec_serde_values))
185205
}
206+
PythonDTO::PyArray(array) => {
207+
let py_list = Python::with_gil(|gil| {
208+
let py_list = postgres_array_to_py(gil, Some(array.clone()));
209+
if let Some(py_list) = py_list {
210+
let mut vec_serde_values: Vec<Value> = vec![];
211+
212+
for py_object in py_list.bind(gil) {
213+
vec_serde_values.push(py_to_rust(&py_object)?.to_serde_value()?);
214+
}
215+
216+
return Ok(json!(vec_serde_values));
217+
}
218+
219+
return Err(RustPSQLDriverError::PyToRustValueConversionError(
220+
"Cannot convert Python sequence into JSON".into(),
221+
));
222+
});
223+
224+
py_list
225+
}
186226
PythonDTO::PyJsonb(py_dict) | PythonDTO::PyJson(py_dict) => Ok(py_dict.clone()),
187227
_ => Err(RustPSQLDriverError::PyToRustValueConversionError(
188228
"Cannot convert your type into Rust type".into(),
@@ -311,6 +351,20 @@ impl ToSql for PythonDTO {
311351
items.to_sql(&items[0].array_type()?, out)?;
312352
}
313353
}
354+
PythonDTO::PyArray(array) => {
355+
if let Some(first_elem) = array.iter().nth(0) {
356+
match first_elem.array_type() {
357+
Ok(ok_type) => {
358+
array.to_sql(&ok_type, out)?;
359+
}
360+
Err(_) => {
361+
return Err(RustPSQLDriverError::PyToRustValueConversionError(
362+
"Cannot define array type.".into(),
363+
))?
364+
}
365+
}
366+
}
367+
}
314368
PythonDTO::PyJsonb(py_dict) | PythonDTO::PyJson(py_dict) => {
315369
<&Value as ToSql>::to_sql(&py_dict, ty, out)?;
316370
}
@@ -358,6 +412,105 @@ pub fn convert_parameters(parameters: Py<PyAny>) -> RustPSQLDriverPyResult<Vec<P
358412
Ok(result_vec)
359413
}
360414

415+
pub fn py_sequence_into_flat_vec(
416+
parameter: &Bound<PyAny>,
417+
) -> RustPSQLDriverPyResult<Vec<PythonDTO>> {
418+
let py_seq = parameter.downcast::<PySequence>().map_err(|_| {
419+
RustPSQLDriverError::PyToRustValueConversionError(
420+
"PostgreSQL ARRAY type can be made only from python Sequence".into(),
421+
)
422+
})?;
423+
424+
let mut final_vec: Vec<PythonDTO> = vec![];
425+
426+
for seq_elem in py_seq.iter()? {
427+
let ok_seq_elem = seq_elem?;
428+
429+
// Check for the string because it's sequence too,
430+
// and in the most cases it should be array type, not new dimension.
431+
if ok_seq_elem.is_instance_of::<PyString>() {
432+
final_vec.push(py_to_rust(&ok_seq_elem)?);
433+
continue;
434+
}
435+
436+
let possible_next_seq = ok_seq_elem.downcast::<PySequence>();
437+
438+
match possible_next_seq {
439+
Ok(next_seq) => {
440+
let mut next_vec = py_sequence_into_flat_vec(next_seq)?;
441+
final_vec.append(&mut next_vec);
442+
}
443+
Err(_) => {
444+
final_vec.push(py_to_rust(&ok_seq_elem)?);
445+
continue;
446+
}
447+
}
448+
}
449+
450+
return Ok(final_vec);
451+
}
452+
453+
pub fn py_sequence_into_postgres_array(
454+
parameter: &Bound<PyAny>,
455+
) -> RustPSQLDriverPyResult<Array<PythonDTO>> {
456+
let mut py_seq = parameter
457+
.downcast::<PySequence>()
458+
.map_err(|_| {
459+
RustPSQLDriverError::PyToRustValueConversionError(
460+
"PostgreSQL ARRAY type can be made only from python Sequence".into(),
461+
)
462+
})?
463+
.clone();
464+
465+
let mut dimensions: Vec<Dimension> = vec![];
466+
let mut continue_iteration = true;
467+
468+
while continue_iteration {
469+
dimensions.push(Dimension {
470+
len: py_seq.len()? as i32,
471+
lower_bound: 1,
472+
});
473+
474+
let first_seq_elem = py_seq.iter()?.nth(0);
475+
match first_seq_elem {
476+
Some(first_seq_elem) => {
477+
if let Ok(first_seq_elem) = first_seq_elem {
478+
// Check for the string because it's sequence too,
479+
// and in the most cases it should be array type, not new dimension.
480+
if first_seq_elem.is_instance_of::<PyString>() {
481+
continue_iteration = false;
482+
continue;
483+
}
484+
let possible_inner_seq = first_seq_elem.downcast::<PySequence>();
485+
486+
match possible_inner_seq {
487+
Ok(possible_inner_seq) => {
488+
py_seq = possible_inner_seq.clone();
489+
}
490+
Err(_) => continue_iteration = false,
491+
}
492+
}
493+
}
494+
None => {
495+
continue_iteration = false;
496+
}
497+
}
498+
}
499+
500+
let array_data = py_sequence_into_flat_vec(parameter)?;
501+
502+
match postgres_array::Array::from_parts_no_panic(array_data, dimensions) {
503+
Ok(result_array) => {
504+
return Ok(result_array);
505+
}
506+
Err(err) => {
507+
return Err(RustPSQLDriverError::PyToRustValueConversionError(format!(
508+
"Cannot convert python sequence to PostgreSQL ARRAY, error - {err}"
509+
)))
510+
}
511+
}
512+
}
513+
361514
/// Convert single python parameter to `PythonDTO` enum.
362515
///
363516
/// # Errors
@@ -467,11 +620,9 @@ pub fn py_to_rust(parameter: &pyo3::Bound<'_, PyAny>) -> RustPSQLDriverPyResult<
467620
}
468621

469622
if parameter.is_instance_of::<PyList>() | parameter.is_instance_of::<PyTuple>() {
470-
let mut items = Vec::new();
471-
for inner in parameter.iter()? {
472-
items.push(py_to_rust(&inner?)?);
473-
}
474-
return Ok(PythonDTO::PyList(items));
623+
return Ok(PythonDTO::PyArray(py_sequence_into_postgres_array(
624+
parameter,
625+
)?));
475626
}
476627

477628
if parameter.is_instance_of::<PyDict>() {
@@ -608,6 +759,69 @@ fn _composite_field_postgres_to_py<'a, T: FromSql<'a>>(
608759
})
609760
}
610761

762+
/// Convert rust array to python list.
763+
///
764+
/// It can convert multidimensional arrays.
765+
fn postgres_array_to_py<T: ToPyObject>(
766+
py: Python<'_>,
767+
array: Option<Array<T>>,
768+
) -> Option<Py<PyList>> {
769+
match array {
770+
Some(array) => {
771+
return Some(_postgres_array_to_py(
772+
py,
773+
array.dimensions(),
774+
array.iter().map(|arg| arg).collect::<Vec<&T>>().as_slice(),
775+
0,
776+
0,
777+
));
778+
}
779+
None => {
780+
return None;
781+
}
782+
}
783+
}
784+
785+
/// Inner postgres array conversion to python list.
786+
fn _postgres_array_to_py<T>(
787+
py: Python<'_>,
788+
dimensions: &[Dimension],
789+
data: &[T],
790+
dimension_index: usize,
791+
mut lower_bound: usize,
792+
) -> Py<PyList>
793+
where
794+
T: ToPyObject,
795+
{
796+
let current_dimension = dimensions.iter().nth(dimension_index).unwrap();
797+
798+
let possible_next_dimension = dimensions.iter().nth(dimension_index + 1);
799+
match possible_next_dimension {
800+
Some(next_dimension) => {
801+
let final_list = PyList::empty_bound(py);
802+
803+
for _ in 0..current_dimension.len as usize {
804+
if dimensions.iter().nth(dimension_index + 1).is_some() {
805+
let inner_pylist = _postgres_array_to_py(
806+
py,
807+
dimensions,
808+
&data[lower_bound..next_dimension.len as usize + lower_bound],
809+
dimension_index + 1,
810+
0,
811+
);
812+
final_list.append(inner_pylist).unwrap();
813+
lower_bound += next_dimension.len as usize;
814+
};
815+
}
816+
817+
return final_list.unbind();
818+
}
819+
None => {
820 10000 +
return PyList::new_bound(py, data).unbind();
821+
}
822+
}
823+
}
824+
611825
#[allow(clippy::too_many_lines)]
612826
fn postgres_bytes_to_py(
613827
py: Python<'_>,
@@ -794,10 +1008,16 @@ fn postgres_bytes_to_py(
7941008
)?
7951009
.to_object(py)),
7961010
// Convert ARRAY of Integer into Vec<i32>, then into list[int]
797-
Type::INT4_ARRAY => Ok(_composite_field_postgres_to_py::<Option<Vec<i32>>>(
798-
type_, buf, is_simple,
799-
)?
800-
.to_object(py)),
1011+
Type::INT4_ARRAY => {
1012+
return Ok(postgres_array_to_py(
1013+
py,
1014+
_composite_field_postgres_to_py::<Option<Array<i32>>>(
1015+
type_,
1016+
buf,
1017+
is_simple,
1018+
)?
1019+
).to_object(py));
1020+
},
8011021
// Convert ARRAY of BigInt into Vec<i64>, then into list[int]
8021022
Type::INT8_ARRAY | Type::MONEY_ARRAY => Ok(_composite_field_postgres_to_py::<Option<Vec<i64>>>(
8031023
type_, buf, is_simple,
@@ -1147,7 +1367,7 @@ pub fn build_serde_value(value: Py<PyAny>) -> RustPSQLDriverPyResult<Value> {
11471367
result_vec.push(serde_value);
11481368
} else {
11491369
return Err(RustPSQLDriverError::PyToRustValueConversionError(
1150-
"PyJSON supports only list of lists or list of dicts.".to_string(),
1370+
"PyJSON must have dicts.".to_string(),
11511371
));
11521372
}
11531373
}
@@ -1156,7 +1376,7 @@ pub fn build_serde_value(value: Py<PyAny>) -> RustPSQLDriverPyResult<Value> {
11561376
return py_to_rust(bind_value)?.to_serde_value();
11571377
} else {
11581378
return Err(RustPSQLDriverError::PyToRustValueConversionError(
1159-
"PyJSON must be list value.".to_string(),
1379+
"PyJSON must be dict value.".to_string(),
11601380
));
11611381
}
11621382
})

0 commit comments

Comments
 (0)
0