JWT auth support

.github/workflows/clickhouse_ci.yml

@@ -34,6 +34,7 @@ jobs:
           CLICKHOUSE_CONNECT_TEST_CLOUD: 'True'
           CLICKHOUSE_CONNECT_TEST_HOST: ${{ secrets.INTEGRATIONS_TEAM_TESTS_CLOUD_HOST_SMT }}
           CLICKHOUSE_CONNECT_TEST_PASSWORD: ${{ secrets.INTEGRATIONS_TEAM_TESTS_CLOUD_PASSWORD_SMT }}
+          CLICKHOUSE_CONNECT_TEST_JWT_SECRET: ${{ secrets.CI_JWT_SIGNING_PRIVATE_KEY }}
         run: pytest tests/integration_tests
 
       - name: Run ClickHouse Container (LATEST)

.github/workflows/on_push.yml

@@ -117,7 +117,6 @@ jobs:
         if: "${{ env.CLOUD_HOST != '' }}"
         run: echo "HAS_SECRETS=true" >> $GITHUB_OUTPUT
 
-
   cloud-tests:
     runs-on: ubuntu-latest
     name: Cloud Tests Py=${{ matrix.python-version }}
@@ -152,5 +151,6 @@ jobs:
           CLICKHOUSE_CONNECT_TEST_PORT: 8443
           CLICKHOUSE_CONNECT_TEST_HOST: ${{ secrets.INTEGRATIONS_TEAM_TESTS_CLOUD_HOST_SMT }}
           CLICKHOUSE_CONNECT_TEST_PASSWORD: ${{ secrets.INTEGRATIONS_TEAM_TESTS_CLOUD_PASSWORD_SMT }}
+          CLICKHOUSE_CONNECT_TEST_JWT_SECRET: ${{ secrets.CI_JWT_SIGNING_PRIVATE_KEY }}
           SQLALCHEMY_SILENCE_UBER_WARNING: 1
         run: pytest tests/integration_tests

CHANGELOG.md

@@ -17,6 +17,12 @@ release (0.9.0), unrecognized arguments/keywords for these methods of creating a
 instead of being passed as ClickHouse server settings. This is in conjunction with some refactoring in Client construction.
 The supported method of passing ClickHouse server settings is to prefix such arguments/query parameters with`ch_`.
 
+## Unreleased
+### Improvement
+- Added support for JWT authentication (ClickHouse Cloud feature). 
+It can be set via the `access_token` client configuration option for both sync and async clients.
+NB: do not mix access token and username/password credentials in the configuration; the client will throw an error if both are set.
+
 ## 0.8.11, 2024-12-17
 ### Improvement
 - Support of ISO8601 strings for inserting values to columns with DateTime64 type was added.

clickhouse_connect/driver/__init__.py

@@ -16,6 +16,7 @@ def create_client(*,
                   host: str = None,
                   username: str = None,
                   password: str = '',
+                  access_token: Optional[str] = None,
                   database: str = '__default__',
                   interface: Optional[str] = None,
                   port: int = 0,
@@ -29,7 +30,11 @@ def create_client(*,
 
     :param host: The hostname or IP address of the ClickHouse server. If not set, localhost will be used.
     :param username: The ClickHouse username. If not set, the default ClickHouse user will be used.
+      Should not be set if `access_token` is used.
     :param password: The password for username.
+      Should not be set if `access_token` is used.
+    :param access_token: JWT access token (ClickHouse Cloud feature).
+      Should not be set if `username`/`password` are used.
     :param database:  The default database for the connection. If not set, ClickHouse Connect will use the
      default database for username.
     :param interface: Must be http or https.  Defaults to http, or to https if port is set to 8443 or 443
@@ -90,6 +95,8 @@ def create_client(*,
     if not interface:
         interface = 'https' if use_tls else 'http'
     port = port or default_port(interface, use_tls)
+    if access_token and (username or password != ''):
+        raise ProgrammingError('Cannot use both access_token and username/password')
     if username is None and 'user' in kwargs:
         username = kwargs.pop('user')
     if username is None and 'user_name' in kwargs:
@@ -112,7 +119,8 @@ def create_client(*,
                     if name.startswith('ch_'):
                         name = name[3:]
                     settings[name] = value
-        return HttpClient(interface, host, port, username, password, database, settings=settings, **kwargs)
+        return HttpClient(interface, host, port, username, password, database, access_token,
+                          settings=settings, **kwargs)
     raise ProgrammingError(f'Unrecognized client type {interface}')
 
 

clickhouse_connect/driver/client.py

@@ -107,7 +107,6 @@ def _init_common_settings(self, apply_server_timezone:Optional[Union[str, bool]]
         if self.min_version('24.8') and not self.min_version('24.10'):
             dynamic_module.json_serialization_format = 0
 
-
     def _validate_settings(self, settings: Optional[Dict[str, Any]]) -> Dict[str, str]:
         """
         This strips any ClickHouse settings that are not recognized or are read only.

clickhouse_connect/driver/httpclient.py

@@ -54,6 +54,7 @@ def __init__(self,
                  username: str,
                  password: str,
                  database: str,
+                 access_token: Optional[str],
                  compress: Union[bool, str] = True,
                  query_limit: int = 0,
                  query_retries: int = 2,
@@ -115,8 +116,11 @@ def __init__(self,
             else:
                 self.http = default_pool_manager()
 
-        if (not client_cert or tls_mode in ('strict', 'proxy')) and username:
+        if access_token:
+            self.headers['Authorization'] = f'Bearer {access_token}'
+        elif (not client_cert or tls_mode in ('strict', 'proxy')) and username:
             self.headers['Authorization'] = 'Basic ' + b64encode(f'{username}:{password}'.encode()).decode()
+
         self.headers['User-Agent'] = common.build_client_name(client_name)
         self._read_format = self._write_format = 'Native'
         self._transform = NativeTransform()

tests/integration_tests/conftest.py

@@ -94,8 +94,8 @@ def test_table_engine_fixture() -> Iterator[str]:
 # pylint: disable=too-many-branches
 @fixture(scope='session', autouse=True, name='test_client')
 def test_client_fixture(test_config: TestConfig, test_create_client: Callable) -> Iterator[Client]:
-    compose_file = f'{PROJECT_ROOT_DIR}/docker-compose.yml'
     if test_config.docker:
+        compose_file = f'{PROJECT_ROOT_DIR}/docker-compose.yml'
         run_cmd(['docker', 'compose', '-f', compose_file, 'down', '-v'])
         sys.stderr.write('Starting docker compose')
         pull_result = run_cmd(['docker', 'compose', '-f', compose_file, 'pull'])
@@ -121,7 +121,6 @@ def test_client_fixture(test_config: TestConfig, test_create_client: Callable) -
     client.command(f'CREATE DATABASE IF NOT EXISTS {test_config.test_database}', use_database=False)
     yield client
 
-    # client.command(f'DROP database IF EXISTS {test_db}', use_database=False)
     if test_config.docker:
         down_result = run_cmd(['docker', 'compose', '-f', compose_file, 'down', '-v'])
         if down_result[0]:

tests/integration_tests/test_jwt_auth.py

@@ -0,0 +1,95 @@
+from datetime import datetime, timezone, timedelta
+from os import environ
+
+import jwt
+import pytest
+
+from clickhouse_connect.driver import create_client, ProgrammingError, create_async_client
+from tests.integration_tests.conftest import TestConfig
+
+
+def test_jwt_auth_sync_client(test_config: TestConfig):
+    if not test_config.cloud:
+        pytest.skip('Skipping JWT test in non-Cloud mode')
+
+    access_token = make_access_token()
+    client = create_client(
+        host=test_config.host,
+        port=test_config.port,
+        access_token=access_token
+    )
+    result = client.query(query=CHECK_CLOUD_MODE_QUERY).result_set
+    assert result == [(True,)]
+
+
+def test_jwt_auth_sync_client_config_errors():
+    with pytest.raises(ProgrammingError):
+        create_client(
+            username='bob',
+            access_token='foobar'
+        )
+    with pytest.raises(ProgrammingError):
+        create_client(
+            username='bob',
+            password='secret',
+            access_token='foo'
+        )
+    with pytest.raises(ProgrammingError):
+        create_client(
+            password='secret',
+            access_token='foo'
+        )
+
+
+@pytest.mark.asyncio
+async def test_jwt_auth_async_client(test_config: TestConfig):
+    if not test_config.cloud:
+        pytest.skip('Skipping JWT test in non-Cloud mode')
+
+    access_token = make_access_token()
+    client = await create_async_client(
+        host=test_config.host,
+        port=test_config.port,
+        access_token=access_token
+    )
+    result = (await client.query(query=CHECK_CLOUD_MODE_QUERY)).result_set
+    assert result == [(True,)]
+
+
+@pytest.mark.asyncio
+async def test_jwt_auth_async_client_config_errors():
+    with pytest.raises(ProgrammingError):
+        await create_async_client(
+            username='bob',
+            access_token='foobar'
+        )
+    with pytest.raises(ProgrammingError):
+        await create_async_client(
+            username='bob',
+            password='secret',
+            access_token='foo'
+        )
+    with pytest.raises(ProgrammingError):
+        await create_async_client(
+            password='secret',
+            access_token='foo'
+        )
+
+
+CHECK_CLOUD_MODE_QUERY = "SELECT value='1' FROM system.settings WHERE name='cloud_mode'"
+JWT_SECRET_ENV_KEY = 'CLICKHOUSE_CONNECT_TEST_JWT_SECRET'
+
+
+def make_access_token():
+    secret = environ.get(JWT_SECRET_ENV_KEY)
+    if not secret:
+        raise ValueError(f'{JWT_SECRET_ENV_KEY} environment variable is not set')
+    payload = {
+        'iss': 'ClickHouse',
+        'sub': 'CI_Test',
+        'aud': '1f7f78b8-da67-480b-8913-726fdd31d2fc',
+        'clickhouse:roles': ['default'],
+        'clickhouse:grants': [],
+        'exp': datetime.now(tz=timezone.utc) + timedelta(minutes=15)
+    }
+    return jwt.encode(payload, secret, algorithm='RS256')

tests/test_requirements.txt

@@ -15,4 +15,5 @@ numpy~=1.26.0; python_version >= '3.11' and python_version <= '3.12'
 numpy~=2.1.0; python_version >= '3.13'
 pandas
 zstandard
-lz4
+lz4
+pyjwt[crypto]==2.10.1