Skip to content

Commit 241dbb7

Browse files
Vladislav YashkovVladislav Yashkov
Vladislav Yashkov
authored and
Vladislav Yashkov
committed
WIP# 1
1 parent cecec4a commit 241dbb7

File tree

3 files changed

+120
-14
lines changed

3 files changed

+120
-14
lines changed

‎python/tests/test_value_converter.py

Lines changed: 32 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
import datetime
2+
import sys
23
import uuid
34
from decimal import Decimal
45
from enum import Enum
@@ -57,6 +58,7 @@
5758

5859
from tests.conftest import DefaultPydanticModel, DefaultPythonModelClass
5960

61+
uuid_ = uuid.uuid4()
6062
pytestmark = pytest.mark.anyio
6163
now_datetime = datetime.datetime.now() # noqa: DTZ005
6264
now_datetime_with_tz = datetime.datetime(
@@ -69,7 +71,30 @@
6971
142574,
7072
tzinfo=datetime.timezone.utc,
7173
)
72-
uuid_ = uuid.uuid4()
74+
75+
now_datetime_with_tz_in_asia_jakarta = datetime.datetime(
76+
2024,
77+
4,
78+
13,
79+
17,
80+
3,
81+
46,
82+
142574,
83+
tzinfo=datetime.timezone.utc,
84+
)
85+
if sys.version_info >= (3, 9):
86+
import zoneinfo
87+
88+
now_datetime_with_tz_in_asia_jakarta = datetime.datetime(
89+
2024,
90+
4,
91+
13,
92+
17,
93+
3,
94+
46,
95+
142574,
96+
tzinfo=zoneinfo.ZoneInfo(key="Asia/Jakarta"),
97+
)
7398

7499

75100
async def test_as_class(
@@ -125,6 +150,7 @@ async def test_as_class(
125150
("TIME", now_datetime.time(), now_datetime.time()),
126151
("TIMESTAMP", now_datetime, now_datetime),
127152
("TIMESTAMPTZ", now_datetime_with_tz, now_datetime_with_tz),
153+
("TIMESTAMPTZ", now_datetime_with_tz_in_asia_jakarta, now_datetime_with_tz_in_asia_jakarta),
128154
("UUID", uuid_, str(uuid_)),
129155
("INET", IPv4Address("192.0.0.1"), IPv4Address("192.0.0.1")),
130156
(
@@ -287,6 +313,11 @@ async def test_as_class(
287313
[now_datetime_with_tz, now_datetime_with_tz],
288314
[now_datetime_with_tz, now_datetime_with_tz],
289315
),
316+
(
317+
"TIMESTAMPTZ ARRAY",
318+
[now_datetime_with_tz, now_datetime_with_tz_in_asia_jakarta],
319+
[now_datetime_with_tz, now_datetime_with_tz_in_asia_jakarta],
320+
),
290321
(
291322
"TIMESTAMPTZ ARRAY",
292323
[[now_datetime_with_tz], [now_datetime_with_tz]],

‎src/driver/connection.rs

Lines changed: 0 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -127,15 +127,6 @@ impl Connection {
127127

128128
#[pymethods]
129129
impl Connection {
130-
#[must_use]
131-
pub fn __aiter__(self_: Py<Self>) -> Py<Self> {
132-
self_
133-
}
134-
135-
fn __await__(self_: Py<Self>) -> Py<Self> {
136-
self_
137-
}
138-
139130
async fn __aenter__<'a>(self_: Py<Self>) -> RustPSQLDriverPyResult<Py<Self>> {
140131
let (db_client, db_pool) = pyo3::Python::with_gil(|gil| {
141132
let self_ = self_.borrow(gil);

‎src/value_converter.rs

Lines changed: 88 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
1-
use chrono::{self, DateTime, FixedOffset, NaiveDate, NaiveDateTime, NaiveTime};
1+
use chrono::{self, DateTime, FixedOffset, NaiveDate, NaiveDateTime, NaiveTime, TimeZone};
2+
use chrono_tz::Tz;
23
use geo_types::{coord, Coord, Line as LineSegment, LineString, Point, Rect};
34
use itertools::Itertools;
45
use macaddr::{MacAddr6, MacAddr8};
@@ -626,8 +627,7 @@ impl ToSql for PythonDTO {
626627
#[allow(clippy::needless_pass_by_value)]
627628
pub fn convert_parameters(parameters: Py<PyAny>) -> RustPSQLDriverPyResult<Vec<PythonDTO>> {
628629
let mut result_vec: Vec<PythonDTO> = vec![];
629-
630-
result_vec = Python::with_gil(|gil| {
630+
Python::with_gil(|gil| {
631631
let params = parameters.extract::<Vec<Py<PyAny>>>(gil).map_err(|_| {
632632
RustPSQLDriverError::PyToRustValueConversionError(
633633
"Cannot convert you parameters argument into Rust type, please use List/Tuple"
@@ -637,8 +637,9 @@ pub fn convert_parameters(parameters: Py<PyAny>) -> RustPSQLDriverPyResult<Vec<P
637637
for parameter in params {
638638
result_vec.push(py_to_rust(parameter.bind(gil))?);
639639
}
640-
Ok::<Vec<PythonDTO>, RustPSQLDriverError>(result_vec)
640+
Ok::<(), RustPSQLDriverError>(())
641641
})?;
642+
642643
Ok(result_vec)
643644
}
644645

@@ -744,6 +745,84 @@ pub fn py_sequence_into_postgres_array(
744745
}
745746
}
746747

748+
/// Extract a value from a Python object, raising an error if missing or invalid
749+
///
750+
/// # Type Parameters
751+
/// - `T`: The type to which the attribute's value will be converted. This type must implement the `FromPyObject` trait
752+
///
753+
/// # Errors
754+
/// This function will return `Err` in the following cases:
755+
/// - The Python object does not have the specified attribute
756+
/// - The attribute exists but cannot be extracted into the specified Rust type
757+
fn extract_value_from_python_object_or_raise<'py, T>(
758+
parameter: &'py pyo3::Bound<'_, PyAny>,
759+
attr_name: &str,
760+
) -> Result<T, RustPSQLDriverError>
761+
where
762+
T: FromPyObject<'py>,
763+
{
764+
parameter
765+
.getattr(attr_name)
766+
.ok()
767+
.and_then(|attr| attr.extract::<T>().ok())
768+
.ok_or_else(|| {
769+
RustPSQLDriverError::PyToRustValueConversionError("Invalid attribute".into())
770+
})
771+
}
772+
773+
/// Extract a timezone-aware datetime from a Python object.
774+
/// This function retrieves various datetime components (`year`, `month`, `day`, etc.)
775+
/// from a Python object and constructs a `DateTime<FixedOffset>`
776+
///
777+
/// # Errors
778+
/// This function will return `Err` in the following cases:
779+
/// - The Python object does not contain or support one or more required datetime attributes
780+
/// - The retrieved values are invalid for constructing a date, time, or datetime (e.g., invalid month or day)
781+
/// - The timezone information (`tzinfo`) is not available or cannot be parsed
782+
/// - The resulting datetime is ambiguous or invalid (e.g., due to DST transitions)
783+
fn extract_datetime_from_python_object_attrs(
784+
parameter: &pyo3::Bound<'_, PyAny>,
785+
) -> Result<DateTime<FixedOffset>, RustPSQLDriverError> {
786+
let year = extract_value_from_python_object_or_raise::<i32>(parameter, "year")?;
787+
let month = extract_value_from_python_object_or_raise::<u32>(parameter, "month")?;
788+
let day = extract_value_from_python_object_or_raise::<u32>(parameter, "day")?;
789+
let hour = extract_value_from_python_object_or_raise::<u32>(parameter, "hour")?;
790+
let minute = extract_value_from_python_object_or_raise::<u32>(parameter, "minute")?;
791+
let second = extract_value_from_python_object_or_raise::<u32>(parameter, "second")?;
792+
let microsecond = extract_value_from_python_object_or_raise::<u32>(parameter, "microsecond")?;
793+
794+
let date = NaiveDate::from_ymd_opt(year, month, day)
795+
.ok_or_else(|| RustPSQLDriverError::PyToRustValueConversionError("Invalid date".into()))?;
796+
let time = NaiveTime::from_hms_micro_opt(hour, minute, second, microsecond)
797+
.ok_or_else(|| RustPSQLDriverError::PyToRustValueConversionError("Invalid time".into()))?;
798+
let naive_datetime = NaiveDateTime::new(date, time);
799+
800+
let raw_timestamp_tz = parameter
801+
.getattr("tzinfo")
802+
.ok()
803+
.and_then(|tzinfo| tzinfo.getattr("key").ok())
804+
.and_then(|key| key.extract::<String>().ok())
805+
.ok_or_else(|| {
806+
RustPSQLDriverError::PyToRustValueConversionError("Invalid timezone info".into())
807+
})?;
808+
809+
let fixed_offset_datetime = raw_timestamp_tz
810+
.parse::<Tz>()
811+
.map_err(|_| {
812+
RustPSQLDriverError::PyToRustValueConversionError("Failed to parse TZ".into())
813+
})?
814+
.from_local_datetime(&naive_datetime)
815+
.single()
816+
.ok_or_else(|| {
817+
RustPSQLDriverError::PyToRustValueConversionError(
818+
"Ambiguous or invalid datetime".into(),
819+
)
820+
})?
821+
.fixed_offset();
822+
823+
Ok(fixed_offset_datetime)
824+
}
825+
747826
/// Convert single python parameter to `PythonDTO` enum.
748827
///
749828
/// # Errors
@@ -849,6 +928,11 @@ pub fn py_to_rust(parameter: &pyo3::Bound<'_, PyAny>) -> RustPSQLDriverPyResult<
849928
return Ok(PythonDTO::PyDateTime(pydatetime_no_tz));
850929
}
851930

931+
let timestamp_tz = extract_datetime_from_python_object_attrs(parameter);
932+
if let Ok(pydatetime_tz) = timestamp_tz {
933+
return Ok(PythonDTO::PyDateTimeTz(pydatetime_tz));
934+
}
935+
852936
return Err(RustPSQLDriverError::PyToRustValueConversionError(
853937
"Can not convert you datetime to rust type".into(),
854938
));

0 commit comments

Comments
 (0)