clickhouse-driver 0.2.10__cp312-cp312-win_amd64.whl

clickhouse_driver/numpy/result.py

@@ -0,0 +1,123 @@
+from itertools import chain
+
+import numpy as np
+import pandas as pd
+from pandas.api.types import union_categoricals
+
+from ..progress import Progress
+from ..result import QueryResult
+
+
+class NumpyQueryResult(QueryResult):
+    """
+    Stores query result from multiple blocks as numpy arrays.
+    """
+
+    def store(self, packet):
+        block = getattr(packet, 'block', None)
+        if block is None:
+            return
+
+        # Header block contains no rows. Pick columns from it.
+        if block.num_rows:
+            if self.columnar:
+                self.data.append(block.get_columns())
+            else:
+                self.data.extend(block.get_rows())
+
+        elif not self.columns_with_types:
+            self.columns_with_types = block.columns_with_types
+
+    def get_result(self):
+        """
+        :return: stored query result.
+        """
+
+        for packet in self.packet_generator:
+            self.store(packet)
+
+        if self.columnar:
+            data = []
+            # Transpose to a list of columns, each column is list of chunks
+            for column_chunks in zip(*self.data):
+                # Concatenate chunks for each column
+                if isinstance(column_chunks[0], np.ndarray):
+                    column = np.concatenate(column_chunks)
+                elif isinstance(column_chunks[0], pd.Categorical):
+                    column = union_categoricals(column_chunks)
+                else:
+                    column = tuple(chain.from_iterable(column_chunks))
+                data.append(column)
+        else:
+            data = self.data
+
+        if self.with_column_types:
+            return data, self.columns_with_types
+        else:
+            return data
+
+
+class NumpyProgressQueryResult(NumpyQueryResult):
+    """
+    Stores query result and progress information from multiple blocks.
+    Provides iteration over query progress.
+    """
+
+    def __init__(self, *args, **kwargs):
+        self.progress_totals = Progress()
+
+        super(NumpyProgressQueryResult, self).__init__(*args, **kwargs)
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        while True:
+            packet = next(self.packet_generator)
+            progress_packet = getattr(packet, 'progress', None)
+            if progress_packet:
+                self.progress_totals.increment(progress_packet)
+                return (
+                    self.progress_totals.rows, self.progress_totals.total_rows
+                )
+            else:
+                self.store(packet)
+
+    def get_result(self):
+        # Read all progress packets.
+        for _ in self:
+            pass
+
+        return super(NumpyProgressQueryResult, self).get_result()
+
+
+class NumpyIterQueryResult(object):
+    """
+    Provides iteration over returned data by chunks (streaming by chunks).
+    """
+
+    def __init__(
+            self, packet_generator,
+            with_column_types=False):
+        self.packet_generator = packet_generator
+        self.with_column_types = with_column_types
+
+        self.first_block = True
+        super(NumpyIterQueryResult, self).__init__()
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        packet = next(self.packet_generator)
+        block = getattr(packet, 'block', None)
+        if block is None:
+            return []
+
+        if self.first_block and self.with_column_types:
+            self.first_block = False
+            rv = [block.columns_with_types]
+            rv.extend(block.get_rows())
+            return rv
+        else:
+            return block.get_rows()

clickhouse_driver/opentelemetry.py

@@ -0,0 +1,43 @@
+
+class OpenTelemetryTraceContext(object):
+    traceparent_tpl = 'xx-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx-xxxxxxxxxxxxxxxx-xx'
+    translation = str.maketrans('1234567890abcdef', 'xxxxxxxxxxxxxxxx')
+
+    def __init__(self, traceparent, tracestate):
+        # xx-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx-xxxxxxxxxxxxxxxx-xx
+        # ^              ^                           ^         ^
+        # version     trace_id                    span_id      flags
+
+        self.trace_id = None  # UUID
+        self.span_id = None  # UInt64
+        self.tracestate = tracestate  # String
+        self.trace_flags = None  # UInt8
+
+        if traceparent is not None:
+            self.parse_traceparent(traceparent)
+
+        super(OpenTelemetryTraceContext, self).__init__()
+
+    def parse_traceparent(self, traceparent):
+        traceparent = traceparent.lower()
+
+        if len(traceparent) != len(self.traceparent_tpl):
+            raise ValueError('unexpected length {}, expected {}'.format(
+                len(traceparent), len(self.traceparent_tpl)
+            ))
+
+        if traceparent.translate(self.translation) != self.traceparent_tpl:
+            raise ValueError(
+                'Malformed traceparant header: {}'.format(traceparent)
+            )
+
+        parts = traceparent.split('-')
+        version = int(parts[0], 16)
+        if version != 0:
+            raise ValueError(
+                'unexpected version {}, expected 00'.format(parts[0])
+            )
+
+        self.trace_id = (int(parts[1][16:], 16) << 64) + int(parts[1][:16], 16)
+        self.span_id = int(parts[2], 16)
+        self.trace_flags = int(parts[3], 16)

clickhouse_driver/progress.py

@@ -0,0 +1,44 @@
+from . import defines
+from .varint import read_varint
+
+
+class Progress(object):
+    def __init__(self):
+        self.rows = 0
+        self.bytes = 0
+        self.total_rows = 0  # total_rows_to_read
+        self.total_bytes = 0  # total_bytes_to_read
+        self.written_rows = 0
+        self.written_bytes = 0
+        self.elapsed_ns = 0
+
+        super(Progress, self).__init__()
+
+    def read(self, server_info, fin):
+        self.rows = read_varint(fin)
+        self.bytes = read_varint(fin)
+
+        revision = server_info.used_revision
+        if revision >= defines.DBMS_MIN_REVISION_WITH_TOTAL_ROWS_IN_PROGRESS:
+            self.total_rows = read_varint(fin)
+
+        if revision >= defines. \
+                DBMS_MIN_PROTOCOL_VERSION_WITH_TOTAL_BYTES_IN_PROGRESS:
+            self.total_bytes = read_varint(fin)
+
+        if revision >= defines.DBMS_MIN_REVISION_WITH_CLIENT_WRITE_INFO:
+            self.written_rows = read_varint(fin)
+            self.written_bytes = read_varint(fin)
+
+        if revision >= defines. \
+                DBMS_MIN_PROTOCOL_VERSION_WITH_SERVER_QUERY_TIME_IN_PROGRESS:
+            self.elapsed_ns = read_varint(fin)
+
+    def increment(self, another_progress):
+        self.rows += another_progress.rows
+        self.bytes += another_progress.bytes
+        self.total_rows += another_progress.total_rows
+        self.total_bytes += another_progress.total_bytes
+        self.written_rows += another_progress.written_rows
+        self.written_bytes += another_progress.written_bytes
+        self.elapsed_ns += another_progress.elapsed_ns

clickhouse_driver/protocol.py

@@ -0,0 +1,130 @@
+
+class ClientPacketTypes(object):
+    """
+    Packet types that client transmits
+    """
+    # Name, version, revision, default DB
+    HELLO = 0
+
+    # Query id, query settings, stage up to which the query must be executed,
+    # whether the compression must be used, query text
+    # (without data for INSERTs).
+    QUERY = 1
+
+    # A block of data (compressed or not).
+    DATA = 2
+
+    # Cancel the query execution.
+    CANCEL = 3
+
+    # Check that connection to the server is alive.
+    PING = 4
+
+    # Check status of tables on the server.
+    TABLES_STATUS_REQUEST = 5
+
+    _types_str = [
+        'Hello', 'Query', 'Data', 'Cancel', 'Ping', 'TablesStatusRequest'
+    ]
+
+    @classmethod
+    def to_str(cls, packet):
+        try:
+            return cls._types_str[packet]
+        except IndexError:
+            return 'Unknown packet'
+
+
+class ServerPacketTypes(object):
+    """
+    Packet types that server transmits.
+    """
+    # Name, version, revision.
+    HELLO = 0
+
+    # A block of data (compressed or not).
+    DATA = 1
+
+    # The exception during query execution.
+    EXCEPTION = 2
+
+    # Query execution progress: rows read, bytes read.
+    PROGRESS = 3
+
+    # Ping response
+    PONG = 4
+
+    # All packets were transmitted
+    END_OF_STREAM = 5
+
+    # Packet with profiling info.
+    PROFILE_INFO = 6
+
+    # A block with totals (compressed or not).
+    TOTALS = 7
+
+    # A block with minimums and maximums (compressed or not).
+    EXTREMES = 8
+
+    # A response to TablesStatus request.
+    TABLES_STATUS_RESPONSE = 9
+
+    # System logs of the query execution
+    LOG = 10
+
+    # Columns' description for default values calculation
+    TABLE_COLUMNS = 11
+
+    # List of unique parts ids.
+    PART_UUIDS = 12
+
+    # String (UUID) describes a request for which next task is needed
+    READ_TASK_REQUEST = 13
+
+    # Packet with profile events from server.
+    PROFILE_EVENTS = 14
+
+    MERGE_TREE_ALL_RANGES_ANNOUNCEMENT = 15
+
+    # Request from a MergeTree replica to a coordinator
+    MERGE_TREE_READ_TASK_REQUEST = 16
+
+    # Receive server's (session-wide) default timezone
+    TIMEZONE_UPDATE = 17
+
+    _types_str = [
+        'Hello', 'Data', 'Exception', 'Progress', 'Pong', 'EndOfStream',
+        'ProfileInfo', 'Totals', 'Extremes', 'TablesStatusResponse', 'Log',
+        'TableColumns', 'PartUUIDs', 'ReadTaskRequest', 'ProfileEvents',
+        'MergeTreeAllRangesAnnouncement', 'MergeTreeReadTaskRequest',
+        'TimezoneUpdate'
+    ]
+
+    @classmethod
+    def to_str(cls, packet):
+        try:
+            return cls._types_str[packet]
+        except IndexError:
+            return 'Unknown packet'
+
+    @classmethod
+    def strings_in_message(cls, packet):
+        if packet == cls.TABLE_COLUMNS:
+            return 2
+        return 0
+
+
+class Compression(object):
+    DISABLED = 0
+    ENABLED = 1
+
+
+class CompressionMethod(object):
+    LZ4 = 1
+    LZ4HC = 2
+    ZSTD = 3
+
+
+class CompressionMethodByte(object):
+    LZ4 = 0x82
+    ZSTD = 0x90

clickhouse_driver/queryprocessingstage.py

@@ -0,0 +1,8 @@
+
+class QueryProcessingStage(object):
+    """
+    Determines till which state SELECT query should be executed.
+    """
+    FETCH_COLUMNS = 0
+    WITH_MERGEABLE_STATE = 1
+    COMPLETE = 2

clickhouse_driver/reader.py

@@ -0,0 +1,69 @@
+from struct import Struct
+
+from .varint import read_varint
+
+
+def read_binary_str(buf):
+    length = read_varint(buf)
+    return read_binary_str_fixed_len(buf, length)
+
+
+def read_binary_bytes(buf):
+    length = read_varint(buf)
+    return read_binary_bytes_fixed_len(buf, length)
+
+
+def read_binary_str_fixed_len(buf, length):
+    return read_binary_bytes_fixed_len(buf, length).decode('utf-8')
+
+
+def read_binary_bytes_fixed_len(buf, length):
+    return buf.read(length)
+
+
+def read_binary_int(buf, fmt):
+    """
+    Reads int from buffer with provided format.
+    """
+    # Little endian.
+    s = Struct('<' + fmt)
+    return s.unpack(buf.read(s.size))[0]
+
+
+def read_binary_int8(buf):
+    return read_binary_int(buf, 'b')
+
+
+def read_binary_int16(buf):
+    return read_binary_int(buf, 'h')
+
+
+def read_binary_int32(buf):
+    return read_binary_int(buf, 'i')
+
+
+def read_binary_int64(buf):
+    return read_binary_int(buf, 'q')
+
+
+def read_binary_uint8(buf):
+    return read_binary_int(buf, 'B')
+
+
+def read_binary_uint16(buf):
+    return read_binary_int(buf, 'H')
+
+
+def read_binary_uint32(buf):
+    return read_binary_int(buf, 'I')
+
+
+def read_binary_uint64(buf):
+    return read_binary_int(buf, 'Q')
+
+
+def read_binary_uint128(buf):
+    hi = read_binary_int(buf, 'Q')
+    lo = read_binary_int(buf, 'Q')
+
+    return (hi << 64) + lo

clickhouse_driver/readhelpers.py

@@ -0,0 +1,26 @@
+from .errors import ServerException
+from .reader import read_binary_str, read_binary_uint8, read_binary_int32
+
+
+def read_exception(buf, additional_message=None):
+    code = read_binary_int32(buf)
+    name = read_binary_str(buf)
+    message = read_binary_str(buf)
+    stack_trace = read_binary_str(buf)
+    has_nested = bool(read_binary_uint8(buf))
+
+    new_message = ''
+
+    if additional_message:
+        new_message += additional_message + '. '
+
+    if name != 'DB::Exception':
+        new_message += name + ". "
+
+    new_message += message + ". Stack trace:\n\n" + stack_trace
+
+    nested = None
+    if has_nested:
+        nested = read_exception(buf)
+
+    return ServerException(new_message, code, nested=nested)

clickhouse_driver/result.py

@@ -0,0 +1,144 @@
+from .blockstreamprofileinfo import BlockStreamProfileInfo
+from .progress import Progress
+
+
+class QueryResult(object):
+    """
+    Stores query result from multiple blocks.
+    """
+
+    def __init__(
+            self, packet_generator,
+            with_column_types=False, columnar=False):
+        self.packet_generator = packet_generator
+        self.with_column_types = with_column_types
+
+        self.data = []
+        self.columns_with_types = []
+        self.columnar = columnar
+
+        super(QueryResult, self).__init__()
+
+    def store(self, packet):
+        block = getattr(packet, 'block', None)
+        if block is None:
+            return
+
+        # Header block contains no rows. Pick columns from it.
+        if block.num_rows:
+            if self.columnar:
+                columns = block.get_columns()
+                if self.data:
+                    # Extend corresponding column.
+                    for i, column in enumerate(columns):
+                        self.data[i].extend(column)
+                else:
+                    # Cast tuples to lists for further extending.
+                    # Concatenating tuples produce new tuple. It's slow.
+                    self.data = [list(c) for c in columns]
+            else:
+                self.data.extend(block.get_rows())
+
+        elif not self.columns_with_types:
+            self.columns_with_types = block.columns_with_types
+
+    def get_result(self):
+        """
+        :return: stored query result.
+        """
+
+        for packet in self.packet_generator:
+            self.store(packet)
+
+        data = self.data
+        if self.columnar:
+            data = [tuple(c) for c in self.data]
+
+        if self.with_column_types:
+            return data, self.columns_with_types
+        else:
+            return data
+
+
+class ProgressQueryResult(QueryResult):
+    """
+    Stores query result and progress information from multiple blocks.
+    Provides iteration over query progress.
+    """
+
+    def __init__(self, *args, **kwargs):
+        self.progress_totals = Progress()
+        super(ProgressQueryResult, self).__init__(*args, **kwargs)
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        while True:
+            packet = next(self.packet_generator)
+            progress_packet = getattr(packet, 'progress', None)
+            if progress_packet:
+                self.progress_totals.increment(progress_packet)
+                return (
+                    self.progress_totals.rows, self.progress_totals.total_rows
+                )
+            else:
+                self.store(packet)
+
+    def get_result(self):
+        # Read all progress packets.
+        for _ in self:
+            pass
+
+        return super(ProgressQueryResult, self).get_result()
+
+
+class IterQueryResult(object):
+    """
+    Provides iteration over returned data by chunks (streaming by chunks).
+    """
+
+    def __init__(
+            self, packet_generator,
+            with_column_types=False):
+        self.packet_generator = packet_generator
+        self.with_column_types = with_column_types
+
+        self.first_block = True
+        super(IterQueryResult, self).__init__()
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        packet = next(self.packet_generator)
+        block = getattr(packet, 'block', None)
+        if block is None:
+            return []
+
+        if self.first_block and self.with_column_types:
+            self.first_block = False
+            rv = [block.columns_with_types]
+            rv.extend(block.get_rows())
+            return rv
+        else:
+            return block.get_rows()
+
+
+class QueryInfo(object):
+    def __init__(self):
+        self.profile_info = BlockStreamProfileInfo()
+        self.progress = Progress()
+        self.elapsed = 0
+
+    def store_profile(self, profile_info):
+        self.profile_info = profile_info
+
+    def store_progress(self, progress):
+        if self.progress:
+            self.progress.increment(progress)
+        else:
+            self.progress = progress
+
+    def store_elapsed(self, elapsed):
+        self.elapsed = elapsed