JWT auth support (#442)

* JWT auth support

* Update JWT secret env var

* Add `set_access_token` method and the related tests

* Bump version for release

---------

Co-authored-by: Geoff Genz <[email protected]>

Author: slvrtrn

.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.INTEGRATIONS_TEAM_TESTS_CLOUD_JWT_PRIVATE_KEY }}
         run: pytest tests/integration_tests
 
       - name: Run ClickHouse Container (LATEST)

.github/workflows/on_push.yml

@@ -118,7 +118,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 }}
@@ -153,5 +152,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.INTEGRATIONS_TEAM_TESTS_CLOUD_JWT_PRIVATE_KEY }}
           SQLALCHEMY_SILENCE_UBER_WARNING: 1
         run: pytest tests/integration_tests

CHANGELOG.md

@@ -17,6 +17,14 @@ 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_`.
 
+## 0.8.12, 2025-01-06
+### 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. 
+The token can also be updated via the `set_access_token` method in the existing client instance.
+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-21
 ### Improvement
 - Support of ISO8601 strings for inserting values to columns with DateTime64 type was added.  If the driver detects

clickhouse_connect/__version__.py

@@ -1 +1 @@
-version = '0.8.11'
+version = '0.8.12'

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/asyncclient.py

@@ -30,7 +30,6 @@ def __init__(self, *, client: Client, executor_threads: int = 0):
             executor_threads = min(32, (os.cpu_count() or 1) + 4)  # Mimic the default behavior
         self.executor = ThreadPoolExecutor(max_workers=executor_threads)
 
-
     def set_client_setting(self, key, value):
         """
         Set a clickhouse setting for the client after initialization.  If a setting is not recognized by ClickHouse,
@@ -48,6 +47,13 @@ def get_client_setting(self, key) -> Optional[str]:
         """
         return self.client.get_client_setting(key=key)
 
+    def set_access_token(self, access_token: str):
+        """
+        Set the ClickHouse access token for the client
+        :param access_token: Access token string
+        """
+        self.client.set_access_token(access_token)
+
     def min_version(self, version_str: str) -> bool:
         """
         Determine whether the connected server is at least the submitted version

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.
@@ -184,6 +183,13 @@ def get_client_setting(self, key) -> Optional[str]:
         :return: The string value of the setting, if it exists, or None
         """
 
+    @abstractmethod
+    def set_access_token(self, access_token: str):
+        """
+        Set the ClickHouse access token for the client
+        :param access_token: Access token string
+        """
+
     # pylint: disable=unused-argument,too-many-locals
     def query(self,
               query: Optional[str] = None,

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()
@@ -180,6 +184,12 @@ def set_client_setting(self, key, value):
     def get_client_setting(self, key) -> Optional[str]:
         return self.params.get(key)
 
+    def set_access_token(self, access_token: str):
+        auth_header = self.headers.get('Authorization')
+        if auth_header and not auth_header.startswith('Bearer'):
+            raise ProgrammingError('Cannot set access token when a different auth type is used')
+        self.headers['Authorization'] = f'Bearer {access_token}'
+
     def _prep_query(self, context: QueryContext):
         final_query = super()._prep_query(context)
         if context.is_insert:

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,168 @@
+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_set_access_token(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,
+    )
+
+    # Should still work after the override
+    access_token = make_access_token()
+    client.set_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'
+        )
+
+
+def test_jwt_auth_sync_client_set_access_token_errors(test_config: TestConfig):
+    if not test_config.cloud:
+        pytest.skip('Skipping JWT test in non-Cloud mode')
+
+    client = create_client(
+        host=test_config.host,
+        port=test_config.port,
+        username=test_config.username,
+        password=test_config.password,
+    )
+
+    # Can't use JWT with username/password
+    access_token = make_access_token()
+    with pytest.raises(ProgrammingError):
+        client.set_access_token(access_token)
+
+
+@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_set_access_token(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,
+    )
+
+    access_token = make_access_token()
+    client.set_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'
+        )
+
+
+@pytest.mark.asyncio
+async def test_jwt_auth_async_client_set_access_token_errors(test_config: TestConfig):
+    if not test_config.cloud:
+        pytest.skip('Skipping JWT test in non-Cloud mode')
+
+    client = await create_async_client(
+        host=test_config.host,
+        port=test_config.port,
+        username=test_config.username,
+        password=test_config.password,
+    )
+
+    # Can't use JWT with username/password
+    access_token = make_access_token()
+    with pytest.raises(ProgrammingError):
+        client.set_access_token(access_token)
+
+
+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