Return tz-aware datetimes for DateTime/DateTime64 columns with a timezone

CLAUDE.md

@@ -114,9 +114,12 @@ supported Python is 3.10 (`python_requires` in `setup.py`).
 - Query `params` use `str.format`, so literal `{`/`}` in SQL (e.g. cluster macros) must be
   escaped as `{{`/`}}`.
 - The TSV streaming backend asserts the response ends on a line boundary (`assert not buffer`).
-- **Native engine specifics**: `decode=False` (raw bytes) is TSV-only; the Native format
-  carries no per-column timezone, so a tz-aware `DateTime`/`DateTime64` comes back as naive
-  UTC. A multi-block native INSERT is **not atomic** on a client-side encoding error (blocks
+- **Timezones**: a `DateTime('TZ')` / `DateTime64(P, 'TZ')` column decodes to a tz-aware
+  `datetime` (`zoneinfo`); a column without a zone stays naive. Writers accept naive (column
+  wall-clock) and aware (any zone) values. One server quirk: the Native header ships a
+  `DateTime('TZ')` column as plain `DateTime` (DateTime64 keeps its zone), so that column is
+  naive UTC on the native engine only.
+- **Native engine specifics**: `decode=False` (raw bytes) is TSV-only. A multi-block native INSERT is **not atomic** on a client-side encoding error (blocks
   already streamed are committed) — `insert_block_size=0` forces a single atomic block. The
   GC is disabled inside the synchronous decode/build sections of a fetch (mass allocation
   otherwise trips the cyclic collector); it is always restored via `try/finally`, and those

README.md

@@ -179,9 +179,9 @@ Because binary encoding is type-specific, a binary INSERT first looks up the
 target column types (one lightweight query) and encodes the rows to match them.
 
 A couple of `native=True` caveats: `decode=False` (raw bytes) is a TSV-only
-feature, and the Native format does not carry a column's timezone, so a tz-aware
-`DateTime`/`DateTime64` comes back as a naive UTC `datetime` (TSV and RowBinary
-apply the timezone).
+feature, and ClickHouse's Native header drops the timezone of a `DateTime('TZ')`
+column (`DateTime64` keeps it), so that one column comes back as a naive UTC
+`datetime` on the native engine (TSV and RowBinary return it tz-aware).
 
 A native INSERT streams its body as several Native blocks so the server can
 insert one while the client encodes the next (`insert_block_size` rows per block,
@@ -259,7 +259,9 @@ vice-versa.
 | `Enum16`             | `str`                   |
 | `Date`               | `datetime.date`         |
 | `DateTime`           | `datetime.datetime`     |
+| `DateTime('TZ')`     | tz-aware `datetime`     |
 | `DateTime64`         | `datetime.datetime`     |
+| `DateTime64(P, 'TZ')`| tz-aware `datetime`     |
 | `Decimal`            | `decimal.Decimal`       |
 | `Decimal32`          | `decimal.Decimal`       |
 | `Decimal64`          | `decimal.Decimal`       |
@@ -273,6 +275,12 @@ vice-versa.
 | `Nullable(T)`        | `None` or `T`           |
 | `LowCardinality(T)`  | `T`                     |
 | `Map(T1, T2)`        | `Dict[T1, T2]`          |
+
+A column declared with a timezone decodes to a `datetime` carrying that zone
+(`zoneinfo.ZoneInfo`); a column without one stays naive. On INSERT a naive
+`datetime` is taken as wall-clock time in the column's zone (UTC for a column
+without one), and a tz-aware `datetime` of any zone is stored as the instant it
+denotes.
 | `Nested(T1, T2, ...)` | `List[Tuple[T1, T2, ...], Tuple[T1, T2, ...]]` |
 
 

aiochclient/_types.pyx

@@ -11,18 +11,12 @@ from uuid import UUID
 from zoneinfo import ZoneInfo
 
 from cpython cimport PyList_Append, PyUnicode_AsEncodedString, PyUnicode_Join
+from cpython.datetime cimport date, date_new, datetime, datetime_new, import_datetime
 from cpython.list cimport PyList_GET_ITEM, PyList_GET_SIZE
+from cpython.mem cimport PyMem_Free, PyMem_Malloc
 from cpython.ref cimport Py_INCREF
 from cpython.tuple cimport PyTuple_New, PyTuple_SET_ITEM
-from cpython.datetime cimport (
-    date,
-    date_new,
-    datetime,
-    datetime_new,
-    import_datetime,
-)
 from cpython.unicode cimport PyUnicode_DecodeUTF8
-from cpython.mem cimport PyMem_Free, PyMem_Malloc
 from libc.stdint cimport (
     int8_t,
     int16_t,
@@ -216,6 +210,21 @@ cdef class Cursor:
 RB_EPOCH_DATE = _dt.date(1970, 1, 1)
 RB_EPOCH_DATETIME = _dt.datetime(1970, 1, 1)
 RB_EPOCH_UTC = _dt.datetime(1970, 1, 1, tzinfo=_dt.timezone.utc)
+_MICROSECOND = _dt.timedelta(microseconds=1)
+
+
+cpdef object to_epoch_micros(object value, object zone):
+    """Microseconds since the epoch for a naive or aware ``value``.
+
+    Naive + column zone -> wall-clock in that zone. Naive + no zone -> UTC
+    wall-clock (the historical behaviour). Aware -> its real instant, whatever
+    the column zone is.
+    """
+    if value.tzinfo is None:
+        if zone is None:
+            return (value - RB_EPOCH_DATETIME) // _MICROSECOND
+        value = value.replace(tzinfo=zone)
+    return (value - RB_EPOCH_UTC) // _MICROSECOND
 _TZ_UNSET = object()
 
 
@@ -884,12 +893,14 @@ cdef class DateTimeType(_RBType):
     cdef object _convert(self, str string):
         string = string.strip("'")
         try:
-            return datetime_parse(string)
+            value = datetime_parse(string)
         except ValueError:
             # In case of 0000-00-00 00:00:00
             if string == "0000-00-00 00:00:00":
                 return None
             raise
+        zone = self._zone()
+        return value.replace(tzinfo=zone) if zone else value
 
     cpdef object p_type(self, str string):
         return self._convert(string)
@@ -902,14 +913,10 @@ cdef class DateTimeType(_RBType):
         zone = self._zone()
         if zone is None:
             return RB_EPOCH_DATETIME + _dt.timedelta(seconds=seconds)
-        return _dt.datetime.fromtimestamp(seconds, zone).replace(tzinfo=None)
+        return _dt.datetime.fromtimestamp(seconds, zone)
 
     cpdef bytes write(self, value):
-        zone = self._zone()
-        if zone is None:
-            seconds = (value - RB_EPOCH_DATETIME) // _dt.timedelta(seconds=1)
-        else:
-            seconds = int(value.replace(tzinfo=zone).timestamp())
+        seconds = to_epoch_micros(value, self._zone()) // 1_000_000
         return int(seconds).to_bytes(4, "little")
 
 
@@ -937,12 +944,14 @@ cdef class DateTime64Type(_RBType):
     cdef object _convert(self, str string):
         string = string.strip("'")
         try:
-            return datetime_parse_f(string)
+            value = datetime_parse_f(string)
         except ValueError:
             # In case of 0000-00-00 00:00:00.000
             if string == "0000-00-00 00:00:00.000":
                 return None
             raise
+        zone = self._zone()
+        return value.replace(tzinfo=zone) if zone else value
 
     cpdef object p_type(self, str string):
         return self._convert(string)
@@ -959,18 +968,10 @@ cdef class DateTime64Type(_RBType):
         zone = self._zone()
         if zone is None:
             return RB_EPOCH_DATETIME + _dt.timedelta(microseconds=micros)
-        return (RB_EPOCH_UTC + _dt.timedelta(microseconds=micros)).astimezone(
-            zone
-        ).replace(tzinfo=None)
+        return (RB_EPOCH_UTC + _dt.timedelta(microseconds=micros)).astimezone(zone)
 
     cpdef bytes write(self, value):
-        zone = self._zone()
-        if zone is None:
-            micros = (value - RB_EPOCH_DATETIME) // _dt.timedelta(microseconds=1)
-        else:
-            micros = (value.replace(tzinfo=zone) - RB_EPOCH_UTC) // _dt.timedelta(
-                microseconds=1
-            )
+        micros = to_epoch_micros(value, self._zone())
         if self._precision <= 6:
             ticks = micros // 10 ** (6 - self._precision)
         else:

aiochclient/native.py

@@ -29,9 +29,9 @@
 
 # Use the compiled Cursor (with bulk column reads) when available.
 try:
-    from aiochclient._types import Cursor  # noqa: F401
+    from aiochclient._types import Cursor, to_epoch_micros  # noqa: F401
 except ImportError:
-    from aiochclient.types import Cursor  # noqa: F401
+    from aiochclient.types import Cursor, to_epoch_micros  # noqa: F401
 
 # Compiled String-column encoder; pure-Python fallback below.
 try:
@@ -213,10 +213,10 @@ def decode_column(cursor, n, ctype):
         return [data[i * width : (i + 1) * width].decode() for i in range(n)]
     if ctype == "Date":
         return cursor.read_date_column(n)
-    if ctype == "DateTime" or ctype.startswith("DateTime("):
-        # The Native format does not carry the per-column timezone (it always
-        # ships the UTC epoch), so DateTime is returned as a naive UTC datetime.
-        # The TSV/RowBinary engines apply the column timezone instead.
+    if ctype == "DateTime":
+        # ClickHouse's Native header drops the timezone of a DateTime('TZ')
+        # column (it ships plain "DateTime"; DateTime64 keeps its zone), so a
+        # timezone DateTime can only come back as naive UTC on this engine.
         return cursor.read_datetime_column(n)
     if ctype.startswith("Nullable("):
         inner = ctype[9:-1]
@@ -341,7 +341,6 @@ async def rows_from_native(
 # Epochs for the null-slot defaults and the bulk Date/DateTime encoders.
 _EPOCH_DATE = dt.date(1970, 1, 1)
 _EPOCH_DATETIME = dt.datetime(1970, 1, 1)
-_ONE_SECOND = dt.timedelta(seconds=1)
 
 
 def _write_varint(value):
@@ -461,10 +460,10 @@ def encode_column(values, ctype):
             column.byteswap()
         return column.tobytes()
     if ctype == "DateTime":
-        # Naive UTC seconds (matches the no-timezone DateTime writer). Timezone
-        # DateTime / DateTime64 keep the per-value writer path below.
+        # Naive UTC / aware seconds (matches the no-timezone DateTime writer).
+        # Timezone DateTime / DateTime64 keep the per-value writer path below.
         column = array.array(
-            "I", [(v - _EPOCH_DATETIME) // _ONE_SECOND for v in values]
+            "I", [to_epoch_micros(v, None) // 1_000_000 for v in values]
         )
         if not _LE:
             column.byteswap()

aiochclient/types.py

@@ -144,6 +144,22 @@ def _parse_tz(name: str) -> Optional[str]:
 # Epoch used to turn day/second/tick offsets into python date/datetime objects.
 RB_EPOCH_DATE = dt.date(1970, 1, 1)
 RB_EPOCH_DATETIME = dt.datetime(1970, 1, 1)
+RB_EPOCH_UTC = dt.datetime(1970, 1, 1, tzinfo=dt.timezone.utc)
+_MICROSECOND = dt.timedelta(microseconds=1)
+
+
+def to_epoch_micros(value: dt.datetime, zone) -> int:
+    """Microseconds since the epoch for a naive or aware ``value``.
+
+    Naive + column zone -> wall-clock in that zone. Naive + no zone -> UTC
+    wall-clock (the historical behaviour). Aware -> its real instant, whatever
+    the column zone is.
+    """
+    if value.tzinfo is None:
+        if zone is None:
+            return (value - RB_EPOCH_DATETIME) // _MICROSECOND
+        value = value.replace(tzinfo=zone)
+    return (value - RB_EPOCH_UTC) // _MICROSECOND
 
 
 def write_varint(value: int) -> bytes:
@@ -455,12 +471,14 @@ def _zone(self):
     def p_type(self, string: str):
         string = string.strip("'")
         try:
-            return datetime_parse(string)
+            value = datetime_parse(string)
         except ValueError:
             # In case of 0000-00-00 00:00:00
             if string == "0000-00-00 00:00:00":
                 return None
             raise
+        zone = self._zone()
+        return value.replace(tzinfo=zone) if zone else value
 
     def convert(self, value: bytes) -> Optional[dt.datetime]:
         return self.p_type(value.decode())
@@ -470,24 +488,18 @@ def read(self, cursor) -> dt.datetime:
         zone = self._zone()
         if zone is None:
             return RB_EPOCH_DATETIME + dt.timedelta(seconds=seconds)
-        # Match the server's TSV output: wall-clock in the column timezone.
-        return dt.datetime.fromtimestamp(seconds, zone).replace(tzinfo=None)
+        return dt.datetime.fromtimestamp(seconds, zone)
 
     def write(self, value: dt.datetime) -> bytes:
-        zone = self._zone()
-        if zone is None:
-            seconds = (value - RB_EPOCH_DATETIME) // dt.timedelta(seconds=1)
-        else:
-            # ``value`` is naive wall-clock in the column timezone.
-            seconds = int(value.replace(tzinfo=zone).timestamp())
+        seconds = to_epoch_micros(value, self._zone()) // 1_000_000
         return seconds.to_bytes(4, "little")
 
     @staticmethod
     def unconvert(value: dt.datetime) -> bytes:
-        if value.microsecond != 0:
-            # In case of 0000-00-00 00:00:00.000 (datetime64)
-            return b"%a" % dt.datetime.strftime(value, '%Y-%m-%d %H:%M:%S.%f')
-        return b"%a" % dt.datetime.strftime(value, '%Y-%m-%d %H:%M:%S')
+        # str() keeps the sub-second part only when present (a DateTime column
+        # rejects it) and the UTC offset of an aware value, which ClickHouse
+        # applies. Same as the compiled ``unconvert_datetime``.
+        return b"'%s'" % str(value).encode()
 
 
 class DateTime64Type(BaseType):
@@ -506,12 +518,14 @@ def _zone(self):
     def p_type(self, string: str):
         string = string.strip("'")
         try:
-            return datetime_parse_f(string)
+            value = datetime_parse_f(string)
         except ValueError:
             # In case of 0000-00-00 00:00:00
             if string == "0000-00-00 00:00:00.000":
                 return None
             raise
+        zone = self._zone()
+        return value.replace(tzinfo=zone) if zone else value
 
     def convert(self, value: bytes) -> Optional[dt.datetime]:
         return self.p_type(value.decode())
@@ -526,25 +540,10 @@ def read(self, cursor) -> dt.datetime:
         zone = self._zone()
         if zone is None:
             return RB_EPOCH_DATETIME + dt.timedelta(microseconds=micros)
-        # Match the server's TSV output: wall-clock in the column timezone.
-        return (
-            (
-                dt.datetime(1970, 1, 1, tzinfo=dt.timezone.utc)
-                + dt.timedelta(microseconds=micros)
-            )
-            .astimezone(zone)
-            .replace(tzinfo=None)
-        )
+        return (RB_EPOCH_UTC + dt.timedelta(microseconds=micros)).astimezone(zone)
 
     def write(self, value: dt.datetime) -> bytes:
-        zone = self._zone()
-        if zone is None:
-            micros = (value - RB_EPOCH_DATETIME) // dt.timedelta(microseconds=1)
-        else:
-            aware = value.replace(tzinfo=zone)
-            micros = (
-                aware - dt.datetime(1970, 1, 1, tzinfo=dt.timezone.utc)
-            ) // dt.timedelta(microseconds=1)
+        micros = to_epoch_micros(value, self._zone())
         if self._precision <= 6:
             ticks = micros // 10 ** (6 - self._precision)
         else: