Files
Joe SpadolaandGitHub bb4496281a rebuild body on connection-reset insert retry (#732)
* rebuild body on connection-reset insert retry

* wait for producer thread on async retry
2026-05-07 13:31:18 -07:00

734 lines
31 KiB
Python

import io
import json
import logging
import re
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:
if isinstance(ex.__context__, ConnectionResetError) and attempts == 1:
# 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. We also only retry this once, as multiple disconnects are unlikely to be
# related to the Keep Alive settings
body = kwargs.get("body")
if retry_body is not None:
kwargs["body"] = retry_body()
logger.debug("Retrying remotely closed connection with rebuilt body")
continue
if body is None or isinstance(body, (bytes, bytearray, str)):
logger.debug("Retrying remotely closed connection")
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).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)
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)