Merge branch 'xzkostyan:master' into master

Author: Simon-Chenzw

clickhouse_sqlalchemy/drivers/base.py

@@ -50,6 +50,7 @@
     'FixedString': types.String,
     'Enum8': types.Enum8,
     'Enum16': types.Enum16,
+    'Object(\'json\')': types.JSON,
     '_array': types.Array,
     '_nullable': types.Nullable,
     '_lowcardinality': types.LowCardinality,
@@ -135,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)
 

clickhouse_sqlalchemy/drivers/compilers/typecompiler.py

@@ -87,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 = [
@@ -121,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):

clickhouse_sqlalchemy/drivers/native/connector.py

@@ -131,7 +131,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/types/__init__.py

@@ -31,6 +31,7 @@
     'Decimal',
     'IPv4',
     'IPv6',
+    'JSON',
     'Nested',
     'Tuple',
     'Map',
@@ -68,6 +69,7 @@
 from .common import Enum8
 from .common import Enum16
 from .common import Decimal
+from .common import JSON
 from .common import Tuple
 from .common import Map
 from .common import AggregateFunction

clickhouse_sqlalchemy/types/common.py

@@ -30,6 +30,10 @@ class Boolean(types.Boolean, ClickHouseTypeEngine):
     pass
 
 
+class JSON(types.JSON, ClickHouseTypeEngine):
+    __visit_name__ = 'json'
+
+
 class Array(ClickHouseTypeEngine):
     __visit_name__ = 'array'
 

tests/drivers/native/test_cursor.py

@@ -1,3 +1,5 @@
+import uuid
+
 from sqlalchemy import text
 
 from tests.testcase import NativeSessionTestCase
@@ -47,3 +49,14 @@ def test_with_settings_in_execution_options(self):
             dict(rv.context.execution_options), {"settings": {"final": 1}}
         )
         self.assertEqual(len(rv.fetchall()), 1)
+
+    def test_set_query_id(self):
+        query_id = str(uuid.uuid4())
+        rv = self.session.execute(
+            text(
+                f"SELECT query_id "
+                f"FROM system.processes "
+                f"WHERE query_id = '{query_id}'"
+            ), execution_options={'query_id': query_id}
+        )
+        self.assertEqual(rv.fetchall()[0][0], query_id)

tests/test_ddl.py

@@ -296,6 +296,26 @@ def test_create_table_tuple(self):
             'ENGINE = Memory'
         )
 
+    def test_create_table_named_tuple(self):
+        table = Table(
+            't1', self.metadata(),
+            Column(
+                'x',
+                types.Tuple(
+                    ('name', types.String),
+                    ('value', types.Float32)
+                )
+            ),
+            engines.Memory()
+        )
+
+        self.assertEqual(
+            self.compile(CreateTable(table)),
+            'CREATE TABLE t1 ('
+            'x Tuple(name String, value Float32)) '
+            'ENGINE = Memory'
+        )
+
     @require_server_version(21, 1, 3)
     def test_create_table_map(self):
         table = Table(

tests/types/test_json.py

@@ -0,0 +1,61 @@
+import json
+from sqlalchemy import Column, text, inspect, func
+from sqlalchemy.sql.ddl import CreateTable
+
+from clickhouse_sqlalchemy import types, engines, Table
+from tests.testcase import BaseTestCase, CompilationTestCase
+from tests.util import class_name_func
+from parameterized import parameterized_class
+from tests.session import native_session
+
+
+class JSONCompilationTestCase(CompilationTestCase):
+    def test_create_table(self):
+        table = Table(
+            'test', CompilationTestCase.metadata(),
+            Column('x', types.JSON),
+            engines.Memory()
+        )
+
+        self.assertEqual(
+            self.compile(CreateTable(table)),
+            'CREATE TABLE test (x JSON) ENGINE = Memory'
+        )
+
+
+@parameterized_class(
+    [{'session': native_session}],
+    class_name_func=class_name_func
+)
+class JSONTestCase(BaseTestCase):
+    required_server_version = (22, 3, 2)
+
+    table = Table(
+        'test', BaseTestCase.metadata(),
+        Column('x', types.JSON),
+        engines.Memory()
+    )
+
+    def test_select_insert(self):
+        data = {'k1': 1, 'k2': '2', 'k3': True}
+
+        self.table.drop(bind=self.session.bind, if_exists=True)
+        try:
+            # http session is unsupport
+            self.session.execute(
+                text('SET allow_experimental_object_type = 1;')
+            )
+            self.session.execute(text(self.compile(CreateTable(self.table))))
+            self.session.execute(self.table.insert(), [{'x': data}])
+            coltype = inspect(self.session.bind).get_columns('test')[0]['type']
+            self.assertIsInstance(coltype, types.JSON)
+            # https://clickhouse.com/docs/en/sql-reference/functions/json-functions#tojsonstring
+            # The json type returns a tuple of values by default,
+            # which needs to be converted to json using the
+            # toJSONString function.
+            res = self.session.query(
+                func.toJSONString(self.table.c.x)
+            ).scalar()
+            self.assertEqual(json.loads(res), data)
+        finally:
+            self.table.drop(bind=self.session.bind, if_exists=True)