def test_when_sample_is_not_provided_without_turbo(self):
query = Query({
"conditions": [],
"aggregations": [],
"groupby": [],
})
request_settings = RequestSettings(turbo=False,
consistent=False,
debug=False)
clickhouse_query = ClickhouseQuery(
dataset=self.dataset,
query=query,
settings=request_settings,
prewhere_conditions=[],
)
assert 'SAMPLE' not in clickhouse_query.format_sql()
def test_provided_sample_should_be_used_with_turbo(self):
query = Query({
"conditions": [],
"aggregations": [],
"groupby": [],
"sample": 0.1
})
request_settings = RequestSettings(turbo=True,
consistent=False,
debug=False)
clickhouse_query = ClickhouseQuery(
dataset=self.dataset,
query=query,
settings=request_settings,
prewhere_conditions=[],
)
assert 'SAMPLE 0.1' in clickhouse_query.format_sql()
def execute(
self,
query: ClickhouseQuery,
# TODO: move Clickhouse specific arguments into DictClickhouseQuery
settings: Optional[Mapping[str, str]] = None,
query_id: Optional[str] = None,
with_totals: bool = False,
) -> Result:
if settings is None:
settings = {}
kwargs = {}
if query_id is not None:
kwargs["query_id"] = query_id
sql = query.format_sql()
return self.__transform_result(
self.__client.execute(
sql, with_column_types=True, settings=settings, **kwargs
),
with_totals=with_totals,
)
def test_when_sample_is_not_provided_with_turbo(self):
source = self.dataset.get_dataset_schemas().get_read_schema(
).get_data_source()
query = Query(
{
"conditions": [],
"aggregations": [],
"groupby": [],
},
source,
)
request_settings = RequestSettings(turbo=True,
consistent=False,
debug=False)
clickhouse_query = ClickhouseQuery(
dataset=self.dataset,
query=query,
settings=request_settings,
prewhere_conditions=[],
)
assert "SAMPLE 0.2" in clickhouse_query.format_sql()
def test_provided_sample_should_be_used(self):
source = self.dataset.get_dataset_schemas().get_read_schema(
).get_data_source()
query = Query(
{
"conditions": [],
"aggregations": [],
"groupby": [],
"sample": 0.1
},
source,
)
request_settings = RequestSettings(turbo=False,
consistent=False,
debug=False)
clickhouse_query = ClickhouseQuery(
dataset=self.dataset,
query=query,
settings=request_settings,
prewhere_conditions=[],
)
assert 'SAMPLE 0.1' in clickhouse_query.format_sql()
def raw_query(
request: Request,
query: ClickhouseQuery,
client: ClickhousePool,
timer: Timer,
stats=None,
) -> QueryResult:
"""
Submit a raw SQL query to clickhouse and do some post-processing on it to
fix some of the formatting issues in the result JSON
"""
from snuba.clickhouse.native import NativeDriverReader
stats = stats or {}
use_cache, use_deduper, uc_max = state.get_configs([
('use_cache', 0),
('use_deduper', 1),
('uncompressed_cache_max_cols', 5),
])
all_confs = state.get_all_configs()
query_settings = {
k.split('/', 1)[1]: v
for k, v in all_confs.items() if k.startswith('query_settings/')
}
# Experiment, if we are going to grab more than X columns worth of data,
# don't use uncompressed_cache in clickhouse, or result cache in snuba.
if len(all_referenced_columns(request.query)) > uc_max:
query_settings['use_uncompressed_cache'] = 0
use_cache = 0
timer.mark('get_configs')
sql = query.format_sql()
query_id = md5(force_bytes(sql)).hexdigest()
with state.deduper(query_id if use_deduper else None) as is_dupe:
timer.mark('dedupe_wait')
result = state.get_result(query_id) if use_cache else None
timer.mark('cache_get')
stats.update({
'is_duplicate': is_dupe,
'query_id': query_id,
'use_cache': bool(use_cache),
'cache_hit': bool(result)
}),
if result:
status = 200
else:
try:
with RateLimitAggregator(
request.settings.get_rate_limit_params(
)) as rate_limit_stats_container:
stats.update(rate_limit_stats_container.to_dict())
timer.mark('rate_limit')
project_rate_limit_stats = rate_limit_stats_container.get_stats(
PROJECT_RATE_LIMIT_NAME)
if 'max_threads' in query_settings and \
project_rate_limit_stats is not None and \
project_rate_limit_stats.concurrent > 1:
maxt = query_settings['max_threads']
query_settings['max_threads'] = max(
1, maxt - project_rate_limit_stats.concurrent + 1)
# Force query to use the first shard replica, which
# should have synchronously received any cluster writes
# before this query is run.
consistent = request.settings.get_consistent()
stats['consistent'] = consistent
if consistent:
query_settings['load_balancing'] = 'in_order'
query_settings['max_threads'] = 1
try:
result = NativeDriverReader(client).execute(
query,
query_settings,
# All queries should already be deduplicated at this point
# But the query_id will let us know if they aren't
query_id=query_id if use_deduper else None,
with_totals=request.query.has_totals(),
)
status = 200
logger.debug(sql)
timer.mark('execute')
stats.update({
'result_rows': len(result['data']),
'result_cols': len(result['meta']),
})
if use_cache:
state.set_result(query_id, result)
timer.mark('cache_set')
except BaseException as ex:
error = str(ex)
status = 500
logger.exception("Error running query: %s\n%s", sql,
error)
if isinstance(ex, ClickHouseError):
result = {
'error': {
'type': 'clickhouse',
'code': ex.code,
'message': error,
}
}
else:
result = {
'error': {
'type': 'unknown',
'message': error,
}
}
except RateLimitExceeded as ex:
error = str(ex)
status = 429
result = {
'error': {
'type': 'ratelimit',
'message': 'rate limit exceeded',
'detail': error
}
}
stats.update(query_settings)
if settings.RECORD_QUERIES:
# send to redis
state.record_query({
'request': request.body,
'sql': sql,
'timing': timer,
'stats': stats,
'status': status,
})
timer.send_metrics_to(metrics,
tags={
'status': str(status),
'referrer': stats.get('referrer', 'none'),
'final': str(stats.get('final', False)),
},
mark_tags={
'final': str(stats.get('final', False)),
})
result['timing'] = timer
if settings.STATS_IN_RESPONSE or request.settings.get_debug():
result['stats'] = stats
result['sql'] = sql
return QueryResult(result, status)
def raw_query(
request: Request,
query: ClickhouseQuery,
timer: Timer,
query_metadata: SnubaQueryMetadata,
stats: MutableMapping[str, Any],
trace_id: Optional[str] = None,
) -> RawQueryResult:
"""
Submits a raw SQL query to the DB and does some post-processing on it to
fix some of the formatting issues in the result JSON.
This function is not supposed to depend on anything higher level than the storage
query (ClickhouseQuery as of now). If this function ends up depending on the
dataset, something is wrong.
TODO: As soon as we have a StorageQuery abstraction remove all the references
to the original query from the request.
"""
use_cache, use_deduper, uc_max = state.get_configs([
("use_cache", settings.USE_RESULT_CACHE),
("use_deduper", 1),
("uncompressed_cache_max_cols", 5),
])
all_confs = state.get_all_configs()
query_settings: MutableMapping[str, Any] = {
k.split("/", 1)[1]: v
for k, v in all_confs.items() if k.startswith("query_settings/")
}
# Experiment, if we are going to grab more than X columns worth of data,
# don't use uncompressed_cache in clickhouse, or result cache in snuba.
if len(request.query.get_all_referenced_columns()) > uc_max:
query_settings["use_uncompressed_cache"] = 0
use_cache = 0
timer.mark("get_configs")
sql = query.format_sql()
query_id = md5(force_bytes(sql)).hexdigest()
with state.deduper(query_id if use_deduper else None) as is_dupe:
timer.mark("dedupe_wait")
result = cache.get(query_id) if use_cache else None
timer.mark("cache_get")
stats.update({
"is_duplicate": is_dupe,
"query_id": query_id,
"use_cache": bool(use_cache),
"cache_hit": bool(result),
}),
update_with_status = partial(
update_query_metadata_and_stats,
request,
sql,
timer,
stats,
query_metadata,
query_settings,
trace_id,
)
if not result:
try:
with RateLimitAggregator(
request.settings.get_rate_limit_params(
)) as rate_limit_stats_container:
stats.update(rate_limit_stats_container.to_dict())
timer.mark("rate_limit")
project_rate_limit_stats = rate_limit_stats_container.get_stats(
PROJECT_RATE_LIMIT_NAME)
if ("max_threads" in query_settings
and project_rate_limit_stats is not None
and project_rate_limit_stats.concurrent > 1):
maxt = query_settings["max_threads"]
query_settings["max_threads"] = max(
1, maxt - project_rate_limit_stats.concurrent + 1)
# Force query to use the first shard replica, which
# should have synchronously received any cluster writes
# before this query is run.
consistent = request.settings.get_consistent()
stats["consistent"] = consistent
if consistent:
query_settings["load_balancing"] = "in_order"
query_settings["max_threads"] = 1
try:
result = reader.execute(
query,
query_settings,
# All queries should already be deduplicated at this point
# But the query_id will let us know if they aren't
query_id=query_id if use_deduper else None,
with_totals=request.query.has_totals(),
)
timer.mark("execute")
stats.update({
"result_rows": len(result["data"]),
"result_cols": len(result["meta"]),
})
if use_cache:
cache.set(query_id, result)
timer.mark("cache_set")
except BaseException as ex:
error = str(ex)
logger.exception("Error running query: %s\n%s", sql,
error)
stats = update_with_status("error")
meta = {}
if isinstance(ex, ClickhouseError):
err_type = "clickhouse"
meta["code"] = ex.code
else:
err_type = "unknown"
raise RawQueryException(
err_type=err_type,
message=error,
stats=stats,
sql=sql,
**meta,
)
except RateLimitExceeded as ex:
stats = update_with_status("rate-limited")
raise RawQueryException(
err_type="rate-limited",
message="rate limit exceeded",
stats=stats,
sql=sql,
detail=str(ex),
)
stats = update_with_status("success")
return RawQueryResult(result, {"stats": stats, "sql": sql})
