def test_format_clickhouse_specific_query() -> None:
"""
Adds a few of the Clickhosue specific fields to the query.
"""
query = ClickhouseQuery(
TableSource("my_table", ColumnSet([])),
selected_columns=[
SelectedExpression("column1", Column(None, None, "column1")),
SelectedExpression("column2", Column(None, "table1", "column2")),
],
condition=binary_condition(
None,
"eq",
lhs=Column(None, None, "column1"),
rhs=Literal(None, "blabla"),
),
groupby=[
Column(None, None, "column1"),
Column(None, "table1", "column2")
],
having=binary_condition(
None,
"eq",
lhs=Column(None, None, "column1"),
rhs=Literal(None, 123),
),
order_by=[
OrderBy(OrderByDirection.ASC, Column(None, None, "column1"))
],
array_join=Column(None, None, "column1"),
sample=0.1,
totals=True,
limitby=(10, "environment"),
)
query.set_final(True)
query.set_offset(50)
query.set_limit(100)
request_settings = HTTPRequestSettings()
clickhouse_query = AstSqlQuery(query, request_settings)
expected = {
"from": "FROM my_table FINAL SAMPLE 0.1",
"group": "GROUP BY (column1, table1.column2) WITH TOTALS",
"having": "HAVING eq(column1, 123)",
"array_join": "ARRAY JOIN column1",
"limit": "LIMIT 100 OFFSET 50",
"limitby": "LIMIT 10 BY environment",
"order": "ORDER BY column1 ASC",
"select": "SELECT column1, table1.column2",
"where": "WHERE eq(column1, 'blabla')",
}
assert clickhouse_query.sql_data() == expected
def test_format_clickhouse_specific_query() -> None:
"""
Adds a few of the Clickhosue specific fields to the query.
"""
query = ClickhouseQuery(
Table("my_table", ColumnSet([]), final=True, sampling_rate=0.1),
selected_columns=[
SelectedExpression("column1", Column(None, None, "column1")),
SelectedExpression("column2", Column(None, "table1", "column2")),
],
condition=binary_condition(
"eq",
lhs=Column(None, None, "column1"),
rhs=Literal(None, "blabla"),
),
groupby=[
Column(None, None, "column1"),
Column(None, "table1", "column2")
],
having=binary_condition(
"eq",
lhs=Column(None, None, "column1"),
rhs=Literal(None, 123),
),
order_by=[
OrderBy(OrderByDirection.ASC, Column(None, None, "column1"))
],
array_join=Column(None, None, "column1"),
totals=True,
limitby=LimitBy(10, Column(None, None, "environment")),
)
query.set_offset(50)
query.set_limit(100)
request_settings = HTTPRequestSettings()
clickhouse_query = format_query(query, request_settings)
expected = [
"SELECT column1, table1.column2",
["FROM", "my_table FINAL SAMPLE 0.1"],
"ARRAY JOIN column1",
"WHERE eq(column1, 'blabla')",
"GROUP BY column1, table1.column2 WITH TOTALS",
"HAVING eq(column1, 123)",
"ORDER BY column1 ASC",
"LIMIT 10 BY environment",
"LIMIT 100 OFFSET 50",
]
assert clickhouse_query.structured() == expected
def execute(
self,
query: Query,
request_settings: RequestSettings,
runner: SplitQueryRunner,
) -> Optional[QueryResult]:
"""
Split query in 2 steps if a large number of columns is being selected.
- First query only selects event_id, project_id and timestamp.
- Second query selects all fields for only those events.
- Shrink the date range.
"""
limit = query.get_limit()
if (limit is None or limit == 0 or query.get_groupby()
or query.get_aggregations()
or not query.get_selected_columns()):
return None
if limit > settings.COLUMN_SPLIT_MAX_LIMIT:
metrics.increment("column_splitter.query_above_limit")
return None
# Do not split if there is already a = or IN condition on an ID column
id_column_matcher = FunctionCall(
Or([String(ConditionFunctions.EQ),
String(ConditionFunctions.IN)]),
(
Column(None, String(self.__id_column)),
AnyExpression(),
),
)
for expr in query.get_condition_from_ast() or []:
match = id_column_matcher.match(expr)
if match:
return None
# We need to count the number of table/column name pairs
# not the number of distinct Column objects in the query
# so to avoid counting aliased columns multiple times.
total_columns = {(col.table_name, col.column_name)
for col in query.get_all_ast_referenced_columns()}
minimal_query = copy.deepcopy(query)
minimal_query.set_selected_columns(
[self.__id_column, self.__project_column, self.__timestamp_column])
# TODO: provide the table alias name to this splitter if we ever use it
# in joins.
minimal_query.set_ast_selected_columns([
SelectedExpression(self.__id_column,
ColumnExpr(None, None, self.__id_column)),
SelectedExpression(self.__project_column,
ColumnExpr(None, None, self.__project_column)),
SelectedExpression(
self.__timestamp_column,
ColumnExpr(None, None, self.__timestamp_column),
),
])
for exp in minimal_query.get_all_expressions():
if exp.alias in (
self.__id_column,
self.__project_column,
self.__timestamp_column,
) and not (isinstance(exp, ColumnExpr)
and exp.column_name == exp.alias):
logger.warning(
"Potential alias shadowing due to column splitter",
extra={"expression": exp},
exc_info=True,
)
minimal_columns = {
(col.table_name, col.column_name)
for col in minimal_query.get_all_ast_referenced_columns()
}
if len(total_columns) <= len(minimal_columns):
return None
# Ensures the AST minimal query is actually runnable on its own.
if not minimal_query.validate_aliases():
return None
legacy_references = set(minimal_query.get_all_referenced_columns())
ast_column_names = {
c.column_name
for c in minimal_query.get_all_ast_referenced_columns()
}
# Ensures the legacy minimal query (which does not expand alias references)
# does not contain alias references we removed when creating minimal_query.
if legacy_references - ast_column_names:
metrics.increment("columns.skip_invalid_legacy_query")
return None
result = runner(minimal_query, request_settings)
del minimal_query
if not result.result["data"]:
return None
# Making a copy just in case runner returned None (which would drive the execution
# strategy to ignore the result of this splitter and try the next one).
query = copy.deepcopy(query)
event_ids = list(
set([event[self.__id_column] for event in result.result["data"]]))
if len(event_ids) > settings.COLUMN_SPLIT_MAX_RESULTS:
# We may be runing a query that is beyond clickhouse maximum query size,
# so we cowardly abandon.
metrics.increment(
"column_splitter.intermediate_results_beyond_limit")
return None
query.add_conditions([(self.__id_column, "IN", event_ids)])
query.add_condition_to_ast(
in_condition(
None,
ColumnExpr(None, None, self.__id_column),
[LiteralExpr(None, e_id) for e_id in event_ids],
))
query.set_offset(0)
# TODO: This is technically wrong. Event ids are unique per project, not globally.
# So, if the minimal query only returned the same event_id from two projects, we
# would be underestimating the limit here.
query.set_limit(len(event_ids))
project_ids = list(
set([
event[self.__project_column] for event in result.result["data"]
]))
_replace_condition(
query,
self.__project_column,
"IN",
project_ids,
)
_replace_ast_condition(
query,
self.__project_column,
"IN",
literals_tuple(None,
[LiteralExpr(None, p_id) for p_id in project_ids]),
)
timestamps = [
event[self.__timestamp_column] for event in result.result["data"]
]
_replace_condition(
query,
self.__timestamp_column,
">=",
util.parse_datetime(min(timestamps)).isoformat(),
)
_replace_ast_condition(
query,
self.__timestamp_column,
">=",
LiteralExpr(None, util.parse_datetime(min(timestamps))),
)
# We add 1 second since this gets translated to ('timestamp', '<', to_date)
# and events are stored with a granularity of 1 second.
_replace_condition(
query,
self.__timestamp_column,
"<",
(util.parse_datetime(max(timestamps)) +
timedelta(seconds=1)).isoformat(),
)
_replace_ast_condition(
query,
self.__timestamp_column,
"<",
LiteralExpr(
None,
(util.parse_datetime(max(timestamps)) + timedelta(seconds=1)),
),
)
return runner(query, request_settings)
def execute(
self,
query: Query,
query_settings: QuerySettings,
runner: SplitQueryRunner,
) -> Optional[QueryResult]:
"""
Split query in 2 steps if a large number of columns is being selected.
- First query only selects event_id, project_id and timestamp.
- Second query selects all fields for only those events.
- Shrink the date range.
"""
limit = query.get_limit()
if (limit is None or limit == 0 or query.get_groupby()
or not query.get_selected_columns()):
return None
if limit > settings.COLUMN_SPLIT_MAX_LIMIT:
metrics.increment("column_splitter.query_above_limit")
return None
# Do not split if there is already a = or IN condition on an ID column
id_column_matcher = FunctionCall(
Or([String(ConditionFunctions.EQ),
String(ConditionFunctions.IN)]),
(
Column(None, String(self.__id_column)),
AnyExpression(),
),
)
for expr in query.get_condition() or []:
match = id_column_matcher.match(expr)
if match:
return None
# We need to count the number of table/column name pairs
# not the number of distinct Column objects in the query
# so to avoid counting aliased columns multiple times.
selected_columns = {
(col.table_name, col.column_name)
for col in query.get_columns_referenced_in_select()
}
if len(selected_columns) < settings.COLUMN_SPLIT_MIN_COLS:
metrics.increment("column_splitter.main_query_min_threshold")
return None
minimal_query = copy.deepcopy(query)
# TODO: provide the table alias name to this splitter if we ever use it
# in joins.
minimal_query.set_ast_selected_columns([
SelectedExpression(
self.__id_column,
ColumnExpr(self.__id_column, None, self.__id_column),
),
SelectedExpression(
self.__project_column,
ColumnExpr(self.__project_column, None, self.__project_column),
),
SelectedExpression(
self.__timestamp_column,
ColumnExpr(self.__timestamp_column, None,
self.__timestamp_column),
),
])
for exp in minimal_query.get_all_expressions():
if exp.alias in (
self.__id_column,
self.__project_column,
self.__timestamp_column,
) and not (isinstance(exp, ColumnExpr)
and exp.column_name == exp.alias):
logger.warning(
"Potential alias shadowing due to column splitter",
extra={"expression": exp},
exc_info=True,
)
# Ensures the AST minimal query is actually runnable on its own.
if not minimal_query.validate_aliases():
return None
# There is a Clickhouse bug where if functions in the ORDER BY clause are not in the SELECT,
# they fail on distributed tables. For that specific case, skip the query splitter.
for orderby in minimal_query.get_orderby():
if isinstance(orderby.expression,
(FunctionCallExpr, CurriedFunctionCallExpr)):
metrics.increment("column_splitter.orderby_has_a_function")
return None
result = runner(minimal_query, query_settings)
del minimal_query
if not result.result["data"]:
metrics.increment("column_splitter.no_data_from_minimal_query")
return None
# Making a copy just in case runner returned None (which would drive the execution
# strategy to ignore the result of this splitter and try the next one).
query = copy.deepcopy(query)
event_ids = list(
set([event[self.__id_column] for event in result.result["data"]]))
if len(event_ids) > settings.COLUMN_SPLIT_MAX_RESULTS:
# We may be runing a query that is beyond clickhouse maximum query size,
# so we cowardly abandon.
metrics.increment(
"column_splitter.intermediate_results_beyond_limit")
return None
query.add_condition_to_ast(
in_condition(
ColumnExpr(None, None, self.__id_column),
[LiteralExpr(None, e_id) for e_id in event_ids],
))
query.set_offset(0)
query.set_limit(len(result.result["data"]))
project_ids = list(
set([
event[self.__project_column] for event in result.result["data"]
]))
_replace_ast_condition(
query,
self.__project_column,
"IN",
literals_tuple(None,
[LiteralExpr(None, p_id) for p_id in project_ids]),
)
timestamps = [
event[self.__timestamp_column] for event in result.result["data"]
]
_replace_ast_condition(
query,
self.__timestamp_column,
">=",
LiteralExpr(None, util.parse_datetime(min(timestamps))),
)
# We add 1 second since this gets translated to ('timestamp', '<', to_date)
# and events are stored with a granularity of 1 second.
_replace_ast_condition(
query,
self.__timestamp_column,
"<",
LiteralExpr(
None,
(util.parse_datetime(max(timestamps)) + timedelta(seconds=1)),
),
)
return runner(query, query_settings)
