def test_translation(
mappers: TranslationMappers, query: SnubaQuery, expected: ClickhouseQuery
) -> None:
translated = QueryTranslator(mappers).translate(query)
# TODO: consider providing an __eq__ method to the Query class. Or turn it into
# a dataclass.
assert expected.get_selected_columns() == translated.get_selected_columns()
assert expected.get_groupby() == translated.get_groupby()
assert expected.get_condition() == translated.get_condition()
assert expected.get_arrayjoin() == translated.get_arrayjoin()
assert expected.get_having() == translated.get_having()
assert expected.get_orderby() == translated.get_orderby()
def get_filtered_mapping_keys(query: Query, column_name: str) -> Sequence[str]:
"""
Identifies the conditions we can apply the arrayFilter optimization
on.
Which means: if the arrayJoin is in the select clause, there
are one or more top level AND condition on the arrayJoin and
there is no OR condition in the query.
"""
array_join_found = any(
array_join_pattern(column_name).match(f) is not None
for selected in query.get_selected_columns() or []
for f in selected.expression
)
if not array_join_found:
return list()
ast_condition = query.get_condition()
cond_keys = (
_get_mapping_keys_in_condition(ast_condition, column_name)
if ast_condition is not None
else set()
)
if cond_keys is None:
# This means we found an OR. Cowardly we give up even though there could
# be cases where this condition is still optimizable.
return []
ast_having = query.get_having()
having_keys = (
_get_mapping_keys_in_condition(ast_having, column_name)
if ast_having is not None
else set()
)
if having_keys is None:
# Same as above
return []
keys = cond_keys | having_keys
return sorted(list(keys))
def test_replace_expression() -> None:
"""
Create a query with the new AST and replaces a function with a different function
replaces f1(...) with tag(f1)
"""
column1 = Column(None, "t1", "c1")
column2 = Column(None, "t1", "c2")
function_1 = FunctionCall("alias", "f1", (column1, column2))
function_2 = FunctionCall("alias", "f2", (column2,))
condition = binary_condition(ConditionFunctions.EQ, function_1, Literal(None, "1"))
prewhere = binary_condition(ConditionFunctions.EQ, function_1, Literal(None, "2"))
orderby = OrderBy(OrderByDirection.ASC, function_2)
query = Query(
Table("my_table", ColumnSet([])),
selected_columns=[SelectedExpression("alias", function_1)],
array_join=None,
condition=condition,
groupby=[function_1],
having=None,
prewhere=prewhere,
order_by=[orderby],
)
def replace(exp: Expression) -> Expression:
if isinstance(exp, FunctionCall) and exp.function_name == "f1":
return FunctionCall(exp.alias, "tag", (Literal(None, "f1"),))
return exp
query.transform_expressions(replace)
expected_query = Query(
Table("my_table", ColumnSet([])),
selected_columns=[
SelectedExpression(
"alias", FunctionCall("alias", "tag", (Literal(None, "f1"),))
)
],
array_join=None,
condition=binary_condition(
ConditionFunctions.EQ,
FunctionCall("alias", "tag", (Literal(None, "f1"),)),
Literal(None, "1"),
),
groupby=[FunctionCall("alias", "tag", (Literal(None, "f1"),))],
prewhere=binary_condition(
ConditionFunctions.EQ,
FunctionCall("alias", "tag", (Literal(None, "f1"),)),
Literal(None, "2"),
),
having=None,
order_by=[orderby],
)
assert query.get_selected_columns() == expected_query.get_selected_columns()
assert query.get_condition() == expected_query.get_condition()
assert query.get_groupby() == expected_query.get_groupby()
assert query.get_having() == expected_query.get_having()
assert query.get_orderby() == expected_query.get_orderby()
assert list(query.get_all_expressions()) == list(
expected_query.get_all_expressions()
)
