Merge pull request #13 from ods/clickhouse-20.10

Adopt to ClickHouse versions up to 20.10

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

@@ -11,7 +11,8 @@ license = MIT
 packages =
     aiochsa
 install_requires =
-    aiohttp>=3.6.2,<4.0.0
+    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
@@ -34,11 +35,7 @@ test = pytest
 [tool:pytest]
 testpaths = tests
 addopts = --strict-markers -r aP --tb=native --cov
-filterwarnings =
-    error
-    # The loop argument is deprecated since Python 3.8, and scheduled for
-    # removal in Python 3.10.
-    ignore::DeprecationWarning:aiohttp.connector
+filterwarnings = error
 
 [coverage:run]
 branch = True

tests/conftest.py

@@ -1,6 +1,7 @@
 import asyncio
 from datetime import datetime
 from decimal import Decimal
+import os
 
 import clickhouse_sqlalchemy
 import pytest
@@ -17,13 +18,34 @@ def pytest_collection_modifyitems(items):
             item.add_marker('asyncio')
 
 
+def pytest_addoption(parser):
+    parser.addoption(
+        '--clickhouse-version', default=None,
+        help='use specified version of ClickHouse',
+    )
+
+
 @pytest.fixture(scope='session')
-def dsn(docker_services):
+def dsn(request, docker_services):
+    ch_version = request.config.getoption('--clickhouse-version')
+    if ch_version is not None:
+        os.environ['CLICKHOUSE_VERSION'] = ch_version
     docker_services.start('clickhouse-server')
     public_port = docker_services.wait_for_service('clickhouse-server', 8123)
     return f'clickhouse://{docker_services.docker_ip}:{public_port}'
 
 
+@pytest.fixture
+async def clickhouse_version(dsn):
+    if not hasattr(pytest, 'clickhouse_version'):
+        async with aiochsa.connect(dsn) as conn:
+            version = await conn.fetchval('SELECT version()')
+            pytest.clickhouse_version = tuple(
+                int(num) for num in version.split('.')
+            )
+    return pytest.clickhouse_version
+
+
 @pytest.fixture
 async def pool(dsn):
     async with aiochsa.create_pool(dsn) as pool:

tests/docker-compose.yml

@@ -3,6 +3,6 @@ version: '2.1'
 services:
 
   clickhouse-server:
-    image: yandex/clickhouse-server
+    image: "yandex/clickhouse-server:${CLICKHOUSE_VERSION:-latest}"
     ports:
       - 8123

tests/test_exc.py

@@ -34,7 +34,7 @@ async def test_exc_stacktrace(dsn, statement):
 
     assert str(error_codes.SYNTAX_ERROR) in str(exc_info.value)
     assert 'Syntax error' in str(exc_info.value)
-    assert len(str(exc_info.value)) < 2000
+    assert len(str(exc_info.value)) < len(exc_info.value.display_text) + 250
 
 
 async def test_exc_row(conn, table_test1):

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(*[
@@ -175,13 +176,29 @@ async def test_as_is_round(conn, value):
     ],
     ids = parametrized_id,
 )
-async def test_zero_dates(conn, sa_type, value):
+async def test_zero_dates(clickhouse_version, conn, sa_type, value):
+    if clickhouse_version >= (20, 7):
+        pytest.skip('Feature is dropped in 20.7')
     result = await conn.fetchval(
         sa.select([sa.func.cast(value, sa_type)])
     )
     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()
@@ -246,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(