Support TzDate, TzDatetime and TzTimestamp types

Register the timezone-carrying primitive types so query results and
parameters map to/from tz-aware datetime via zoneinfo (raw string when
native result-set conversion is disabled). Reading them previously
crashed with AttributeError.

Author: vgvoleg

CHANGELOG.md

@@ -1,3 +1,4 @@
+* Support the timezone-carrying types `TzDate`, `TzDatetime` and `TzTimestamp` in query results and parameters (reading them previously raised `AttributeError`)
 * Drop support for Python 3.8 and 3.9; the minimum supported version is now Python 3.10
 * Accept native `datetime.datetime` values for `Datetime` and `Datetime64` query parameters (previously only integer seconds since the epoch were accepted)
 

docs/types.rst

@@ -183,6 +183,12 @@ Date and time
      - extended-range timestamp (microsecond precision); maps to/from ``datetime.datetime``
    * - ``PrimitiveType.Interval64``
      - extended-range interval; maps to/from ``datetime.timedelta``
+   * - ``PrimitiveType.TzDate``
+     - date with a timezone name; maps to/from a tz-aware ``datetime.datetime`` at midnight
+   * - ``PrimitiveType.TzDatetime``
+     - date + time with a timezone name (second precision); maps to/from a tz-aware ``datetime.datetime``
+   * - ``PrimitiveType.TzTimestamp``
+     - date + time with a timezone name (microsecond precision); maps to/from a tz-aware ``datetime.datetime``
 
 .. note::
 
@@ -192,6 +198,16 @@ Date and time
    Integer inputs (seconds for ``Datetime``/``Datetime64``, microseconds for
    ``Timestamp``/``Timestamp64`` since the Unix epoch) are still accepted.
 
+.. note::
+
+   ``TzDate``, ``TzDatetime`` and ``TzTimestamp`` carry an IANA timezone name.
+   On write, pass a tz-aware ``datetime.datetime`` whose ``tzinfo`` is a
+   ``zoneinfo.ZoneInfo`` — its local wall-clock and zone name are stored. On
+   read, they are returned as tz-aware ``datetime.datetime`` values (``TzDate``
+   at midnight in that zone). When the matching
+   ``with_native_*_in_result_sets`` option is disabled, the raw
+   ``"<value>,<zone>"`` string is returned instead.
+
 Other
 ^^^^^
 

tests/aio/test_types.py

@@ -4,6 +4,7 @@
 from datetime import date, datetime, timedelta, timezone
 from decimal import Decimal
 from uuid import uuid4
+from zoneinfo import ZoneInfo
 
 
 @pytest.mark.parametrize(
@@ -39,6 +40,9 @@
         (-100, "Interval"),
         (100, "Timestamp"),
         (1511789040123456, "Timestamp"),
+        ("2019-09-17,Europe/Moscow", "TzDate"),
+        ("2019-09-16T18:24:00,Europe/Moscow", "TzDatetime"),
+        ("2019-09-16T18:24:00.123456,Europe/Moscow", "TzTimestamp"),
     ],
 )
 @pytest.mark.asyncio
@@ -54,6 +58,7 @@ async def test_types(driver, database, value, ydb_type):
 test_today = test_now.date()
 test_dt_today = datetime.today()
 tz4h = timezone(timedelta(hours=4))
+tzmsk = ZoneInfo("Europe/Moscow")
 
 
 @pytest.mark.parametrize(
@@ -76,6 +81,13 @@ async def test_types(driver, database, value, ydb_type):
         ),
         ('{"foo": "bar"}', "Json", {"foo": "bar"}),
         ('{"foo": "bar"}', "JsonDocument", {"foo": "bar"}),
+        (datetime(2019, 9, 17, tzinfo=tzmsk), "TzDate", datetime(2019, 9, 17, tzinfo=tzmsk)),
+        (datetime(2019, 9, 16, 18, 24, tzinfo=tzmsk), "TzDatetime", datetime(2019, 9, 16, 18, 24, tzinfo=tzmsk)),
+        (
+            datetime(2019, 9, 16, 18, 24, 0, 123456, tzinfo=tzmsk),
+            "TzTimestamp",
+            datetime(2019, 9, 16, 18, 24, 0, 123456, tzinfo=tzmsk),
+        ),
     ],
 )
 @pytest.mark.asyncio

tests/query/test_types.py

@@ -4,6 +4,7 @@
 from datetime import date, datetime, timedelta, timezone
 from decimal import Decimal
 from uuid import uuid4
+from zoneinfo import ZoneInfo
 
 
 @pytest.mark.parametrize(
@@ -56,6 +57,9 @@
         (1511789040123456, ydb.PrimitiveType.Timestamp),
         (1511789040123456, ydb.PrimitiveType.Timestamp64),
         (-1511789040123456, ydb.PrimitiveType.Timestamp64),
+        ("2019-09-17,Europe/Moscow", ydb.PrimitiveType.TzDate),
+        ("2019-09-16T18:24:00,Europe/Moscow", ydb.PrimitiveType.TzDatetime),
+        ("2019-09-16T18:24:00.123456,Europe/Moscow", ydb.PrimitiveType.TzTimestamp),
     ],
 )
 def test_types(driver_sync: ydb.Driver, value, ydb_type):
@@ -81,6 +85,7 @@ def test_types(driver_sync: ydb.Driver, value, ydb_type):
 test_today = test_now.date()
 test_dt_today = datetime.today()
 tz4h = timezone(timedelta(hours=4))
+tzmsk = ZoneInfo("Europe/Moscow")
 
 
 @pytest.mark.parametrize(
@@ -108,6 +113,17 @@ def test_types(driver_sync: ydb.Driver, value, ydb_type):
         ),
         ('{"foo": "bar"}', ydb.PrimitiveType.Json, {"foo": "bar"}),
         ('{"foo": "bar"}', ydb.PrimitiveType.JsonDocument, {"foo": "bar"}),
+        (datetime(2019, 9, 17, tzinfo=tzmsk), ydb.PrimitiveType.TzDate, datetime(2019, 9, 17, tzinfo=tzmsk)),
+        (
+            datetime(2019, 9, 16, 18, 24, tzinfo=tzmsk),
+            ydb.PrimitiveType.TzDatetime,
+            datetime(2019, 9, 16, 18, 24, tzinfo=tzmsk),
+        ),
+        (
+            datetime(2019, 9, 16, 18, 24, 0, 123456, tzinfo=tzmsk),
+            ydb.PrimitiveType.TzTimestamp,
+            datetime(2019, 9, 16, 18, 24, 0, 123456, tzinfo=tzmsk),
+        ),
     ],
 )
 def test_types_native(driver_sync, value, ydb_type, result_value):
@@ -138,6 +154,24 @@ def test_types_native(driver_sync, value, ydb_type, result_value):
         (test_td, ydb.PrimitiveType.Interval, "-PT0.0001S", test_td),
         (test_td, ydb.PrimitiveType.Interval64, "-PT0.0001S", test_td),
         (test_old_date, ydb.PrimitiveType.Timestamp64, "1221-01-01T00:00:00Z", test_old_date),
+        (
+            datetime(2019, 9, 17, tzinfo=tzmsk),
+            ydb.PrimitiveType.TzDate,
+            "2019-09-17,Europe/Moscow",
+            datetime(2019, 9, 17, tzinfo=tzmsk),
+        ),
+        (
+            datetime(2019, 9, 16, 18, 24, tzinfo=tzmsk),
+            ydb.PrimitiveType.TzDatetime,
+            "2019-09-16T18:24:00,Europe/Moscow",
+            datetime(2019, 9, 16, 18, 24, tzinfo=tzmsk),
+        ),
+        (
+            datetime(2019, 9, 16, 18, 24, 0, 123456, tzinfo=tzmsk),
+            ydb.PrimitiveType.TzTimestamp,
+            "2019-09-16T18:24:00.123456,Europe/Moscow",
+            datetime(2019, 9, 16, 18, 24, 0, 123456, tzinfo=tzmsk),
+        ),
     ],
 )
 def test_type_str_repr(driver_sync, value, ydb_type, str_repr, result_value):

ydb/types.py

@@ -7,6 +7,7 @@
 import json
 from . import _utilities, _apis
 from datetime import date, datetime, timedelta, timezone
+from zoneinfo import ZoneInfo, ZoneInfoNotFoundError
 import typing
 import uuid
 import struct
@@ -80,6 +81,67 @@ def _to_datetime64(pb: ydb_value_pb2.Value, value: typing.Union[datetime, int])
         pb.int64_value = value
 
 
+def _tz_name(value: datetime) -> str:
+    tz = value.tzinfo
+    if isinstance(tz, ZoneInfo):
+        return tz.key
+    return str(tz)
+
+
+def _parse_tz(value: str) -> typing.Union[datetime, str]:
+    wall_clock, _, tz_name = value.partition(",")
+    try:
+        tz = ZoneInfo(tz_name)
+    except ZoneInfoNotFoundError:
+        # No IANA database available for this zone; return the raw text as-is.
+        return value
+    if "T" in wall_clock:
+        naive = datetime.fromisoformat(wall_clock)
+    else:
+        parsed_date = date.fromisoformat(wall_clock)
+        naive = datetime(parsed_date.year, parsed_date.month, parsed_date.day)
+    return naive.replace(tzinfo=tz)
+
+
+def _from_tz_date(x: str, table_client_settings: table.TableClientSettings) -> typing.Union[datetime, str]:
+    if table_client_settings is not None and table_client_settings._native_date_in_result_sets:
+        return _parse_tz(x)
+    return x
+
+
+def _to_tz_date(pb: ydb_value_pb2.Value, value: typing.Union[datetime, str]) -> None:
+    if isinstance(value, datetime):
+        pb.text_value = value.strftime("%Y-%m-%d") + "," + _tz_name(value)
+    else:
+        pb.text_value = value
+
+
+def _from_tz_datetime(x: str, table_client_settings: table.TableClientSettings) -> typing.Union[datetime, str]:
+    if table_client_settings is not None and table_client_settings._native_datetime_in_result_sets:
+        return _parse_tz(x)
+    return x
+
+
+def _to_tz_datetime(pb: ydb_value_pb2.Value, value: typing.Union[datetime, str]) -> None:
+    if isinstance(value, datetime):
+        pb.text_value = value.strftime("%Y-%m-%dT%H:%M:%S") + "," + _tz_name(value)
+    else:
+        pb.text_value = value
+
+
+def _from_tz_timestamp(x: str, table_client_settings: table.TableClientSettings) -> typing.Union[datetime, str]:
+    if table_client_settings is not None and table_client_settings._native_timestamp_in_result_sets:
+        return _parse_tz(x)
+    return x
+
+
+def _to_tz_timestamp(pb: ydb_value_pb2.Value, value: typing.Union[datetime, str]) -> None:
+    if isinstance(value, datetime):
+        pb.text_value = value.strftime("%Y-%m-%dT%H:%M:%S.%f") + "," + _tz_name(value)
+    else:
+        pb.text_value = value
+
+
 def _from_json(x: typing.Union[str, bytearray, bytes], table_client_settings: table.TableClientSettings) -> typing.Any:
     if table_client_settings is not None and table_client_settings._native_json_in_result_sets:
         return json.loads(x)
@@ -214,6 +276,24 @@ class PrimitiveType(enum.Enum):
         _from_timestamp64,
         _to_timestamp64,
     )
+    TzDate = (
+        _apis.primitive_types.TZ_DATE,
+        "text_value",
+        _from_tz_date,
+        _to_tz_date,
+    )
+    TzDatetime = (
+        _apis.primitive_types.TZ_DATETIME,
+        "text_value",
+        _from_tz_datetime,
+        _to_tz_datetime,
+    )
+    TzTimestamp = (
+        _apis.primitive_types.TZ_TIMESTAMP,
+        "text_value",
+        _from_tz_timestamp,
+        _to_tz_timestamp,
+    )
     Interval = (
         _apis.primitive_types.INTERVAL,
         None,