Skip to content

Commit 2e1fc3d

Browse files
oriangenzgd
orian
and
genzgd
authored
Add param extra_http_headers to query/command methods (#489)
* Add param extra_http_headers to query/command methods * add test, fix dict copy --------- Co-authored-by: Geoff Genz <[email protected]>
1 parent 076f1c6 commit 2e1fc3d

File tree

4 files changed

+83
-38
lines changed

4 files changed

+83
-38
lines changed

‎clickhouse_connect/driver/client.py

Lines changed: 59 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -212,7 +212,8 @@ def query(self,
212212
context: QueryContext = None,
213213
query_tz: Optional[Union[str, tzinfo]] = None,
214214
column_tzs: Optional[Dict[str, Union[str, tzinfo]]] = None,
215-
external_data: Optional[ExternalData] = None) -> QueryResult:
215+
external_data: Optional[ExternalData] = None,
216+
extra_http_headers: Optional[Dict[str, str]] = None) -> QueryResult:
216217
"""
217218
Main query method for SELECT, DESCRIBE and other SQL statements that return a result matrix. For
218219
parameters, see the create_query_context method
@@ -228,7 +229,8 @@ def query(self,
228229
response = self.command(query,
229230
parameters=query_context.parameters,
230231
settings=query_context.settings,
231-
external_data=query_context.external_data)
232+
external_data=query_context.external_data,
233+
extra_http_headers=query_context.extra_http_headers)
232234
if isinstance(response, QuerySummary):
233235
return response.as_query_result()
234236
return QueryResult([response] if isinstance(response, list) else [[response]])
@@ -245,7 +247,8 @@ def query_column_block_stream(self,
245247
context: QueryContext = None,
246248
query_tz: Optional[Union[str, tzinfo]] = None,
247249
column_tzs: Optional[Dict[str, Union[str, tzinfo]]] = None,
248-
external_data: Optional[ExternalData] = None) -> StreamContext:
250+
external_data: Optional[ExternalData] = None,
251+
extra_http_headers: Optional[Dict[str, str]] = None) -> StreamContext:
249252
"""
250253
Variation of main query method that returns a stream of column oriented blocks. For
251254
parameters, see the create_query_context method.
@@ -264,7 +267,8 @@ def query_row_block_stream(self,
264267
context: QueryContext = None,
265268
query_tz: Optional[Union[str, tzinfo]] = None,
266269
column_tzs: Optional[Dict[str, Union[str, tzinfo]]] = None,
267-
external_data: Optional[ExternalData] = None) -> StreamContext:
270+
external_data: Optional[ExternalData] = None,
271+
extra_http_headers: Optional[Dict[str, str]] = None) -> StreamContext:
268272
"""
269273
Variation of main query method that returns a stream of row oriented blocks. For
270274
parameters, see the create_query_context method.
@@ -283,7 +287,8 @@ def query_rows_stream(self,
283287
context: QueryContext = None,
284288
query_tz: Optional[Union[str, tzinfo]] = None,
285289
column_tzs: Optional[Dict[str, Union[str, tzinfo]]] = None,
286-
external_data: Optional[ExternalData] = None) -> StreamContext:
290+
external_data: Optional[ExternalData] = None,
291+
extra_http_headers: Optional[Dict[str, str]] = None) -> StreamContext:
287292
"""
288293
Variation of main query method that returns a stream of row oriented blocks. For
289294
parameters, see the create_query_context method.
@@ -297,16 +302,19 @@ def raw_query(self, query: str,
297302
settings: Optional[Dict[str, Any]] = None,
298303
fmt: str = None,
299304
use_database: bool = True,
300-
external_data: Optional[ExternalData] = None) -> bytes:
305+
external_data: Optional[ExternalData] = None,
306+
extra_http_headers: Optional[Dict[str, str]] = None) -> bytes:
301307
"""
302308
Query method that simply returns the raw ClickHouse format bytes
303309
:param query: Query statement/format string
304310
:param parameters: Optional dictionary used to format the query
305311
:param settings: Optional dictionary of ClickHouse settings (key/string values)
306312
:param fmt: ClickHouse output format
307-
:param use_database Send the database parameter to ClickHouse so the command will be executed in the client
313+
:param use_database: Send the database parameter to ClickHouse so the command will be executed in the client
308314
database context.
309-
:param external_data External data to send with the query
315+
:param external_data: External data to send with the query
316+
:param extra_http_headers: Optional dictionary of extra HTTP headers to pass to ClickHouse,
317+
useful if using Proxy.
310318
:return: bytes representing raw ClickHouse return value based on format
311319
"""
312320

@@ -316,7 +324,8 @@ def raw_stream(self, query: str,
316324
settings: Optional[Dict[str, Any]] = None,
317325
fmt: str = None,
318326
use_database: bool = True,
319-
external_data: Optional[ExternalData] = None) -> io.IOBase:
327+
external_data: Optional[ExternalData] = None,
328+
extra_http_headers: Optional[Dict[str, str]] = None) -> io.IOBase:
320329
"""
321330
Query method that returns the result as an io.IOBase iterator
322331
:param query: Query statement/format string
@@ -325,7 +334,9 @@ def raw_stream(self, query: str,
325334
:param fmt: ClickHouse output format
326335
:param use_database Send the database parameter to ClickHouse so the command will be executed in the client
327336
database context.
328-
:param external_data External data to send with the query
337+
:param external_data: External data to send with the query.
338+
:param extra_http_headers: Optional dictionary of extra HTTP headers to pass to ClickHouse,
339+
useful if using Proxy.
329340
:return: io.IOBase stream/iterator for the result
330341
"""
331342

@@ -340,7 +351,8 @@ def query_np(self,
340351
use_none: Optional[bool] = None,
341352
max_str_len: Optional[int] = None,
342353
context: QueryContext = None,
343-
external_data: Optional[ExternalData] = None):
354+
external_data: Optional[ExternalData] = None,
355+
extra_http_headers: Optional[Dict[str, str]] = None):
344356
"""
345357
Query method that returns the results as a numpy array. For parameter values, see the
346358
create_query_context method
@@ -360,7 +372,8 @@ def query_np_stream(self,
360372
use_none: Optional[bool] = None,
361373
max_str_len: Optional[int] = None,
362374
context: QueryContext = None,
363-
external_data: Optional[ExternalData] = None) -> StreamContext:
375+
external_data: Optional[ExternalData] = None,
376+
extra_http_headers: Optional[Dict[str, str]] = None) -> StreamContext:
364377
"""
365378
Query method that returns the results as a stream of numpy arrays. For parameter values, see the
366379
create_query_context method
@@ -384,7 +397,8 @@ def query_df(self,
384397
column_tzs: Optional[Dict[str, Union[str, tzinfo]]] = None,
385398
context: QueryContext = None,
386399
external_data: Optional[ExternalData] = None,
387-
use_extended_dtypes: Optional[bool] = None):
400+
use_extended_dtypes: Optional[bool] = None,
401+
extra_http_headers: Optional[Dict[str, str]] = None):
388402
"""
389403
Query method that results the results as a pandas dataframe. For parameter values, see the
390404
create_query_context method
@@ -408,7 +422,8 @@ def query_df_stream(self,
408422
column_tzs: Optional[Dict[str, Union[str, tzinfo]]] = None,
409423
context: QueryContext = None,
410424
external_data: Optional[ExternalData] = None,
411-
use_extended_dtypes: Optional[bool] = None) -> StreamContext:
425+
use_extended_dtypes: Optional[bool] = None,
426+
extra_http_headers: Optional[Dict[str, str]] = None) -> StreamContext:
412427
"""
413428
Query method that returns the results as a StreamContext. For parameter values, see the
414429
create_query_context method
@@ -437,7 +452,8 @@ def create_query_context(self,
437452
streaming: bool = False,
438453
as_pandas: bool = False,
439454
external_data: Optional[ExternalData] = None,
440-
use_extended_dtypes: Optional[bool] = None) -> QueryContext:
455+
use_extended_dtypes: Optional[bool] = None,
456+
extra_http_headers: Optional[Dict[str, str]] = None) -> QueryContext:
441457
"""
442458
Creates or updates a reusable QueryContext object
443459
:param query: Query statement/format string
@@ -455,10 +471,10 @@ def create_query_context(self,
455471
structured array even with ClickHouse variable length String columns. If 0, Numpy arrays for
456472
String columns will always be object arrays
457473
:param context: An existing QueryContext to be updated with any provided parameter values
458-
:param query_tz Either a string or a pytz tzinfo object. (Strings will be converted to tzinfo objects).
474+
:param query_tz: Either a string or a pytz tzinfo object. (Strings will be converted to tzinfo objects).
459475
Values for any DateTime or DateTime64 column in the query will be converted to Python datetime.datetime
460476
objects with the selected timezone.
461-
:param column_tzs A dictionary of column names to tzinfo objects (or strings that will be converted to
477+
:param column_tzs: A dictionary of column names to tzinfo objects (or strings that will be converted to
462478
tzinfo objects). The timezone will be applied to datetime objects returned in the query
463479
:param use_na_values: Deprecated alias for use_advanced_dtypes
464480
:param as_pandas Return the result columns as pandas.Series objects
@@ -467,6 +483,8 @@ def create_query_context(self,
467483
:param use_extended_dtypes: Only relevant to Pandas Dataframe queries. Use Pandas "missing types", such as
468484
pandas.NA and pandas.NaT for ClickHouse NULL values, as well as extended Pandas dtypes such as IntegerArray
469485
and StringArray. Defaulted to True for query_df methods
486+
:param extra_http_headers: Optional dictionary of extra HTTP headers to pass to ClickHouse,
487+
useful if using Proxy.
470488
:return: Reusable QueryContext
471489
"""
472490
if context:
@@ -510,21 +528,25 @@ def create_query_context(self,
510528
as_pandas=as_pandas,
511529
streaming=streaming,
512530
apply_server_tz=self.apply_server_timezone,
513-
external_data=external_data)
531+
external_data=external_data,
532+
extra_http_headers=extra_http_headers)
514533

515534
def query_arrow(self,
516535
query: str,
517536
parameters: Optional[Union[Sequence, Dict[str, Any]]] = None,
518537
settings: Optional[Dict[str, Any]] = None,
519538
use_strings: Optional[bool] = None,
520-
external_data: Optional[ExternalData] = None):
539+
external_data: Optional[ExternalData] = None,
540+
extra_http_headers: Optional[Dict[str, str]] = None):
521541
"""
522542
Query method using the ClickHouse Arrow format to return a PyArrow table
523543
:param query: Query statement/format string
524544
:param parameters: Optional dictionary used to format the query
525545
:param settings: Optional dictionary of ClickHouse settings (key/string values)
526-
:param use_strings: Convert ClickHouse String type to Arrow string type (instead of binary)
527-
:param external_data ClickHouse "external data" to send with query
546+
:param use_strings: Convert ClickHouse String type to Arrow string type (instead of binary)
547+
:param external_data: ClickHouse "external data" to send with query
548+
:param extra_http_headers: Optional dictionary of extra HTTP headers to pass to ClickHouse,
549+
useful if using Proxy.
528550
:return: PyArrow.Table
529551
"""
530552
check_arrow()
@@ -533,21 +555,25 @@ def query_arrow(self,
533555
parameters,
534556
settings,
535557
fmt='Arrow',
536-
external_data=external_data))
558+
external_data=external_data,
559+
extra_http_headers=extra_http_headers))
537560

538561
def query_arrow_stream(self,
539562
query: str,
540563
parameters: Optional[Union[Sequence, Dict[str, Any]]] = None,
541564
settings: Optional[Dict[str, Any]] = None,
542565
use_strings: Optional[bool] = None,
543-
external_data: Optional[ExternalData] = None) -> StreamContext:
566+
external_data: Optional[ExternalData] = None,
567+
extra_http_headers: Optional[Dict[str, str]] = None) -> StreamContext:
544568
"""
545569
Query method that returns the results as a stream of Arrow tables
546570
:param query: Query statement/format string
547571
:param parameters: Optional dictionary used to format the query
548572
:param settings: Optional dictionary of ClickHouse settings (key/string values)
549-
:param use_strings: Convert ClickHouse String type to Arrow string type (instead of binary)
550-
:param external_data ClickHouse "external data" to send with query
573+
:param use_strings: Convert ClickHouse String type to Arrow string type (instead of binary)
574+
:param external_data: ClickHouse "external data" to send with query
575+
:param extra_http_headers: Optional dictionary of extra HTTP headers to pass to ClickHouse,
576+
useful if using Proxy.
551577
:return: Generator that yields a PyArrow.Table for per block representing the result set
552578
"""
553579
check_arrow()
@@ -556,7 +582,8 @@ def query_arrow_stream(self,
556582
parameters,
557583
settings,
558584
fmt='ArrowStream',
559-
external_data=external_data))
585+
external_data=external_data,
586+
extra_http_headers=extra_http_headers))
560587

561588
def _update_arrow_settings(self,
562589
settings: Optional[Dict[str, Any]],
@@ -581,17 +608,20 @@ def command(self,
581608
data: Union[str, bytes] = None,
582609
settings: Dict[str, Any] = None,
583610
use_database: bool = True,
584-
external_data: Optional[ExternalData] = None) -> Union[str, int, Sequence[str], QuerySummary]:
611+
external_data: Optional[ExternalData] = None,
612+
extra_http_headers: Optional[Dict[str, str]] = None) -> Union[str, int, Sequence[str], QuerySummary]:
585613
"""
586614
Client method that returns a single value instead of a result set
587615
:param cmd: ClickHouse query/command as a python format string
588616
:param parameters: Optional dictionary of key/values pairs to be formatted
589617
:param data: Optional 'data' for the command (for INSERT INTO in particular)
590618
:param settings: Optional dictionary of ClickHouse settings (key/string values)
591619
:param use_database: Send the database parameter to ClickHouse so the command will be executed in the client
592-
database context. Otherwise, no database will be specified with the command. This is useful for determining
620+
database context. Otherwise, no database will be specified with the command. This is useful for determining
593621
the default user database
594-
:param external_data ClickHouse "external data" to send with command/query
622+
:param external_data: ClickHouse "external data" to send with command/query
623+
:param extra_http_headers: Optional dictionary of extra HTTP headers to pass to ClickHouse,
624+
useful if using Proxy.
595625
:return: Decoded response from ClickHouse as either a string, int, or sequence of strings, or QuerySummary
596626
if no data returned
597627
"""

‎clickhouse_connect/driver/httpclient.py

Lines changed: 10 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -204,7 +204,7 @@ def _prep_query(self, context: QueryContext):
204204
return final_query + fmt
205205

206206
def _query_with_context(self, context: QueryContext) -> QueryResult:
207-
headers = {}
207+
headers = dict_copy(context.extra_http_headers)
208208
params = {}
209209
if self.database:
210210
params['database'] = self.database
@@ -332,7 +332,8 @@ def command(self,
332332
data: Union[str, bytes] = None,
333333
settings: Optional[Dict] = None,
334334
use_database: int = True,
335-
external_data: Optional[ExternalData] = None) -> Union[str, int, Sequence[str], QuerySummary]:
335+
external_data: Optional[ExternalData] = None,
336+
extra_http_headers: Optional[Dict[str, str]] = None) -> Union[str, int, Sequence[str], QuerySummary]:
336337
"""
337338
See BaseClient doc_string for this method
338339
"""
@@ -482,24 +483,27 @@ def raw_query(self, query: str,
482483
settings: Optional[Dict[str, Any]] = None,
483484
fmt: str = None,
484485
use_database: bool = True,
485-
external_data: Optional[ExternalData] = None) -> bytes:
486+
external_data: Optional[ExternalData] = None,
487+
extra_http_headers: Optional[Dict[str, str]] = None) -> bytes:
486488
"""
487489
See BaseClient doc_string for this method
488490
"""
489491
body, params, fields = self._prep_raw_query(query, parameters, settings, fmt, use_database, external_data)
490-
return self._raw_request(body, params, fields=fields).data
492+
return self._raw_request(body, params, fields=fields, headers=extra_http_headers).data
491493

492494
def raw_stream(self, query: str,
493495
parameters: Optional[Union[Sequence, Dict[str, Any]]] = None,
494496
settings: Optional[Dict[str, Any]] = None,
495497
fmt: str = None,
496498
use_database: bool = True,
497-
external_data: Optional[ExternalData] = None) -> io.IOBase:
499+
external_data: Optional[ExternalData] = None,
500+
extra_http_headers: Optional[Dict[str, str]] = None) -> io.IOBase:
498501
"""
499502
See BaseClient doc_string for this method
500503
"""
501504
body, params, fields = self._prep_raw_query(query, parameters, settings, fmt, use_database, external_data)
502-
return self._raw_request(body, params, fields=fields, stream=True, server_wait=False)
505+
return self._raw_request(body, params, fields=fields, stream=True, server_wait=False,
506+
headers=extra_http_headers)
503507

504508
def _prep_raw_query(self, query: str,
505509
parameters: Optional[Union[Sequence, Dict[str, Any]]],

0 commit comments

Comments
 (0)