def _prepare_query(client: SyncClient, query: str, args: QueryArgs):
"""
Given a string query with placeholders we do one of two things:
1. for a insert query we just format, and remove comments
2. for non-insert queries, we return the sql with placeholders
evaluated with the contents of `args`
We also return `tags` which contains some detail around the context
within which the query was executed e.g. the django view name
NOTE: `client.execute` would normally handle substitution, but
because we want to strip the comments to make it easier to copy
and past queries from the `system.query_log` easily with metabase
(metabase doesn't show new lines, so with comments, you can't get
a working query without exporting to csv or similar), we need to
do it manually.
We only want to try to substitue for SELECT queries, which
clickhouse_driver at this moment in time decides based on the
below predicate.
"""
prepared_args: Any = QueryArgs
if isinstance(args, (list, tuple, types.GeneratorType)):
# If we get one of these it means we have an insert, let the clickhouse
# client handle substitution here.
rendered_sql = query
prepared_args = args
elif not args:
# If `args` is not truthy then make prepared_args `None`, which the
# clickhouse client uses to signal no substitution is desired. Expected
# args balue are `None` or `{}` for instance
rendered_sql = query
prepared_args = None
else:
# Else perform the substitution so we can perform operations on the raw
# non-templated SQL
rendered_sql = client.substitute_params(query, args)
prepared_args = None
formatted_sql = sqlparse.format(rendered_sql, strip_comments=True)
annotated_sql, tags = _annotate_tagged_query(formatted_sql, args)
if app_settings.SHELL_PLUS_PRINT_SQL:
print()
print(format_sql(formatted_sql))
return annotated_sql, prepared_args, tags
class Check(object):
def __init__(self, *args, **kwargs):
self.conn = Client(*args, **kwargs)
self.code = Code()
def execute(self, *args, **kwargs) -> List[tuple]:
"""
Wrapper to execute ClickHouse SQL
"""
if len(args) == 1:
logger.info('Execute query: {}'.format(args[0]))
elif len(args) >= 2:
logger.info('Execute query: {}'.format(
self.conn.substitute_params(args[0], args[1])))
return self.conn.execute(*args, **kwargs)
def execute_dict(self, *args, **kwargs) -> List[dict]:
"""
Wrapper around execute() to return list of rows as dict
"""
kwargs['with_column_types'] = True
rows, columns = self.execute(*args, **kwargs)
result = [{columns[i][0]: v for i, v in enumerate(r)} for r in rows]
return result
def exit(self, message: str) -> ExitStruct:
message = self.code.name + ': ' + message
return self.code.current, message
def check_config(self, config: dict, keys: set):
"""
Checks if all mandatory keys are presented in the config dict
"""
keys_in_config = config.keys() & keys
if keys_in_config != keys:
raise KeyError('Not all of {} presented in config: {}'.format(
keys, config))
def _prepare_query(client: SyncClient, query: str, args: QueryArgs):
"""
Given a string query with placeholders we do one of two things:
1. for a insert query we just format, and remove comments
2. for non-insert queries, we return the sql with placeholders
evaluated with the contents of `args`
We also return `tags` which contains some detail around the context
within which the query was executed e.g. the django view name
NOTE: `client.execute` would normally handle substitution, but
because we want to strip the comments to make it easier to copy
and past queries from the `system.query_log` easily with metabase
(metabase doesn't show new lines, so with comments, you can't get
a working query without exporting to csv or similar), we need to
do it manually.
We only want to try to substitue for SELECT queries, which
clickhouse_driver at this moment in time decides based on the
below predicate.
"""
if isinstance(args, (list, tuple, types.GeneratorType)):
rendered_sql = query
else:
rendered_sql = client.substitute_params(query, args or {})
args = None
formatted_sql = sqlparse.format(rendered_sql, strip_comments=True)
annotated_sql, tags = _annotate_tagged_query(formatted_sql, args)
if app_settings.SHELL_PLUS_PRINT_SQL:
print()
print(format_sql(formatted_sql))
return annotated_sql, args, tags
class Client(object):
def __init__(self, *args, **kwargs):
self._loop = kwargs.pop('loop', None) or asyncio.get_event_loop()
self._executor = kwargs.pop('executor', None)
if '_client' not in kwargs:
self._client = BlockingClient(*args, **kwargs)
else:
self._client = kwargs.pop('_client')
super(Client, self).__init__()
@classmethod
def from_url(cls, url, loop=None, executor=None):
"""
*New in version 0.0.2.*
"""
_client = BlockingClient.from_url(url)
return cls(_client=_client, loop=loop, executor=executor)
def run_in_executor(self, *args, **kwargs):
return run_in_executor(self._executor, self._loop, *args, **kwargs)
async def disconnect(self):
return await self.run_in_executor(self._client.disconnect)
async def execute(self, *args, **kwargs):
return await self.run_in_executor(self._client.execute, *args,
**kwargs)
async def execute_with_progress(
self, query, params=None, with_column_types=False,
external_tables=None, query_id=None, settings=None,
types_check=False, columnar=False):
self._client.make_query_settings(settings)
await self.run_in_executor(self._client.connection.force_connect)
self._client.last_query = QueryInfo()
return await self.process_ordinary_query_with_progress(
query, params=params, with_column_types=with_column_types,
external_tables=external_tables,
query_id=query_id, types_check=types_check, columnar=columnar
)
async def execute_iter(
self, query, params=None, with_column_types=False,
external_tables=None, query_id=None, settings=None,
types_check=False):
"""
*New in version 0.0.2.*
"""
self._client.make_query_settings(settings)
await self.run_in_executor(self._client.connection.force_connect)
self._client.last_query = QueryInfo()
return await self.iter_process_ordinary_query(
query, params=params, with_column_types=with_column_types,
external_tables=external_tables,
query_id=query_id, types_check=types_check
)
async def process_ordinary_query_with_progress(
self, query, params=None, with_column_types=False,
external_tables=None, query_id=None,
types_check=False, columnar=False):
if params is not None:
query = self._client.substitute_params(query, params)
await self.run_in_executor(
self._client.connection.send_query, query, query_id=query_id
)
await self.run_in_executor(
self._client.connection.send_external_tables, external_tables,
types_check=types_check
)
return await self.receive_result(
with_column_types=with_column_types, progress=True,
columnar=columnar
)
async def iter_process_ordinary_query(
self, query, params=None, with_column_types=False,
external_tables=None, query_id=None,
types_check=False):
if params is not None:
query = self._client.substitute_params(query, params)
await self.run_in_executor(
self._client.connection.send_query, query, query_id=query_id
)
await self.run_in_executor(
self._client.connection.send_external_tables, external_tables,
types_check=types_check
)
return self.iter_receive_result(with_column_types=with_column_types)
async def cancel(self, with_column_types=False):
# TODO: Add warning if already cancelled.
await self.run_in_executor(self._client.connection.send_cancel)
# Client must still read until END_OF_STREAM packet.
return await self.receive_result(with_column_types=with_column_types)
async def iter_receive_result(self, with_column_types=False):
gen = self.packet_generator()
rows_gen = IterQueryResult(gen, with_column_types=with_column_types)
async for rows in rows_gen:
for row in rows:
yield row
async def packet_generator(self):
receive_packet = self._client.receive_packet
while True:
try:
packet = await self.run_in_executor(receive_packet)
if not packet:
break
if packet is True:
continue
yield packet
except (Exception, KeyboardInterrupt):
await self.disconnect()
raise
async def receive_result(
self, with_column_types=False, progress=False, columnar=False):
gen = self.packet_generator()
if progress:
return ProgressQueryResult(
gen, with_column_types=with_column_types, columnar=columnar
)
else:
result = QueryResult(
gen, with_column_types=with_column_types, columnar=columnar
)
return await result.get_result()
