Add support for timezones (fixes #12)

Author: ods

aiochsa/parser.py

@@ -1,3 +1,4 @@
+from collections import namedtuple
 import pkgutil
 import simplejson as json
 from typing import Iterable
@@ -21,6 +22,9 @@
 )
 
 
+EnumOption = namedtuple('EnumOption', ['label', 'value'])
+
+
 @v_args(inline=True)
 class TypeTransformer(Transformer):
 
@@ -37,7 +41,17 @@ def aggregate_type(self, name, func, type_):
         return self._types[name](type_)
 
     def simple_type(self, name, *params):
-        return self._types[name]()
+        return self._types[name](*params)
+
+    def enum_param(self, label, value):
+        return EnumOption(label, value)
+
+    def STRING(self, value):
+        assert value[0] == value[-1] == "'"
+        return value[1:-1]
+
+    def INT(self, value):
+        return int(value)
 
 
 def parse_type(types: TypeRegistry, type_str):

aiochsa/type.lark

@@ -3,12 +3,15 @@ start: _type
 _type: composite_type | aggregate_type | simple_type
 composite_type: COMPOSITE_NAME "(" _type ("," _type)* ")"
 aggregate_type: AGGREGATE_NAME "(" SIMPLE_NAME "," _type ")"
-simple_type: SIMPLE_NAME [params]
-params: "(" PARAM ("," PARAM)* ")"
+simple_type: SIMPLE_NAME [_params]
+_params: "(" _param ("," _param)* ")"
+_param: enum_param | STRING | INT
+enum_param: STRING "=" INT
 
 COMPOSITE_NAME: /Tuple|Array|Nullable|LowCardinality/
 AGGREGATE_NAME: /AggregateFunction|SimpleAggregateFunction/
 SIMPLE_NAME: /(?!Tuple|Array|Nullable|LowCardinality|AggregateFunction|SimpleAggregateFunction)\w+/
-PARAM: /('([^\\']|\\.)*'\s*=\s)?-?\d+/
+STRING: /'([^\\']|\\.)*'/
+INT: /-?\d+/
 
 %ignore /[ \t\f\r\n]+/ // whitespace

aiochsa/types.py

@@ -6,6 +6,11 @@
 )
 from uuid import UUID
 
+try:
+    import zoneinfo
+except ImportError:
+    from backports import zoneinfo
+
 
 NoneType = type(None)
 PyType = TypeVar('PyType')
@@ -47,6 +52,9 @@ def __repr__(self):
 class StrType(BaseType):
     py_type = str
 
+    def __init__(self, *params):
+        pass
+
     @classmethod
     def escape(cls, value: str, escape=None) -> str:
         value = value.replace('\\', '\\\\').replace("'", "\\'")
@@ -70,6 +78,9 @@ class FloatType(BaseType):
 class DecimalType(BaseType):
     py_type = Decimal
 
+    def __init__(self, *params):
+        pass
+
     @classmethod
     def escape(cls, value: PyType, escape: Callable) -> str:
         return f"'{value}'"
@@ -101,6 +112,14 @@ def from_json(self, value: str) -> Optional[date]:
 class DateTimeType(BaseType):
     py_type = datetime
 
+    __slots__ = ('_tzinfo',)
+
+    def __init__(self, tz_name=None):
+        if tz_name is None:
+            self._tzinfo = None
+        else:
+            self._tzinfo = zoneinfo.ZoneInfo(tz_name)
+
     @classmethod
     def escape(cls, value: datetime, escape=None) -> str:
         value = value.replace(tzinfo=None, microsecond=0)
@@ -114,7 +133,10 @@ def to_json(cls, value: datetime, to_json: Callable) -> JsonType:
     def from_json(self, value: str) -> Optional[datetime]:
         if value == '0000-00-00 00:00:00':
             return None
-        return datetime.fromisoformat(value)
+        result = datetime.fromisoformat(value)
+        if self._tzinfo is not None:
+            result = result.replace(tzinfo=self._tzinfo)
+        return result
 
 
 class DateTimeUTCType(DateTimeType):
@@ -144,7 +166,11 @@ def to_json(cls, value: datetime, to_json: Callable) -> JsonType:
         return value.isoformat()
 
     def from_json(self, value: str) -> datetime:
-        return datetime.fromisoformat(value).replace(tzinfo=timezone.utc)
+        result = datetime.fromisoformat(value)
+        if self._tzinfo is None:
+            return result.replace(tzinfo=timezone.utc)
+        else:
+            return result.replace(tzinfo=self._tzinfo).astimezone(timezone.utc)
 
 
 class UUIDType(BaseType):

setup.cfg

@@ -12,6 +12,7 @@ packages =
     aiochsa
 install_requires =
     aiohttp>=3.7.2,<4.0.0
+    backports.zoneinfo;python_version<"3.9"
     clickhouse_sqlalchemy>=0.1.4
     lark-parser>=0.7.7
     simplejson>=3.16.0

tests/test_types.py

@@ -1,4 +1,4 @@
-from datetime import date, datetime, timezone
+from datetime import date, datetime, timedelta, timezone
 from decimal import Decimal
 import enum
 from ipaddress import IPv4Address, IPv6Address
@@ -54,6 +54,7 @@ class CustomStr(str):
 
 SaType = Union[sa.types.TypeEngine, Type[sa.types.TypeEngine]]
 
+
 def combine_typed_rapameters(spec_seq: Iterable[Tuple[SaType, Iterable]]):
     return list( # Wrap into list to make it reusable (iterator is one-off)
         itertools.chain(*[
@@ -184,6 +185,20 @@ async def test_zero_dates(clickhouse_version, conn, sa_type, value):
     assert result is None
 
 
+@pytest.mark.parametrize('tz_name,tz_offset', [
+    ('UTC', 0),
+    ('EST', -18_000),
+    ('Europe/Moscow', 10_800),
+])
+async def test_timezones(conn, tz_name, tz_offset):
+    dt = datetime(2020, 1, 1)
+    result = await conn.fetchval(
+        sa.func.toTimeZone(sa.func.toDateTime(dt), tz_name).select()
+    )
+    assert result.utcoffset().total_seconds() == tz_offset
+    assert result.astimezone(timezone.utc).replace(tzinfo=None) == dt
+
+
 @pytest.fixture
 async def conn_utc(dsn):
     types = TypeRegistry()
@@ -248,6 +263,19 @@ async def test_datetime_utc_insert_naive(conn_utc, table_for_type):
         )
 
 
+@pytest.mark.parametrize('tz_name,tz_offset', [
+    ('UTC', 0),
+    ('EST', -18_000),
+    ('Europe/Moscow', 10_800),
+])
+async def test_timezones_with_utc(conn_utc, tz_name, tz_offset):
+    dt = datetime(2020, 1, 1, tzinfo=timezone.utc)
+    result = await conn_utc.fetchval(
+        sa.func.toTimeZone(sa.func.toDateTime(dt, 'UTC'), tz_name).select()
+    )
+    assert result == dt
+
+
 @pytest.mark.parametrize('value', [0, 4294967295])
 async def test_simple_aggregate_function(conn, recreate_table_for_type, value):
     table_name = await recreate_table_for_type(