Change http_headers to transport settings, add transport settings to async client and insert methods (#490)

Author: genzgd

CHANGELOG.md

@@ -22,16 +22,21 @@ instead of being passed as ClickHouse server settings. This is in conjunction wi
 The supported method of passing ClickHouse server settings is to prefix such arguments/query parameters with`ch_`.
 
 ## 0.8.17, 2025-04-10
-### Bug Fix
-- Version 0.8.16 introduced a bug where changing a Client setting value and then changing that setting value back to the
-original server value would fail to restore the original setting.  This has been fixed.  Closes
-https://github.com/ClickHouse/clickhouse-connect/issues/487
 
-### Improvement
+### Improvements
+- The parameter `transport_settings` has been added to the Client query and insert methods.  For the HTTP client (currently
+the only  option), this dictionary of string is directly translated into additional HTTP headers at a query level.  This can
+be used to provide additional proxy directives or other extra 'non-ClickHouse' information that is passed via headers.
+Thanks to [Paweł Szczur](https://github.com/orian) of PostHog for the original PR!
 - There was previously no way to add a path to the ClickHouse server host in cases where the ClickHouse server was
 behind a proxy that used path based routing (such as `https://big_proxy:8080/clickhouse).  The new `proxy_path`
 `get_client` argument can now be used to set that path.  Closes https://github.com/ClickHouse/clickhouse-connect/issues/486
 
+### Bug Fix
+- Version 0.8.16 introduced a bug where changing a Client setting value and then changing that setting value back to the
+  original server value would fail to restore the original setting.  This has been fixed.  Closes
+  https://github.com/ClickHouse/clickhouse-connect/issues/487
+
 ## 0.8.16, 2025-03-28
 ### Bug Fixes
 - Don't send a setting value if the setting is already correct according to the `system.settings` table. 

clickhouse_connect/driver/asyncclient.py

@@ -16,7 +16,8 @@ def __init__(self,
                  column_formats: Optional[Dict[str, Union[str, Dict[str, str]]]] = None,
                  encoding: Optional[str] = None,
                  use_extended_dtypes: bool = False,
-                 use_numpy: bool = False):
+                 use_numpy: bool = False,
+                 transport_settings: Optional[Dict[str, str]] = None):
         self.settings = settings or {}
         if query_formats is None:
             self.type_formats = _empty_map
@@ -36,6 +37,7 @@ def __init__(self,
                                                        for type_name, fmt in fmt.items()}
         self.query_formats = query_formats or {}
         self.column_formats = column_formats or {}
+        self.transport_settings = transport_settings
         self.column_name = None
         self.encoding = encoding
         self.use_numpy = use_numpy

clickhouse_connect/driver/client.py

@@ -204,7 +204,7 @@ def _prep_query(self, context: QueryContext):
         return final_query + fmt
 
     def _query_with_context(self, context: QueryContext) -> QueryResult:
-        headers = dict_copy(context.extra_http_headers)
+        headers = {}
         params = {}
         if self.database:
             params['database'] = self.database
@@ -242,7 +242,7 @@ def _query_with_context(self, context: QueryContext) -> QueryResult:
             headers['Content-Type'] = 'text/plain; charset=utf-8'
         response = self._raw_request(body,
                                      params,
-                                     headers,
+                                     dict_copy(headers, context.transport_settings),
                                      stream=True,
                                      retries=self.query_retries,
                                      fields=fields,
@@ -280,7 +280,7 @@ def error_handler(resp: HTTPResponse):
         if self.database:
             params['database'] = self.database
         params.update(self._validate_settings(context.settings))
-
+        headers = dict_copy(headers, context.transport_settings)
         response = self._raw_request(block_gen, params, headers, error_handler=error_handler, server_wait=False)
         logger.debug('Context insert response code: %d, content: %s', response.status, response.data)
         context.data = None
@@ -291,7 +291,8 @@ def raw_insert(self, table: str = None,
                    insert_block: Union[str, bytes, Generator[bytes, None, None], BinaryIO] = None,
                    settings: Optional[Dict] = None,
                    fmt: Optional[str] = None,
-                   compression: Optional[str] = None) -> QuerySummary:
+                   compression: Optional[str] = None,
+                   transport_settings: Optional[Dict[str, str]] = None) -> QuerySummary:
         """
         See BaseClient doc_string for this method
         """
@@ -311,6 +312,7 @@ def raw_insert(self, table: str = None,
         if self.database:
             params['database'] = self.database
         params.update(self._validate_settings(settings or {}))
+        headers = dict_copy(headers, transport_settings)
         response = self._raw_request(insert_block, params, headers, server_wait=False)
         logger.debug('Raw insert response code: %d, content: %s', response.status, response.data)
         return QuerySummary(self._summary(response))
@@ -333,7 +335,7 @@ def command(self,
                 settings: Optional[Dict] = None,
                 use_database: int = True,
                 external_data: Optional[ExternalData] = None,
-                extra_http_headers: Optional[Dict[str, str]] = None) -> Union[str, int, Sequence[str], QuerySummary]:
+                transport_settings: Optional[Dict[str, str]] = None) -> Union[str, int, Sequence[str], QuerySummary]:
         """
         See BaseClient doc_string for this method
         """
@@ -361,7 +363,7 @@ def command(self,
         if use_database and self.database:
             params['database'] = self.database
         params.update(self._validate_settings(settings or {}))
-
+        headers = dict_copy(headers, transport_settings)
         method = 'POST' if payload or fields else 'GET'
         response = self._raw_request(payload, params, headers, method, fields=fields, server_wait=False)
         if response.data:
@@ -484,26 +486,26 @@ def raw_query(self, query: str,
                   fmt: str = None,
                   use_database: bool = True,
                   external_data: Optional[ExternalData] = None,
-                  extra_http_headers: Optional[Dict[str, str]] = None) -> bytes:
+                  transport_settings: Optional[Dict[str, str]] = None) -> bytes:
         """
         See BaseClient doc_string for this method
         """
         body, params, fields = self._prep_raw_query(query, parameters, settings, fmt, use_database, external_data)
-        return self._raw_request(body, params, fields=fields, headers=extra_http_headers).data
+        return self._raw_request(body, params, fields=fields, headers=transport_settings).data
 
     def raw_stream(self, query: str,
                    parameters: Optional[Union[Sequence, Dict[str, Any]]] = None,
                    settings: Optional[Dict[str, Any]] = None,
                    fmt: str = None,
                    use_database: bool = True,
                    external_data: Optional[ExternalData] = None,
-                   extra_http_headers: Optional[Dict[str, str]] = None) -> io.IOBase:
+                   transport_settings: Optional[Dict[str, str]] = None) -> io.IOBase:
         """
         See BaseClient doc_string for this method
         """
         body, params, fields = self._prep_raw_query(query, parameters, settings, fmt, use_database, external_data)
         return self._raw_request(body, params, fields=fields, stream=True, server_wait=False,
-                                 headers=extra_http_headers)
+                                 headers=transport_settings)
 
     def _prep_raw_query(self, query: str,
                         parameters: Optional[Union[Sequence, Dict[str, Any]]],

clickhouse_connect/driver/context.py

@@ -42,8 +42,9 @@ def __init__(self,
                  compression: Optional[Union[str, bool]] = None,
                  query_formats: Optional[Dict[str, str]] = None,
                  column_formats: Optional[Dict[str, Union[str, Dict[str, str]]]] = None,
-                 block_size: Optional[int] = None):
-        super().__init__(settings, query_formats, column_formats)
+                 block_size: Optional[int] = None,
+                 transport_settings: Optional[Dict[str, str]] = None):
+        super().__init__(settings, query_formats, column_formats, transport_settings=transport_settings)
         self.table = table
         self.column_names = column_names
         self.column_types = column_types

clickhouse_connect/driver/httpclient.py

@@ -53,7 +53,7 @@ def __init__(self,
                  streaming: bool = False,
                  apply_server_tz: bool = False,
                  external_data: Optional[ExternalData] = None,
-                 extra_http_headers: Optional[Dict[str, str]] = None):
+                 transport_settings: Optional[Dict[str, str]] = None):
         """
         Initializes various configuration settings for the query context
 
@@ -86,7 +86,8 @@ def __init__(self,
                          column_formats,
                          encoding,
                          use_extended_dtypes if use_extended_dtypes is not None else False,
-                         use_numpy if use_numpy is not None else False)
+                         use_numpy if use_numpy is not None else False,
+                         transport_settings=transport_settings)
         self.query = query
         self.parameters = parameters or {}
         self.use_none = True if use_none is None else use_none
@@ -117,7 +118,6 @@ def __init__(self,
         self.as_pandas = as_pandas
         self.use_pandas_na = as_pandas and pd_extended_dtypes
         self.streaming = streaming
-        self.extra_http_headers = extra_http_headers
         self._update_query()
 
     @property
@@ -192,7 +192,7 @@ def updated_copy(self,
                      as_pandas: bool = False,
                      streaming: bool = False,
                      external_data: Optional[ExternalData] = None,
-                     extra_http_headers: Optional[Dict[str, str]] = None) -> 'QueryContext':
+                     transport_settings: Optional[Dict[str, str]] = None) -> 'QueryContext':
         """
         Creates Query context copy with parameters overridden/updated as appropriate.
         """
@@ -214,7 +214,7 @@ def updated_copy(self,
                             streaming,
                             self.apply_server_tz,
                             self.external_data if external_data is None else external_data,
-                            self.extra_http_headers if extra_http_headers is None else extra_http_headers)
+                            self.transport_settings if transport_settings is None else transport_settings)
 
     def _update_query(self):
         self.final_query, self.bind_params = bind_query(self.query, self.parameters, self.server_tz)

clickhouse_connect/driver/insert.py

@@ -40,9 +40,9 @@ def test_client_name(test_client: Client):
     assert 'py/' in user_agent
 
 
-def test_extra_http_headers(test_client: Client):
+def test_transport_settings(test_client: Client):
     result = test_client.query('SELECT name,database FROM system.tables',
-                               extra_http_headers={'X-Workload': 'ONLINE'})
+                               transport_settings={'X-Workload': 'ONLINE'})
     assert result.column_names == ('name', 'database')
     assert len(result.result_set) > 0