mirror of
https://github.com/hanzo-ds/connect.git
synced 2026-10-30 20:31:24 +00:00
* async client: retry consecutive ServerDisconnectedErrors up to query_retries * - Add CHANGELOG.md - Add 0.1 * attempts delay between retries - add handling for httpclient as well and test * pass query_retries to raw_query and raw_stream for parity with async * linting --------- Co-authored-by: Joe S <[email protected]>
749 lines
32 KiB
Python
749 lines
32 KiB
Python
import io
|
|
import json
|
|
import logging
|
|
import re
|
|
import time
|
|
import uuid
|
|
from base64 import b64encode
|
|
from collections.abc import Callable, Generator, Sequence
|
|
from importlib import import_module
|
|
from importlib.metadata import version as dist_version
|
|
from typing import Any, BinaryIO
|
|
from urllib.parse import urlencode
|
|
|
|
from urllib3 import Timeout
|
|
from urllib3.exceptions import HTTPError
|
|
from urllib3.poolmanager import PoolManager
|
|
from urllib3.response import HTTPResponse
|
|
|
|
from clickhouse_connect import common
|
|
from clickhouse_connect.datatypes import registry
|
|
from clickhouse_connect.datatypes.base import ClickHouseType
|
|
from clickhouse_connect.driver.binding import bind_query, quote_identifier
|
|
from clickhouse_connect.driver.client import Client
|
|
from clickhouse_connect.driver.common import coerce_bool, coerce_int, dict_add, dict_copy
|
|
from clickhouse_connect.driver.compression import available_compression
|
|
from clickhouse_connect.driver.ctypes import RespBuffCls
|
|
from clickhouse_connect.driver.exceptions import DatabaseError, OperationalError, ProgrammingError
|
|
from clickhouse_connect.driver.external import ExternalData
|
|
from clickhouse_connect.driver.httputil import (
|
|
ResponseSource,
|
|
all_managers,
|
|
check_conn_expiration,
|
|
check_env_proxy,
|
|
default_pool_manager,
|
|
get_pool_manager,
|
|
get_proxy_manager,
|
|
get_response_data,
|
|
)
|
|
from clickhouse_connect.driver.insert import InsertContext
|
|
from clickhouse_connect.driver.query import QueryContext, QueryResult, TzSource
|
|
from clickhouse_connect.driver.summary import QuerySummary
|
|
from clickhouse_connect.driver.transform import NativeTransform
|
|
|
|
logger = logging.getLogger(__name__)
|
|
columns_only_re = re.compile(r"LIMIT 0\s*$", re.IGNORECASE)
|
|
ex_header = "X-ClickHouse-Exception-Code"
|
|
ex_tag_header = "X-ClickHouse-Exception-Tag"
|
|
|
|
|
|
class HttpClient(Client):
|
|
params = {}
|
|
valid_transport_settings = {
|
|
"database",
|
|
"buffer_size",
|
|
"session_id",
|
|
"compress",
|
|
"decompress",
|
|
"session_timeout",
|
|
"session_check",
|
|
"query_id",
|
|
"quota_key",
|
|
"wait_end_of_query",
|
|
"client_protocol_version",
|
|
"role",
|
|
}
|
|
optional_transport_settings = {"send_progress_in_http_headers", "http_headers_progress_interval_ms", "enable_http_compression"}
|
|
_owns_pool_manager = False
|
|
|
|
# R0917: too-many-positional-arguments
|
|
|
|
def __init__(
|
|
self,
|
|
interface: str,
|
|
host: str,
|
|
port: int,
|
|
username: str,
|
|
password: str,
|
|
database: str,
|
|
access_token: str | None = None,
|
|
compress: bool | str = True,
|
|
query_limit: int = 0,
|
|
query_retries: int = 2,
|
|
connect_timeout: int = 10,
|
|
send_receive_timeout: int = 300,
|
|
client_name: str | None = None,
|
|
verify: bool | str = True,
|
|
ca_cert: str | None = None,
|
|
client_cert: str | None = None,
|
|
client_cert_key: str | None = None,
|
|
session_id: str | None = None,
|
|
settings: dict[str, Any] | None = None,
|
|
pool_mgr: PoolManager | None = None,
|
|
http_proxy: str | None = None,
|
|
https_proxy: str | None = None,
|
|
server_host_name: str | None = None,
|
|
tz_source: TzSource | None = None,
|
|
tz_mode: str | None = None,
|
|
show_clickhouse_errors: bool | None = None,
|
|
autogenerate_session_id: bool | None = None,
|
|
autogenerate_query_id: bool | None = None,
|
|
tls_mode: str | None = None,
|
|
proxy_path: str = "",
|
|
form_encode_query_params: bool = False,
|
|
rename_response_column: str | None = None,
|
|
):
|
|
"""
|
|
Create an HTTP ClickHouse Connect client
|
|
See clickhouse_connect.get_client for parameters
|
|
"""
|
|
proxy_path = proxy_path.lstrip("/")
|
|
if proxy_path:
|
|
proxy_path = "/" + proxy_path
|
|
self.url = f"{interface}://{host}:{port}{proxy_path}"
|
|
self.headers = {}
|
|
self.form_encode_query_params = form_encode_query_params
|
|
self.params = dict_copy(HttpClient.params)
|
|
ch_settings = dict_copy(settings, self.params)
|
|
self.http = pool_mgr
|
|
if interface == "https":
|
|
if isinstance(verify, str) and verify.lower() == "proxy":
|
|
verify = True
|
|
tls_mode = tls_mode or "proxy"
|
|
if not https_proxy:
|
|
https_proxy = check_env_proxy("https", host, port)
|
|
verify = coerce_bool(verify)
|
|
if client_cert and (tls_mode is None or tls_mode == "mutual"):
|
|
if not username:
|
|
raise ProgrammingError("username parameter is required for Mutual TLS authentication")
|
|
self.headers["X-ClickHouse-User"] = username
|
|
self.headers["X-ClickHouse-SSL-Certificate-Auth"] = "on"
|
|
|
|
if not self.http and (server_host_name or ca_cert or client_cert or not verify or https_proxy):
|
|
options = {"verify": verify}
|
|
dict_add(options, "ca_cert", ca_cert)
|
|
dict_add(options, "client_cert", client_cert)
|
|
dict_add(options, "client_cert_key", client_cert_key)
|
|
if server_host_name:
|
|
if options["verify"]:
|
|
options["assert_hostname"] = server_host_name
|
|
options["server_hostname"] = server_host_name
|
|
self.http = get_pool_manager(https_proxy=https_proxy, **options)
|
|
self._owns_pool_manager = True
|
|
if not self.http:
|
|
if not http_proxy:
|
|
http_proxy = check_env_proxy("http", host, port)
|
|
if http_proxy:
|
|
self.http = get_proxy_manager(host, http_proxy)
|
|
else:
|
|
self.http = default_pool_manager()
|
|
|
|
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._reported_libs = set()
|
|
self.headers["User-Agent"] = common.build_client_name(client_name)
|
|
self._read_format = self._write_format = "Native"
|
|
self._transform = NativeTransform()
|
|
|
|
# There are use cases when the client needs to disable timeouts.
|
|
if connect_timeout is not None:
|
|
connect_timeout = coerce_int(connect_timeout)
|
|
if send_receive_timeout is not None:
|
|
send_receive_timeout = coerce_int(send_receive_timeout)
|
|
self.timeout = Timeout(connect=connect_timeout, read=send_receive_timeout)
|
|
self.http_retries = 1
|
|
self._send_progress = None
|
|
self._send_comp_setting = False
|
|
self._progress_interval = None
|
|
self._active_session = None
|
|
self._rename_response_column = rename_response_column
|
|
|
|
# allow to override the global autogenerate_session_id setting via the constructor params
|
|
_autogenerate_session_id = (
|
|
common.get_setting("autogenerate_session_id") if autogenerate_session_id is None else autogenerate_session_id
|
|
)
|
|
|
|
if session_id:
|
|
ch_settings["session_id"] = session_id
|
|
elif "session_id" not in ch_settings and _autogenerate_session_id:
|
|
ch_settings["session_id"] = str(uuid.uuid4())
|
|
|
|
# allow to override the global autogenerate_query_id setting via the constructor params
|
|
self._autogenerate_query_id = (
|
|
common.get_setting("autogenerate_query_id") if autogenerate_query_id is None else autogenerate_query_id
|
|
)
|
|
|
|
if coerce_bool(compress):
|
|
compression = ",".join(available_compression)
|
|
self.write_compression = available_compression[0]
|
|
elif compress and compress not in ("False", "false", "0"):
|
|
if compress not in available_compression:
|
|
raise ProgrammingError(f"Unsupported compression method {compress}")
|
|
compression = compress
|
|
self.write_compression = compress
|
|
else:
|
|
compression = None
|
|
|
|
super().__init__(
|
|
database=database,
|
|
uri=self.url,
|
|
query_limit=query_limit,
|
|
query_retries=query_retries,
|
|
server_host_name=server_host_name,
|
|
tz_source=tz_source,
|
|
tz_mode=tz_mode,
|
|
show_clickhouse_errors=show_clickhouse_errors,
|
|
autoconnect=True,
|
|
)
|
|
self.params = dict_copy(self.params, self._validate_settings(ch_settings))
|
|
cancel_setting = self._setting_status("cancel_http_readonly_queries_on_client_close")
|
|
if (
|
|
cancel_setting.is_writable
|
|
and not cancel_setting.is_set
|
|
and "cancel_http_readonly_queries_on_client_close" not in (settings or {})
|
|
):
|
|
self.params["cancel_http_readonly_queries_on_client_close"] = "1"
|
|
comp_setting = self._setting_status("enable_http_compression")
|
|
self._send_comp_setting = not comp_setting.is_set and comp_setting.is_writable
|
|
if comp_setting.is_set or comp_setting.is_writable:
|
|
self.compression = compression
|
|
send_setting = self._setting_status("send_progress_in_http_headers")
|
|
self._send_progress = not send_setting.is_set and send_setting.is_writable
|
|
if (send_setting.is_set or send_setting.is_writable) and self._setting_status("http_headers_progress_interval_ms").is_writable:
|
|
self._progress_interval = str(min(120000, max(10000, (send_receive_timeout - 5) * 1000)))
|
|
|
|
def set_client_setting(self, key: str, value: Any) -> None:
|
|
str_value = self._validate_setting(key, value, common.get_setting("invalid_setting_action"))
|
|
if str_value is not None:
|
|
self.params[key] = str_value
|
|
|
|
def get_client_setting(self, key: str) -> str | None:
|
|
return self.params.get(key)
|
|
|
|
def set_access_token(self, access_token: str) -> None:
|
|
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:
|
|
return final_query
|
|
fmt = f"\n FORMAT {self._read_format}"
|
|
if isinstance(final_query, bytes):
|
|
return final_query + fmt.encode()
|
|
return final_query + fmt
|
|
|
|
def _query_with_context(self, context: QueryContext) -> QueryResult:
|
|
headers = {}
|
|
params = {}
|
|
if self.database:
|
|
params["database"] = self.database
|
|
if self.protocol_version:
|
|
params["client_protocol_version"] = self.protocol_version
|
|
context.block_info = True
|
|
params.update(self._validate_settings(context.settings))
|
|
context.rename_response_column = self._rename_response_column
|
|
if not context.is_insert and columns_only_re.search(context.uncommented_query):
|
|
# Mirror normal query behavior for form encoding and external data
|
|
fmt_json_query = f"{context.final_query}\n FORMAT JSON"
|
|
if self.form_encode_query_params:
|
|
fields = {"query": fmt_json_query}
|
|
fields.update(context.bind_params)
|
|
if context.external_data: # Deal with form encoding + external data
|
|
params.update(context.external_data.query_params)
|
|
fields.update(context.external_data.form_data)
|
|
response = self._raw_request(b"", params, headers, retries=self.query_retries, fields=fields)
|
|
elif context.external_data: # Deal with external data without form encoding
|
|
fields = context.external_data.form_data
|
|
params.update(context.bind_params)
|
|
params.update(context.external_data.query_params)
|
|
params["query"] = fmt_json_query
|
|
response = self._raw_request(b"", params, headers, retries=self.query_retries, fields=fields)
|
|
else: # Legacy behavior (plain body, bind params in URL)
|
|
params.update(context.bind_params)
|
|
response = self._raw_request(fmt_json_query, params, headers, retries=self.query_retries)
|
|
json_result = json.loads(response.data)
|
|
# ClickHouse will respond with a JSON object of meta, data, and some other objects
|
|
# We just grab the column names and column types from the metadata sub object
|
|
names: list[str] = []
|
|
types: list[ClickHouseType] = []
|
|
renamer = context.column_renamer
|
|
for col in json_result["meta"]:
|
|
name = col["name"]
|
|
if renamer is not None:
|
|
try:
|
|
name = renamer(name)
|
|
except Exception as e:
|
|
logger.debug("Failed to rename col '%s'. Skipping rename. Error: %s", name, e)
|
|
names.append(name)
|
|
types.append(registry.get_from_name(col["type"]))
|
|
return QueryResult([], None, tuple(names), tuple(types))
|
|
|
|
if self.compression:
|
|
headers["Accept-Encoding"] = self.compression
|
|
if self._send_comp_setting:
|
|
params["enable_http_compression"] = "1"
|
|
final_query = self._prep_query(context)
|
|
fields = {}
|
|
# Setup additional query parameters and body
|
|
if self.form_encode_query_params:
|
|
body = b""
|
|
fields["query"] = final_query
|
|
fields.update(context.bind_params)
|
|
if context.external_data:
|
|
params.update(context.external_data.query_params)
|
|
fields.update(context.external_data.form_data)
|
|
elif context.external_data:
|
|
params.update(context.bind_params)
|
|
body = b""
|
|
params["query"] = final_query
|
|
params.update(context.external_data.query_params)
|
|
fields = context.external_data.form_data
|
|
else:
|
|
params.update(context.bind_params)
|
|
body = final_query
|
|
fields = None
|
|
headers["Content-Type"] = "text/plain; charset=utf-8"
|
|
response = self._raw_request(
|
|
body,
|
|
params,
|
|
dict_copy(headers, context.transport_settings),
|
|
stream=True,
|
|
retries=self.query_retries,
|
|
fields=fields,
|
|
server_wait=not context.streaming,
|
|
)
|
|
exception_tag = response.headers.get(ex_tag_header)
|
|
byte_source = RespBuffCls(ResponseSource(response, exception_tag=exception_tag))
|
|
context.set_response_tz(self._check_tz_change(response.headers.get("X-ClickHouse-Timezone")))
|
|
query_result = self._transform.parse_response(byte_source, context)
|
|
query_result.summary = self._summary(response)
|
|
return query_result
|
|
|
|
def data_insert(self, context: InsertContext) -> QuerySummary:
|
|
"""
|
|
See BaseClient doc_string for this method
|
|
"""
|
|
if context.empty:
|
|
logger.debug("No data included in insert, skipping")
|
|
return QuerySummary()
|
|
|
|
def error_handler(resp: HTTPResponse):
|
|
# If we actually had a local exception when building the insert, throw that instead
|
|
if context.insert_exception:
|
|
ex = context.insert_exception
|
|
context.insert_exception = None
|
|
raise ex
|
|
self._error_handler(resp)
|
|
|
|
headers = {"Content-Type": "application/octet-stream"}
|
|
if context.compression is None:
|
|
context.compression = self.write_compression
|
|
if context.compression:
|
|
headers["Content-Encoding"] = context.compression
|
|
block_gen = self._transform.build_insert(context)
|
|
|
|
def rebuild_block_gen():
|
|
context.current_row = 0
|
|
context.current_block = 0
|
|
return self._transform.build_insert(context)
|
|
|
|
params = {}
|
|
if self.database:
|
|
params["database"] = self.database
|
|
params.update(self._validate_settings(context.settings))
|
|
headers = dict_copy(headers, context.transport_settings)
|
|
try:
|
|
response = self._raw_request(
|
|
block_gen,
|
|
params,
|
|
headers,
|
|
error_handler=error_handler,
|
|
server_wait=False,
|
|
retry_body=rebuild_block_gen,
|
|
)
|
|
logger.debug("Context insert response code: %d, content: %s", response.status, response.data)
|
|
return QuerySummary(self._summary(response))
|
|
finally:
|
|
context.data = None
|
|
|
|
def raw_insert(
|
|
self,
|
|
table: str = None,
|
|
column_names: Sequence[str] | None = None,
|
|
insert_block: str | bytes | Generator[bytes, None, None] | BinaryIO = None,
|
|
settings: dict | None = None,
|
|
fmt: str | None = None,
|
|
compression: str | None = None,
|
|
transport_settings: dict[str, str] | None = None,
|
|
) -> QuerySummary:
|
|
"""
|
|
See BaseClient doc_string for this method
|
|
"""
|
|
params = {}
|
|
headers = {"Content-Type": "application/octet-stream"}
|
|
if compression:
|
|
headers["Content-Encoding"] = compression
|
|
if table:
|
|
cols = f" ({', '.join([quote_identifier(x) for x in column_names])})" if column_names is not None else ""
|
|
query = f"INSERT INTO {table}{cols} FORMAT {fmt if fmt else self._write_format}"
|
|
if not compression and isinstance(insert_block, str):
|
|
insert_block = query + "\n" + insert_block
|
|
elif not compression and isinstance(insert_block, (bytes, bytearray, BinaryIO)):
|
|
insert_block = (query + "\n").encode() + insert_block
|
|
else:
|
|
params["query"] = query
|
|
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))
|
|
|
|
@staticmethod
|
|
def _summary(response: HTTPResponse):
|
|
summary = {}
|
|
if "X-ClickHouse-Summary" in response.headers:
|
|
try:
|
|
summary = json.loads(response.headers["X-ClickHouse-Summary"])
|
|
except json.JSONDecodeError:
|
|
pass
|
|
summary["query_id"] = response.headers.get("X-ClickHouse-Query-Id", "")
|
|
return summary
|
|
|
|
def command(
|
|
self,
|
|
cmd: str,
|
|
parameters: Sequence | dict[str, Any] | None = None,
|
|
data: str | bytes = None,
|
|
settings: dict | None = None,
|
|
use_database: bool = True,
|
|
external_data: ExternalData | None = None,
|
|
transport_settings: dict[str, str] | None = None,
|
|
) -> str | int | Sequence[str] | QuerySummary:
|
|
"""
|
|
See BaseClient doc_string for this method
|
|
"""
|
|
cmd, params = bind_query(cmd, parameters, self.server_tz)
|
|
headers = {}
|
|
payload = None
|
|
fields = None
|
|
if external_data:
|
|
if data:
|
|
raise ProgrammingError("Cannot combine command data with external data") from None
|
|
fields = external_data.form_data
|
|
params.update(external_data.query_params)
|
|
elif isinstance(data, str):
|
|
headers["Content-Type"] = "text/plain; charset=utf-8"
|
|
payload = data.encode()
|
|
elif isinstance(data, bytes):
|
|
headers["Content-Type"] = "application/octet-stream"
|
|
payload = data
|
|
if payload is None and not cmd:
|
|
raise ProgrammingError("Command sent without query or recognized data") from None
|
|
if payload or fields:
|
|
params["query"] = cmd
|
|
else:
|
|
payload = cmd
|
|
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:
|
|
try:
|
|
result = response.data.decode()[:-1].split("\t")
|
|
if len(result) == 1:
|
|
try:
|
|
return int(result[0])
|
|
except ValueError:
|
|
return result[0]
|
|
return result
|
|
except UnicodeDecodeError:
|
|
return str(response.data)
|
|
return QuerySummary(self._summary(response))
|
|
|
|
def _error_handler(self, response: HTTPResponse, retried: bool = False) -> None:
|
|
"""
|
|
Handles HTTP errors. Tries to be robust and provide maximum context.
|
|
"""
|
|
try:
|
|
body = ""
|
|
try:
|
|
raw_body = get_response_data(response)
|
|
body = common.format_error(raw_body.decode(errors="backslashreplace")).strip()
|
|
except Exception:
|
|
logger.warning("Failed to read error response body", exc_info=True)
|
|
|
|
if self.show_clickhouse_errors:
|
|
err_code = response.headers.get(ex_header)
|
|
if err_code:
|
|
err_str = f"Received ClickHouse exception, code: {err_code}"
|
|
else:
|
|
err_str = f"HTTP driver received HTTP status {response.status}"
|
|
|
|
if body:
|
|
err_str = f"{err_str}, server response: {body}"
|
|
else:
|
|
err_str = "The ClickHouse server returned an error"
|
|
|
|
err_str = f"{err_str} (for url {self.url})"
|
|
|
|
finally:
|
|
response.close()
|
|
|
|
raise OperationalError(err_str) if retried else DatabaseError(err_str) from None
|
|
|
|
def _raw_request(
|
|
self,
|
|
data,
|
|
params: dict[str, str],
|
|
headers: dict[str, Any] | None = None,
|
|
method: str = "POST",
|
|
retries: int = 0,
|
|
stream: bool = False,
|
|
server_wait: bool = True,
|
|
fields: dict[str, tuple] | None = None,
|
|
error_handler: Callable = None,
|
|
retry_body: Callable[[], Any] | None = None,
|
|
) -> HTTPResponse:
|
|
if isinstance(data, str):
|
|
data = data.encode()
|
|
headers = dict_copy(self.headers, headers)
|
|
attempts = 0
|
|
final_params = {}
|
|
if server_wait:
|
|
final_params["wait_end_of_query"] = "1"
|
|
# We can't actually read the progress headers, but we enable them so ClickHouse sends something
|
|
# to keep the connection alive when waiting for long-running queries and (2) to get summary information
|
|
# if not streaming
|
|
if self._send_progress:
|
|
final_params["send_progress_in_http_headers"] = "1"
|
|
if self._progress_interval:
|
|
final_params["http_headers_progress_interval_ms"] = self._progress_interval
|
|
final_params = dict_copy(self.params, final_params)
|
|
final_params = dict_copy(final_params, params)
|
|
|
|
if self._autogenerate_query_id and "query_id" not in final_params:
|
|
final_params["query_id"] = str(uuid.uuid4())
|
|
|
|
url = f"{self.url}?{urlencode(final_params)}"
|
|
kwargs = {"headers": headers, "timeout": self.timeout, "retries": self.http_retries, "preload_content": not stream}
|
|
if self.server_host_name:
|
|
kwargs["assert_same_host"] = False
|
|
kwargs["headers"].update({"Host": self.server_host_name})
|
|
if fields:
|
|
kwargs["fields"] = fields
|
|
else:
|
|
kwargs["body"] = data
|
|
check_conn_expiration(self.http)
|
|
query_session = final_params.get("session_id")
|
|
while True:
|
|
attempts += 1
|
|
if query_session:
|
|
if query_session == self._active_session:
|
|
raise ProgrammingError(
|
|
"Attempt to execute concurrent queries within the same session. "
|
|
+ "Please use a separate client instance per thread/process."
|
|
)
|
|
# There is a race condition here when using multiprocessing -- in that case the server will
|
|
# throw an error instead, but in most cases this more helpful error will be thrown first
|
|
self._active_session = query_session
|
|
try:
|
|
response = self.http.request(method, url, **kwargs)
|
|
except HTTPError as ex:
|
|
# Always allow at least one retry on a clean connection error so a single stale
|
|
# keep-alive socket doesn't surface to the caller, and additionally honor the
|
|
# retries budget when it is larger (e.g. query_retries for reads), so that
|
|
# bursts of stale pooled connections can be drained before giving up.
|
|
max_attempts = max(2, retries + 1)
|
|
if isinstance(ex.__context__, ConnectionResetError) and attempts < max_attempts:
|
|
# The server closed the connection, probably because the Keep Alive has expired.
|
|
# We should be safe to retry, as ClickHouse should not have processed anything on
|
|
# a connection that it killed.
|
|
body = kwargs.get("body")
|
|
if retry_body is not None:
|
|
kwargs["body"] = retry_body()
|
|
logger.debug("Retrying remotely closed connection with rebuilt body (attempt %s/%s)", attempts, max_attempts)
|
|
time.sleep(0.1 * attempts)
|
|
continue
|
|
if body is None or isinstance(body, (bytes, bytearray, str)):
|
|
logger.debug("Retrying remotely closed connection (attempt %s/%s)", attempts, max_attempts)
|
|
time.sleep(0.1 * attempts)
|
|
continue
|
|
logger.warning("Unexpected Http Driver Exception")
|
|
err_url = f" ({self.url})" if self.show_clickhouse_errors else ""
|
|
raise OperationalError(f"Error {ex} executing HTTP request attempt {attempts}{err_url}") from ex
|
|
finally:
|
|
if query_session:
|
|
self._active_session = None # Make sure we always clear this
|
|
if 200 <= response.status < 300 and not response.headers.get(ex_header):
|
|
return response
|
|
if response.status in (429, 503, 504):
|
|
if attempts > retries:
|
|
self._error_handler(response, True)
|
|
logger.debug("Retrying requests with status code %d", response.status)
|
|
elif error_handler:
|
|
error_handler(response)
|
|
else:
|
|
self._error_handler(response)
|
|
|
|
def raw_query(
|
|
self,
|
|
query: str,
|
|
parameters: Sequence | dict[str, Any] | None = None,
|
|
settings: dict[str, Any] | None = None,
|
|
fmt: str = None,
|
|
use_database: bool = True,
|
|
external_data: ExternalData | None = None,
|
|
transport_settings: dict[str, str] | None = 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=transport_settings, retries=self.query_retries).data
|
|
|
|
def raw_stream(
|
|
self,
|
|
query: str,
|
|
parameters: Sequence | dict[str, Any] | None = None,
|
|
settings: dict[str, Any] | None = None,
|
|
fmt: str = None,
|
|
use_database: bool = True,
|
|
external_data: ExternalData | None = None,
|
|
transport_settings: dict[str, str] | None = 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=transport_settings,
|
|
retries=self.query_retries,
|
|
)
|
|
|
|
def _prep_raw_query(
|
|
self,
|
|
query: str,
|
|
parameters: Sequence | dict[str, Any] | None,
|
|
settings: dict[str, Any] | None,
|
|
fmt: str,
|
|
use_database: bool,
|
|
external_data: ExternalData | None,
|
|
):
|
|
if fmt:
|
|
query += f"\n FORMAT {fmt}"
|
|
final_query, bind_params = bind_query(query, parameters, self.server_tz)
|
|
params = self._validate_settings(settings or {})
|
|
if use_database and self.database:
|
|
params["database"] = self.database
|
|
fields = {}
|
|
# Setup query body
|
|
if external_data and not self.form_encode_query_params and isinstance(final_query, bytes):
|
|
raise ProgrammingError("Binary query cannot be placed in URL when using External Data; enable form encoding.")
|
|
# Setup additional query parameters and body
|
|
if self.form_encode_query_params:
|
|
body = b""
|
|
fields["query"] = final_query
|
|
fields.update(bind_params)
|
|
if external_data:
|
|
params.update(external_data.query_params)
|
|
fields.update(external_data.form_data)
|
|
elif external_data:
|
|
params.update(bind_params)
|
|
body = b""
|
|
params["query"] = final_query
|
|
params.update(external_data.query_params)
|
|
fields = external_data.form_data
|
|
else:
|
|
params.update(bind_params)
|
|
body = final_query
|
|
fields = None
|
|
return body, params, fields
|
|
|
|
def _add_integration_tag(self, name: str):
|
|
"""
|
|
Dynamically adds a product (like pandas or sqlalchemy) to the User-Agent string details section.
|
|
"""
|
|
if not common.get_setting("send_integration_tags") or name in self._reported_libs:
|
|
return
|
|
|
|
try:
|
|
ver = "unknown"
|
|
try:
|
|
ver = dist_version(name)
|
|
except Exception:
|
|
try:
|
|
mod = import_module(name)
|
|
ver = getattr(mod, "__version__", "unknown")
|
|
except Exception:
|
|
pass
|
|
|
|
product_info = f"{name}/{ver}"
|
|
|
|
ua = self.headers.get("User-Agent", "")
|
|
start = ua.find("(")
|
|
if start == -1:
|
|
return
|
|
end = ua.find(")", start + 1)
|
|
if end == -1:
|
|
return
|
|
|
|
details = ua[start + 1 : end].strip()
|
|
|
|
if product_info in details:
|
|
self._reported_libs.add(name)
|
|
return
|
|
|
|
new_details = f"{product_info}; {details}" if details else product_info
|
|
new_ua = f"{ua[: start + 1]}{new_details}{ua[end:]}"
|
|
self.headers["User-Agent"] = new_ua.strip()
|
|
|
|
self._reported_libs.add(name)
|
|
logger.debug("Added '%s' to User-Agent", product_info)
|
|
|
|
except Exception as e:
|
|
logger.debug("Problem adding '%s' to User-Agent: %s", name, e)
|
|
|
|
def ping(self) -> bool:
|
|
"""
|
|
See BaseClient doc_string for this method
|
|
"""
|
|
try:
|
|
response = self.http.request("GET", f"{self.url}/ping", timeout=3, preload_content=True)
|
|
return 200 <= response.status < 300
|
|
except HTTPError:
|
|
logger.debug("ping failed", exc_info=True)
|
|
return False
|
|
|
|
def close_connections(self) -> None:
|
|
self.http.clear()
|
|
|
|
def close(self) -> None:
|
|
if self._owns_pool_manager:
|
|
self.http.clear()
|
|
all_managers.pop(self.http, None)
|