From 2abb79358f91a16ab43e840f2ad2412698357cba Mon Sep 17 00:00:00 2001 From: simozzy Date: Fri, 19 Jun 2026 03:17:50 +0100 Subject: [PATCH] fix: keep timezone-aware datetimes aware and UTC-normalise naive coercion DatabendDateTime.result_processor unconditionally did value.replace(tzinfo=None), which has two problems: - It strips tzinfo even for timezone=True columns, so a DateTime(timezone=True) / TIMESTAMP WITH TIME ZONE column loses its awareness. - It drops tzinfo without converting to UTC first, so a tz-aware value carrying a non-UTC offset yields the wrong naive wall-clock. Only coerce to naive for timezone=False columns, and convert to UTC before dropping tzinfo so the wall-clock is correct regardless of the driver's offset; timezone=True columns are returned unchanged. None and string parsing are preserved. Adds a server-less unit test. --- databend_sqlalchemy/databend_dialect.py | 10 +++++++-- tests/unit/test_databend_dialect.py | 30 +++++++++++++++++++++++++ 2 files changed, 38 insertions(+), 2 deletions(-) diff --git a/databend_sqlalchemy/databend_dialect.py b/databend_sqlalchemy/databend_dialect.py index cfb5482..9af3ef3 100644 --- a/databend_sqlalchemy/databend_dialect.py +++ b/databend_sqlalchemy/databend_dialect.py @@ -737,8 +737,14 @@ def process(value): "could not parse %r as a datetime value" % (value,) ) return datetime.datetime(*[int(x or 0) for x in m.groups()]) - else: - return value.replace(tzinfo=None) + if value.tzinfo is not None and not self.timezone: + # The driver returns tz-aware (UTC) datetimes, but a + # timezone=False column must yield naive values to honour + # SQLAlchemy's typing contract. Normalise to UTC before + # dropping tzinfo so a non-UTC offset doesn't shift the + # wall-clock; a timezone=True column keeps its awareness. + return value.astimezone(datetime.timezone.utc).replace(tzinfo=None) + return value return process diff --git a/tests/unit/test_databend_dialect.py b/tests/unit/test_databend_dialect.py index 889bfc5..122e037 100644 --- a/tests/unit/test_databend_dialect.py +++ b/tests/unit/test_databend_dialect.py @@ -1,3 +1,4 @@ +import datetime import os from unittest import mock @@ -228,3 +229,32 @@ def test_extract_nullable_string(): assert expected_types[i] == true_type i += 1 print(true_type) + + +def test_datetime_result_processor_timezone_handling(): + # The driver returns tz-aware (UTC) datetimes; a plain DateTime()/TIMESTAMP + # column (timezone=False) must yield naive values so the compliance suite's + # naive round-trips hold, while a timezone=True column keeps its awareness. + DatabendDateTime = databend_sqlalchemy.databend_dialect.DatabendDateTime + aware_utc = datetime.datetime( + 2012, 10, 15, 12, 57, 18, 396, tzinfo=datetime.timezone.utc + ) + naive = datetime.datetime(2012, 10, 15, 12, 57, 18, 396) + + process = DatabendDateTime().result_processor(None, None) + # tz-aware input -> naive output for a timezone=False column + assert process(aware_utc) == naive + assert process(aware_utc).tzinfo is None + # a non-UTC offset is normalised to UTC before tzinfo is dropped, so the + # wall-clock is converted rather than merely truncated + aware_plus2 = aware_utc.astimezone(datetime.timezone(datetime.timedelta(hours=2))) + assert process(aware_plus2) == naive + # naive passes through unchanged; None is preserved; strings still parse + assert process(naive) == naive + assert process(None) is None + assert process("2012-10-15 12:57:18.000396") == naive + + # timezone=True columns keep their tzinfo untouched + process_tz = DatabendDateTime(timezone=True).result_processor(None, None) + assert process_tz(aware_utc) == aware_utc + assert process_tz(aware_utc).tzinfo is not None