Keep the time zone of TIMESTAMP results (non-UTC session time zones)

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]:

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