From 241dbb785b409ef17c86ce0b2ee8744d94009d64 Mon Sep 17 00:00:00 2001 From: Vladislav Yashkov Date: Tue, 3 Dec 2024 12:12:51 +0300 Subject: [PATCH 1/2] WIP# 1 --- python/tests/test_value_converter.py | 33 +++++++++- src/driver/connection.rs | 9 --- src/value_converter.rs | 92 ++++++++++++++++++++++++++-- 3 files changed, 120 insertions(+), 14 deletions(-) diff --git a/python/tests/test_value_converter.py b/python/tests/test_value_converter.py index f8f8ade..cb53bad 100644 --- a/python/tests/test_value_converter.py +++ b/python/tests/test_value_converter.py @@ -1,4 +1,5 @@ import datetime +import sys import uuid from decimal import Decimal from enum import Enum @@ -57,6 +58,7 @@ from tests.conftest import DefaultPydanticModel, DefaultPythonModelClass +uuid_ = uuid.uuid4() pytestmark = pytest.mark.anyio now_datetime = datetime.datetime.now() # noqa: DTZ005 now_datetime_with_tz = datetime.datetime( @@ -69,7 +71,30 @@ 142574, tzinfo=datetime.timezone.utc, ) -uuid_ = uuid.uuid4() + +now_datetime_with_tz_in_asia_jakarta = datetime.datetime( + 2024, + 4, + 13, + 17, + 3, + 46, + 142574, + tzinfo=datetime.timezone.utc, +) +if sys.version_info >= (3, 9): + import zoneinfo + + now_datetime_with_tz_in_asia_jakarta = datetime.datetime( + 2024, + 4, + 13, + 17, + 3, + 46, + 142574, + tzinfo=zoneinfo.ZoneInfo(key="Asia/Jakarta"), + ) async def test_as_class( @@ -125,6 +150,7 @@ async def test_as_class( ("TIME", now_datetime.time(), now_datetime.time()), ("TIMESTAMP", now_datetime, now_datetime), ("TIMESTAMPTZ", now_datetime_with_tz, now_datetime_with_tz), + ("TIMESTAMPTZ", now_datetime_with_tz_in_asia_jakarta, now_datetime_with_tz_in_asia_jakarta), ("UUID", uuid_, str(uuid_)), ("INET", IPv4Address("192.0.0.1"), IPv4Address("192.0.0.1")), ( @@ -287,6 +313,11 @@ async def test_as_class( [now_datetime_with_tz, now_datetime_with_tz], [now_datetime_with_tz, now_datetime_with_tz], ), + ( + "TIMESTAMPTZ ARRAY", + [now_datetime_with_tz, now_datetime_with_tz_in_asia_jakarta], + [now_datetime_with_tz, now_datetime_with_tz_in_asia_jakarta], + ), ( "TIMESTAMPTZ ARRAY", [[now_datetime_with_tz], [now_datetime_with_tz]], diff --git a/src/driver/connection.rs b/src/driver/connection.rs index 23e86a4..97dc66a 100644 --- a/src/driver/connection.rs +++ b/src/driver/connection.rs @@ -127,15 +127,6 @@ impl Connection { #[pymethods] impl Connection { - #[must_use] - pub fn __aiter__(self_: Py) -> Py { - self_ - } - - fn __await__(self_: Py) -> Py { - self_ - } - async fn __aenter__<'a>(self_: Py) -> RustPSQLDriverPyResult> { let (db_client, db_pool) = pyo3::Python::with_gil(|gil| { let self_ = self_.borrow(gil); diff --git a/src/value_converter.rs b/src/value_converter.rs index 7402023..fc19e2f 100644 --- a/src/value_converter.rs +++ b/src/value_converter.rs @@ -1,4 +1,5 @@ -use chrono::{self, DateTime, FixedOffset, NaiveDate, NaiveDateTime, NaiveTime}; +use chrono::{self, DateTime, FixedOffset, NaiveDate, NaiveDateTime, NaiveTime, TimeZone}; +use chrono_tz::Tz; use geo_types::{coord, Coord, Line as LineSegment, LineString, Point, Rect}; use itertools::Itertools; use macaddr::{MacAddr6, MacAddr8}; @@ -626,8 +627,7 @@ impl ToSql for PythonDTO { #[allow(clippy::needless_pass_by_value)] pub fn convert_parameters(parameters: Py) -> RustPSQLDriverPyResult> { let mut result_vec: Vec = vec![]; - - result_vec = Python::with_gil(|gil| { + Python::with_gil(|gil| { let params = parameters.extract::>>(gil).map_err(|_| { RustPSQLDriverError::PyToRustValueConversionError( "Cannot convert you parameters argument into Rust type, please use List/Tuple" @@ -637,8 +637,9 @@ pub fn convert_parameters(parameters: Py) -> RustPSQLDriverPyResult, RustPSQLDriverError>(result_vec) + Ok::<(), RustPSQLDriverError>(()) })?; + Ok(result_vec) } @@ -744,6 +745,84 @@ pub fn py_sequence_into_postgres_array( } } +/// Extract a value from a Python object, raising an error if missing or invalid +/// +/// # Type Parameters +/// - `T`: The type to which the attribute's value will be converted. This type must implement the `FromPyObject` trait +/// +/// # Errors +/// This function will return `Err` in the following cases: +/// - The Python object does not have the specified attribute +/// - The attribute exists but cannot be extracted into the specified Rust type +fn extract_value_from_python_object_or_raise<'py, T>( + parameter: &'py pyo3::Bound<'_, PyAny>, + attr_name: &str, +) -> Result +where + T: FromPyObject<'py>, +{ + parameter + .getattr(attr_name) + .ok() + .and_then(|attr| attr.extract::().ok()) + .ok_or_else(|| { + RustPSQLDriverError::PyToRustValueConversionError("Invalid attribute".into()) + }) +} + +/// Extract a timezone-aware datetime from a Python object. +/// This function retrieves various datetime components (`year`, `month`, `day`, etc.) +/// from a Python object and constructs a `DateTime` +/// +/// # Errors +/// This function will return `Err` in the following cases: +/// - The Python object does not contain or support one or more required datetime attributes +/// - The retrieved values are invalid for constructing a date, time, or datetime (e.g., invalid month or day) +/// - The timezone information (`tzinfo`) is not available or cannot be parsed +/// - The resulting datetime is ambiguous or invalid (e.g., due to DST transitions) +fn extract_datetime_from_python_object_attrs( + parameter: &pyo3::Bound<'_, PyAny>, +) -> Result, RustPSQLDriverError> { + let year = extract_value_from_python_object_or_raise::(parameter, "year")?; + let month = extract_value_from_python_object_or_raise::(parameter, "month")?; + let day = extract_value_from_python_object_or_raise::(parameter, "day")?; + let hour = extract_value_from_python_object_or_raise::(parameter, "hour")?; + let minute = extract_value_from_python_object_or_raise::(parameter, "minute")?; + let second = extract_value_from_python_object_or_raise::(parameter, "second")?; + let microsecond = extract_value_from_python_object_or_raise::(parameter, "microsecond")?; + + let date = NaiveDate::from_ymd_opt(year, month, day) + .ok_or_else(|| RustPSQLDriverError::PyToRustValueConversionError("Invalid date".into()))?; + let time = NaiveTime::from_hms_micro_opt(hour, minute, second, microsecond) + .ok_or_else(|| RustPSQLDriverError::PyToRustValueConversionError("Invalid time".into()))?; + let naive_datetime = NaiveDateTime::new(date, time); + + let raw_timestamp_tz = parameter + .getattr("tzinfo") + .ok() + .and_then(|tzinfo| tzinfo.getattr("key").ok()) + .and_then(|key| key.extract::().ok()) + .ok_or_else(|| { + RustPSQLDriverError::PyToRustValueConversionError("Invalid timezone info".into()) + })?; + + let fixed_offset_datetime = raw_timestamp_tz + .parse::() + .map_err(|_| { + RustPSQLDriverError::PyToRustValueConversionError("Failed to parse TZ".into()) + })? + .from_local_datetime(&naive_datetime) + .single() + .ok_or_else(|| { + RustPSQLDriverError::PyToRustValueConversionError( + "Ambiguous or invalid datetime".into(), + ) + })? + .fixed_offset(); + + Ok(fixed_offset_datetime) +} + /// Convert single python parameter to `PythonDTO` enum. /// /// # Errors @@ -849,6 +928,11 @@ pub fn py_to_rust(parameter: &pyo3::Bound<'_, PyAny>) -> RustPSQLDriverPyResult< return Ok(PythonDTO::PyDateTime(pydatetime_no_tz)); } + let timestamp_tz = extract_datetime_from_python_object_attrs(parameter); + if let Ok(pydatetime_tz) = timestamp_tz { + return Ok(PythonDTO::PyDateTimeTz(pydatetime_tz)); + } + return Err(RustPSQLDriverError::PyToRustValueConversionError( "Can not convert you datetime to rust type".into(), )); From 5cba8584bcfdac3983f62a4b32661908d1e0b2c5 Mon Sep 17 00:00:00 2001 From: Vladislav Yashkov Date: Sat, 7 Dec 2024 18:11:38 +0300 Subject: [PATCH 2/2] WIP# 2 --- src/value_converter.rs | 3 --- 1 file changed, 3 deletions(-) diff --git a/src/value_converter.rs b/src/value_converter.rs index fc19e2f..8f34bd9 100644 --- a/src/value_converter.rs +++ b/src/value_converter.rs @@ -747,9 +747,6 @@ pub fn py_sequence_into_postgres_array( /// Extract a value from a Python object, raising an error if missing or invalid /// -/// # Type Parameters -/// - `T`: The type to which the attribute's value will be converted. This type must implement the `FromPyObject` trait -/// /// # Errors /// This function will return `Err` in the following cases: /// - The Python object does not have the specified attribute