Merge pull request #72 from ClickHouse/database_fix

Ensure database exists before configuring driver

Author: genzgd

Diff for: dbt/adapters/clickhouse/__version__.py

@@ -1 +1 @@
-version = '1.1.4'
+version = '1.1.5'

Diff for: dbt/adapters/clickhouse/connections.py

@@ -51,9 +51,6 @@ class ClickhouseCredentials(Credentials):
     sync_request_timeout: int = 5
     compress_block_size: int = 1048576
     compression: str = ''
-    use_default_schema: bool = (
-        False  # This is used in tests to make sure we connect always to the default database.
-    )
     custom_settings: Optional[Dict[str, Any]] = None
 
     @property
@@ -89,7 +86,6 @@ def _connection_keys(self):
             'sync_request_timeout',
             'compress_block_size',
             'compression',
-            'use_default_schema',
             'custom_settings',
         )
 
@@ -132,63 +128,22 @@ def open(cls, connection):
                 driver = 'http'
             elif clickhouse_driver and port in (9000, 9440):
                 driver = 'native'
-            else:
-                driver = 'unspecified'
-        custom_settings = credentials.custom_settings or {}
-        connection.state = 'fail'
-        db_err = None
+        client = None
         if clickhouse_connect and driver == 'http':
-            try:
-                connection.handle = clickhouse_connect.get_client(
-                    host=credentials.host,
-                    port=credentials.port,
-                    database='default' if credentials.use_default_schema else credentials.schema,
-                    username=credentials.user,
-                    password=credentials.password,
-                    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,
-                    http_user_agent=f'cc-dbt-{dbt_version}',
-                    verify=credentials.verify,
-                    query_limit=0,
-                    session_id='dbt::' + str(uuid.uuid4()),
-                    **custom_settings,
-                )
-            except clickhouse_connect.driver.exceptions.DatabaseError as exp:
-                db_err = exp
+            client, db_err = _connect_http(credentials)
         elif clickhouse_driver and driver == 'native':
-            try:
-                client = clickhouse_driver.Client(
-                    host=credentials.host,
-                    port=credentials.port,
-                    database='default',
-                    user=credentials.user,
-                    password=credentials.password,
-                    client_name=f'dbt-{dbt_version}',
-                    secure=credentials.secure,
-                    verify=credentials.verify,
-                    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,
-                    **custom_settings,
-                )
-                connection.handle = ChNativeAdapter(client)
-            except clickhouse_driver.errors as exp:
-                db_err = exp
+            client, db_err = _connect_native(credentials)
         else:
-            raise dbt.exceptions.FailedToConnectException(
-                f'Library for ClickHouse driver type {driver} not found'
-            )
+            db_err = f'Library for ClickHouse driver type {driver} not found'
         if db_err:
+            connection.state = 'fail'
             logger.debug(
                 'Got an error when attempting to open a clickhouse connection: \'{}\'', str(db_err)
             )
+            if client:
+                client.close()
             raise dbt.exceptions.FailedToConnectException(str(db_err))
+        connection.handle = client
         connection.state = 'open'
         return connection
 
@@ -286,3 +241,61 @@ def begin(self):
 
     def commit(self):
         pass
+
+
+def _connect_http(credentials):
+    try:
+        client = clickhouse_connect.get_client(
+            host=credentials.host,
+            port=credentials.port,
+            username=credentials.user,
+            password=credentials.password,
+            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,
+            client_name=f'cc-dbt-{dbt_version}',
+            verify=credentials.verify,
+            query_limit=0,
+            session_id='dbt::' + str(uuid.uuid4()),
+            **(credentials.custom_settings or {}),
+        )
+        return client, _ensure_database(client, credentials.schema)
+    except clickhouse_connect.driver.exceptions.DatabaseError as exp:
+        return None, exp
+
+
+def _connect_native(credentials):
+    try:
+        client = clickhouse_driver.Client(
+            host=credentials.host,
+            port=credentials.port,
+            user=credentials.user,
+            password=credentials.password,
+            client_name=f'dbt-{dbt_version}',
+            secure=credentials.secure,
+            verify=credentials.verify,
+            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,
+            **(credentials.custom_settings or {}),
+        )
+        client = ChNativeAdapter(client)
+        return client, _ensure_database(client, credentials.schema)
+    except clickhouse_driver.errors.Error as exp:
+        return None, exp
+
+
+def _ensure_database(client, database):
+    if database:
+        check_db = f'EXISTS DATABASE {database}'
+        db_exists = client.command(check_db)
+        if not db_exists:
+            client.command(f'CREATE DATABASE {database}')
+            db_exists = client.command(check_db)
+            if not db_exists:
+                return f'Unable to create DBT database {database}'
+        client.database = database
+    return None

Diff for: dbt/adapters/clickhouse/impl.py

@@ -8,7 +8,7 @@
 import dbt.exceptions
 from dbt.adapters.base import AdapterConfig, available
 from dbt.adapters.base.impl import catch_as_completed
-from dbt.adapters.base.relation import InformationSchema
+from dbt.adapters.base.relation import BaseRelation, InformationSchema
 from dbt.adapters.sql import SQLAdapter
 from dbt.clients.agate_helper import table_from_rows
 from dbt.contracts.graph.manifest import Manifest
@@ -82,6 +82,12 @@ def check_schema_exists(self, database, schema):
         exists = True if schema in [row[0] for row in results] else False
         return exists
 
+    def drop_schema(self, relation: BaseRelation) -> None:
+        super().drop_schema(relation)
+        conn = self.connections.get_if_exists()
+        if conn:
+            conn.handle.database = None
+
     def list_relations_without_caching(
         self, schema_relation: ClickhouseRelation
     ) -> List[ClickhouseRelation]:

Diff for: dbt/adapters/clickhouse/nativeadapter.py

@@ -5,15 +5,25 @@ class ChNativeAdapter:
     def __init__(self, client: clickhouse_driver.Client):
         self.client = client
 
-    def query(self, sql):
-        return NativeAdapterResult(self.client.execute(sql, with_column_types=True))
+    def query(self, sql, **kwargs):
+        return NativeAdapterResult(self.client.execute(sql, with_column_types=True, **kwargs))
 
-    def command(self, sql):
-        self.client.execute(sql)
+    def command(self, sql, **kwargs):
+        result = self.client.execute(sql, **kwargs)
+        if len(result) and len(result[0]):
+            return result[0][0]
 
     def close(self):
         self.client.disconnect()
 
+    @property
+    def database(self):
+        return self.client.connection.database
+
+    @database.setter
+    def database(self, database):
+        self.client.connection.database = database
+
 
 class NativeAdapterResult:
     def __init__(self, native_result):

Diff for: dev_requirements.txt

@@ -1,5 +1,5 @@
 dbt-core==1.1.0
-clickhouse-connect>=0.1.4
+clickhouse-connect>=0.1.5
 clickhouse-driver>=0.2.3
 pytest==7.0.0
 pytest-dotenv==0.5.2

Diff for: setup.py

@@ -54,7 +54,7 @@ def _dbt_clickhouse_version():
     },
     install_requires=[
         f'dbt-core~={dbt_version}',
-        'clickhouse-connect>=0.1.4',
+        'clickhouse-connect>=0.1.5',
         'clickhouse-driver>=0.2.3',
     ],
     python_requires=">=3.7",

Diff for: tests/conftest.py

@@ -58,7 +58,6 @@ def dbt_profile_target():
         'port': int(os.environ.get('PORT_ENV_VAR_NAME', 8123)),  # docker client port
         'secure': False,
         'driver': driver,
-        'use_default_schema': True,  # In tests we always use default schema (a.k.a default database)
     }