Add set_access_token method and the related tests

CHANGELOG.md

@@ -20,8 +20,10 @@ The supported method of passing ClickHouse server settings is to prefix such arg
 ## 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.
+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

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

@@ -183,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

@@ -184,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/test_jwt_auth.py

@@ -22,6 +22,25 @@ def test_jwt_auth_sync_client(test_config: TestConfig):
     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(
@@ -41,6 +60,23 @@ def test_jwt_auth_sync_client_config_errors():
         )
 
 
+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:
@@ -56,6 +92,25 @@ async def test_jwt_auth_async_client(test_config: TestConfig):
     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):
@@ -76,6 +131,24 @@ async def test_jwt_auth_async_client_config_errors():
         )
 
 
+@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'