mirror of
https://github.com/hanzo-ds/connect.git
synced 2026-10-30 20:31:24 +00:00
d19ef5e0be
* refactor(driver): accept None for port and database as canonical 'not specified' * docs: track internal driver signature default change * keep legacy port and database sentinels working * allow trailing slash --------- Co-authored-by: Joe S <[email protected]>
1280 lines
57 KiB
Python
1280 lines
57 KiB
Python
from __future__ import annotations
|
|
|
|
import io
|
|
import logging
|
|
from abc import ABC, abstractmethod
|
|
from collections.abc import Generator, Sequence
|
|
from datetime import timezone, tzinfo
|
|
from typing import (
|
|
TYPE_CHECKING,
|
|
Any,
|
|
BinaryIO,
|
|
cast,
|
|
)
|
|
from zoneinfo import ZoneInfoNotFoundError
|
|
|
|
from clickhouse_connect import common
|
|
from clickhouse_connect.common import version
|
|
from clickhouse_connect.datatypes import dynamic as dynamic_module
|
|
from clickhouse_connect.datatypes.base import ClickHouseType
|
|
from clickhouse_connect.datatypes.registry import get_from_name
|
|
from clickhouse_connect.driver import options, tzutil
|
|
from clickhouse_connect.driver.binding import quote_identifier
|
|
from clickhouse_connect.driver.common import (
|
|
StreamContext,
|
|
coerce_bool,
|
|
coerce_int,
|
|
dict_copy,
|
|
)
|
|
from clickhouse_connect.driver.constants import (
|
|
CH_VERSION_WITH_PROTOCOL,
|
|
PROTOCOL_VERSION_WITH_LOW_CARD,
|
|
)
|
|
from clickhouse_connect.driver.exceptions import (
|
|
DataError,
|
|
OperationalError,
|
|
ProgrammingError,
|
|
)
|
|
from clickhouse_connect.driver.external import ExternalData
|
|
from clickhouse_connect.driver.insert import InsertContext
|
|
from clickhouse_connect.driver.models import ColumnDef, SettingDef, SettingStatus
|
|
from clickhouse_connect.driver.options import (
|
|
check_arrow,
|
|
check_numpy,
|
|
check_pandas,
|
|
check_polars,
|
|
)
|
|
from clickhouse_connect.driver.query import (
|
|
_VALID_TZ_MODES,
|
|
_VALID_TZ_SOURCES,
|
|
QueryContext,
|
|
QueryResult,
|
|
TzMode,
|
|
TzSource,
|
|
arrow_buffer,
|
|
to_arrow,
|
|
to_arrow_batches,
|
|
)
|
|
from clickhouse_connect.driver.summary import QuerySummary
|
|
from clickhouse_connect.driver.types import Closable
|
|
|
|
if TYPE_CHECKING:
|
|
import numpy
|
|
import pandas
|
|
import polars
|
|
import pyarrow
|
|
|
|
io.DEFAULT_BUFFER_SIZE = 1024 * 256 # type: ignore[misc] # override module default buffer size
|
|
logger = logging.getLogger(__name__)
|
|
arrow_str_setting = "output_format_arrow_string_as_string"
|
|
|
|
|
|
def _strip_utc_timezone_from_arrow(table: pyarrow.Table) -> pyarrow.Table:
|
|
"""Strip UTC timezone from timestamp columns in Arrow table.
|
|
|
|
This ensures naive datetimes are returned when the server timezone is UTC
|
|
and tz_mode is 'naive_utc' (the default).
|
|
|
|
Only UTC-equivalent timezones (UTC, Etc/UTC, GMT, etc.) are stripped.
|
|
Non-UTC timezones carry important offset information and are always
|
|
preserved regardless of tz_mode setting.
|
|
"""
|
|
new_fields = []
|
|
needs_cast = False
|
|
for field in table.schema:
|
|
if options.arrow.types.is_timestamp(field.type) and tzutil.is_utc_timezone(field.type.tz):
|
|
new_fields.append(options.arrow.field(field.name, options.arrow.timestamp(field.type.unit)))
|
|
needs_cast = True
|
|
else:
|
|
new_fields.append(field)
|
|
if needs_cast:
|
|
return table.cast(options.arrow.schema(new_fields))
|
|
return table
|
|
|
|
|
|
def _apply_arrow_tz_policy(table: pyarrow.Table, tz_mode: str) -> pyarrow.Table:
|
|
"""Apply the tz_mode policy to an Arrow table before conversion.
|
|
|
|
Handles UTC stripping when tz_mode is "naive_utc" and warns when
|
|
tz_mode is "schema" since that mode is not yet implemented for
|
|
Arrow-based queries.
|
|
"""
|
|
if tz_mode == "schema":
|
|
logger.warning(
|
|
'tz_mode="schema" is not yet supported for Arrow-based query methods. '
|
|
"It would require a separate schema lookup since ClickHouse attaches the server "
|
|
"timezone to all DateTime columns in Arrow format. Use query/query_df for "
|
|
"schema-matching behavior or open an issue if you need Arrow support."
|
|
)
|
|
if tz_mode == "naive_utc":
|
|
table = _strip_utc_timezone_from_arrow(table)
|
|
return table
|
|
|
|
|
|
class Client(ABC):
|
|
"""
|
|
Base ClickHouse Connect client
|
|
"""
|
|
|
|
compression: str | None = None
|
|
write_compression: str | None = None
|
|
protocol_version = 0
|
|
valid_transport_settings: set[str] = set()
|
|
optional_transport_settings: set[str] = set()
|
|
database = None
|
|
max_error_message = 0
|
|
_tz_source: TzSource = "auto"
|
|
_apply_server_tz = False
|
|
tz_mode: TzMode = "naive_utc"
|
|
show_clickhouse_errors = True
|
|
|
|
@property
|
|
def tz_source(self) -> TzSource:
|
|
return self._tz_source
|
|
|
|
@tz_source.setter
|
|
def tz_source(self, value: TzSource):
|
|
if value not in _VALID_TZ_SOURCES:
|
|
raise ProgrammingError(f'tz_source must be "auto", "server", or "local", got "{value}"')
|
|
self._tz_source = value
|
|
if value == "auto":
|
|
self._apply_server_tz = self._dst_safe
|
|
else:
|
|
self._apply_server_tz = value == "server"
|
|
|
|
def __init__(
|
|
self,
|
|
database: str | None,
|
|
query_limit: int,
|
|
uri: str,
|
|
query_retries: int,
|
|
server_host_name: str | None,
|
|
tz_source: TzSource | None = None,
|
|
tz_mode: TzMode | None = None,
|
|
show_clickhouse_errors: bool | None = None,
|
|
autoconnect: bool = True,
|
|
):
|
|
"""
|
|
Shared initialization of ClickHouse Connect client
|
|
:param database: database name
|
|
:param query_limit: default LIMIT for queries
|
|
:param uri: uri for error messages
|
|
:param tz_source: Controls how the client determines the fallback timezone for DateTime columns without an
|
|
explicit timezone. "auto" (default) auto-detects based on DST safety of server timezone. "server" always
|
|
uses the server timezone. "local" always uses the local timezone.
|
|
:param tz_mode: Controls timezone-aware behavior for UTC DateTime columns. "naive_utc" (default) returns
|
|
naive UTC timestamps. "aware" forces timezone-aware UTC datetimes. "schema" returns datetimes that
|
|
match the server's column definition which means timezone-aware when the column defines a timezone and naive
|
|
for bare DateTime columns.
|
|
:param autoconnect: If True, immediately connect to server and fetch settings. If False,
|
|
defer connection to _connect() method. Used by async clients to avoid blocking I/O in __init__.
|
|
"""
|
|
self.query_limit = coerce_int(query_limit)
|
|
self.query_retries = coerce_int(query_retries)
|
|
if database and database != "__default__":
|
|
self.database = database
|
|
if show_clickhouse_errors is not None:
|
|
self.show_clickhouse_errors = coerce_bool(show_clickhouse_errors)
|
|
self.server_host_name = server_host_name
|
|
self.uri = uri
|
|
self.tz_mode = tz_mode if tz_mode is not None else "naive_utc"
|
|
if self.tz_mode not in _VALID_TZ_MODES:
|
|
raise ProgrammingError(f'tz_mode must be "naive_utc", "aware", or "schema", got "{self.tz_mode}"')
|
|
resolved_tz_source = tz_source if tz_source is not None else "auto"
|
|
if resolved_tz_source not in _VALID_TZ_SOURCES:
|
|
raise ProgrammingError(f'tz_source must be "auto", "server", or "local", got "{resolved_tz_source}"')
|
|
self._tz_source = resolved_tz_source
|
|
|
|
# Initialize attributes that will be set during connection
|
|
self.server_version: str | None = None
|
|
self.server_tz: tzinfo = timezone.utc
|
|
self.server_settings: dict[str, SettingDef] = {}
|
|
|
|
if autoconnect:
|
|
self._init_common_settings(resolved_tz_source)
|
|
else:
|
|
# Store for deferred async initialization
|
|
self._deferred_tz_source = resolved_tz_source
|
|
|
|
def _init_common_settings(self, tz_source: TzSource):
|
|
self.server_tz, self._dst_safe = timezone.utc, True
|
|
version_result = self.command("SELECT version(), timezone()", use_database=False)
|
|
if not isinstance(version_result, Sequence) or isinstance(version_result, str):
|
|
raise OperationalError(f"Unexpected response to server version query: {version_result!r}")
|
|
self.server_version, server_tz = version_result[0], version_result[1]
|
|
try:
|
|
server_tz_info = tzutil.resolve_zone(server_tz)
|
|
server_tz_info, self._dst_safe = tzutil.normalize_timezone(server_tz_info, trust_fixed_offset=True)
|
|
self.server_tz = server_tz_info
|
|
except ZoneInfoNotFoundError:
|
|
logger.warning(
|
|
"Server timezone %s could not be resolved, falling back to UTC; %s",
|
|
server_tz,
|
|
tzutil.TZDATA_HINT,
|
|
)
|
|
if tz_source == "auto":
|
|
self._apply_server_tz = self._dst_safe
|
|
else:
|
|
self._apply_server_tz = tz_source == "server"
|
|
|
|
if not self._apply_server_tz and not tzutil.local_tz_dst_safe:
|
|
logger.warning(
|
|
"local timezone %s may return unexpected times due to Daylight Savings Time/" + "Summer Time differences",
|
|
tzutil.local_tz.tzname(None),
|
|
)
|
|
readonly = "readonly"
|
|
if not self.min_version("19.17"):
|
|
readonly = common.get_setting("readonly")
|
|
server_settings = self.query(f"SELECT name, value, {readonly} as readonly FROM system.settings LIMIT 10000")
|
|
self.server_settings = {row["name"]: SettingDef(**row) for row in server_settings.named_results()}
|
|
|
|
if self.min_version(CH_VERSION_WITH_PROTOCOL) and common.get_setting("use_protocol_version"):
|
|
# Unfortunately we have to validate that the client protocol version is actually used by ClickHouse
|
|
# since the query parameter could be stripped off (in particular, by CHProxy)
|
|
test_data = self.raw_query(
|
|
"SELECT 1 AS check", fmt="Native", settings={"client_protocol_version": PROTOCOL_VERSION_WITH_LOW_CARD}
|
|
)
|
|
if test_data[8:16] == b"\x01\x01\x05check":
|
|
self.protocol_version = PROTOCOL_VERSION_WITH_LOW_CARD
|
|
if self._setting_status("date_time_input_format").is_writable:
|
|
self.set_client_setting("date_time_input_format", "best_effort")
|
|
if (
|
|
self._setting_status("allow_experimental_json_type").is_set
|
|
and self._setting_status("cast_string_to_dynamic_use_inference").is_writable
|
|
):
|
|
self.set_client_setting("cast_string_to_dynamic_use_inference", "1")
|
|
if self.min_version("24.8") and not self.min_version("24.10"):
|
|
dynamic_module.json_serialization_format = 0
|
|
|
|
def _validate_settings(self, settings: dict[str, Any] | None) -> dict[str, str]:
|
|
"""
|
|
This strips any ClickHouse settings that are not recognized or are read only.
|
|
:param settings: Dictionary of setting name and values
|
|
:return: A filtered dictionary of settings with values rendered as strings
|
|
"""
|
|
validated: dict[str, str] = {}
|
|
invalid_action = common.get_setting("invalid_setting_action")
|
|
if not settings:
|
|
return validated
|
|
for key, value in settings.items():
|
|
str_value = self._validate_setting(key, value, invalid_action)
|
|
if str_value is not None:
|
|
validated[key] = value
|
|
return validated
|
|
|
|
def _validate_setting(self, key: str, value: Any, invalid_action: str) -> str | None:
|
|
str_value = str(value)
|
|
if value is True:
|
|
str_value = "1"
|
|
elif value is False:
|
|
str_value = "0"
|
|
if key not in self.valid_transport_settings:
|
|
setting_def = self.server_settings.get(key)
|
|
current_setting = self.get_client_setting(key)
|
|
# Skip if the requested value matches the server's value and either the setting
|
|
# is readonly i.e. there's nothing to change or already explicitly stored on the client
|
|
if setting_def and setting_def.value == str_value:
|
|
if setting_def.readonly or (current_setting is not None and current_setting == setting_def.value):
|
|
return None
|
|
if setting_def is None or setting_def.readonly:
|
|
if key in self.optional_transport_settings:
|
|
return None
|
|
if invalid_action == "send":
|
|
logger.warning("Attempting to send unrecognized or readonly setting %s", key)
|
|
elif invalid_action == "drop":
|
|
logger.warning("Dropping unrecognized or readonly settings %s", key)
|
|
return None
|
|
else:
|
|
raise ProgrammingError(f"Setting {key} is unknown or readonly") from None
|
|
return str_value
|
|
|
|
def _setting_status(self, key: str) -> SettingStatus:
|
|
comp_setting = self.server_settings.get(key)
|
|
if not comp_setting:
|
|
return SettingStatus(False, False)
|
|
return SettingStatus(comp_setting.value != "0", comp_setting.readonly != 1)
|
|
|
|
def _prep_query(self, context: QueryContext):
|
|
if context.is_select and not context.has_limit and self.query_limit:
|
|
limit = f"\n LIMIT {self.query_limit}"
|
|
if isinstance(context.query, bytes):
|
|
return context.final_query + limit.encode()
|
|
return context.final_query + limit
|
|
return context.final_query
|
|
|
|
def _check_tz_change(self, new_tz) -> tzinfo | None:
|
|
if new_tz:
|
|
try:
|
|
new_tzinfo = tzutil.resolve_zone(new_tz)
|
|
if new_tzinfo != self.server_tz:
|
|
return new_tzinfo
|
|
except ZoneInfoNotFoundError:
|
|
logger.warning(
|
|
"Unrecognized timezone %s received from ClickHouse; %s",
|
|
new_tz,
|
|
tzutil.TZDATA_HINT,
|
|
)
|
|
return None
|
|
|
|
@abstractmethod
|
|
def _query_with_context(self, context: QueryContext) -> QueryResult:
|
|
pass
|
|
|
|
@abstractmethod
|
|
def set_client_setting(self, key: str, value: Any) -> None:
|
|
"""
|
|
Set a clickhouse setting for the client after initialization. If a setting is not recognized by ClickHouse,
|
|
or the setting is identified as "read_only", this call will either throw a Programming exception or attempt
|
|
to send the setting anyway based on the common setting 'invalid_setting_action'
|
|
:param key: ClickHouse setting name
|
|
:param value: ClickHouse setting value
|
|
"""
|
|
|
|
@abstractmethod
|
|
def get_client_setting(self, key: str) -> str | None:
|
|
"""
|
|
:param key: The setting key
|
|
:return: The string value of the setting, if it exists, or None
|
|
"""
|
|
|
|
@abstractmethod
|
|
def set_access_token(self, access_token: str) -> None:
|
|
"""
|
|
Set the ClickHouse access token for the client
|
|
:param access_token: Access token string
|
|
"""
|
|
|
|
def query(
|
|
self,
|
|
query: str | None = None,
|
|
parameters: Sequence | dict[str, Any] | None = None,
|
|
settings: dict[str, Any] | None = None,
|
|
query_formats: dict[str, str] | None = None,
|
|
column_formats: dict[str, str | dict[str, str]] | None = None,
|
|
encoding: str | None = None,
|
|
use_none: bool | None = None,
|
|
column_oriented: bool | None = None,
|
|
use_numpy: bool | None = None,
|
|
max_str_len: int | None = None,
|
|
context: QueryContext | None = None,
|
|
query_tz: str | tzinfo | None = None,
|
|
column_tzs: dict[str, str | tzinfo] | None = None,
|
|
external_data: ExternalData | None = None,
|
|
transport_settings: dict[str, str] | None = None,
|
|
tz_mode: TzMode | None = None,
|
|
) -> QueryResult:
|
|
"""
|
|
Main query method for SELECT, DESCRIBE and other SQL statements that return a result matrix. For
|
|
parameters, see the create_query_context method
|
|
:return: QueryResult -- data and metadata from response
|
|
"""
|
|
if query and query.lower().strip().startswith("select __connect_version__"):
|
|
return QueryResult(
|
|
[[f"ClickHouse Connect v.{version()} â“’ ClickHouse Inc."]],
|
|
None,
|
|
("connect_version",),
|
|
(get_from_name("String"),),
|
|
)
|
|
kwargs = locals().copy()
|
|
del kwargs["self"]
|
|
query_context = self.create_query_context(**kwargs)
|
|
if query_context.is_command:
|
|
response = self.command(
|
|
cast(str, query),
|
|
parameters=query_context.parameters,
|
|
settings=query_context.settings,
|
|
external_data=query_context.external_data,
|
|
transport_settings=query_context.transport_settings,
|
|
)
|
|
if isinstance(response, QuerySummary):
|
|
return response.as_query_result()
|
|
return QueryResult([response] if isinstance(response, list) else [[response]])
|
|
return self._query_with_context(query_context)
|
|
|
|
def query_column_block_stream(
|
|
self,
|
|
query: str | None = None,
|
|
parameters: Sequence | dict[str, Any] | None = None,
|
|
settings: dict[str, Any] | None = None,
|
|
query_formats: dict[str, str] | None = None,
|
|
column_formats: dict[str, str | dict[str, str]] | None = None,
|
|
encoding: str | None = None,
|
|
use_none: bool | None = None,
|
|
context: QueryContext | None = None,
|
|
query_tz: str | tzinfo | None = None,
|
|
column_tzs: dict[str, str | tzinfo] | None = None,
|
|
external_data: ExternalData | None = None,
|
|
transport_settings: dict[str, str] | None = None,
|
|
tz_mode: TzMode | None = None,
|
|
) -> StreamContext:
|
|
"""
|
|
Variation of main query method that returns a stream of column oriented blocks. For
|
|
parameters, see the create_query_context method.
|
|
:return: StreamContext -- Iterable stream context that returns column oriented blocks
|
|
"""
|
|
return self._context_query(locals(), use_numpy=False, streaming=True).column_block_stream
|
|
|
|
def query_row_block_stream(
|
|
self,
|
|
query: str | None = None,
|
|
parameters: Sequence | dict[str, Any] | None = None,
|
|
settings: dict[str, Any] | None = None,
|
|
query_formats: dict[str, str] | None = None,
|
|
column_formats: dict[str, str | dict[str, str]] | None = None,
|
|
encoding: str | None = None,
|
|
use_none: bool | None = None,
|
|
context: QueryContext | None = None,
|
|
query_tz: str | tzinfo | None = None,
|
|
column_tzs: dict[str, str | tzinfo] | None = None,
|
|
external_data: ExternalData | None = None,
|
|
transport_settings: dict[str, str] | None = None,
|
|
tz_mode: TzMode | None = None,
|
|
) -> StreamContext:
|
|
"""
|
|
Variation of main query method that returns a stream of row oriented blocks. For
|
|
parameters, see the create_query_context method.
|
|
:return: StreamContext -- Iterable stream context that returns blocks of rows
|
|
"""
|
|
return self._context_query(locals(), use_numpy=False, streaming=True).row_block_stream
|
|
|
|
def query_rows_stream(
|
|
self,
|
|
query: str | None = None,
|
|
parameters: Sequence | dict[str, Any] | None = None,
|
|
settings: dict[str, Any] | None = None,
|
|
query_formats: dict[str, str] | None = None,
|
|
column_formats: dict[str, str | dict[str, str]] | None = None,
|
|
encoding: str | None = None,
|
|
use_none: bool | None = None,
|
|
context: QueryContext | None = None,
|
|
query_tz: str | tzinfo | None = None,
|
|
column_tzs: dict[str, str | tzinfo] | None = None,
|
|
external_data: ExternalData | None = None,
|
|
transport_settings: dict[str, str] | None = None,
|
|
tz_mode: TzMode | None = None,
|
|
) -> StreamContext:
|
|
"""
|
|
Variation of main query method that returns a stream of row oriented blocks. For
|
|
parameters, see the create_query_context method.
|
|
:return: StreamContext -- Iterable stream context that returns blocks of rows
|
|
"""
|
|
return self._context_query(locals(), use_numpy=False, streaming=True).rows_stream
|
|
|
|
@abstractmethod
|
|
def raw_query(
|
|
self,
|
|
query: str,
|
|
parameters: Sequence | dict[str, Any] | None = None,
|
|
settings: dict[str, Any] | None = None,
|
|
fmt: str | None = None,
|
|
use_database: bool = True,
|
|
external_data: ExternalData | None = None,
|
|
transport_settings: dict[str, str] | None = None,
|
|
) -> bytes:
|
|
"""
|
|
Query method that simply returns the raw ClickHouse format bytes
|
|
:param query: Query statement/format string
|
|
:param parameters: Optional dictionary used to format the query
|
|
:param settings: Optional dictionary of ClickHouse settings (key/string values)
|
|
:param fmt: ClickHouse output format
|
|
:param use_database: Send the database parameter to ClickHouse so the command will be executed in the client
|
|
database context.
|
|
:param external_data: External data to send with the query
|
|
:param transport_settings: Optional dictionary of transport level settings (HTTP headers, etc.)
|
|
:return: bytes representing raw ClickHouse return value based on format
|
|
"""
|
|
|
|
@abstractmethod
|
|
def raw_stream(
|
|
self,
|
|
query: str,
|
|
parameters: Sequence | dict[str, Any] | None = None,
|
|
settings: dict[str, Any] | None = None,
|
|
fmt: str | None = None,
|
|
use_database: bool = True,
|
|
external_data: ExternalData | None = None,
|
|
transport_settings: dict[str, str] | None = None,
|
|
) -> io.IOBase | StreamContext:
|
|
"""
|
|
Query method that returns the result as a stream iterator.
|
|
:param query: Query statement/format string
|
|
:param parameters: Optional dictionary used to format the query
|
|
:param settings: Optional dictionary of ClickHouse settings (key/string values)
|
|
:param fmt: ClickHouse output format
|
|
:param use_database Send the database parameter to ClickHouse so the command will be executed in the client
|
|
database context.
|
|
:param external_data: External data to send with the query.
|
|
:param transport_settings: Optional dictionary of transport level settings (HTTP headers, etc.)
|
|
:return: io.IOBase (sync) or StreamContext (async) - both support iteration over raw bytes
|
|
"""
|
|
|
|
def query_np(
|
|
self,
|
|
query: str | None = None,
|
|
parameters: Sequence | dict[str, Any] | None = None,
|
|
settings: dict[str, Any] | None = None,
|
|
query_formats: dict[str, str] | None = None,
|
|
column_formats: dict[str, str] | None = None,
|
|
encoding: str | None = None,
|
|
use_none: bool | None = None,
|
|
max_str_len: int | None = None,
|
|
context: QueryContext | None = None,
|
|
external_data: ExternalData | None = None,
|
|
transport_settings: dict[str, str] | None = None,
|
|
) -> numpy.ndarray:
|
|
"""
|
|
Query method that returns the results as a numpy array. For parameter values, see the
|
|
create_query_context method
|
|
:return: Numpy array representing the result set
|
|
"""
|
|
check_numpy()
|
|
self._add_integration_tag("numpy")
|
|
return self._context_query(locals(), use_numpy=True).np_result
|
|
|
|
def query_np_stream(
|
|
self,
|
|
query: str | None = None,
|
|
parameters: Sequence | dict[str, Any] | None = None,
|
|
settings: dict[str, Any] | None = None,
|
|
query_formats: dict[str, str] | None = None,
|
|
column_formats: dict[str, str] | None = None,
|
|
encoding: str | None = None,
|
|
use_none: bool | None = None,
|
|
max_str_len: int | None = None,
|
|
context: QueryContext | None = None,
|
|
external_data: ExternalData | None = None,
|
|
transport_settings: dict[str, str] | None = None,
|
|
) -> StreamContext:
|
|
"""
|
|
Query method that returns the results as a stream of numpy arrays. For parameter values, see the
|
|
create_query_context method
|
|
:return: Generator that yield a numpy array per block representing the result set
|
|
"""
|
|
check_numpy()
|
|
self._add_integration_tag("numpy")
|
|
return self._context_query(locals(), use_numpy=True, streaming=True).np_stream
|
|
|
|
def query_df(
|
|
self,
|
|
query: str | None = None,
|
|
parameters: Sequence | dict[str, Any] | None = None,
|
|
settings: dict[str, Any] | None = None,
|
|
query_formats: dict[str, str] | None = None,
|
|
column_formats: dict[str, str] | None = None,
|
|
encoding: str | None = None,
|
|
use_none: bool | None = None,
|
|
max_str_len: int | None = None,
|
|
use_na_values: bool | None = None,
|
|
query_tz: str | None = None,
|
|
column_tzs: dict[str, str | tzinfo] | None = None,
|
|
context: QueryContext | None = None,
|
|
external_data: ExternalData | None = None,
|
|
use_extended_dtypes: bool | None = None,
|
|
transport_settings: dict[str, str] | None = None,
|
|
tz_mode: TzMode | None = None,
|
|
) -> pandas.DataFrame:
|
|
"""
|
|
Query method that results the results as a pandas dataframe. For parameter values, see the
|
|
create_query_context method
|
|
:return: Pandas dataframe representing the result set
|
|
"""
|
|
check_pandas()
|
|
self._add_integration_tag("pandas")
|
|
return self._context_query(locals(), use_numpy=True, as_pandas=True).df_result
|
|
|
|
def query_df_stream(
|
|
self,
|
|
query: str | None = None,
|
|
parameters: Sequence | dict[str, Any] | None = None,
|
|
settings: dict[str, Any] | None = None,
|
|
query_formats: dict[str, str] | None = None,
|
|
column_formats: dict[str, str] | None = None,
|
|
encoding: str | None = None,
|
|
use_none: bool | None = None,
|
|
max_str_len: int | None = None,
|
|
use_na_values: bool | None = None,
|
|
query_tz: str | None = None,
|
|
column_tzs: dict[str, str | tzinfo] | None = None,
|
|
context: QueryContext | None = None,
|
|
external_data: ExternalData | None = None,
|
|
use_extended_dtypes: bool | None = None,
|
|
transport_settings: dict[str, str] | None = None,
|
|
tz_mode: TzMode | None = None,
|
|
) -> StreamContext:
|
|
"""
|
|
Query method that returns the results as a StreamContext. For parameter values, see the
|
|
create_query_context method
|
|
:return: Generator that yields a Pandas dataframe per block representing the result set
|
|
"""
|
|
check_pandas()
|
|
self._add_integration_tag("pandas")
|
|
return self._context_query(locals(), use_numpy=True, as_pandas=True, streaming=True).df_stream
|
|
|
|
def create_query_context(
|
|
self,
|
|
query: str | bytes | None = None,
|
|
parameters: Sequence | dict[str, Any] | None = None,
|
|
settings: dict[str, Any] | None = None,
|
|
query_formats: dict[str, str] | None = None,
|
|
column_formats: dict[str, str | dict[str, str]] | None = None,
|
|
encoding: str | None = None,
|
|
use_none: bool | None = None,
|
|
column_oriented: bool | None = None,
|
|
use_numpy: bool | None = False,
|
|
max_str_len: int | None = 0,
|
|
context: QueryContext | None = None,
|
|
query_tz: str | tzinfo | None = None,
|
|
column_tzs: dict[str, str | tzinfo] | None = None,
|
|
use_na_values: bool | None = None,
|
|
streaming: bool = False,
|
|
as_pandas: bool = False,
|
|
external_data: ExternalData | None = None,
|
|
use_extended_dtypes: bool | None = None,
|
|
transport_settings: dict[str, str] | None = None,
|
|
tz_mode: TzMode | None = None,
|
|
) -> QueryContext:
|
|
"""
|
|
Creates or updates a reusable QueryContext object
|
|
:param query: Query statement/format string
|
|
:param parameters: Optional dictionary used to format the query
|
|
:param settings: Optional dictionary of ClickHouse settings (key/string values)
|
|
:param query_formats: See QueryContext __init__ docstring
|
|
:param column_formats: See QueryContext __init__ docstring
|
|
:param encoding: See QueryContext __init__ docstring
|
|
:param use_none: Use None for ClickHouse NULL instead of default values. Note that using None in Numpy
|
|
arrays will force the numpy array dtype to 'object', which is often inefficient. This effect also
|
|
will impact the performance of Pandas dataframes.
|
|
:param column_oriented: Deprecated. Controls orientation of the QueryResult result_set property
|
|
:param use_numpy: Return QueryResult columns as one-dimensional numpy arrays
|
|
:param max_str_len: Limit returned ClickHouse String values to this length, which allows a Numpy
|
|
structured array even with ClickHouse variable length String columns. If 0, Numpy arrays for
|
|
String columns will always be object arrays
|
|
:param context: An existing QueryContext to be updated with any provided parameter values
|
|
:param query_tz: Either a string IANA timezone name or a tzinfo object (strings are resolved via zoneinfo).
|
|
Values for any DateTime or DateTime64 column in the query will be converted to Python datetime.datetime
|
|
objects with the selected timezone.
|
|
:param column_tzs: A dictionary of column names to tzinfo objects (or strings that will be converted to
|
|
tzinfo objects). The timezone will be applied to datetime objects returned in the query
|
|
:param tz_mode: Override the client default for handling UTC results. "aware" forces timezone-aware
|
|
UTC datetimes, "naive_utc" returns naive UTC datetimes, and "schema" returns datetimes matching the
|
|
server's column definition.
|
|
:param use_na_values: Deprecated alias for use_advanced_dtypes
|
|
:param as_pandas Return the result columns as pandas.Series objects
|
|
:param streaming Marker used to correctly configure streaming queries
|
|
:param external_data ClickHouse "external data" to send with query
|
|
:param use_extended_dtypes: Only relevant to Pandas Dataframe queries. Use Pandas "missing types", such as
|
|
pandas.NA and pandas.NaT for ClickHouse NULL values, as well as extended Pandas dtypes such as IntegerArray
|
|
and StringArray. Defaulted to True for query_df methods
|
|
:param transport_settings: Optional dictionary of transport level settings (HTTP headers, etc.)
|
|
:return: Reusable QueryContext
|
|
"""
|
|
resolved_tz_mode = tz_mode if tz_mode is not None else self.tz_mode
|
|
if context:
|
|
return context.updated_copy(
|
|
query=query,
|
|
parameters=parameters,
|
|
settings=settings,
|
|
query_formats=query_formats,
|
|
column_formats=column_formats,
|
|
encoding=encoding,
|
|
server_tz=self.server_tz,
|
|
use_none=use_none,
|
|
column_oriented=column_oriented,
|
|
use_numpy=use_numpy,
|
|
max_str_len=max_str_len,
|
|
query_tz=query_tz,
|
|
column_tzs=column_tzs,
|
|
tz_mode=resolved_tz_mode,
|
|
as_pandas=as_pandas,
|
|
use_extended_dtypes=use_extended_dtypes,
|
|
streaming=streaming,
|
|
external_data=external_data,
|
|
transport_settings=transport_settings,
|
|
)
|
|
if use_numpy and max_str_len is None:
|
|
max_str_len = 0
|
|
if use_extended_dtypes is None:
|
|
use_extended_dtypes = use_na_values
|
|
if as_pandas and use_extended_dtypes is None:
|
|
use_extended_dtypes = True
|
|
return QueryContext(
|
|
query=cast(str | bytes, query),
|
|
parameters=parameters,
|
|
settings=settings,
|
|
query_formats=query_formats,
|
|
column_formats=column_formats,
|
|
encoding=encoding,
|
|
server_tz=self.server_tz,
|
|
use_none=use_none,
|
|
column_oriented=column_oriented,
|
|
use_numpy=use_numpy,
|
|
max_str_len=max_str_len,
|
|
query_tz=query_tz,
|
|
column_tzs=column_tzs,
|
|
tz_mode=resolved_tz_mode,
|
|
use_extended_dtypes=use_extended_dtypes,
|
|
as_pandas=as_pandas,
|
|
streaming=streaming,
|
|
apply_server_tz=self._apply_server_tz,
|
|
external_data=external_data,
|
|
transport_settings=transport_settings,
|
|
)
|
|
|
|
def query_arrow(
|
|
self,
|
|
query: str,
|
|
parameters: Sequence | dict[str, Any] | None = None,
|
|
settings: dict[str, Any] | None = None,
|
|
use_strings: bool | None = None,
|
|
external_data: ExternalData | None = None,
|
|
transport_settings: dict[str, str] | None = None,
|
|
) -> pyarrow.Table:
|
|
"""
|
|
Query method using the ClickHouse Arrow format to return a PyArrow table
|
|
:param query: Query statement/format string
|
|
:param parameters: Optional dictionary used to format the query
|
|
:param settings: Optional dictionary of ClickHouse settings (key/string values)
|
|
:param use_strings: Convert ClickHouse String type to Arrow string type (instead of binary)
|
|
:param external_data: ClickHouse "external data" to send with query
|
|
:param transport_settings: Optional dictionary of transport level settings (HTTP headers, etc.)
|
|
:return: PyArrow.Table
|
|
"""
|
|
check_arrow()
|
|
self._add_integration_tag("arrow")
|
|
settings = self._update_arrow_settings(settings, use_strings)
|
|
return to_arrow(
|
|
self.raw_query(
|
|
query,
|
|
parameters,
|
|
settings,
|
|
fmt="Arrow",
|
|
external_data=external_data,
|
|
transport_settings=transport_settings,
|
|
)
|
|
)
|
|
|
|
def query_arrow_stream(
|
|
self,
|
|
query: str,
|
|
parameters: Sequence | dict[str, Any] | None = None,
|
|
settings: dict[str, Any] | None = None,
|
|
use_strings: bool | None = None,
|
|
external_data: ExternalData | None = None,
|
|
transport_settings: dict[str, str] | None = None,
|
|
) -> StreamContext:
|
|
"""
|
|
Query method that returns the results as a stream of Arrow tables
|
|
:param query: Query statement/format string
|
|
:param parameters: Optional dictionary used to format the query
|
|
:param settings: Optional dictionary of ClickHouse settings (key/string values)
|
|
:param use_strings: Convert ClickHouse String type to Arrow string type (instead of binary)
|
|
:param external_data: ClickHouse "external data" to send with query
|
|
:param transport_settings: Optional dictionary of transport level settings (HTTP headers, etc.)
|
|
:return: Generator that yields a PyArrow.Table for per block representing the result set
|
|
"""
|
|
check_arrow()
|
|
self._add_integration_tag("arrow")
|
|
settings = self._update_arrow_settings(settings, use_strings)
|
|
return to_arrow_batches(
|
|
cast(
|
|
io.IOBase,
|
|
self.raw_stream(
|
|
query,
|
|
parameters,
|
|
settings,
|
|
fmt="ArrowStream",
|
|
external_data=external_data,
|
|
transport_settings=transport_settings,
|
|
),
|
|
)
|
|
)
|
|
|
|
def query_df_arrow(
|
|
self,
|
|
query: str,
|
|
parameters: Sequence | dict[str, Any] | None = None,
|
|
settings: dict[str, Any] | None = None,
|
|
use_strings: bool | None = None,
|
|
external_data: ExternalData | None = None,
|
|
transport_settings: dict[str, str] | None = None,
|
|
dataframe_library: str = "pandas",
|
|
) -> pandas.DataFrame | polars.DataFrame:
|
|
"""
|
|
Query method using the ClickHouse Arrow format to return a DataFrame
|
|
with PyArrow dtype backend. This provides better performance and memory efficiency
|
|
compared to the standard query_df method, though fewer output formatting options.
|
|
|
|
:param query: Query statement/format string
|
|
:param parameters: Optional dictionary used to format the query
|
|
:param settings: Optional dictionary of ClickHouse settings (key/string values)
|
|
:param use_strings: Convert ClickHouse String type to Arrow string type (instead of binary)
|
|
:param external_data: ClickHouse "external data" to send with query
|
|
:param transport_settings: Optional dictionary of transport level settings (HTTP headers, etc.)
|
|
:param dataframe_library: Library to use for DataFrame creation ("pandas" or "polars")
|
|
:return: DataFrame (pandas or polars based on dataframe_library parameter)
|
|
"""
|
|
check_arrow()
|
|
|
|
if dataframe_library == "pandas":
|
|
check_pandas()
|
|
self._add_integration_tag("pandas")
|
|
|
|
def converter(table: pyarrow.Table) -> pandas.DataFrame:
|
|
table = _apply_arrow_tz_policy(table, self.tz_mode)
|
|
return table.to_pandas(types_mapper=options.pd.ArrowDtype, safe=False)
|
|
|
|
elif dataframe_library == "polars":
|
|
check_polars()
|
|
self._add_integration_tag("polars")
|
|
|
|
def converter(table: pyarrow.Table) -> polars.DataFrame: # type: ignore[misc]
|
|
table = _apply_arrow_tz_policy(table, self.tz_mode)
|
|
return options.pl.from_arrow(table)
|
|
|
|
else:
|
|
raise ValueError(f"dataframe_library must be 'pandas' or 'polars', got '{dataframe_library}'")
|
|
|
|
arrow_table = self.query_arrow(
|
|
query=query,
|
|
parameters=parameters,
|
|
settings=settings,
|
|
use_strings=use_strings,
|
|
external_data=external_data,
|
|
transport_settings=transport_settings,
|
|
)
|
|
|
|
return converter(arrow_table)
|
|
|
|
def query_df_arrow_stream(
|
|
self,
|
|
query: str,
|
|
parameters: Sequence | dict[str, Any] | None = None,
|
|
settings: dict[str, Any] | None = None,
|
|
use_strings: bool | None = None,
|
|
external_data: ExternalData | None = None,
|
|
transport_settings: dict[str, str] | None = None,
|
|
dataframe_library: str = "pandas",
|
|
) -> StreamContext:
|
|
"""
|
|
Query method that returns the results as a stream of DataFrames with PyArrow dtype backend.
|
|
Each DataFrame represents a block from the ClickHouse response.
|
|
|
|
:param query: Query statement/format string
|
|
:param parameters: Optional dictionary used to format the query
|
|
:param settings: Optional dictionary of ClickHouse settings (key/string values)
|
|
:param use_strings: Convert ClickHouse String type to Arrow string type (instead of binary)
|
|
:param external_data: ClickHouse "external data" to send with query
|
|
:param transport_settings: Optional dictionary of transport level settings (HTTP headers, etc.)
|
|
:param dataframe_library: Library to use for DataFrame creation ("pandas" or "polars")
|
|
:return: StreamContext that yields DataFrames (pandas or polars based on dataframe_library parameter)
|
|
"""
|
|
check_arrow()
|
|
if dataframe_library == "pandas":
|
|
check_pandas()
|
|
self._add_integration_tag("pandas")
|
|
|
|
def converter(table: pyarrow.Table) -> pandas.DataFrame:
|
|
table = _apply_arrow_tz_policy(table, self.tz_mode)
|
|
return table.to_pandas(types_mapper=options.pd.ArrowDtype, safe=False)
|
|
elif dataframe_library == "polars":
|
|
check_polars()
|
|
self._add_integration_tag("polars")
|
|
|
|
def converter(table: pyarrow.Table) -> polars.DataFrame: # type: ignore[misc]
|
|
table = _apply_arrow_tz_policy(table, self.tz_mode)
|
|
return options.pl.from_arrow(table)
|
|
else:
|
|
raise ValueError(f"dataframe_library must be 'pandas' or 'polars', got '{dataframe_library}'")
|
|
settings = self._update_arrow_settings(settings, use_strings)
|
|
raw_stream = self.raw_stream(
|
|
query, parameters, settings, fmt="ArrowStream", external_data=external_data, transport_settings=transport_settings
|
|
)
|
|
reader = options.arrow.ipc.open_stream(raw_stream)
|
|
|
|
def df_generator():
|
|
for batch in reader:
|
|
yield converter(batch)
|
|
|
|
return StreamContext(cast(Closable, raw_stream), df_generator())
|
|
|
|
def _update_arrow_settings(self, settings: dict[str, Any] | None, use_strings: bool | None) -> dict[str, Any]:
|
|
settings = dict_copy(settings)
|
|
if self.database:
|
|
settings["database"] = self.database
|
|
str_status = self._setting_status(arrow_str_setting)
|
|
if use_strings is None:
|
|
if str_status.is_writable and not str_status.is_set:
|
|
settings[arrow_str_setting] = "1" # Default to returning strings if possible
|
|
elif use_strings != str_status.is_set:
|
|
if not str_status.is_writable:
|
|
raise OperationalError(f"Cannot change readonly {arrow_str_setting} to {use_strings}")
|
|
settings[arrow_str_setting] = "1" if use_strings else "0"
|
|
return settings
|
|
|
|
@abstractmethod
|
|
def command(
|
|
self,
|
|
cmd: str,
|
|
parameters: Sequence | dict[str, Any] | None = None,
|
|
data: str | bytes | None = None,
|
|
settings: dict[str, Any] | None = None,
|
|
use_database: bool = True,
|
|
external_data: ExternalData | None = None,
|
|
transport_settings: dict[str, str] | None = None,
|
|
) -> str | int | Sequence[str] | QuerySummary:
|
|
"""
|
|
Client method that returns a single value instead of a result set
|
|
:param cmd: ClickHouse query/command as a python format string
|
|
:param parameters: Optional dictionary of key/values pairs to be formatted
|
|
:param data: Optional 'data' for the command (for INSERT INTO in particular)
|
|
:param settings: Optional dictionary of ClickHouse settings (key/string values)
|
|
:param use_database: Send the database parameter to ClickHouse so the command will be executed in the client
|
|
database context. Otherwise, no database will be specified with the command. This is useful for determining
|
|
the default user database
|
|
:param external_data: ClickHouse "external data" to send with command/query
|
|
:param transport_settings: Optional dictionary of transport level settings (HTTP headers, etc.)
|
|
:return: Decoded response from ClickHouse as either a string, int, or sequence of strings, or QuerySummary
|
|
if no data returned. Explicitly handled read-style commands can return an empty string for empty results.
|
|
"""
|
|
|
|
@abstractmethod
|
|
def ping(self) -> bool:
|
|
"""
|
|
Validate the connection, does not throw an Exception (see debug logs)
|
|
:return: ClickHouse server is up and reachable
|
|
"""
|
|
|
|
def insert(
|
|
self,
|
|
table: str | None = None,
|
|
data: Sequence[Sequence[Any]] | None = None,
|
|
column_names: str | Sequence[str] | None = "*",
|
|
database: str | None = None,
|
|
column_types: Sequence[ClickHouseType] | None = None,
|
|
column_type_names: Sequence[str] | None = None,
|
|
column_oriented: bool = False,
|
|
settings: dict[str, Any] | None = None,
|
|
context: InsertContext | None = None,
|
|
transport_settings: dict[str, str] | None = None,
|
|
) -> QuerySummary:
|
|
"""
|
|
Method to insert multiple rows/data matrix of native Python objects. If context is specified arguments
|
|
other than data are ignored
|
|
:param table: Target table
|
|
:param data: Sequence of sequences of Python data
|
|
:param column_names: Ordered list of column names or '*' if column types should be retrieved from the
|
|
ClickHouse table definition
|
|
:param database: Target database -- will use client default database if not specified.
|
|
:param column_types: ClickHouse column types. If set then column data does not need to be retrieved from
|
|
the server
|
|
:param column_type_names: ClickHouse column type names. If set then column data does not need to be
|
|
retrieved from the server
|
|
:param column_oriented: If true the data is already "pivoted" in column form
|
|
:param settings: Optional dictionary of ClickHouse settings (key/string values)
|
|
:param context: Optional reusable insert context to allow repeated inserts into the same table with
|
|
different data batches
|
|
:param transport_settings: Optional dictionary of transport level settings (HTTP headers, etc.)
|
|
:return: QuerySummary with summary information, throws exception if insert fails
|
|
"""
|
|
if (context is None or context.empty) and data is None:
|
|
raise ProgrammingError("No data specified for insert") from None
|
|
if context is None:
|
|
if table is None:
|
|
raise ProgrammingError("No table specified for insert") from None
|
|
context = self.create_insert_context(
|
|
table,
|
|
column_names,
|
|
database,
|
|
column_types,
|
|
column_type_names,
|
|
column_oriented,
|
|
settings,
|
|
transport_settings=transport_settings,
|
|
)
|
|
if data is not None:
|
|
if not context.empty:
|
|
raise ProgrammingError("Attempting to insert new data with non-empty insert context") from None
|
|
context.data = data
|
|
return self.data_insert(context)
|
|
|
|
def insert_df(
|
|
self,
|
|
table: str | None = None,
|
|
df=None,
|
|
database: str | None = None,
|
|
settings: dict | None = None,
|
|
column_names: Sequence[str] | None = None,
|
|
column_types: Sequence[ClickHouseType] | None = None,
|
|
column_type_names: Sequence[str] | None = None,
|
|
context: InsertContext | None = None,
|
|
transport_settings: dict[str, str] | None = None,
|
|
) -> QuerySummary:
|
|
"""
|
|
Insert a pandas DataFrame into ClickHouse. If context is specified arguments other than df are ignored
|
|
:param table: ClickHouse table
|
|
:param df: two-dimensional pandas dataframe
|
|
:param database: Optional ClickHouse database
|
|
:param settings: Optional dictionary of ClickHouse settings (key/string values)
|
|
:param column_names: An optional list of ClickHouse column names. If not set, the DataFrame column names
|
|
will be used
|
|
:param column_types: ClickHouse column types. If set then column data does not need to be retrieved from
|
|
the server
|
|
:param column_type_names: ClickHouse column type names. If set then column data does not need to be
|
|
retrieved from the server
|
|
:param context: Optional reusable insert context to allow repeated inserts into the same table with
|
|
different data batches
|
|
:param transport_settings: Optional dictionary of transport level settings (HTTP headers, etc.)
|
|
:return: QuerySummary with summary information, throws exception if insert fails
|
|
"""
|
|
check_pandas()
|
|
self._add_integration_tag("pandas")
|
|
if context is None:
|
|
if column_names is None:
|
|
column_names = df.columns
|
|
elif len(column_names) != len(df.columns):
|
|
raise ProgrammingError("DataFrame column count does not match insert_columns") from None
|
|
return self.insert(
|
|
table,
|
|
df,
|
|
column_names,
|
|
database,
|
|
column_types=column_types,
|
|
column_type_names=column_type_names,
|
|
settings=settings,
|
|
transport_settings=transport_settings,
|
|
context=context,
|
|
)
|
|
|
|
def insert_arrow(
|
|
self,
|
|
table: str,
|
|
arrow_table,
|
|
database: str | None = None,
|
|
settings: dict | None = None,
|
|
transport_settings: dict[str, str] | None = None,
|
|
) -> QuerySummary:
|
|
"""
|
|
Insert a PyArrow table DataFrame into ClickHouse using raw Arrow format
|
|
:param table: ClickHouse table
|
|
:param arrow_table: PyArrow Table object
|
|
:param database: Optional ClickHouse database
|
|
:param settings: Optional dictionary of ClickHouse settings (key/string values)
|
|
:param transport_settings: Optional dictionary of transport level settings (HTTP headers, etc.)
|
|
"""
|
|
check_arrow()
|
|
self._add_integration_tag("arrow")
|
|
full_table = table if "." in table or not database else f"{database}.{table}"
|
|
compression = self.write_compression if self.write_compression in ("zstd", "lz4") else None
|
|
column_names, insert_block = arrow_buffer(arrow_table, compression)
|
|
return self.raw_insert(full_table, column_names, insert_block, settings, "Arrow", transport_settings=transport_settings)
|
|
|
|
def insert_df_arrow(
|
|
self,
|
|
table: str,
|
|
df: pandas.DataFrame | polars.DataFrame,
|
|
database: str | None = None,
|
|
settings: dict | None = None,
|
|
transport_settings: dict[str, str] | None = None,
|
|
) -> QuerySummary:
|
|
"""
|
|
Insert a pandas DataFrame with PyArrow backend or a polars DataFrame into ClickHouse using Arrow format.
|
|
This method is optimized for DataFrames that already use Arrow format, providing
|
|
better performance than the standard insert_df method.
|
|
|
|
Validation is performed and an exception will be raised if this requirement is not met.
|
|
Polars DataFrames are natively Arrow-based and don't require additional validation.
|
|
|
|
:param table: ClickHouse table name
|
|
:param df: Pandas DataFrame with PyArrow dtype backend or Polars DataFrame
|
|
:param database: Optional ClickHouse database name
|
|
:param settings: Optional dictionary of ClickHouse settings (key/string values)
|
|
:param transport_settings: Optional dictionary of transport level settings (HTTP headers, etc.)
|
|
:return: QuerySummary with summary information, throws exception if insert fails
|
|
"""
|
|
check_arrow()
|
|
|
|
if options.pd is not None and isinstance(df, options.pd.DataFrame):
|
|
df_lib = "pandas"
|
|
elif options.pl is not None and isinstance(df, options.pl.DataFrame):
|
|
df_lib = "polars"
|
|
else:
|
|
if options.pd is None and options.pl is None:
|
|
raise ImportError("A DataFrame library (pandas or polars) must be installed to use insert_df_arrow.")
|
|
raise TypeError(f"df must be either a pandas DataFrame or polars DataFrame, got {type(df).__name__}")
|
|
|
|
if df_lib == "pandas":
|
|
non_arrow_cols = [col for col, dtype in df.dtypes.items() if not isinstance(dtype, options.pd.ArrowDtype)]
|
|
if non_arrow_cols:
|
|
raise ProgrammingError(
|
|
f"insert_df_arrow requires all columns to use PyArrow dtypes. Non-Arrow columns found: [{', '.join(non_arrow_cols)}]. "
|
|
)
|
|
try:
|
|
arrow_table = options.arrow.Table.from_pandas(df, preserve_index=False)
|
|
except Exception as e:
|
|
raise DataError(f"Failed to convert pandas DataFrame to Arrow table: {e}") from e
|
|
else:
|
|
try:
|
|
arrow_table = df.to_arrow()
|
|
except Exception as e:
|
|
raise DataError(f"Failed to convert polars DataFrame to Arrow table: {e}") from e
|
|
|
|
self._add_integration_tag(df_lib)
|
|
return self.insert_arrow(
|
|
table=table,
|
|
arrow_table=arrow_table,
|
|
database=database,
|
|
settings=settings,
|
|
transport_settings=transport_settings,
|
|
)
|
|
|
|
def create_insert_context(
|
|
self,
|
|
table: str,
|
|
column_names: str | Sequence[str] | None = None,
|
|
database: str | None = None,
|
|
column_types: Sequence[ClickHouseType] | None = None,
|
|
column_type_names: Sequence[str] | None = None,
|
|
column_oriented: bool = False,
|
|
settings: dict[str, Any] | None = None,
|
|
data: Sequence[Sequence[Any]] | None = None,
|
|
transport_settings: dict[str, str] | None = None,
|
|
) -> InsertContext:
|
|
"""
|
|
Builds a reusable insert context to hold state for a duration of an insert
|
|
:param table: Target table
|
|
:param database: Target database. If not set, uses the client default database
|
|
:param column_names: Optional ordered list of column names. If not set, all columns ('*') will be assumed
|
|
in the order specified by the table definition
|
|
:param database: Target database -- will use client default database if not specified
|
|
:param column_types: ClickHouse column types. Optional Sequence of ClickHouseType objects. If neither column
|
|
types nor column type names are set, actual column types will be retrieved from the server.
|
|
:param column_type_names: ClickHouse column type names. Specified column types by name string
|
|
:param column_oriented: If true the data is already "pivoted" in column form
|
|
:param settings: Optional dictionary of ClickHouse settings (key/string values)
|
|
:param data: Initial dataset for insert
|
|
:param transport_settings: Optional dictionary of transport level settings (HTTP headers, etc.)
|
|
:return: Reusable insert context
|
|
"""
|
|
full_table = table
|
|
if "." not in table:
|
|
if database:
|
|
full_table = f"{quote_identifier(database)}.{quote_identifier(table)}"
|
|
else:
|
|
full_table = quote_identifier(table)
|
|
column_defs: list[ColumnDef] = []
|
|
if column_types is None and column_type_names is None:
|
|
describe_result = self.query(f"DESCRIBE TABLE {full_table}", settings=settings)
|
|
column_defs = [
|
|
ColumnDef(**row) for row in describe_result.named_results() if row["default_type"] not in ("ALIAS", "MATERIALIZED")
|
|
]
|
|
if column_names is None or isinstance(column_names, str) and column_names == "*":
|
|
column_names = [cd.name for cd in column_defs]
|
|
column_types = [cd.ch_type for cd in column_defs]
|
|
elif isinstance(column_names, str):
|
|
column_names = [column_names]
|
|
if len(column_names) == 0:
|
|
raise ValueError("Column names must be specified for insert")
|
|
if not column_types:
|
|
if column_type_names:
|
|
column_types = [get_from_name(name) for name in column_type_names]
|
|
else:
|
|
column_map = {d.name: d for d in column_defs}
|
|
try:
|
|
column_types = [column_map[name].ch_type for name in column_names]
|
|
except KeyError as ex:
|
|
raise ProgrammingError(f"Unrecognized column {ex} in table {table}") from None
|
|
if len(column_names) != len(column_types):
|
|
raise ProgrammingError("Column names do not match column types") from None
|
|
return InsertContext(
|
|
full_table,
|
|
column_names,
|
|
column_types,
|
|
column_oriented=column_oriented,
|
|
settings=settings,
|
|
transport_settings=transport_settings,
|
|
data=data,
|
|
)
|
|
|
|
def min_version(self, version_str: str) -> bool:
|
|
"""
|
|
Determine whether the connected server is at least the submitted version
|
|
For Altinity Stable versions like 22.8.15.25.altinitystable
|
|
the last condition in the first list comprehension expression is added
|
|
:param version_str: A version string consisting of up to 4 integers delimited by dots
|
|
:return: True if version_str is greater than the server_version, False if less than
|
|
"""
|
|
try:
|
|
server_parts = [int(x) for x in (self.server_version or "").split(".") if x.isnumeric()]
|
|
server_parts.extend([0] * (4 - len(server_parts)))
|
|
version_parts = [int(x) for x in version_str.split(".")]
|
|
version_parts.extend([0] * (4 - len(version_parts)))
|
|
except ValueError:
|
|
logger.warning(
|
|
"Server %s or requested version %s does not match format of numbers separated by dots", self.server_version, version_str
|
|
)
|
|
return False
|
|
for x, y in zip(server_parts, version_parts):
|
|
if x > y:
|
|
return True
|
|
if x < y:
|
|
return False
|
|
return True
|
|
|
|
def _add_integration_tag(self, name: str) -> None:
|
|
"""Transport hook to surface 3rd party lib integration info (default: no-op)."""
|
|
return
|
|
|
|
@abstractmethod
|
|
def data_insert(self, context: InsertContext) -> QuerySummary:
|
|
"""
|
|
Subclass implementation of the data insert
|
|
:context: InsertContext parameter object
|
|
:return: No return, throws an exception if the insert fails
|
|
"""
|
|
|
|
@abstractmethod
|
|
def raw_insert(
|
|
self,
|
|
table: str,
|
|
column_names: Sequence[str] | None = None,
|
|
insert_block: str | bytes | Generator[bytes, None, None] | BinaryIO | None = None,
|
|
settings: dict | None = None,
|
|
fmt: str | None = None,
|
|
compression: str | None = None,
|
|
transport_settings: dict[str, str] | None = None,
|
|
) -> QuerySummary:
|
|
"""
|
|
Insert data already formatted in a bytes object
|
|
:param table: Table name (whether qualified with the database name or not)
|
|
:param column_names: Sequence of column names
|
|
:param insert_block: Binary or string data already in a recognized ClickHouse format
|
|
:param settings: Optional dictionary of ClickHouse settings (key/string values)
|
|
:param fmt: Valid clickhouse format
|
|
:param compression: Recognized ClickHouse `Accept-Encoding` header compression value
|
|
:param transport_settings: Optional dictionary of transport level settings (HTTP headers, etc.)
|
|
"""
|
|
|
|
@abstractmethod
|
|
def close(self) -> None:
|
|
"""
|
|
Subclass implementation to close the connection to the server/deallocate the client
|
|
"""
|
|
|
|
@abstractmethod
|
|
def close_connections(self) -> None:
|
|
"""
|
|
Subclass implementation to disconnect all "re-used" client connections
|
|
"""
|
|
|
|
def _context_query(self, lcls: dict, **overrides):
|
|
kwargs = lcls.copy()
|
|
kwargs.pop("self")
|
|
kwargs.update(overrides)
|
|
return self._query_with_context(self.create_query_context(**kwargs))
|
|
|
|
def __enter__(self) -> Client:
|
|
return self
|
|
|
|
def __exit__(self, exc_type, exc_value, exc_traceback) -> None:
|
|
self.close()
|