clickhouse-driver 0.2.1__cp39-cp39-win_amd64.whl → 0.2.10__cp39-cp39-win_amd64.whl

clickhouse_driver/clientinfo.py

@@ -1,80 +1,119 @@
-import socket
-import getpass
-
-from . import defines
-from . import errors
-from .varint import write_varint
-from .writer import write_binary_str, write_binary_uint8
-
-
-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
-    client_revision = defines.CLIENT_REVISION
-    interface = Interface.TCP
-
-    initial_user = ''
-    initial_query_id = ''
-    initial_address = '0.0.0.0:0'
-
-    quota_key = ''
-
-    def __init__(self, client_name):
-        self.query_kind = ClientInfo.QueryKind.NO_QUERY
-
-        try:
-            self.os_user = getpass.getuser()
-        except KeyError:
-            self.os_user = ''
-        self.client_hostname = socket.gethostname()
-        self.client_name = client_name
-
-        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)
-
-        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_REVISION_WITH_VERSION_PATCH:
-            write_varint(self.client_version_patch, fout)
+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

@@ -1,150 +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.size_column = UInt64Column()
-        self.nested_column = nested_column
-        self._write_depth_0_size = True
-        super(ArrayColumn, self).__init__(**kwargs)
-
-    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.nested_column.nullable = self.nullable
-        self.nullable = False
-        self._write_depth_0_size = False
-        self._write(data, buf)
-
-    def read_data(self, rows, buf):
-        self.nested_column = ArrayColumn(self.nested_column)
-        self.nested_column.nullable = self.nullable
-        self.nullable = False
-        return self._read(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):
-        self._write_sizes(value, buf)
-        self._write_nulls_data(value, buf)
-        self._write_data(value, buf)
-
-    def read_state_prefix(self, buf):
-        return self.nested_column.read_state_prefix(buf)
-
-    def write_state_prefix(self, 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):
-    inner = spec[6:-1]
-    return ArrayColumn(column_by_spec_getter(inner))
+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)