diff --git a/clickhouse/setup.py b/clickhouse/setup.py index 004a151c7..ec8c48a7d 100644 --- a/clickhouse/setup.py +++ b/clickhouse/setup.py @@ -12,6 +12,7 @@ url="https://github.com/testcontainers/testcontainers-python", install_requires=[ "testcontainers-core", + "clickhouse-connect", "clickhouse-driver", ], python_requires=">=3.7", diff --git a/clickhouse/testcontainers/clickhouse/__init__.py b/clickhouse/testcontainers/clickhouse/__init__.py index b6ffce609..70a46f185 100644 --- a/clickhouse/testcontainers/clickhouse/__init__.py +++ b/clickhouse/testcontainers/clickhouse/__init__.py @@ -13,8 +13,10 @@ import os from typing import Optional +import clickhouse_connect +from clickhouse_connect.driver.exceptions import Error as ClickhouseConnectError import clickhouse_driver -from clickhouse_driver.errors import Error +from clickhouse_driver.errors import Error as ClickhouseDriverError from testcontainers.core.generic import DbContainer from testcontainers.core.waiting_utils import wait_container_is_ready @@ -25,19 +27,25 @@ class ClickHouseContainer(DbContainer): ClickHouse database container. Example: - - The example spins up a ClickHouse database and connects to it using the - :code:`clickhouse-driver`. + The example spins up a ClickHouse database and connects to it + using the :code:`clickhouse-driver`. .. doctest:: - >>> import clickhouse_driver >>> from testcontainers.clickhouse import ClickHouseContainer + >>> import clickhouse_driver >>> with ClickHouseContainer("clickhouse/clickhouse-server:21.8") as clickhouse: ... client = clickhouse_driver.Client.from_url(clickhouse.get_connection_url()) ... client.execute("select 'working'") [('working',)] + + >>> import clickhouse_connect + >>> with ClickHouseContainer("clickhouse/clickhouse-server:21.8", + ... port=8123) as clickhouse: + ... client = clickhouse_connect.get_client(dsn=clickhouse.get_connection_url()) + ... client.command("select 'working'") + 'working' """ CLICKHOUSE_USER = os.environ.get("CLICKHOUSE_USER", "test") @@ -60,10 +68,14 @@ def __init__( self.port_to_expose = port self.with_exposed_ports(self.port_to_expose) - @wait_container_is_ready(Error, EOFError) + @wait_container_is_ready(ClickhouseDriverError, ClickhouseConnectError, EOFError) def _connect(self) -> None: - with clickhouse_driver.Client.from_url(self.get_connection_url()) as client: - client.execute("SELECT version()") + if self.port_to_expose == 8123: + with clickhouse_connect.get_client(dsn=self.get_connection_url()) as client: + client.command("SELECT version()") + else: + with clickhouse_driver.Client.from_url(self.get_connection_url()) as client: + client.execute("SELECT version()") def _configure(self) -> None: self.with_env("CLICKHOUSE_USER", self.CLICKHOUSE_USER)