Merge branch 'master' into github_actions

.github/workflows/actions.yml

@@ -13,6 +13,7 @@ jobs:
           - "3.11"
           - "3.12"
         clickhouse-version:
+          - 24.1.8.22
           - 23.8.4.69
           - 22.5.1.2079
           - 19.3.5
@@ -77,7 +78,7 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - name: Finished
-        uses: coverallsapp/github-action@v2.2.3
+        uses: coverallsapp/github-action@v2.3.0
         with:
           github-token: ${{ secrets.GITHUB_TOKEN }}
           parallel-finished: true

CHANGELOG.md

@@ -2,6 +2,15 @@
 
 ## [Unreleased]
 
+## [0.3.2] - 2024-06-12
+### Added
+- ``quantile`` and ``quantileIf`` functions. Pull request [#303](https://github.com/xzkostyan/clickhouse-sqlalchemy/pull/303) by [aronbierbaum](https://github.com/aronbierbaum).
+- ``AggregateFunction`` and ``SimpleAggregateFunction`` aggregate types. Pull request [#297](https://github.com/xzkostyan/clickhouse-sqlalchemy/pull/297) by [aronbierbaum](https://github.com/aronbierbaum).
+- Date32 type. Pull request [#307](https://github.com/xzkostyan/clickhouse-sqlalchemy/pull/307) by [BTheunissen](https://github.com/BTheunissen). Solves issue [#302](https://github.com/xzkostyan/clickhouse-sqlalchemy/issues/302).
+
+### Fixed
+- Broken nested Map types. Pull request [#315](https://github.com/xzkostyan/clickhouse-sqlalchemy/pull/315) by [aksenof](https://github.com/aksenof). Solves issue [#314](https://github.com/xzkostyan/clickhouse-sqlalchemy/issues/314).
+
 ## [0.3.1] - 2024-03-14
 ### Added
 - ``SETTINGS`` clause. Pull request [#292](https://github.com/xzkostyan/clickhouse-sqlalchemy/pull/292) by [limonyellow](https://github.com/limonyellow).
@@ -344,7 +353,8 @@ Log, TinyLog, Null.
 - Chunked `INSERT INTO` in one request.
 - Engines: MergeTree, CollapsingMergeTree, SummingMergeTree, Buffer, Memory. 
 
-[Unreleased]: https://github.com/xzkostyan/clickhouse-sqlalchemy/compare/0.3.1...HEAD
+[Unreleased]: https://github.com/xzkostyan/clickhouse-sqlalchemy/compare/0.3.2...HEAD
+[0.3.2]: https://github.com/xzkostyan/clickhouse-sqlalchemy/compare/0.3.1...0.3.2
 [0.3.1]: https://github.com/xzkostyan/clickhouse-sqlalchemy/compare/0.3.0...0.3.1
 [0.3.0]: https://github.com/xzkostyan/clickhouse-sqlalchemy/compare/0.2.5...0.3.0
 [0.2.5]: https://github.com/xzkostyan/clickhouse-sqlalchemy/compare/0.2.4...0.2.5

clickhouse_sqlalchemy/__init__.py

@@ -4,7 +4,7 @@
 from .sql import Table, MaterializedView, select
 
 
-VERSION = (0, 3, 1)
+VERSION = (0, 3, 2)
 __version__ = '.'.join(str(x) for x in VERSION)
 
 

clickhouse_sqlalchemy/drivers/asynch/base.py

@@ -4,17 +4,24 @@
 from sqlalchemy.pool import AsyncAdaptedQueuePool
 
 from .connector import AsyncAdapt_asynch_dbapi
-from ..native.base import ClickHouseDialect_native
+from ..native.base import ClickHouseDialect_native, ClickHouseExecutionContext
 
 # Export connector version
 VERSION = (0, 0, 1, None)
 
 
+class ClickHouseAsynchExecutionContext(ClickHouseExecutionContext):
+    def create_server_side_cursor(self):
+        return self.create_default_cursor()
+
+
 class ClickHouseDialect_asynch(ClickHouseDialect_native):
     driver = 'asynch'
+    execution_ctx_cls = ClickHouseAsynchExecutionContext
 
     is_async = True
     supports_statement_cache = True
+    supports_server_side_cursors = True
 
     @classmethod
     def import_dbapi(cls):

clickhouse_sqlalchemy/drivers/base.py

@@ -13,7 +13,7 @@
 from .compilers.sqlcompiler import ClickHouseSQLCompiler
 from .compilers.typecompiler import ClickHouseTypeCompiler
 from .reflection import ClickHouseInspector
-from .util import get_inner_spec
+from .util import get_inner_spec, parse_arguments
 from .. import types
 
 # Column specifications
@@ -35,6 +35,7 @@
     'UInt16': types.UInt16,
     'UInt8': types.UInt8,
     'Date': types.Date,
+    'Date32': types.Date32,
     'DateTime': types.DateTime,
     'DateTime64': types.DateTime64,
     'Float64': types.Float64,
@@ -49,11 +50,14 @@
     'FixedString': types.String,
     'Enum8': types.Enum8,
     'Enum16': types.Enum16,
+    'Object(\'json\')': types.JSON,
     '_array': types.Array,
     '_nullable': types.Nullable,
     '_lowcardinality': types.LowCardinality,
     '_tuple': types.Tuple,
     '_map': types.Map,
+    '_aggregatefunction': types.AggregateFunction,
+    '_simpleaggregatefunction': types.SimpleAggregateFunction,
 }
 
 
@@ -132,6 +136,16 @@ class ClickHouseDialect(default.DefaultDialect):
 
     inspector = ClickHouseInspector
 
+    def __init__(
+        self,
+        json_serializer=None,
+        json_deserializer=None,
+        **kwargs,
+    ):
+        default.DefaultDialect.__init__(self, **kwargs)
+        self._json_deserializer = json_deserializer
+        self._json_serializer = json_serializer
+
     def initialize(self, connection):
         super(ClickHouseDialect, self).initialize(connection)
 
@@ -230,6 +244,32 @@ def _get_column_type(self, name, spec):
             coltype = self.ischema_names['_lowcardinality']
             return coltype(self._get_column_type(name, inner))
 
+        elif spec.startswith('AggregateFunction'):
+            params = spec[18:-1]
+
+            arguments = parse_arguments(params)
+            agg_func, inner = arguments[0], arguments[1:]
+
+            inner_types = [
+                self._get_column_type(name, param)
+                for param in inner
+            ]
+            coltype = self.ischema_names['_aggregatefunction']
+            return coltype(agg_func, *inner_types)
+
+        elif spec.startswith('SimpleAggregateFunction'):
+            params = spec[24:-1]
+
+            arguments = parse_arguments(params)
+            agg_func, inner = arguments[0], arguments[1:]
+
+            inner_types = [
+                self._get_column_type(name, param)
+                for param in inner
+            ]
+            coltype = self.ischema_names['_simpleaggregatefunction']
+            return coltype(agg_func, *inner_types)
+
         elif spec.startswith('Tuple'):
             inner = spec[6:-1]
             coltype = self.ischema_names['_tuple']
@@ -244,7 +284,7 @@ def _get_column_type(self, name, spec):
             coltype = self.ischema_names['_map']
             inner_types = [
                 self._get_column_type(name, t.strip())
-                for t in inner.split(',')
+                for t in inner.split(',', 1)
             ]
             return coltype(*inner_types)
 

clickhouse_sqlalchemy/drivers/compilers/sqlcompiler.py

@@ -4,6 +4,8 @@
 from sqlalchemy.sql import type_api
 from sqlalchemy.util import inspect_getfullargspec
 
+import clickhouse_sqlalchemy.sql.functions  # noqa:F401
+
 from ... import types
 
 

clickhouse_sqlalchemy/drivers/compilers/typecompiler.py

@@ -60,6 +60,9 @@ def visit_uint256(self, type_, **kw):
     def visit_date(self, type_, **kw):
         return 'Date'
 
+    def visit_date32(self, type_, **kw):
+        return 'Date32'
+
     def visit_datetime(self, type_, **kw):
         if type_.timezone:
             return "DateTime('%s')" % type_.timezone
@@ -84,6 +87,9 @@ def visit_numeric(self, type_, **kw):
     def visit_boolean(self, type_, **kw):
         return 'Bool'
 
+    def visit_json(self, type_, **kw):
+        return 'JSON'
+
     def visit_nested(self, nested, **kwargs):
         ddl_compiler = self.dialect.ddl_compiler(self.dialect, None)
         cols_create = [
@@ -118,10 +124,26 @@ def visit_ipv6(self, type_, **kw):
         return 'IPv6'
 
     def visit_tuple(self, type_, **kw):
-        cols = (
-            self.process(type_api.to_instance(nested_type), **kw)
+        cols = []
+        is_named_type = all([
+            isinstance(nested_type, tuple) and len(nested_type) == 2
             for nested_type in type_.nested_types
-        )
+        ])
+        if is_named_type:
+            for nested_type in type_.nested_types:
+                name = nested_type[0]
+                name_type = nested_type[1]
+                inner_type = self.process(
+                    type_api.to_instance(name_type),
+                    **kw
+                )
+                cols.append(
+                    f'{name} {inner_type}')
+        else:
+            cols = (
+                self.process(type_api.to_instance(nested_type), **kw)
+                for nested_type in type_.nested_types
+            )
         return 'Tuple(%s)' % ', '.join(cols)
 
     def visit_map(self, type_, **kw):
@@ -131,3 +153,29 @@ def visit_map(self, type_, **kw):
             self.process(key_type, **kw),
             self.process(value_type, **kw)
         )
+
+    def visit_aggregatefunction(self, type_, **kw):
+        types = [type_api.to_instance(val) for val in type_.nested_types]
+        type_strings = [self.process(val, **kw) for val in types]
+
+        if isinstance(type_.agg_func, str):
+            agg_str = type_.agg_func
+        else:
+            agg_str = str(type_.agg_func.compile(dialect=self.dialect))
+
+        return "AggregateFunction(%s, %s)" % (
+            agg_str, ", ".join(type_strings)
+        )
+
+    def visit_simpleaggregatefunction(self, type_, **kw):
+        types = [type_api.to_instance(val) for val in type_.nested_types]
+        type_strings = [self.process(val, **kw) for val in types]
+
+        if isinstance(type_.agg_func, str):
+            agg_str = type_.agg_func
+        else:
+            agg_str = str(type_.agg_func.compile(dialect=self.dialect))
+
+        return "SimpleAggregateFunction(%s, %s)" % (
+            agg_str, ", ".join(type_strings)
+        )

clickhouse_sqlalchemy/drivers/http/escaper.py

@@ -1,6 +1,7 @@
 from datetime import date, datetime
 from decimal import Decimal
 import enum
+import uuid
 
 
 class Escaper(object):
@@ -49,6 +50,9 @@ def escape_datetime64(self, item):
     def escape_decimal(self, item):
         return float(item)
 
+    def escape_uuid(self, item):
+        return str(item)
+
     def escape_item(self, item):
         if item is None:
             return 'NULL'
@@ -75,5 +79,7 @@ def escape_item(self, item):
             ) + "}"
         elif isinstance(item, enum.Enum):
             return self.escape_string(item.name)
+        elif isinstance(item, uuid.UUID):
+            return self.escape_uuid(item)
         else:
             raise Exception("Unsupported object {}".format(item))

clickhouse_sqlalchemy/drivers/native/connector.py

@@ -138,7 +138,8 @@ def _prepare(self, context=None):
         execute_kwargs = {
             'settings': settings,
             'external_tables': external_tables,
-            'types_check': execution_options.get('types_check', False)
+            'types_check': execution_options.get('types_check', False),
+            'query_id': execution_options.get('query_id', None)
         }
 
         return execute, execute_kwargs

clickhouse_sqlalchemy/drivers/util.py

@@ -16,3 +16,30 @@ def get_inner_spec(spec):
             break
 
     return spec[offset + 1:i]
+
+
+def parse_arguments(param_string):
+    """
+    Given a string of function arguments, parse them into a tuple.
+    """
+    params = []
+    bracket_level = 0
+    current_param = ''
+
+    for char in param_string:
+        if char == '(':
+            bracket_level += 1
+        elif char == ')':
+            bracket_level -= 1
+        elif char == ',' and bracket_level == 0:
+            params.append(current_param.strip())
+            current_param = ''
+            continue
+
+        current_param += char
+
+    # Append the last parameter
+    if current_param:
+        params.append(current_param.strip())
+
+    return tuple(params)

clickhouse_sqlalchemy/sql/functions.py

@@ -0,0 +1,61 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, TypeVar
+
+from sqlalchemy.ext.compiler import compiles
+from sqlalchemy.sql import coercions, roles
+from sqlalchemy.sql.elements import ColumnElement
+from sqlalchemy.sql.functions import GenericFunction
+
+from clickhouse_sqlalchemy import types
+
+if TYPE_CHECKING:
+    from sqlalchemy.sql._typing import _ColumnExpressionArgument
+
+_T = TypeVar('_T', bound=Any)
+
+
+class quantile(GenericFunction[_T]):
+    inherit_cache = True
+
+    def __init__(
+        self, level: float, expr: _ColumnExpressionArgument[Any],
+        condition: _ColumnExpressionArgument[Any] = None, **kwargs: Any
+    ):
+        arg: ColumnElement[Any] = coercions.expect(
+            roles.ExpressionElementRole, expr, apply_propagate_attrs=self
+        )
+
+        args = [arg]
+        if condition is not None:
+            condition = coercions.expect(
+                roles.ExpressionElementRole, condition,
+                apply_propagate_attrs=self
+            )
+            args.append(condition)
+
+        self.level = level
+
+        if isinstance(arg.type, (types.Decimal, types.Float, types.Int)):
+            return_type = types.Float64
+        elif isinstance(arg.type, types.DateTime):
+            return_type = types.DateTime
+        elif isinstance(arg.type, types.Date):
+            return_type = types.Date
+        else:
+            return_type = types.Float64
+
+        kwargs['type_'] = return_type
+        kwargs['_parsed_args'] = args
+        super().__init__(arg, **kwargs)
+
+
+class quantileIf(quantile[_T]):
+    inherit_cache = True
+
+
+@compiles(quantile, 'clickhouse')
+@compiles(quantileIf, 'clickhouse')
+def compile_quantile(element, compiler, **kwargs):
+    args_str = compiler.function_argspec(element, **kwargs)
+    return f'{element.name}({element.level}){args_str}'