diff --git a/src/databricks/sqlalchemy/_types.py b/src/databricks/sqlalchemy/_types.py index c180404..59c83ea 100644 --- a/src/databricks/sqlalchemy/_types.py +++ b/src/databricks/sqlalchemy/_types.py @@ -191,6 +191,20 @@ def process_result_value(self, value: Union[None, datetime], dialect): return value.replace(tzinfo=timezone.utc) return value + def result_processor(self, dialect, coltype): + """Use only process_result_value. + + The impl (DateTime) is adapted through the dialect's colspecs to + TIMESTAMP_NTZ, whose result processing drops tzinfo. Composing it would + relabel the session-local wall clock as UTC, shifting every value when + the session time zone is not UTC. + """ + + def process(value): + return self.process_result_value(value, dialect) + + return process + def process_bind_param( self, value: Union[datetime, None], dialect ) -> Optional[datetime]: diff --git a/tests/test_local/test_types.py b/tests/test_local/test_types.py index 7036c8d..73ba2c2 100644 --- a/tests/test_local/test_types.py +++ b/tests/test_local/test_types.py @@ -366,3 +366,26 @@ def test_as_uuid_false_round_trip_normalizes_hex_input(self): assert bind(self.HYPHENATED) == self.HYPHENATED assert result(self.HYPHENATED) == self.HYPHENATED assert result(self.HEX) == self.HYPHENATED + + +def test_timestamp_result_keeps_session_time_zone(): + """A TIMESTAMP read in a non-UTC session keeps its instant.""" + import datetime + import zoneinfo + + import sqlalchemy + + from databricks.sqlalchemy import TIMESTAMP + + dialect = sqlalchemy.create_engine( + "databricks://token:x@host?http_path=p" + ).dialect + process = TIMESTAMP().dialect_impl(dialect).result_processor(dialect, None) + ist = datetime.datetime( + 2024, 3, 1, 12, 29, 58, tzinfo=zoneinfo.ZoneInfo("Asia/Kolkata") + ) + assert process(ist) == ist + assert process(ist).utcoffset() == ist.utcoffset() + naive = datetime.datetime(2024, 3, 1) + assert process(naive) == naive.replace(tzinfo=datetime.timezone.utc) + assert process(None) is None