clickhouse-driver 0.2.10__pp311-pypy311_pp73-win_amd64.whl

clickhouse_driver/clientinfo.py

@@ -0,0 +1,119 @@
+import socket
+import getpass
+from time import time
+
+from . import defines
+from . import errors
+from .opentelemetry import OpenTelemetryTraceContext
+from .varint import write_varint
+from .writer import write_binary_str, write_binary_uint8, \
+    write_binary_uint64, write_binary_uint128
+
+
+class ClientInfo(object):
+    class Interface(object):
+        TCP = 1
+        HTTP = 2
+
+    class QueryKind(object):
+        # Uninitialized object.
+        NO_QUERY = 0
+
+        INITIAL_QUERY = 1
+
+        # Query that was initiated by another query for distributed query
+        # execution.
+        SECONDARY_QUERY = 2
+
+    client_version_major = defines.CLIENT_VERSION_MAJOR
+    client_version_minor = defines.CLIENT_VERSION_MINOR
+    client_version_patch = defines.CLIENT_VERSION_PATCH
+    interface = Interface.TCP
+
+    initial_user = ''
+    initial_query_id = ''
+    initial_address = '0.0.0.0:0'
+
+    def __init__(self, client_name, context, client_revision):
+        self.query_kind = ClientInfo.QueryKind.NO_QUERY
+
+        try:
+            self.os_user = getpass.getuser()
+        except (KeyError, OSError):
+            self.os_user = ''
+        self.client_hostname = socket.gethostname()
+        self.client_name = client_name
+        self.client_revision = client_revision
+
+        self.client_trace_context = OpenTelemetryTraceContext(
+            context.client_settings['opentelemetry_traceparent'],
+            context.client_settings['opentelemetry_tracestate']
+        )
+
+        self.quota_key = context.client_settings['quota_key']
+        self.distributed_depth = 0
+        self.initial_query_start_time_microseconds = int(time() * 1000000)
+
+        super(ClientInfo, self).__init__()
+
+    @property
+    def empty(self):
+        return self.query_kind == ClientInfo.QueryKind.NO_QUERY
+
+    def write(self, server_revision, fout):
+        revision = server_revision
+        if server_revision < defines.DBMS_MIN_REVISION_WITH_CLIENT_INFO:
+            raise errors.LogicalError('Method ClientInfo.write is called '
+                                      'for unsupported server revision')
+
+        write_binary_uint8(self.query_kind, fout)
+        if self.empty:
+            return
+
+        write_binary_str(self.initial_user, fout)
+        write_binary_str(self.initial_query_id, fout)
+        write_binary_str(self.initial_address, fout)
+
+        if (
+            revision >=
+            defines.DBMS_MIN_PROTOCOL_VERSION_WITH_INITIAL_QUERY_START_TIME
+        ):
+            write_binary_uint64(
+                self.initial_query_start_time_microseconds, fout
+            )
+
+        write_binary_uint8(self.interface, fout)
+
+        write_binary_str(self.os_user, fout)
+        write_binary_str(self.client_hostname, fout)
+        write_binary_str(self.client_name, fout)
+        write_varint(self.client_version_major, fout)
+        write_varint(self.client_version_minor, fout)
+        write_varint(self.client_revision, fout)
+
+        if revision >= defines.DBMS_MIN_REVISION_WITH_QUOTA_KEY_IN_CLIENT_INFO:
+            write_binary_str(self.quota_key, fout)
+
+        if revision >= \
+                defines.DBMS_MIN_PROTOCOL_VERSION_WITH_DISTRIBUTED_DEPTH:
+            write_varint(self.distributed_depth, fout)
+
+        if revision >= defines.DBMS_MIN_REVISION_WITH_VERSION_PATCH:
+            write_varint(self.client_version_patch, fout)
+
+        if revision >= defines.DBMS_MIN_REVISION_WITH_OPENTELEMETRY:
+            if self.client_trace_context.trace_id is not None:
+                # Have OpenTelemetry header.
+                write_binary_uint8(1, fout)
+                write_binary_uint128(self.client_trace_context.trace_id, fout)
+                write_binary_uint64(self.client_trace_context.span_id, fout)
+                write_binary_str(self.client_trace_context.tracestate, fout)
+                write_binary_uint8(self.client_trace_context.trace_flags, fout)
+            else:
+                # Don't have OpenTelemetry header.
+                write_binary_uint8(0, fout)
+
+        if revision >= defines.DBMS_MIN_REVISION_WITH_PARALLEL_REPLICAS:
+            write_varint(0, fout)  # collaborate_with_initiator
+            write_varint(0, fout)  # count_participating_replicas
+            write_varint(0, fout)  # number_of_current_replica

clickhouse_driver/columns/arraycolumn.py

@@ -0,0 +1,161 @@
+from itertools import chain
+from struct import Struct
+
+from .base import Column
+from .intcolumn import UInt64Column
+from ..util.helpers import pairwise
+
+
+class ArrayColumn(Column):
+    """
+    Nested arrays written in flatten form after information about their
+    sizes (offsets really).
+    One element of array of arrays can be represented as tree:
+    (0 depth)          [[3, 4], [5, 6]]
+                      |               |
+    (1 depth)      [3, 4]           [5, 6]
+                   |    |           |    |
+    (leaf)        3     4          5     6
+
+    Offsets (sizes) written in breadth-first search order. In example above
+    following sequence of offset will be written: 4 -> 2 -> 4
+    1) size of whole array: 4
+    2) size of array 1 in depth=1: 2
+    3) size of array 2 plus size of all array before in depth=1: 2 + 2 = 4
+
+    After sizes info comes flatten data: 3 -> 4 -> 5 -> 6
+    """
+    py_types = (list, tuple)
+
+    def __init__(self, nested_column, **kwargs):
+        self.init_kwargs = kwargs
+        self.size_column = UInt64Column(**kwargs)
+        self.nested_column = nested_column
+        self._write_depth_0_size = True
+        super(ArrayColumn, self).__init__(**kwargs)
+        self.null_value = []
+
+    def write_data(self, data, buf):
+        # Column of Array(T) is stored in "compact" format and passed to server
+        # wrapped into another Array without size of wrapper array.
+        self.nested_column = ArrayColumn(
+            self.nested_column, **self.init_kwargs
+        )
+        self.nested_column.nullable = self.nullable
+        self.nullable = False
+        self._write_depth_0_size = False
+        self._write(data, buf)
+
+    def read_data(self, n_rows, buf):
+        self.nested_column = ArrayColumn(
+            self.nested_column, **self.init_kwargs
+        )
+        self.nested_column.nullable = self.nullable
+        self.nullable = False
+        return self._read(n_rows, buf)[0]
+
+    def _write_sizes(self, value, buf):
+        nulls_map = []
+
+        column = self
+        sizes = [len(value)] if self._write_depth_0_size else []
+
+        while True:
+            nested_column = column.nested_column
+            if not isinstance(nested_column, ArrayColumn):
+                if column.nullable:
+                    nulls_map = [x is None for x in value]
+                break
+
+            offset = 0
+            new_value = []
+            for x in value:
+                offset += len(x)
+                sizes.append(offset)
+                new_value.extend(x)
+
+            value = new_value
+            column = nested_column
+
+        if nulls_map:
+            self._write_nulls_map(nulls_map, buf)
+
+        ns = Struct('<{}Q'.format(len(sizes)))
+        buf.write(ns.pack(*sizes))
+
+    def _write_data(self, value, buf):
+        if self.nullable:
+            value = value or []
+
+        if isinstance(self.nested_column, ArrayColumn):
+            value = list(chain.from_iterable(value))
+
+        if value:
+            self.nested_column._write_data(value, buf)
+
+    def _write_nulls_data(self, value, buf):
+        if self.nullable:
+            value = value or []
+
+        if isinstance(self.nested_column, ArrayColumn):
+            value = list(chain.from_iterable(value))
+            self.nested_column._write_nulls_data(value, buf)
+        else:
+            if self.nested_column.nullable:
+                self.nested_column._write_nulls_map(value, buf)
+
+    def _write(self, value, buf):
+        value = self.prepare_items(value)
+        self._write_sizes(value, buf)
+        self._write_nulls_data(value, buf)
+        self._write_data(value, buf)
+
+    def read_state_prefix(self, buf):
+        super(ArrayColumn, self).read_state_prefix(buf)
+
+        self.nested_column.read_state_prefix(buf)
+
+    def write_state_prefix(self, buf):
+        super(ArrayColumn, self).write_state_prefix(buf)
+
+        self.nested_column.write_state_prefix(buf)
+
+    def _read(self, size, buf):
+        slices_series = [[0, size]]
+        nested_column = self.nested_column
+
+        cur_level_slice_size = size
+        cur_level_slice = None
+        while (isinstance(nested_column, ArrayColumn)):
+            if cur_level_slice is None:
+                cur_level_slice = [0]
+            ns = Struct('<{}Q'.format(cur_level_slice_size))
+            nested_sizes = ns.unpack(buf.read(ns.size))
+            cur_level_slice.extend(nested_sizes)
+            slices_series.append(cur_level_slice)
+            cur_level_slice = None
+            cur_level_slice_size = nested_sizes[-1] if len(nested_sizes) > 0 \
+                else 0
+            nested_column = nested_column.nested_column
+
+        n_items = cur_level_slice_size if size > 0 else 0
+        nulls_map = None
+        if nested_column.nullable:
+            nulls_map = self._read_nulls_map(n_items, buf)
+
+        data = []
+        if n_items:
+            data = list(nested_column._read_data(
+                n_items, buf, nulls_map=nulls_map
+            ))
+
+        # Build nested structure.
+        for slices in reversed(slices_series):
+            data = [data[begin:end] for begin, end in pairwise(slices)]
+
+        return tuple(data)
+
+
+def create_array_column(spec, column_by_spec_getter, column_options):
+    inner = spec[6:-1]
+    return ArrayColumn(column_by_spec_getter(inner), **column_options)

clickhouse_driver/columns/base.py

@@ -0,0 +1,221 @@
+from struct import Struct, error as struct_error
+
+from . import exceptions
+from ..varint import read_varint
+
+
+class CommonSerialization(object):
+    def __init__(self, column):
+        self.column = column
+        super(CommonSerialization, self).__init__()
+
+    def read_sparse(self, n_items, buf):
+        return n_items
+
+    def apply_sparse(self, items):
+        return items
+
+
+class SparseSerialization(CommonSerialization):
+
+    def __init__(self, *args, **kwargs):
+        self.sparse_indexes = []
+        self.items_total = None
+        super(SparseSerialization, self).__init__(*args, **kwargs)
+
+    def read_sparse(self, n_items, buf):
+        sparse_indexes = []
+        items_total = 0
+        non_default_items = 0
+
+        END_OF_GRANULE_FLAG = 1 << 62
+        end_of_granule = False
+
+        while not end_of_granule:
+            group_size = read_varint(buf)
+            end_of_granule = group_size & END_OF_GRANULE_FLAG
+            group_size &= ~END_OF_GRANULE_FLAG
+
+            items_total += group_size + 1
+            if not end_of_granule:
+                non_default_items += 1
+                sparse_indexes.append(items_total)
+
+        self.sparse_indexes = sparse_indexes
+        self.items_total = items_total
+
+        return non_default_items
+
+    def apply_sparse(self, items):
+        default = self.column.null_value
+        if self.column.after_read_items:
+            default = self.column.after_read_items([default])[0]
+
+        rv = [default] * (self.items_total - 1)
+        for item_number, i in enumerate(self.sparse_indexes):
+            rv[i - 1] = items[item_number]
+
+        return rv
+
+
+class Column(object):
+    ch_type = None
+    py_types = None
+
+    check_item = None
+    after_read_items = None
+    before_write_items = None
+
+    types_check_enabled = False
+
+    null_value = 0
+
+    def __init__(self, types_check=False, has_custom_serialization=False,
+                 **kwargs):
+        self.nullable = False
+        self.types_check_enabled = types_check
+        self.has_custom_serialization = has_custom_serialization
+        self.serialization = CommonSerialization(self)
+        self.input_null_as_default = False
+
+        self.context = kwargs['context']
+        self.input_null_as_default = self.context.client_settings \
+            .get('input_format_null_as_default', False)
+
+        super(Column, self).__init__()
+
+    def make_null_struct(self, n_items):
+        return Struct('<{}B'.format(n_items))
+
+    def _read_nulls_map(self, n_items, buf):
+        s = self.make_null_struct(n_items)
+        return s.unpack(buf.read(s.size))
+
+    def _write_nulls_map(self, items, buf):
+        s = self.make_null_struct(len(items))
+        items = [x is None for x in items]
+        buf.write(s.pack(*items))
+
+    def check_item_type(self, value):
+        if not isinstance(value, self.py_types):
+            raise exceptions.ColumnTypeMismatchException(value)
+
+    def prepare_items(self, items):
+        nullable = self.nullable
+        null_value = self.null_value
+        null_as_default = self.input_null_as_default
+
+        check_item = self.check_item
+        if self.types_check_enabled:
+            check_item_type = self.check_item_type
+        else:
+            check_item_type = False
+
+        if (not (self.nullable or null_as_default) and not check_item_type and
+                not check_item and not self.before_write_items):
+            return items
+
+        nulls_map = [False] * len(items) if self.nullable else None
+        for i, x in enumerate(items):
+            if x is None:
+                if nullable:
+                    nulls_map[i] = True
+                    x = null_value
+                elif null_as_default:
+                    x = null_value
+
+            else:
+                if check_item_type:
+                    check_item_type(x)
+
+                if check_item:
+                    check_item(x)
+
+            items[i] = x
+
+        if self.before_write_items:
+            self.before_write_items(items, nulls_map=nulls_map)
+
+        return items
+
+    def write_data(self, items, buf):
+        if self.nullable:
+            self._write_nulls_map(items, buf)
+
+        self._write_data(items, buf)
+
+    def _write_data(self, items, buf):
+        prepared = self.prepare_items(items)
+        self.write_items(prepared, buf)
+
+    def write_items(self, items, buf):
+        raise NotImplementedError
+
+    def read_data(self, n_items, buf):
+        n_items = self.serialization.read_sparse(n_items, buf)
+
+        if self.nullable:
+            nulls_map = self._read_nulls_map(n_items, buf)
+        else:
+            nulls_map = None
+
+        items = self._read_data(n_items, buf, nulls_map=nulls_map)
+        return self.serialization.apply_sparse(items)
+
+    def _read_data(self, n_items, buf, nulls_map=None):
+        items = self.read_items(n_items, buf)
+
+        if self.after_read_items:
+            return self.after_read_items(items, nulls_map)
+        elif nulls_map is not None:
+            return tuple(
+                (None if is_null else items[i])
+                for i, is_null in enumerate(nulls_map)
+            )
+        return items
+
+    def read_items(self, n_items, buf):
+        raise NotImplementedError
+
+    def read_state_prefix(self, buf):
+        if self.has_custom_serialization:
+            use_custom_serialization = read_varint(buf)
+            if use_custom_serialization:
+                self.serialization = SparseSerialization(self)
+
+    def write_state_prefix(self, buf):
+        pass
+
+
+class FormatColumn(Column):
+    """
+    Uses struct.pack for bulk items writing.
+    """
+
+    format = None
+
+    def make_struct(self, n_items):
+        return Struct('<{}{}'.format(n_items, self.format))
+
+    def write_items(self, items, buf):
+        s = self.make_struct(len(items))
+        try:
+            buf.write(s.pack(*items))
+
+        except struct_error as e:
+            raise exceptions.StructPackException(e)
+
+    def read_items(self, n_items, buf):
+        s = self.make_struct(n_items)
+        return s.unpack(buf.read(s.size))
+
+
+# How to write new column?
+# - Check ClickHouse documentation for column
+# - Wireshark and tcpdump are your friends.
+# - Use `clickhouse-client --compression 0` to see what's going on data
+#   transmission.
+# - Check for similar existing columns and tests.
+# - Use `FormatColumn` for columns that use "simple" types under the hood.
+# - Some columns have before_write and after_read hooks.
+#   Use them to convert items in column into "simple" types.

clickhouse_driver/columns/boolcolumn.py

@@ -0,0 +1,7 @@
+from .base import FormatColumn
+
+
+class BoolColumn(FormatColumn):
+    ch_type = 'Bool'
+    py_types = (bool, )
+    format = '?'

clickhouse_driver/columns/datecolumn.py

@@ -0,0 +1,108 @@
+from os import getenv
+from datetime import date, timedelta
+
+from .base import FormatColumn
+
+
+epoch_start = date(1970, 1, 1)
+epoch_end = date(2149, 6, 6)
+
+epoch_start_date32 = date(1900, 1, 1)
+epoch_end_date32 = date(2299, 12, 31)
+
+
+class LazyLUT(dict):
+    def __init__(self, *args, _factory, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._default_factory = _factory
+
+    def __missing__(self, key):
+        return self.setdefault(key, self._default_factory(key))
+
+
+def make_date_lut_range(date_start, date_end):
+    return range(
+        (date_start - epoch_start).days,
+        (date_end - epoch_start).days + 1,
+    )
+
+
+enable_lazy_date_lut = getenv('CLICKHOUSE_DRIVER_LASY_DATE_LUT', False)
+if enable_lazy_date_lut:
+    try:
+        start, end = enable_lazy_date_lut.split(':')
+        start_date = date.fromisoformat(start)
+        end_date = date.fromisoformat(end)
+
+        date_range = make_date_lut_range(start_date, end_date)
+    except ValueError:
+        date_range = ()
+
+    # Since we initialize lazy lut with some initially warmed values,
+    # we use iterator and not dict comprehension for memory & time optimization
+    _date_lut = LazyLUT(
+        ((x, epoch_start + timedelta(days=x)) for x in date_range),
+        _factory=lambda x: epoch_start + timedelta(days=x),
+    )
+    _date_lut_reverse = LazyLUT(
+        ((value, key) for key, value in _date_lut.items()),
+        _factory=lambda x: (x - epoch_start).days,
+    )
+else:
+    # If lazy lut is not enabled, we fallback to static dict initialization
+    # In both cases, we use same lut for both data types,
+    # since one encompasses the other and we can avoid duplicating overlap
+    date_range = make_date_lut_range(epoch_start_date32, epoch_end_date32)
+    _date_lut = {x: epoch_start + timedelta(days=x) for x in date_range}
+    _date_lut_reverse = {value: key for key, value in _date_lut.items()}
+
+
+class DateColumn(FormatColumn):
+    ch_type = 'Date'
+    py_types = (date, )
+    format = 'H'
+
+    min_value = epoch_start
+    max_value = epoch_end
+
+    date_lut = _date_lut
+    date_lut_reverse = _date_lut_reverse
+
+    def before_write_items(self, items, nulls_map=None):
+        null_value = self.null_value
+
+        date_lut_reverse = self.date_lut_reverse
+        min_value = self.min_value
+        max_value = self.max_value
+
+        for i, item in enumerate(items):
+            if nulls_map and nulls_map[i]:
+                items[i] = null_value
+                continue
+
+            if item is not date:
+                item = date(item.year, item.month, item.day)
+
+            if min_value <= item <= max_value:
+                items[i] = date_lut_reverse[item]
+            else:
+                items[i] = 0
+
+    def after_read_items(self, items, nulls_map=None):
+        date_lut = self.date_lut
+
+        if nulls_map is None:
+            return tuple(date_lut[item] for item in items)
+        else:
+            return tuple(
+                (None if is_null else date_lut[items[i]])
+                for i, is_null in enumerate(nulls_map)
+            )
+
+
+class Date32Column(DateColumn):
+    ch_type = 'Date32'
+    format = 'i'
+
+    min_value = epoch_start_date32
+    max_value = epoch_end_date32