Merge pull request #95 from psqlpy-python/feature/add_interval_type_support

Added INTERVAL type support

Author: chandr-andr

Cargo.lock

@@ -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" }

Cargo.toml

@@ -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:
 

docs/usage/types/array_types.md

@@ -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 |

docs/usage/types/supported_types.md

@@ -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
@@ -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.
         """

python/psqlpy/_internal/extra_types.pyi

@@ -14,6 +14,7 @@
     Int32Array,
     Int64Array,
     Integer,
+    IntervalArray,
     IpAddressArray,
     JSONArray,
     JSONBArray,
@@ -96,4 +97,5 @@
     "LineArray",
     "LsegArray",
     "CircleArray",
+    "IntervalArray",
 ]

python/psqlpy/extra_types.py

@@ -27,6 +27,7 @@
     Int32Array,
     Int64Array,
     Integer,
+    IntervalArray,
     IpAddressArray,
     JSONArray,
     JSONBArray,
@@ -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"],
@@ -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(
@@ -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(

python/tests/test_value_converter.py

@@ -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)]
@@ -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(())
 }

src/extra_types.rs

@@ -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};
@@ -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,
@@ -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);
 
@@ -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
@@ -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),
@@ -164,6 +208,7 @@ pub enum PythonDTO {
     PyTime(NaiveTime),
     PyDateTime(NaiveDateTime),
     PyDateTimeTz(DateTime<FixedOffset>),
+    PyInterval(Interval),
     PyIpAddress(IpAddr),
     PyList(Vec<PythonDTO>),
     PyArray(Array<PythonDTO>),
@@ -180,6 +225,7 @@ pub enum PythonDTO {
     PyLine(Line),
     PyLineSegment(LineSegment),
     PyCircle(Circle),
+    // Arrays
     PyBoolArray(Array<PythonDTO>),
     PyUuidArray(Array<PythonDTO>),
     PyVarCharArray(Array<PythonDTO>),
@@ -206,6 +252,7 @@ pub enum PythonDTO {
     PyLineArray(Array<PythonDTO>),
     PyLsegArray(Array<PythonDTO>),
     PyCircleArray(Array<PythonDTO>),
+    PyIntervalArray(Array<PythonDTO>),
 }
 
 impl ToPyObject for PythonDTO {
@@ -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(),
             )),
@@ -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)?;
             }
@@ -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 {
@@ -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,
@@ -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));
     }
@@ -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",
     )))
@@ -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,
@@ -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.")
         )),