Skip to content
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

Added INTERVAL type support #95

Merged
merged 2 commits into from
Oct 6, 2024
Merged
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
11 changes: 11 additions & 0 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

1 change: 1 addition & 0 deletions Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -51,3 +51,4 @@ openssl = { version = "0.10.64", features = ["vendored"] }
itertools = "0.12.1"
openssl-src = "300.2.2"
openssl-sys = "0.9.102"
pg_interval = { git = "https://github.com/chandr-andr/rust-postgres-interval.git", branch = "psqlpy" }
1 change: 1 addition & 0 deletions docs/usage/types/array_types.md
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,7 @@ For type safety and better performance we have predefined array types.
| LineArray | LINE ARRAY |
| LsegArray | LSEG ARRAY |
| CircleArray | CIRCLE ARRAY |
| IntervalArray | INTERVAL ARRAY |

### Example:

Expand Down
1 change: 1 addition & 0 deletions docs/usage/types/supported_types.md
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@ Here you can find all types supported by `PSQLPy`. If PSQLPy isn't `-`, you can
| datetime.time | - | TIME |
| datetime.datetime | - | TIMESTAMP |
| datetime.datetime | - | TIMESTAMPTZ |
| datetime.timedelta | - | INTERVAL |
| UUID | - | UUID |
| dict | - | JSONB |
| dict | PyJSONB | JSONB |
Expand Down
23 changes: 21 additions & 2 deletions python/psqlpy/_internal/extra_types.pyi
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
import typing
from datetime import date, datetime, time
from datetime import date, datetime, time, timedelta
from decimal import Decimal
from ipaddress import IPv4Address, IPv6Address
from uuid import UUID
Expand Down Expand Up @@ -753,5 +753,24 @@ class CircleArray:
"""Create new instance of CircleArray.

### Parameters:
- `inner`: inner value, sequence of PyLineSegment values.
- `inner`: inner value, sequence of PyCircle values.
"""

class IntervalArray:
"""Represent INTERVAL ARRAY in PostgreSQL."""

def __init__(
self: Self,
inner: typing.Sequence[
typing.Union[
timedelta,
typing.Sequence[timedelta],
typing.Any,
],
],
) -> None:
"""Create new instance of IntervalArray.

### Parameters:
- `inner`: inner value, sequence of timedelta values.
"""
2 changes: 2 additions & 0 deletions python/psqlpy/extra_types.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@
Int32Array,
Int64Array,
Integer,
IntervalArray,
IpAddressArray,
JSONArray,
JSONBArray,
Expand Down Expand Up @@ -96,4 +97,5 @@
"LineArray",
"LsegArray",
"CircleArray",
"IntervalArray",
]
18 changes: 18 additions & 0 deletions python/tests/test_value_converter.py
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@
Int32Array,
Int64Array,
Integer,
IntervalArray,
IpAddressArray,
JSONArray,
JSONBArray,
Expand Down Expand Up @@ -198,6 +199,11 @@ async def test_as_class(
PyCircle([1, 2.8, 3]),
((1.0, 2.8), 3.0),
),
(
"INTERVAL",
datetime.timedelta(days=100, microseconds=100),
datetime.timedelta(days=100, microseconds=100),
),
(
"VARCHAR ARRAY",
["Some String", "Some String"],
Expand Down Expand Up @@ -598,6 +604,11 @@ async def test_as_class(
[((5.0, 1.8), 10.0)],
],
),
(
"INTERVAL ARRAY",
[datetime.timedelta(days=100, microseconds=100), datetime.timedelta(days=100, microseconds=100)],
[datetime.timedelta(days=100, microseconds=100), datetime.timedelta(days=100, microseconds=100)],
),
),
)
async def test_deserialization_simple_into_python(
Expand Down Expand Up @@ -1501,6 +1512,13 @@ async def test_empty_array(
[((5.0, 1.8), 10.0)],
],
),
(
"INTERVAL ARRAY",
IntervalArray(
[[datetime.timedelta(days=100, microseconds=100)], [datetime.timedelta(days=100, microseconds=100)]],
),
[[datetime.timedelta(days=100, microseconds=100)], [datetime.timedelta(days=100, microseconds=100)]],
),
),
)
async def test_array_types(
Expand Down
2 changes: 2 additions & 0 deletions src/extra_types.rs
Original file line number Diff line number Diff line change
Expand Up @@ -361,6 +361,7 @@ build_array_type!(PathArray, PythonDTO::PyPathArray);
build_array_type!(LineArray, PythonDTO::PyLineArray);
build_array_type!(LsegArray, PythonDTO::PyLsegArray);
build_array_type!(CircleArray, PythonDTO::PyCircleArray);
build_array_type!(IntervalArray, PythonDTO::PyIntervalArray);

#[allow(clippy::module_name_repetitions)]
#[allow(clippy::missing_errors_doc)]
Expand Down Expand Up @@ -410,5 +411,6 @@ pub fn extra_types_module(_py: Python<'_>, pymod: &Bound<'_, PyModule>) -> PyRes
pymod.add_class::<LineArray>()?;
pymod.add_class::<LsegArray>()?;
pymod.add_class::<CircleArray>()?;
pymod.add_class::<IntervalArray>()?;
Ok(())
}
89 changes: 84 additions & 5 deletions src/value_converter.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ use chrono::{self, DateTime, FixedOffset, NaiveDate, NaiveDateTime, NaiveTime};
use geo_types::{coord, Coord, Line as LineSegment, LineString, Point, Rect};
use itertools::Itertools;
use macaddr::{MacAddr6, MacAddr8};
use pg_interval::Interval;
use postgres_types::{Field, FromSql, Kind, ToSql};
use rust_decimal::Decimal;
use serde_json::{json, Map, Value};
Expand All @@ -13,8 +14,8 @@ use postgres_protocol::types;
use pyo3::{
sync::GILOnceCell,
types::{
PyAnyMethods, PyBool, PyBytes, PyDate, PyDateTime, PyDict, PyDictMethods, PyFloat, PyInt,
PyIterator, PyList, PyListMethods, PySequence, PySet, PyString, PyTime, PyTuple, PyType,
PyAnyMethods, PyBool, PyBytes, PyDate, PyDateTime, PyDelta, PyDict, PyDictMethods, PyFloat,
PyInt, PyList, PyListMethods, PySequence, PySet, PyString, PyTime, PyTuple, PyType,
PyTypeMethods,
},
Bound, FromPyObject, IntoPy, Py, PyAny, PyObject, PyResult, Python, ToPyObject,
Expand All @@ -35,6 +36,7 @@ use crate::{
use postgres_array::{array::Array, Dimension};

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

pub type QueryParameter = (dyn ToSql + Sync);

Expand All @@ -50,6 +52,18 @@ fn get_decimal_cls(py: Python<'_>) -> PyResult<&Bound<'_, PyType>> {
.map(|ty| ty.bind(py))
}

fn get_timedelta_cls(py: Python<'_>) -> PyResult<&Bound<'_, PyType>> {
TIMEDELTA_CLS
.get_or_try_init(py, || {
let type_object = py
.import_bound("datetime")?
.getattr("timedelta")?
.downcast_into()?;
Ok(type_object.unbind())
})
.map(|ty| ty.bind(py))
}

/// Struct for Uuid.
///
/// We use custom struct because we need to implement external traits
Expand Down Expand Up @@ -138,13 +152,43 @@ impl<'a> FromSql<'a> for InnerDecimal {
}
}

struct InnerInterval(Interval);

impl ToPyObject for InnerInterval {
fn to_object(&self, py: Python<'_>) -> PyObject {
let td_cls = get_timedelta_cls(py).expect("failed to load datetime.timedelta");
let pydict = PyDict::new_bound(py);
let months = self.0.months * 30;
let _ = pydict.set_item("days", self.0.days + months);
let _ = pydict.set_item("microseconds", self.0.microseconds);
let ret = td_cls
.call((), Some(&pydict))
.expect("failed to call datetime.timedelta(days=<>, microseconds=<>)");
ret.to_object(py)
}
}

impl<'a> FromSql<'a> for InnerInterval {
fn from_sql(
ty: &Type,
raw: &'a [u8],
) -> Result<Self, Box<dyn std::error::Error + Sync + Send>> {
Ok(InnerInterval(<Interval as FromSql>::from_sql(ty, raw)?))
}

fn accepts(_ty: &Type) -> bool {
true
}
}

/// Additional type for types come from Python.
///
/// It's necessary because we need to pass this
/// enum into `to_sql` method of `ToSql` trait from
/// `postgres` crate.
#[derive(Debug, Clone, PartialEq)]
pub enum PythonDTO {
// Primitive
PyNone,
PyBytes(Vec<u8>),
PyBool(bool),
Expand All @@ -164,6 +208,7 @@ pub enum PythonDTO {
PyTime(NaiveTime),
PyDateTime(NaiveDateTime),
PyDateTimeTz(DateTime<FixedOffset>),
PyInterval(Interval),
PyIpAddress(IpAddr),
PyList(Vec<PythonDTO>),
PyArray(Array<PythonDTO>),
Expand All @@ -180,6 +225,7 @@ pub enum PythonDTO {
PyLine(Line),
PyLineSegment(LineSegment),
PyCircle(Circle),
// Arrays
PyBoolArray(Array<PythonDTO>),
PyUuidArray(Array<PythonDTO>),
PyVarCharArray(Array<PythonDTO>),
Expand All @@ -206,6 +252,7 @@ pub enum PythonDTO {
PyLineArray(Array<PythonDTO>),
PyLsegArray(Array<PythonDTO>),
PyCircleArray(Array<PythonDTO>),
PyIntervalArray(Array<PythonDTO>),
}

impl ToPyObject for PythonDTO {
Expand Down Expand Up @@ -267,6 +314,7 @@ impl PythonDTO {
PythonDTO::PyLine(_) => Ok(tokio_postgres::types::Type::LINE_ARRAY),
PythonDTO::PyLineSegment(_) => Ok(tokio_postgres::types::Type::LSEG_ARRAY),
PythonDTO::PyCircle(_) => Ok(tokio_postgres::types::Type::CIRCLE_ARRAY),
PythonDTO::PyInterval(_) => Ok(tokio_postgres::types::Type::INTERVAL_ARRAY),
_ => Err(RustPSQLDriverError::PyToRustValueConversionError(
"Can't process array type, your type doesn't have support yet".into(),
)),
Expand Down Expand Up @@ -385,6 +433,9 @@ impl ToSql for PythonDTO {
PythonDTO::PyDateTimeTz(pydatetime_tz) => {
<&DateTime<FixedOffset> as ToSql>::to_sql(&pydatetime_tz, ty, out)?;
}
PythonDTO::PyInterval(pyinterval) => {
<&Interval as ToSql>::to_sql(&pyinterval, ty, out)?;
}
PythonDTO::PyIpAddress(pyidaddress) => {
<&IpAddr as ToSql>::to_sql(&pyidaddress, ty, out)?;
}
Expand Down Expand Up @@ -525,6 +576,9 @@ impl ToSql for PythonDTO {
PythonDTO::PyCircleArray(array) => {
array.to_sql(&Type::CIRCLE_ARRAY, out)?;
}
PythonDTO::PyIntervalArray(array) => {
array.to_sql(&Type::INTERVAL_ARRAY, out)?;
}
}

if return_is_null_true {
Expand Down Expand Up @@ -787,6 +841,16 @@ pub fn py_to_rust(parameter: &pyo3::Bound<'_, PyAny>) -> RustPSQLDriverPyResult<
return Ok(PythonDTO::PyTime(parameter.extract::<NaiveTime>()?));
}

if parameter.is_instance_of::<PyDelta>() {
let duration = parameter.extract::<chrono::Duration>()?;
if let Some(interval) = Interval::from_duration(duration) {
return Ok(PythonDTO::PyInterval(interval));
}
return Err(RustPSQLDriverError::PyToRustValueConversionError(
"Cannot convert timedelta from Python to inner Rust type.".to_string(),
));
}

if parameter.is_instance_of::<PyList>() | parameter.is_instance_of::<PyTuple>() {
return Ok(PythonDTO::PyArray(py_sequence_into_postgres_array(
parameter,
Expand Down Expand Up @@ -1052,6 +1116,12 @@ pub fn py_to_rust(parameter: &pyo3::Bound<'_, PyAny>) -> RustPSQLDriverPyResult<
._convert_to_python_dto();
}

if parameter.is_instance_of::<extra_types::IntervalArray>() {
return parameter
.extract::<extra_types::IntervalArray>()?
._convert_to_python_dto();
}

if let Ok(id_address) = parameter.extract::<IpAddr>() {
return Ok(PythonDTO::PyIpAddress(id_address));
}
Expand All @@ -1065,9 +1135,6 @@ pub fn py_to_rust(parameter: &pyo3::Bound<'_, PyAny>) -> RustPSQLDriverPyResult<
}
}

let a = parameter.downcast::<PyIterator>();
println!("{:?}", a.iter());

Err(RustPSQLDriverError::PyToRustValueConversionError(format!(
"Can not covert you type {parameter} into inner one",
)))
Expand Down Expand Up @@ -1387,6 +1454,13 @@ fn postgres_bytes_to_py(
None => Ok(py.None().to_object(py)),
}
}
Type::INTERVAL => {
let interval = _composite_field_postgres_to_py::<Option<Interval>>(type_, buf, is_simple)?;
if let Some(interval) = interval {
return Ok(InnerInterval(interval).to_object(py));
}
Ok(py.None())
}
// ---------- Array Text Types ----------
Type::BOOL_ARRAY => Ok(postgres_array_to_py(py, _composite_field_postgres_to_py::<Option<Array<bool>>>(
type_, buf, is_simple,
Expand Down Expand Up @@ -1505,6 +1579,11 @@ fn postgres_bytes_to_py(

Ok(postgres_array_to_py(py, circle_array_).to_object(py))
}
Type::INTERVAL_ARRAY => {
let interval_array_ = _composite_field_postgres_to_py::<Option<Array<InnerInterval>>>(type_, buf, is_simple)?;

Ok(postgres_array_to_py(py, interval_array_).to_object(py))
}
_ => Err(RustPSQLDriverError::RustToPyValueConversionError(
format!("Cannot convert {type_} into Python type, please look at the custom_decoders functionality.")
)),
Expand Down
Loading