826 lines
29 KiB
Python
826 lines
29 KiB
Python
import logging
|
|
import socket
|
|
import ssl
|
|
from collections import deque
|
|
from contextlib import contextmanager
|
|
from sys import platform
|
|
from time import time
|
|
from urllib.parse import urlparse
|
|
|
|
try:
|
|
import certifi
|
|
except ImportError:
|
|
certifi = None
|
|
|
|
from . import defines
|
|
from . import errors
|
|
from .block import RowOrientedBlock
|
|
from .blockstreamprofileinfo import BlockStreamProfileInfo
|
|
from .bufferedreader import BufferedSocketReader
|
|
from .bufferedwriter import BufferedSocketWriter
|
|
from .clientinfo import ClientInfo
|
|
from .compression import get_compressor_cls
|
|
from .context import Context
|
|
from .log import log_block
|
|
from .progress import Progress
|
|
from .protocol import Compression, ClientPacketTypes, ServerPacketTypes
|
|
from .queryprocessingstage import QueryProcessingStage
|
|
from .reader import read_binary_str, read_binary_uint64
|
|
from .readhelpers import read_exception
|
|
from .settings.writer import write_settings, SettingsFlags
|
|
from .streams.native import BlockInputStream, BlockOutputStream
|
|
from .util.compat import threading
|
|
from .util.escape import escape_params
|
|
from .varint import write_varint, read_varint
|
|
from .writer import write_binary_str
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class Packet(object):
|
|
def __init__(self):
|
|
self.type = None
|
|
self.block = None
|
|
self.exception = None
|
|
self.progress = None
|
|
self.profile_info = None
|
|
self.multistring_message = None
|
|
|
|
super(Packet, self).__init__()
|
|
|
|
|
|
class ServerInfo(object):
|
|
def __init__(self, name, version_major, version_minor, version_patch,
|
|
revision, timezone, display_name, used_revision):
|
|
self.name = name
|
|
self.version_major = version_major
|
|
self.version_minor = version_minor
|
|
self.version_patch = version_patch
|
|
self.revision = revision
|
|
self.timezone = timezone
|
|
self.session_timezone = None
|
|
self.display_name = display_name
|
|
self.used_revision = used_revision
|
|
|
|
super(ServerInfo, self).__init__()
|
|
|
|
def get_timezone(self):
|
|
return self.session_timezone or self.timezone
|
|
|
|
def version_tuple(self):
|
|
return self.version_major, self.version_minor, self.version_patch
|
|
|
|
def __repr__(self):
|
|
version = '%s.%s.%s' % (
|
|
self.version_major, self.version_minor, self.version_patch
|
|
)
|
|
items = [
|
|
('name', self.name),
|
|
('version', version),
|
|
('revision', self.revision),
|
|
('used revision', self.used_revision),
|
|
('timezone', self.timezone),
|
|
('display_name', self.display_name)
|
|
]
|
|
|
|
params = ', '.join('{}={}'.format(key, value) for key, value in items)
|
|
return '<ServerInfo(%s)>' % (params)
|
|
|
|
|
|
class Connection(object):
|
|
"""
|
|
Represents connection between client and ClickHouse server.
|
|
|
|
:param host: host with running ClickHouse server.
|
|
:param port: port ClickHouse server is bound to.
|
|
Defaults to ``9000`` if connection is not secured and
|
|
to ``9440`` if connection is secured.
|
|
:param database: database connect to. Defaults to ``'default'``.
|
|
:param user: database user. Defaults to ``'default'``.
|
|
:param password: user's password. Defaults to ``''`` (no password).
|
|
:param client_name: this name will appear in server logs.
|
|
Defaults to ``'python-driver'``.
|
|
:param connect_timeout: timeout for establishing connection.
|
|
Defaults to ``10`` seconds.
|
|
:param send_receive_timeout: timeout for sending and receiving data.
|
|
Defaults to ``300`` seconds.
|
|
:param sync_request_timeout: timeout for server ping.
|
|
Defaults to ``5`` seconds.
|
|
:param compress_block_size: size of compressed block to send.
|
|
Defaults to ``1048576``.
|
|
:param compression: specifies whether or not use compression.
|
|
Defaults to ``False``. Possible choices:
|
|
|
|
* ``True`` is equivalent to ``'lz4'``.
|
|
* ``'lz4'``.
|
|
* ``'lz4hc'`` high-compression variant of
|
|
``'lz4'``.
|
|
* ``'zstd'``.
|
|
|
|
:param secure: establish secure connection. Defaults to ``False``.
|
|
:param verify: specifies whether a certificate is required and whether it
|
|
will be validated after connection.
|
|
Defaults to ``True``.
|
|
:param ssl_version: see :func:`ssl.wrap_socket` docs.
|
|
:param ca_certs: see :func:`ssl.wrap_socket` docs.
|
|
:param ciphers: see :func:`ssl.wrap_socket` docs.
|
|
:param keyfile: see :func:`ssl.wrap_socket` docs.
|
|
:param keypass: see :func:`ssl.wrap_socket` docs.
|
|
:param certfile: see :func:`ssl.wrap_socket` docs.
|
|
:param check_hostname: see :func:`ssl.wrap_socket` docs.
|
|
Defaults to ``True``.
|
|
:param server_hostname: Hostname to use in SSL Wrapper construction.
|
|
Defaults to `None` which will send the passed
|
|
host param during SSL initialization. This param
|
|
may be used when connecting over an SSH tunnel
|
|
to correctly identify the desired server via SNI.
|
|
:param alt_hosts: list of alternative hosts for connection.
|
|
Example: alt_hosts=host1:port1,host2:port2.
|
|
:param settings_is_important: ``False`` means unknown settings will be
|
|
ignored, ``True`` means that the query will
|
|
fail with UNKNOWN_SETTING error.
|
|
Defaults to ``False``.
|
|
:param tcp_keepalive: enables `TCP keepalive <https://tldp.org/HOWTO/
|
|
TCP-Keepalive-HOWTO/overview.html>`_ on established
|
|
connection. If is set to ``True``` system keepalive
|
|
settings are used. You can also specify custom
|
|
keepalive setting with tuple:
|
|
``(idle_time_sec, interval_sec, probes)``.
|
|
Defaults to ``False``.
|
|
:param client_revision: can be used for client version downgrading.
|
|
Defaults to ``None``.
|
|
:param disable_reconnect: disable automatic reconnect in case of
|
|
failed ``ping``, helpful when every reconnect
|
|
need to be caught in calling code.
|
|
Defaults to ``False``.
|
|
"""
|
|
|
|
def __init__(
|
|
self, host, port=None,
|
|
database=defines.DEFAULT_DATABASE,
|
|
user=defines.DEFAULT_USER, password=defines.DEFAULT_PASSWORD,
|
|
client_name=defines.CLIENT_NAME,
|
|
connect_timeout=defines.DBMS_DEFAULT_CONNECT_TIMEOUT_SEC,
|
|
send_receive_timeout=defines.DBMS_DEFAULT_TIMEOUT_SEC,
|
|
sync_request_timeout=defines.DBMS_DEFAULT_SYNC_REQUEST_TIMEOUT_SEC,
|
|
compress_block_size=defines.DEFAULT_COMPRESS_BLOCK_SIZE,
|
|
compression=False,
|
|
secure=False,
|
|
# Secure socket parameters.
|
|
verify=True, ssl_version=None, ca_certs=None, ciphers=None,
|
|
keyfile=None, keypass=None, certfile=None, check_hostname=True,
|
|
server_hostname=None,
|
|
alt_hosts=None,
|
|
settings_is_important=False,
|
|
tcp_keepalive=False,
|
|
client_revision=None,
|
|
disable_reconnect=False,
|
|
):
|
|
if secure:
|
|
default_port = defines.DEFAULT_SECURE_PORT
|
|
else:
|
|
default_port = defines.DEFAULT_PORT
|
|
|
|
self.hosts = deque([(host, port or default_port)])
|
|
|
|
if alt_hosts:
|
|
for host in alt_hosts.split(','):
|
|
url = urlparse('clickhouse://' + host)
|
|
self.hosts.append((url.hostname, url.port or default_port))
|
|
|
|
self.database = database
|
|
self.user = user
|
|
self.password = password
|
|
self.client_name = defines.DBMS_NAME + ' ' + client_name
|
|
self.connect_timeout = connect_timeout
|
|
self.send_receive_timeout = send_receive_timeout
|
|
self.sync_request_timeout = sync_request_timeout
|
|
self.settings_is_important = settings_is_important
|
|
self.tcp_keepalive = tcp_keepalive
|
|
self.client_revision = min(
|
|
client_revision or defines.CLIENT_REVISION, defines.CLIENT_REVISION
|
|
)
|
|
self.disable_reconnect = disable_reconnect
|
|
|
|
self.secure_socket = secure
|
|
self.verify_cert = verify
|
|
|
|
if certifi is not None:
|
|
ca_certs = ca_certs or certifi.where()
|
|
|
|
ssl_options = {}
|
|
if ssl_version is not None:
|
|
ssl_options['ssl_version'] = ssl_version
|
|
if ca_certs is not None:
|
|
ssl_options['ca_certs'] = ca_certs
|
|
if ciphers is not None:
|
|
ssl_options['ciphers'] = ciphers
|
|
if keyfile is not None:
|
|
ssl_options['keyfile'] = keyfile
|
|
if keypass is not None:
|
|
ssl_options['keypass'] = keypass
|
|
if certfile is not None:
|
|
ssl_options['certfile'] = certfile
|
|
|
|
self.ssl_options = ssl_options
|
|
|
|
self.check_hostname = check_hostname if self.verify_cert else False
|
|
self.server_hostname = server_hostname
|
|
|
|
# Use LZ4 compression by default.
|
|
if compression is True:
|
|
compression = 'lz4'
|
|
|
|
if compression is False:
|
|
self.compression = Compression.DISABLED
|
|
self.compressor_cls = None
|
|
self.compress_block_size = None
|
|
else:
|
|
self.compression = Compression.ENABLED
|
|
self.compressor_cls = get_compressor_cls(compression)
|
|
self.compress_block_size = compress_block_size
|
|
|
|
self.socket = None
|
|
self.fin = None
|
|
self.fout = None
|
|
|
|
self.connected = False
|
|
|
|
self.client_trace_context = None
|
|
self.server_info = None
|
|
self.context = Context()
|
|
|
|
# Block writer/reader
|
|
self.block_in = None
|
|
self.block_out = None
|
|
self.block_in_raw = None # log blocks are always not compressed
|
|
|
|
self._lock = threading.Lock()
|
|
self.is_query_executing = False
|
|
|
|
super(Connection, self).__init__()
|
|
|
|
def __repr__(self):
|
|
dsn = '%s://%s:***@%s:%s/%s' % (
|
|
'clickhouses' if self.secure_socket else 'clickhouse',
|
|
self.user, self.host, self.port, self.database
|
|
) if self.connected else '(not connected)'
|
|
|
|
return '<Connection(dsn=%s, compression=%s)>' % (dsn, self.compression)
|
|
|
|
def get_description(self):
|
|
return '{}:{}'.format(self.host, self.port)
|
|
|
|
def force_connect(self):
|
|
self.check_query_execution()
|
|
|
|
if not self.connected:
|
|
self.connect()
|
|
|
|
elif not self.ping():
|
|
if self.disable_reconnect:
|
|
raise errors.NetworkError(
|
|
"Connection was closed, reconnect is disabled."
|
|
)
|
|
|
|
logger.warning('Connection was closed, reconnecting.')
|
|
self.connect()
|
|
|
|
def _create_socket(self, host, port):
|
|
"""
|
|
Acts like socket.create_connection, but wraps socket with SSL
|
|
if connection is secure.
|
|
"""
|
|
ssl_options = {}
|
|
if self.secure_socket:
|
|
if self.verify_cert:
|
|
cert_reqs = ssl.CERT_REQUIRED
|
|
else:
|
|
cert_reqs = ssl.CERT_NONE
|
|
|
|
ssl_options = self.ssl_options.copy()
|
|
ssl_options['cert_reqs'] = cert_reqs
|
|
|
|
err = None
|
|
for res in socket.getaddrinfo(host, port, 0, socket.SOCK_STREAM):
|
|
af, socktype, proto, canonname, sa = res
|
|
sock = None
|
|
try:
|
|
sock = socket.socket(af, socktype, proto)
|
|
sock.settimeout(self.connect_timeout)
|
|
|
|
if self.secure_socket:
|
|
ssl_context = self._create_ssl_context(ssl_options)
|
|
sock = ssl_context.wrap_socket(
|
|
sock, server_hostname=self.server_hostname or host)
|
|
|
|
sock.connect(sa)
|
|
return sock
|
|
|
|
except socket.error as _:
|
|
err = _
|
|
if sock is not None:
|
|
sock.close()
|
|
|
|
if err is not None:
|
|
raise err
|
|
else:
|
|
raise socket.error("getaddrinfo returns an empty list")
|
|
|
|
def _create_ssl_context(self, ssl_options):
|
|
purpose = ssl.Purpose.SERVER_AUTH
|
|
|
|
version = ssl_options.get('ssl_version', ssl.PROTOCOL_TLS_CLIENT)
|
|
context = ssl.SSLContext(version)
|
|
context.check_hostname = self.check_hostname
|
|
|
|
if 'ca_certs' in ssl_options:
|
|
context.load_verify_locations(ssl_options['ca_certs'])
|
|
elif ssl_options.get('cert_reqs') != ssl.CERT_NONE:
|
|
context.load_default_certs(purpose)
|
|
if 'ciphers' in ssl_options:
|
|
context.set_ciphers(ssl_options['ciphers'])
|
|
|
|
if 'cert_reqs' in ssl_options:
|
|
context.verify_mode = ssl_options['cert_reqs']
|
|
|
|
if 'certfile' in ssl_options:
|
|
keyfile = ssl_options.get('keyfile')
|
|
keypass = ssl_options.get('keypass')
|
|
context.load_cert_chain(
|
|
ssl_options['certfile'],
|
|
keyfile=keyfile, password=keypass
|
|
)
|
|
|
|
return context
|
|
|
|
def _init_connection(self, host, port):
|
|
self.socket = self._create_socket(host, port)
|
|
self.connected = True
|
|
self.host, self.port = host, port
|
|
self.socket.settimeout(self.send_receive_timeout)
|
|
|
|
# performance tweak
|
|
self.socket.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
|
|
if self.tcp_keepalive:
|
|
self._set_keepalive()
|
|
|
|
self.fin = BufferedSocketReader(self.socket, defines.BUFFER_SIZE)
|
|
self.fout = BufferedSocketWriter(self.socket, defines.BUFFER_SIZE)
|
|
|
|
self.send_hello()
|
|
self.receive_hello()
|
|
|
|
revision = self.server_info.used_revision
|
|
if revision >= defines.DBMS_MIN_PROTOCOL_VERSION_WITH_ADDENDUM:
|
|
self.send_addendum()
|
|
|
|
self.block_in = self.get_block_in_stream()
|
|
self.block_in_raw = BlockInputStream(self.fin, self.context)
|
|
self.block_out = self.get_block_out_stream()
|
|
|
|
def _set_keepalive(self):
|
|
self.socket.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)
|
|
|
|
if not isinstance(self.tcp_keepalive, tuple):
|
|
return
|
|
|
|
idle_time_sec, interval_sec, probes = self.tcp_keepalive
|
|
|
|
if platform == 'linux' or platform == 'win32':
|
|
# This should also work for Windows
|
|
# starting with Windows 10, version 1709.
|
|
self.socket.setsockopt(
|
|
socket.IPPROTO_TCP, socket.TCP_KEEPIDLE, idle_time_sec
|
|
)
|
|
self.socket.setsockopt(
|
|
socket.IPPROTO_TCP, socket.TCP_KEEPINTVL, interval_sec
|
|
)
|
|
self.socket.setsockopt(
|
|
socket.IPPROTO_TCP, socket.TCP_KEEPCNT, probes
|
|
)
|
|
|
|
elif platform == 'darwin':
|
|
TCP_KEEPALIVE = 0x10
|
|
# Only interval is available in mac os.
|
|
self.socket.setsockopt(
|
|
socket.IPPROTO_TCP, TCP_KEEPALIVE, interval_sec
|
|
)
|
|
|
|
def _format_connection_error(self, e, host, port):
|
|
err = (e.strerror + ' ') if e.strerror else ''
|
|
return err + '({}:{})'.format(host, port)
|
|
|
|
def connect(self):
|
|
if self.connected:
|
|
self.disconnect()
|
|
|
|
logger.debug(
|
|
'Connecting. Database: %s. User: %s', self.database, self.user
|
|
)
|
|
|
|
err = None
|
|
for i in range(len(self.hosts)):
|
|
host, port = self.hosts[0]
|
|
logger.debug('Connecting to %s:%s', host, port)
|
|
|
|
try:
|
|
return self._init_connection(host, port)
|
|
|
|
except socket.timeout as e:
|
|
self.disconnect()
|
|
logger.warning(
|
|
'Failed to connect to %s:%s', host, port, exc_info=True
|
|
)
|
|
err_str = self._format_connection_error(e, host, port)
|
|
err = errors.SocketTimeoutError(err_str)
|
|
|
|
except socket.error as e:
|
|
self.disconnect()
|
|
logger.warning(
|
|
'Failed to connect to %s:%s', host, port, exc_info=True
|
|
)
|
|
err_str = self._format_connection_error(e, host, port)
|
|
err = errors.NetworkError(err_str)
|
|
|
|
self.hosts.rotate(-1)
|
|
|
|
if err is not None:
|
|
raise err
|
|
|
|
def reset_state(self):
|
|
self.host = None
|
|
self.port = None
|
|
self.socket = None
|
|
self.fin = None
|
|
self.fout = None
|
|
|
|
self.connected = False
|
|
|
|
self.client_trace_context = None
|
|
self.server_info = None
|
|
|
|
self.block_in = None
|
|
self.block_in_raw = None
|
|
self.block_out = None
|
|
|
|
self.is_query_executing = False
|
|
|
|
def disconnect(self):
|
|
"""
|
|
Closes connection between server and client.
|
|
Frees resources: e.g. closes socket.
|
|
"""
|
|
|
|
if self.connected:
|
|
# There can be errors on shutdown.
|
|
# We need to close socket and reset state even if it happens.
|
|
try:
|
|
self.socket.shutdown(socket.SHUT_RDWR)
|
|
|
|
except socket.error as e:
|
|
logger.warning('Error on socket shutdown: %s', e)
|
|
|
|
self.socket.close()
|
|
|
|
# Socket can be constructed but not connected.
|
|
elif self.socket:
|
|
self.socket.close()
|
|
|
|
self.reset_state()
|
|
|
|
def send_hello(self):
|
|
write_varint(ClientPacketTypes.HELLO, self.fout)
|
|
write_binary_str(self.client_name, self.fout)
|
|
write_varint(defines.CLIENT_VERSION_MAJOR, self.fout)
|
|
write_varint(defines.CLIENT_VERSION_MINOR, self.fout)
|
|
# NOTE For backward compatibility of the protocol,
|
|
# client cannot send its version_patch.
|
|
write_varint(self.client_revision, self.fout)
|
|
write_binary_str(self.database, self.fout)
|
|
write_binary_str(self.user, self.fout)
|
|
write_binary_str(self.password, self.fout)
|
|
|
|
self.fout.flush()
|
|
|
|
def receive_hello(self):
|
|
packet_type = read_varint(self.fin)
|
|
|
|
if packet_type == ServerPacketTypes.HELLO:
|
|
server_name = read_binary_str(self.fin)
|
|
server_version_major = read_varint(self.fin)
|
|
server_version_minor = read_varint(self.fin)
|
|
server_revision = read_varint(self.fin)
|
|
|
|
used_revision = min(self.client_revision, server_revision)
|
|
|
|
server_timezone = None
|
|
if used_revision >= \
|
|
defines.DBMS_MIN_REVISION_WITH_SERVER_TIMEZONE:
|
|
server_timezone = read_binary_str(self.fin)
|
|
|
|
server_display_name = ''
|
|
if used_revision >= \
|
|
defines.DBMS_MIN_REVISION_WITH_SERVER_DISPLAY_NAME:
|
|
server_display_name = read_binary_str(self.fin)
|
|
|
|
server_version_patch = server_revision
|
|
if used_revision >= \
|
|
defines.DBMS_MIN_REVISION_WITH_VERSION_PATCH:
|
|
server_version_patch = read_varint(self.fin)
|
|
|
|
if used_revision >= defines. \
|
|
DBMS_MIN_PROTOCOL_VERSION_WITH_PASSWORD_COMPLEXITY_RULES:
|
|
rules_size = read_varint(self.fin)
|
|
for _i in range(rules_size):
|
|
read_binary_str(self.fin) # original_pattern
|
|
read_binary_str(self.fin) # exception_message
|
|
|
|
if used_revision >= defines. \
|
|
DBMS_MIN_REVISION_WITH_INTERSERVER_SECRET_V2:
|
|
read_binary_uint64(self.fin) # read_nonce
|
|
|
|
self.server_info = ServerInfo(
|
|
server_name, server_version_major, server_version_minor,
|
|
server_version_patch, server_revision,
|
|
server_timezone, server_display_name, used_revision
|
|
)
|
|
self.context.server_info = self.server_info
|
|
|
|
logger.debug(
|
|
'Connected to %s server version %s.%s.%s, revision: %s',
|
|
server_name, server_version_major, server_version_minor,
|
|
server_version_patch, server_revision
|
|
)
|
|
|
|
elif packet_type == ServerPacketTypes.EXCEPTION:
|
|
raise self.receive_exception()
|
|
|
|
else:
|
|
message = self.unexpected_packet_message('Hello or Exception',
|
|
packet_type)
|
|
self.disconnect()
|
|
raise errors.UnexpectedPacketFromServerError(message)
|
|
|
|
def send_addendum(self):
|
|
revision = self.server_info.used_revision
|
|
|
|
if revision >= defines.DBMS_MIN_PROTOCOL_VERSION_WITH_QUOTA_KEY:
|
|
write_binary_str(
|
|
self.context.client_settings['quota_key'], self.fout
|
|
)
|
|
|
|
def ping(self):
|
|
if not self.socket:
|
|
return None
|
|
|
|
timeout = self.sync_request_timeout
|
|
|
|
with self.timeout_setter(timeout):
|
|
try:
|
|
write_varint(ClientPacketTypes.PING, self.fout)
|
|
self.fout.flush()
|
|
|
|
packet_type = read_varint(self.fin)
|
|
while packet_type == ServerPacketTypes.PROGRESS:
|
|
self.receive_progress()
|
|
packet_type = read_varint(self.fin)
|
|
|
|
if packet_type != ServerPacketTypes.PONG:
|
|
msg = self.unexpected_packet_message('Pong', packet_type)
|
|
raise errors.UnexpectedPacketFromServerError(msg)
|
|
|
|
except errors.Error:
|
|
raise
|
|
|
|
except (socket.error, EOFError) as e:
|
|
# It's just a warning now.
|
|
# Current connection will be closed, new will be established.
|
|
logger.warning(
|
|
'Error on %s ping: %s', self.get_description(), e
|
|
)
|
|
return False
|
|
|
|
return True
|
|
|
|
def receive_packet(self):
|
|
packet = Packet()
|
|
|
|
packet.type = packet_type = read_varint(self.fin)
|
|
|
|
if packet_type == ServerPacketTypes.DATA:
|
|
packet.block = self.receive_data(may_be_use_numpy=True)
|
|
|
|
elif packet_type == ServerPacketTypes.EXCEPTION:
|
|
packet.exception = self.receive_exception()
|
|
|
|
elif packet_type == ServerPacketTypes.PROGRESS:
|
|
packet.progress = self.receive_progress()
|
|
|
|
elif packet_type == ServerPacketTypes.PROFILE_INFO:
|
|
packet.profile_info = self.receive_profile_info()
|
|
|
|
elif packet_type == ServerPacketTypes.TOTALS:
|
|
packet.block = self.receive_data()
|
|
|
|
elif packet_type == ServerPacketTypes.EXTREMES:
|
|
packet.block = self.receive_data()
|
|
|
|
elif packet_type == ServerPacketTypes.LOG:
|
|
packet.block = self.receive_data(may_be_compressed=False)
|
|
log_block(packet.block)
|
|
|
|
elif packet_type == ServerPacketTypes.END_OF_STREAM:
|
|
self.is_query_executing = False
|
|
pass
|
|
|
|
elif packet_type == ServerPacketTypes.TABLE_COLUMNS:
|
|
packet.multistring_message = self.receive_multistring_message(
|
|
packet_type
|
|
)
|
|
|
|
elif packet_type == ServerPacketTypes.PART_UUIDS:
|
|
packet.block = self.receive_data()
|
|
|
|
elif packet_type == ServerPacketTypes.READ_TASK_REQUEST:
|
|
packet.block = self.receive_data()
|
|
|
|
elif packet_type == ServerPacketTypes.PROFILE_EVENTS:
|
|
packet.block = self.receive_data(may_be_compressed=False)
|
|
|
|
elif packet_type == ServerPacketTypes.TIMEZONE_UPDATE:
|
|
timezone = read_binary_str(self.fin)
|
|
if timezone:
|
|
logger.info('Server timezone changed to %s', timezone)
|
|
self.server_info.session_timezone = timezone
|
|
|
|
else:
|
|
message = 'Unknown packet {} from server {}'.format(
|
|
packet_type, self.get_description()
|
|
)
|
|
self.disconnect()
|
|
raise errors.UnknownPacketFromServerError(message)
|
|
|
|
return packet
|
|
|
|
def get_block_in_stream(self):
|
|
if self.compression:
|
|
from .streams.compressed import CompressedBlockInputStream
|
|
|
|
return CompressedBlockInputStream(self.fin, self.context)
|
|
else:
|
|
return BlockInputStream(self.fin, self.context)
|
|
|
|
def get_block_out_stream(self):
|
|
if self.compression:
|
|
from .streams.compressed import CompressedBlockOutputStream
|
|
|
|
return CompressedBlockOutputStream(
|
|
self.compressor_cls, self.compress_block_size,
|
|
self.fout, self.context
|
|
)
|
|
else:
|
|
return BlockOutputStream(self.fout, self.context)
|
|
|
|
def receive_data(self, may_be_compressed=True, may_be_use_numpy=False):
|
|
revision = self.server_info.used_revision
|
|
|
|
if revision >= defines.DBMS_MIN_REVISION_WITH_TEMPORARY_TABLES:
|
|
read_binary_str(self.fin)
|
|
|
|
reader = self.block_in if may_be_compressed else self.block_in_raw
|
|
use_numpy = False if not may_be_use_numpy else None
|
|
return reader.read(use_numpy=use_numpy)
|
|
|
|
def receive_exception(self):
|
|
return read_exception(self.fin)
|
|
|
|
def receive_progress(self):
|
|
progress = Progress()
|
|
progress.read(self.server_info, self.fin)
|
|
return progress
|
|
|
|
def receive_profile_info(self):
|
|
profile_info = BlockStreamProfileInfo()
|
|
profile_info.read(self.fin)
|
|
return profile_info
|
|
|
|
def receive_multistring_message(self, packet_type):
|
|
num = ServerPacketTypes.strings_in_message(packet_type)
|
|
return [read_binary_str(self.fin) for _i in range(num)]
|
|
|
|
def send_data(self, block, table_name=''):
|
|
start = time()
|
|
write_varint(ClientPacketTypes.DATA, self.fout)
|
|
|
|
revision = self.server_info.used_revision
|
|
if revision >= defines.DBMS_MIN_REVISION_WITH_TEMPORARY_TABLES:
|
|
write_binary_str(table_name, self.fout)
|
|
|
|
self.block_out.write(block)
|
|
logger.debug('Block "%s" send time: %f', table_name, time() - start)
|
|
|
|
def send_query(self, query, query_id=None, params=None):
|
|
if not self.connected:
|
|
self.connect()
|
|
|
|
write_varint(ClientPacketTypes.QUERY, self.fout)
|
|
|
|
write_binary_str(query_id or '', self.fout)
|
|
|
|
revision = self.server_info.used_revision
|
|
if revision >= defines.DBMS_MIN_REVISION_WITH_CLIENT_INFO:
|
|
client_info = ClientInfo(self.client_name, self.context,
|
|
client_revision=self.client_revision)
|
|
client_info.query_kind = ClientInfo.QueryKind.INITIAL_QUERY
|
|
|
|
client_info.write(revision, self.fout)
|
|
|
|
settings_as_strings = (
|
|
revision >= defines
|
|
.DBMS_MIN_REVISION_WITH_SETTINGS_SERIALIZED_AS_STRINGS
|
|
)
|
|
settings_flags = 0
|
|
if self.settings_is_important:
|
|
settings_flags |= SettingsFlags.IMPORTANT
|
|
write_settings(self.context.settings, self.fout, settings_as_strings,
|
|
settings_flags)
|
|
|
|
if revision >= defines.DBMS_MIN_REVISION_WITH_INTERSERVER_SECRET:
|
|
write_binary_str('', self.fout)
|
|
|
|
write_varint(QueryProcessingStage.COMPLETE, self.fout)
|
|
write_varint(self.compression, self.fout)
|
|
|
|
write_binary_str(query, self.fout)
|
|
|
|
if revision >= defines.DBMS_MIN_PROTOCOL_VERSION_WITH_PARAMETERS:
|
|
if self.context.client_settings['server_side_params']:
|
|
# Always settings_as_strings = True
|
|
escaped = escape_params(
|
|
params or {}, self.context, for_server=True
|
|
)
|
|
else:
|
|
escaped = {}
|
|
write_settings(escaped, self.fout, True, SettingsFlags.CUSTOM)
|
|
|
|
logger.debug('Query: %s', query)
|
|
|
|
self.fout.flush()
|
|
|
|
def send_cancel(self):
|
|
write_varint(ClientPacketTypes.CANCEL, self.fout)
|
|
|
|
self.fout.flush()
|
|
|
|
def send_external_tables(self, tables, types_check=False):
|
|
for table in tables or []:
|
|
if not table['structure']:
|
|
raise ValueError(
|
|
'Empty table "{}" structure'.format(table['name'])
|
|
)
|
|
|
|
data = table['data']
|
|
block_cls = RowOrientedBlock
|
|
|
|
if self.context.client_settings['use_numpy']:
|
|
from .numpy.block import NumpyColumnOrientedBlock
|
|
|
|
columns = [x[0] for x in table['structure']]
|
|
data = [data[column].values for column in columns]
|
|
|
|
block_cls = NumpyColumnOrientedBlock
|
|
|
|
block = block_cls(table['structure'], data,
|
|
types_check=types_check)
|
|
self.send_data(block, table_name=table['name'])
|
|
|
|
# Empty block, end of data transfer.
|
|
self.send_data(RowOrientedBlock())
|
|
|
|
@contextmanager
|
|
def timeout_setter(self, new_timeout):
|
|
old_timeout = self.socket.gettimeout()
|
|
self.socket.settimeout(new_timeout)
|
|
|
|
yield
|
|
|
|
self.socket.settimeout(old_timeout)
|
|
|
|
def unexpected_packet_message(self, expected, packet_type):
|
|
packet_type = ServerPacketTypes.to_str(packet_type)
|
|
|
|
return (
|
|
'Unexpected packet from server {} (expected {}, got {})'
|
|
.format(self.get_description(), expected, packet_type)
|
|
)
|
|
|
|
def check_query_execution(self):
|
|
self._lock.acquire(blocking=False)
|
|
|
|
if self.is_query_executing:
|
|
raise errors.PartiallyConsumedQueryError()
|
|
|
|
self.is_query_executing = True
|
|
self._lock.release()
|