From b05a6c0a75ea8d31e4ba1f3f7d06d3242f085e83 Mon Sep 17 00:00:00 2001 From: John Bodley Date: Sat, 29 Apr 2023 10:11:55 +1200 Subject: [PATCH] fix(dbapi): Address timestamp with time zone localization issue --- tests/integration/test_dbapi_integration.py | 48 +++++++++++++++++---- trino/client.py | 4 +- 2 files changed, 43 insertions(+), 9 deletions(-) diff --git a/tests/integration/test_dbapi_integration.py b/tests/integration/test_dbapi_integration.py index 0f442bb1..60d14142 100644 --- a/tests/integration/test_dbapi_integration.py +++ b/tests/integration/test_dbapi_integration.py @@ -364,7 +364,7 @@ def test_datetime_with_numeric_offset_time_zone_query_param(trino_connection): def test_datetime_with_named_time_zone_query_param(trino_connection): cur = trino_connection.cursor() - params = datetime(2020, 1, 1, 16, 43, 22, 320000, tzinfo=pytz.timezone('America/Los_Angeles')) + params = pytz.timezone('America/Los_Angeles').localize(datetime(2020, 1, 1, 16, 43, 22, 320000)) cur.execute("SELECT ?", params=(params,)) rows = cur.fetchall() @@ -393,13 +393,44 @@ def test_null_datetime_with_time_zone(trino_connection): assert_cursor_description(cur, trino_type="timestamp(3) with time zone", precision=3) -def test_datetime_with_time_zone_numeric_offset(trino_connection): +@pytest.mark.parametrize( + 'operation,result,precision', + [ + ( + "SELECT TIMESTAMP '2001-08-22 03:04:05.321 -08:00'", + datetime(2001, 8, 22, 3, 4, 5, 321000, tzinfo=timezone(-timedelta(hours=8))), + 3, + + ), + ( + "SELECT TIMESTAMP '2001-08-22 03:04:05.321 America/Los_Angeles'", + pytz.timezone("America/Los_Angeles").localize(datetime(2001, 8, 22, 3, 4, 5, 321000)), + 3, + ), + ], +) +def test_datetime_with_time_zone(trino_connection, operation, result, precision): + cur = trino_connection.cursor() + + cur.execute(operation) + rows = cur.fetchall() + + assert rows[0][0] == result + + assert_cursor_description( + cur, + trino_type=f"timestamp({precision}) with time zone", + precision=precision, + ) + + +def test_datetime_with_at_time_zone(trino_connection): cur = trino_connection.cursor() - cur.execute("SELECT TIMESTAMP '2001-08-22 03:04:05.321 -08:00'") + cur.execute("SELECT TIMESTAMP '2001-08-22 03:04:05.321 UTC' AT TIME ZONE 'America/Los_Angeles'") rows = cur.fetchall() - assert rows[0][0] == datetime.strptime("2001-08-22 03:04:05.321 -08:00", "%Y-%m-%d %H:%M:%S.%f %z") + assert rows[0][0] == pytz.timezone('America/Los_Angeles').localize(datetime(2001, 8, 21, 20, 4, 5, 321000)) assert_cursor_description(cur, trino_type="timestamp(3) with time zone", precision=3) @@ -416,23 +447,24 @@ def test_datetimes_with_time_zone_in_dst_gap_query_param(trino_connection): def test_doubled_datetimes(trino_connection): # Trino doesn't distinguish between doubled datetimes that lie within a DST transition. See also # See also https://github.com/trinodb/trino/issues/5781 + dt = datetime(2002, 10, 27, 1, 30, 0) cur = trino_connection.cursor() - params = pytz.timezone('US/Eastern').localize(datetime(2002, 10, 27, 1, 30, 0), is_dst=True) + params = pytz.timezone('US/Eastern').localize(dt, is_dst=True) cur.execute("SELECT ?", params=(params,)) rows = cur.fetchall() - assert rows[0][0] == datetime(2002, 10, 27, 1, 30, 0, tzinfo=pytz.timezone('US/Eastern')) + assert rows[0][0] == pytz.timezone('US/Eastern').localize(dt) cur = trino_connection.cursor() - params = pytz.timezone('US/Eastern').localize(datetime(2002, 10, 27, 1, 30, 0), is_dst=False) + params = pytz.timezone('US/Eastern').localize(dt, is_dst=False) cur.execute("SELECT ?", params=(params,)) rows = cur.fetchall() - assert rows[0][0] == datetime(2002, 10, 27, 1, 30, 0, tzinfo=pytz.timezone('US/Eastern')) + assert rows[0][0] == pytz.timezone('US/Eastern').localize(dt) def test_date_query_param(trino_connection): diff --git a/trino/client.py b/trino/client.py index 9f9b75f8..28ebc151 100644 --- a/trino/client.py +++ b/trino/client.py @@ -1108,8 +1108,10 @@ def map(self, value) -> Optional[datetime]: datetime_with_fraction, timezone_part = value.rsplit(' ', 1) whole_python_temporal_value = datetime_with_fraction[:self.datetime_default_size] remaining_fractional_seconds = datetime_with_fraction[self.datetime_default_size + 1:] + tz = _create_tzinfo(timezone_part) + dt = datetime.fromisoformat(whole_python_temporal_value) return TimestampWithTimeZone( - datetime.fromisoformat(whole_python_temporal_value).replace(tzinfo=_create_tzinfo(timezone_part)), + dt.replace(tzinfo=tz) if isinstance(tz, timezone) else tz.localize(dt), _fraction_to_decimal(remaining_fractional_seconds), ).round_to(self.precision).to_python_type()