Files
0537253223 async client: retry ServerDisconnectedErrors up to retries (#755)
* 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]>
2026-05-21 14:52:02 -07:00

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)