JWT auth support (ClickHouse#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 <geoff@clickhouse.com>

.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

tests/integration_tests/conftest.py

@@ -90,8 +90,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'])
@@ -118,7 +118,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/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

timeplus_connect/__version__.py

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

timeplus_connect/driver/__init__.py

@@ -17,7 +17,7 @@ def create_client(*,
                   username: str = None,
                   password: str = '',
                   access_token: Optional[str] = None,
-                  database: str = 'default',
+                  database: str = '__default__',
                   interface: Optional[str] = None,
                   port: int = 0,
                   secure: Union[bool, str] = False,

timeplus_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

timeplus_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,

timeplus_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()
@@ -179,6 +183,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: