Merge pull request #51 from ClickHouse/update_clickhouse_driver

Update clickhouse driver

Author: guykoh

.github/workflows/build.yml

@@ -25,9 +25,10 @@ jobs:
 
     services:
       clickhouse:
-        image: yandex/clickhouse-server:${{ matrix.clickhouse-version}}
+        image: clickhouse/clickhouse-server:${{ matrix.clickhouse-version}}
         ports:
           - 9000:9000
+          - 8123:8123
 
     steps:
       - name: Checkout

README.md

@@ -67,7 +67,6 @@ your_profile_name:
       connect_timeout: [10] # default 10
       send_receive_timeout: [300] # default 300
       sync_request_timeout: [5] # default 5
-      compress_block_size: [1048576] # default 1048576
       compression: ['lz4'] # default '' (disable)
 ```
 

dbt/adapters/clickhouse/__version__.py

@@ -1 +1 @@
-version = '1.0.5'
+version = '1.1.0'

dbt/adapters/clickhouse/connections.py

@@ -4,19 +4,25 @@
 from typing import Any, Optional, Tuple
 
 import agate
+import clickhouse_connect
 import dbt.exceptions
-from clickhouse_driver import Client, errors
+from clickhouse_connect.driver.client import Client as ChClient
+from clickhouse_connect.driver.exceptions import DatabaseError, Error
 from dbt.adapters.base import Credentials
 from dbt.adapters.sql import SQLConnectionManager
 from dbt.contracts.connection import Connection
 from dbt.events import AdapterLogger
-from dbt.version import __version__ as dbt_version
 
 logger = AdapterLogger('clickhouse')
 
 
 @dataclass
 class ClickhouseCredentials(Credentials):
+    """
+    ClickHouse connectio credentials data class.
+    """
+
+    # pylint: disable=too-many-instance-attributes
     host: str = 'localhost'
     port: Optional[int] = None
     user: Optional[str] = 'default'
@@ -52,37 +58,38 @@ def __post_init__(self):
         self.database = None
 
     def _connection_keys(self):
-        return ('host', 'port', 'user', 'schema', 'secure', 'verify')
+        return 'host', 'port', 'user', 'schema', 'secure', 'verify'
 
 
 class ClickhouseConnectionManager(SQLConnectionManager):
+    """
+    ClickHouse Connector connection manager.
+    """
+
     TYPE = 'clickhouse'
 
     @contextmanager
     def exception_handler(self, sql):
         try:
             yield
-
-        except errors.ServerException as e:
-            logger.debug('Clickhouse error: {}', str(e))
-
+        except DatabaseError as err:
+            logger.debug('Clickhouse error: {}', str(err))
             try:
                 # attempt to release the connection
                 self.release()
-            except errors.Error:
+            except Error:
                 logger.debug('Failed to release connection!')
-                pass
 
-            raise dbt.exceptions.DatabaseException(str(e).strip()) from e
+            raise dbt.exceptions.DatabaseException(str(err).strip()) from err
 
-        except Exception as e:
+        except Exception as exp:
             logger.debug('Error running SQL: {}', sql)
             logger.debug('Rolling back transaction.')
             self.release()
-            if isinstance(e, dbt.exceptions.RuntimeException):
+            if isinstance(exp, dbt.exceptions.RuntimeException):
                 raise
 
-            raise dbt.exceptions.RuntimeException(e) from e
+            raise dbt.exceptions.RuntimeException(exp) from exp
 
     @classmethod
     def open(cls, connection):
@@ -94,34 +101,32 @@ def open(cls, connection):
         kwargs = {}
 
         try:
-            handle = Client(
+            handle = clickhouse_connect.get_client(
                 host=credentials.host,
                 port=credentials.port,
                 database='default',
-                user=credentials.user,
+                username=credentials.user,
                 password=credentials.password,
-                client_name=f'dbt-{dbt_version}',
-                secure=credentials.secure,
-                verify=credentials.verify,
+                interface='https' if credentials.secure else 'http',
+                compress=False if credentials.compression == '' else bool(credentials.compression),
                 connect_timeout=credentials.connect_timeout,
                 send_receive_timeout=credentials.send_receive_timeout,
-                sync_request_timeout=credentials.sync_request_timeout,
-                compress_block_size=credentials.compress_block_size,
-                compression=False if credentials.compression == '' else credentials.compression,
+                verify=credentials.verify,
+                query_limit=0,
                 **kwargs,
             )
             connection.handle = handle
             connection.state = 'open'
-        except errors.ServerException as e:
+        except DatabaseError as err:
             logger.debug(
                 'Got an error when attempting to open a clickhouse connection: \'{}\'',
-                str(e),
+                str(err),
             )
 
             connection.handle = None
             connection.state = 'fail'
 
-            raise dbt.exceptions.FailedToConnectException(str(e))
+            raise dbt.exceptions.FailedToConnectException(str(err))
 
         return connection
 
@@ -135,9 +140,12 @@ def cancel(self, connection):
         logger.debug('Cancel query \'{}\'', connection_name)
 
     @classmethod
-    def get_table_from_response(cls, response, columns) -> agate.Table:
-        column_names = [x[0] for x in columns]
-
+    def get_table_from_response(cls, response, column_names) -> agate.Table:
+        """
+        Build agate tabel from response.
+        :param response: ClickHouse query result
+        :param column_names: Table column names
+        """
         data = []
         for row in response:
             data.append(dict(zip(column_names, row)))
@@ -152,28 +160,37 @@ def execute(
         client = conn.handle
 
         with self.exception_handler(sql):
-            logger.debug(
-                'On {connection_name}: {sql}'.format(connection_name=conn.name, sql=f'{sql}...'),
-            )
+            logger.debug(f'On {conn.name}: {sql}...')
 
             pre = time.time()
 
-            response, columns = client.execute(sql, with_column_types=True)
+            if fetch:
+                query_result = client.query(sql)
+            else:
+                query_result = client.command(sql)
 
             status = self.get_status(client)
 
-            logger.debug(
-                'SQL status: {status} in {elapsed:0.2f} seconds'.format(
-                    status=status, elapsed=(time.time() - pre)
-                ),
-            )
+            logger.debug(f'SQL status: {status} in {(time.time() - pre):.2f} seconds')
 
             if fetch:
-                table = self.get_table_from_response(response, columns)
+                table = self.get_table_from_response(
+                    query_result.result_set, query_result.column_names
+                )
             else:
                 table = dbt.clients.agate_helper.empty_table()
             return status, table
 
+    def insert_table_data(self, table_name, table: agate.Table):
+        """
+        Insert data into ClickHouse table
+        :param table_name: Target table name
+        :param table: Data to be inserted
+        """
+        client: ChClient = self.get_thread_connection().handle
+        with self.exception_handler(f'INSERT INTO {table_name}'):
+            client.insert(table_name, table.rows, table.column_names)
+
     def add_query(
         self,
         sql: str,
@@ -186,33 +203,33 @@ def add_query(
         client = conn.handle
 
         with self.exception_handler(sql):
-            logger.debug(
-                'On {connection_name}: {sql}'.format(connection_name=conn.name, sql=f'{sql}...')
-            )
+            logger.debug(f'On {conn.name}: {sql}...')
 
             pre = time.time()
-            client.execute(sql)
+            client.query(sql)
 
             status = self.get_status(client)
 
-            logger.debug(
-                'SQL status: {status} in {elapsed:0.2f} seconds'.format(
-                    status=status, elapsed=(time.time() - pre)
-                )
-            )
+            logger.debug(f'SQL status: {status} in {(time.time() - pre):0.2f} seconds')
 
             return conn, None
 
     @classmethod
     def get_credentials(cls, credentials):
+        """
+        Returns ClickHouse credentials
+        """
         return credentials
 
     @classmethod
-    def get_status(cls, cursor):
+    def get_status(cls, _):
+        """
+        Returns connection status
+        """
         return 'OK'
 
     @classmethod
-    def get_response(cls, cursor):
+    def get_response(cls, _):
         return 'OK'
 
     def begin(self):

dbt/adapters/clickhouse/impl.py

@@ -235,16 +235,21 @@ def get_csv_data(self, table):
 
         return buf.getvalue()
 
+    @available
+    def insert_table_data(self, table_name, table):
+        self.connections.insert_table_data(table_name, table)
+
     def run_sql_for_tests(self, sql, fetch, conn):
-        cursor = conn.handle
+        client = conn.handle
         try:
-            result = cursor.execute(sql)
+            if fetch:
+                result = client.query(sql).result_set
+            else:
+                result = client.command(sql)
             if fetch == "one" and len(result) > 0:
                 return result[0]
-            elif fetch == "all":
+            if fetch == "all":
                 return result
-            else:
-                return
         except BaseException as e:
             self.logger.error(sql)
             self.logger.error(e)

dbt/include/clickhouse/macros/catalog.sql

@@ -20,7 +20,7 @@
       {%- if not loop.last %} or {% endif -%}
     {%- endfor -%}
     )
-    order by columns.database, columns.table, columns.position;
+    order by columns.database, columns.table, columns.position
   {%- endcall -%}
   {{ return(load_result('catalog').table) }}
 {%- endmacro %}

dbt/include/clickhouse/macros/materializations/seed.sql

@@ -1,13 +1,5 @@
 {% macro clickhouse__load_csv_rows(model, agate_table) %}
-  {% set cols_sql = get_seed_column_quoted_csv(model, agate_table.column_names) %}
-  {% set data_sql = adapter.get_csv_data(agate_table) %}
-
-  {% set sql -%}
-      insert into {{ this.render() }} ({{ cols_sql }}) format CSV
-      {{ data_sql }}
-  {%- endset %}
-
-  {% do adapter.add_query(sql, bindings=agate_table, abridge_sql_log=True) %}
+  {% do adapter.insert_table_data(this.render(), agate_table) %}
 {% endmacro %}
 
 {% macro clickhouse__create_csv_table(model, agate_table) %}

dev_requirements.txt

@@ -1,7 +1,7 @@
 dbt-core==1.1.0
-clickhouse-driver==0.2.2
+clickhouse-connect>=0.0.14
 pytest==7.0.0
-pytest-dotenv
+pytest-dotenv==0.5.2
 dbt-tests-adapter==1.1
 black==22.3.0
 isort==5.10.1

pyproject.toml

@@ -9,3 +9,7 @@ line_length = 100
 profile = "black"
 use_parentheses = true
 skip = '.eggs/,.mypy_cache/,.venv/,venv/,env/'
+
+[tool.pytest.ini_options]
+log_cli = true
+log_cli_level = "WARNING"

setup.py

@@ -54,17 +54,16 @@ def _dbt_clickhouse_version():
     },
     install_requires=[
         f'dbt-core=={dbt_version}',
-        'clickhouse-driver>=0.2.2',
+        'clickhouse-connect>=0.0.14',
     ],
-    python_requires=">=3.6",
+    python_requires=">=3.7",
     platforms='any',
     classifiers=[
         'Development Status :: 5 - Production/Stable',
         'License :: OSI Approved :: Apache Software License',
         'Operating System :: Microsoft :: Windows',
         'Operating System :: MacOS :: MacOS X',
         'Operating System :: POSIX :: Linux',
-        'Programming Language :: Python :: 3.6',
         'Programming Language :: Python :: 3.7',
         'Programming Language :: Python :: 3.8',
     ],

tests/conftest.py

@@ -27,7 +27,7 @@ def test_config():
             up_result = run_cmd(['docker-compose', '-f', compose_file, 'up', '-d'])
             if up_result[0]:
                 raise Exception(f'Failed to start docker: {up_result[2]}')
-            url = "http://{}:{}".format(os.environ.get('HOST_ENV_VAR_NAME'), 10723)
+            url = "http://{}:{}".format(os.environ.get('HOST_ENV_VAR_NAME', 'localhost'), 10723)
             wait_until_responsive(timeout=30.0, pause=0.5, check=lambda: is_responsive(url))
         except Exception as e:
             raise Exception('Failed to run docker-compose: {}', str(e))
@@ -50,7 +50,7 @@ def dbt_profile_target():
         'host': os.environ.get('HOST_ENV_VAR_NAME', 'localhost'),
         'user': os.environ.get('USER_ENV_VAR_NAME', 'default'),
         'password': os.environ.get('PASSWORD_ENV_VAR_NAME', ''),
-        'port': int(os.environ.get('PORT_ENV_VAR_NAME', 9000)),  # docker client port
+        'port': int(os.environ.get('PORT_ENV_VAR_NAME', 8123)),  # docker client port
         'secure': False,
     }