def test_visit_expression():
col1 = Column("al", "t1", "c1")
literal1 = Literal("al2", "test")
mapping = Column("al2", "t1", "tags")
key = Literal(None, "myTag")
tag = SubscriptableReference(None, mapping, key)
f1 = FunctionCall("al3", "f1", [col1, literal1, tag])
col2 = Column("al4", "t1", "c2")
literal2 = Literal("al5", "test2")
f2 = FunctionCall("al6", "f2", [col2, literal2])
curried = CurriedFunctionCall("al7", f2, [f1])
visitor = DummyVisitor()
ret = curried.accept(visitor)
expected = [
curried,
f2,
col2,
literal2,
f1,
col1,
literal1,
tag,
mapping,
key,
]
# Tests the state changes on the Visitor
assert visitor.get_visited_nodes() == expected
# Tests the return value of the visitor
assert ret == expected
def attempt_map(
self,
expression: CurriedFunctionCall,
children_translator: SnubaClickhouseStrictTranslator,
) -> Optional[CurriedFunctionCall]:
internal_function = expression.internal_function.accept(
children_translator)
assert isinstance(internal_function, FunctionCall) # mypy
parameters = tuple(
p.accept(children_translator) for p in expression.parameters)
for param in parameters:
# All impossible columns that have been converted to NULL will be the identity function.
# So we know that if a function has the identity function as a parameter, we can
# collapse the entire expression.
fmatch = self.function_match.match(param)
if fmatch is not None:
return CurriedFunctionCall(
alias=expression.alias,
internal_function=FunctionCall(
None,
f"{internal_function.function_name}OrNull",
internal_function.parameters,
),
parameters=tuple(Literal(None, None) for p in parameters),
)
return None
def builder(alias: Optional[str], params: List[Expression]) -> Expression:
# The first element in the sequence is the inner function.
# Unfortunately I could not find a better way to reuse this
# between FunctionCall and CurriedFunctionCall.
inner_function = params.pop(0)
assert isinstance(inner_function, FunctionCall)
return CurriedFunctionCall(alias, inner_function, tuple(params))
def test_column_curried_function_translation() -> None:
assert ColumnToCurriedFunction(
None,
"duration_quantiles",
FunctionCall(
None,
"quantilesIfMerge",
(
Literal(None, 0.5),
Literal(None, 0.9),
),
),
(Column(None, None, "duration_quantiles"),),
).attempt_map(
Column("duration_quantiles", None, "duration_quantiles"),
SnubaClickhouseMappingTranslator(TranslationMappers()),
) == CurriedFunctionCall(
"duration_quantiles",
FunctionCall(
None,
"quantilesIfMerge",
(
Literal(None, 0.5),
Literal(None, 0.9),
),
),
(Column(None, None, "duration_quantiles"),),
)
def attempt_map(
self,
expression: CurriedFunctionCall,
children_translator: SnubaClickhouseStrictTranslator,
) -> Optional[CurriedFunctionCall]:
if not _should_transform_aggregation(
expression.internal_function.function_name,
self.from_name,
self.column_to_map,
expression,
):
return None
return CurriedFunctionCall(
expression.alias,
FunctionCall(
None,
self.to_name,
tuple(
p.accept(children_translator)
for p in expression.internal_function.parameters),
),
_build_parameters(expression, children_translator,
self.aggr_col_name),
)
def query_runner(query: Query, settings: RequestSettings,
reader: Reader[SqlQuery]) -> QueryResult:
assert query.get_selected_columns_from_ast() == [
SelectedExpression(
"duration_quantiles",
CurriedFunctionCall(
"duration_quantiles",
FunctionCall(
None,
"quantilesIfMerge",
(Literal(None, 0.5), Literal(None, 0.9)),
),
(Column(None, None, "duration_quantiles"), ),
),
),
SelectedExpression(
"sessions",
FunctionCall("sessions", "countIfMerge",
(Column(None, None, "sessions"), )),
),
SelectedExpression(
"users",
FunctionCall("users", "uniqIfMerge",
(Column(None, None, "users"), )),
),
]
return QueryResult({}, {})
def visit_function_call(
self,
node: Node,
visited_children: Tuple[
str,
Any,
List[Expression],
Any,
Union[Node, List[Expression]],
Union[Node, List[Any]],
],
) -> Expression:
name, _, params1, _, params2, alias = visited_children
if isinstance(alias, Node) or len(alias) == 0:
alias = None
else:
_, _, _, alias = alias
alias = alias.text
param_list1 = tuple(params1)
if isinstance(params2, Node) and params2.text == "":
# params2.text == "" means empty node.
return FunctionCall(alias, name, param_list1)
internal_f = FunctionCall(None, name, param_list1)
_, param_list2, _ = params2
if isinstance(param_list2, (list, tuple)) and len(param_list2) > 0:
param_list2 = tuple(param_list2)
else:
# This happens when the second parameter list is empty. Somehow
# it does not turn into an empty list.
param_list2 = ()
return CurriedFunctionCall(alias, internal_f, param_list2)
def test_iterate() -> None:
"""
Test iteration over a subtree. The subtree is a function call in the form
f2(c3, f1(c1, c2))
"""
column1 = Column(None, "c1", "t1")
column2 = Column(None, "c2", "t1")
function_1 = FunctionCall(None, "f1", (column1, column2))
column3 = Column(None, "c2", "t1")
column4 = Column(None, "c3", "t1")
literal = Literal(None, "blablabla")
function_2i = FunctionCall(None, "f2", (column3, function_1, literal))
function_2 = CurriedFunctionCall(None, function_2i, (column4, ))
expected = [
column3,
column1,
column2,
function_1,
literal,
function_2i,
column4,
function_2,
]
assert list(function_2) == expected
def test_visit_expression():
col1 = Column("al", "c1", "t1")
literal1 = Literal("al2", "test")
f1 = FunctionCall("al3", "f1", [col1, literal1])
col2 = Column("al4", "c2", "t1")
literal2 = Literal("al5", "test2")
f2 = FunctionCall("al6", "f2", [col2, literal2])
curried = CurriedFunctionCall("al7", f2, [f1])
visitor = DummyVisitor()
ret = curried.accept(visitor)
expected = [curried, f2, col2, literal2, f1, col1, literal1]
# Tests the state changes on the Visitor
assert visitor.get_visited_nodes() == expected
# Tests the return value of the visitor
assert ret == expected
def attempt_map(
self,
expression: CurriedFunctionCall,
children_translator: SnubaClickhouseStrictTranslator,
) -> Optional[CurriedFunctionCall]:
return CurriedFunctionCall(
alias=expression.alias,
internal_function=children_translator.translate_function_strict(
expression.internal_function),
parameters=tuple(
p.accept(children_translator) for p in expression.parameters),
)
def parse_aggregation(
aggregation_function: str,
column: Any,
alias: Optional[str],
dataset_columns: ColumnSet,
array_join_cols: Set[str],
) -> Expression:
"""
Aggregations, unfortunately, support both Snuba syntax and a subset
of Clickhouse syntax. In order to preserve this behavior and still build
a meaningful AST when parsing the query, we need to do some parsing of
the clickhouse expression. (not that we should support this, but it is
used in production).
"""
if not isinstance(column, (list, tuple)):
columns: Iterable[Any] = (column, )
else:
columns = column
columns_expr = [
parse_expression(column, dataset_columns, array_join_cols)
for column in columns if column
]
matched = FUNCTION_NAME_RE.fullmatch(aggregation_function)
if matched is not None:
return FunctionCall(alias, aggregation_function, tuple(columns_expr))
parsed_expression = parse_clickhouse_function(aggregation_function)
if (
# Simple Clickhouse expression with no snuba syntax
# ["ifNull(count(somthing), something)", None, None]
isinstance(parsed_expression, (FunctionCall, CurriedFunctionCall))
and not columns_expr):
return replace(parsed_expression, alias=alias)
elif isinstance(parsed_expression, FunctionCall) and columns_expr:
# Mix of clickhouse syntax and snuba syntax that generates a CurriedFunction
# ["f(a)", "b", None]
return CurriedFunctionCall(
alias,
parsed_expression,
tuple(columns_expr),
)
else:
raise ParsingException(
f"Invalid aggregation format {aggregation_function} {column}",
report=False)
def test_mapping_curried_function() -> None:
c1 = Column(None, "c1", "t1")
f1 = FunctionCall(None, "f1", (c1, ))
c2 = Column(None, "c1", "t1")
f2 = CurriedFunctionCall(None, f1, (c2, ))
def replace_col(e: Expression) -> Expression:
if isinstance(e, Column) and e.column_name == "c1":
return Column(None, "c2", "t1")
return e
f3 = f2.transform(replace_col)
replaced_col = Column(None, "c2", "t1")
replaced_function = FunctionCall(None, "f1", (replaced_col, ))
expected = [
replaced_col,
replaced_function,
replaced_col,
CurriedFunctionCall(None, replaced_function, (replaced_col, )),
]
assert list(f3) == expected
def test_hash() -> None:
"""
Ensures expressions are hashable
"""
column1 = Column(None, "c1", "t1")
column2 = Column(None, "c2", "t1")
function_1 = FunctionCall(None, "f1", (column1, column2))
literal = Literal(None, "blablabla")
function_2 = CurriedFunctionCall(None, function_1, (column1, ))
lm = Lambda(None, ("x", "y"),
FunctionCall(None, "test", (Argument(None, "x"))))
s = set()
s.add(column1)
s.add(column2)
s.add(function_1)
s.add(literal)
s.add(function_2)
s.add(lm)
assert len(s) == 6
def query_runner(query: Query, settings: RequestSettings,
reader: Reader) -> QueryResult:
quantiles = tuple(
Literal(None, quant) for quant in [0.5, 0.75, 0.9, 0.95, 0.99, 1])
assert query.get_selected_columns() == [
SelectedExpression(
"duration_quantiles",
CurriedFunctionCall(
"_snuba_duration_quantiles",
FunctionCall(
None,
"quantilesIfMerge",
quantiles,
),
(Column(None, None, "duration_quantiles"), ),
),
),
SelectedExpression(
"sessions",
FunctionCall(
"_snuba_sessions",
"plus",
(
FunctionCall(None, "countIfMerge",
(Column(None, None, "sessions"), )),
FunctionCall(
None,
"sumIfMerge",
(Column(None, None, "sessions_preaggr"), ),
),
),
),
),
SelectedExpression(
"users",
FunctionCall("_snuba_users", "uniqIfMerge",
(Column(None, None, "users"), )),
),
]
return QueryResult({}, {})
def parse_aggregation(aggregation_function: str, column: Any,
alias: Optional[str]) -> Expression:
"""
Aggregations, unfortunately, support both Snuba syntax and a subset
of ClickHosue syntax. In order to preserve this behavior and still build
a meaningful AST when parsing the query, we need to do some parsing of
the clickhouse expression. (not that we should support this, but it is
used in production).
"""
expression_tree = minimal_clickhouse_grammar.parse(aggregation_function)
parsed_expression = ClickhouseVisitor().visit(expression_tree)
if not isinstance(column, (list, tuple)):
columns: Iterable[Any] = (column, )
else:
columns = column
columns_expr = [parse_expression(column) for column in columns if column]
if isinstance(parsed_expression, str):
# Simple aggregation with snuba syntax ["count", ["c1", "c2"]]
return FunctionCall(alias, parsed_expression, tuple(columns_expr))
elif (
# Simple Clickhouse expression with no snuba syntax
# ["ifNull(count(somthing), something)", None, None]
isinstance(parsed_expression, (FunctionCall, CurriedFunctionCall))
and not columns_expr):
return replace(parsed_expression, alias=alias)
elif isinstance(parsed_expression, FunctionCall) and columns_expr:
# Mix of clickhouse syntax and snuba syntax that generates a CurriedFunction
# ["f(a)", "b", None]
return CurriedFunctionCall(
alias,
parsed_expression,
tuple(columns_expr),
)
else:
raise ValueError(
f"Invalid aggregation format {aggregation_function} {column}")
def attempt_map(
self,
expression: CurriedFunctionCall,
children_translator: SnubaClickhouseStrictTranslator,
) -> Optional[CurriedFunctionCall]:
internal_function = expression.internal_function.accept(
children_translator)
assert isinstance(internal_function, FunctionCall) # mypy
parameters = tuple(
p.accept(children_translator) for p in expression.parameters)
all_null = True
for param in parameters:
# Handle wrapped functions that have been converted to ifNull(NULL, NULL)
fmatch = self.function_match.match(param)
if fmatch is None:
if isinstance(param, Literal):
if param.value is not None:
all_null = False
break
else:
all_null = False
break
if all_null and len(parameters) > 0:
# Currently curried function mappers require returning other curried functions.
# So return this to keep the mapper happy.
return CurriedFunctionCall(
alias=expression.alias,
internal_function=FunctionCall(
None,
f"{internal_function.function_name}OrNull",
internal_function.parameters,
),
parameters=tuple(Literal(None, None) for p in parameters),
)
return None
["f(123e+06)", None, None],
FunctionCall(None, "f", (Literal(None, float(123e06)), )),
), # Int literal exp
(
["f(123.2)", None, None],
FunctionCall(None, "f", (Literal(None, 123.2), )),
), # Float literal
(
["f(123.2e-06)", None, None],
FunctionCall(None, "f", (Literal(None, 123.2e-06), )),
), # Float literal exp
(
["count()", "event_id", None],
CurriedFunctionCall(
None,
FunctionCall(None, "count", ()),
(Column(None, None, "event_id"), ),
),
), # This is probably wrong, but we cannot disambiguate it at this level
(
["uniq", "platform", "uniq_platforms"],
FunctionCall("uniq_platforms", "uniq",
(Column(None, None, "platform"), )),
), # Use the columns provided as parameters
(
["topK(1)", "platform", "top_platforms"],
CurriedFunctionCall(
"top_platforms",
FunctionCall(None, "topK", (Literal(None, 1), )),
(Column(None, None, "platform"), ),
),
"ORDER BY column1 ASC, table1.column2 DESC"),
id="Simple query with aliases and multiple tables",
),
pytest.param(
Query(
Table("my_table", ColumnSet([])),
selected_columns=[
SelectedExpression(
"my_complex_math",
CurriedFunctionCall(
"my_complex_math",
FunctionCall(
None,
"doSomething",
(
Column(None, None, "column1"),
Column(None, "table1", "column2"),
Column("al", None, "column3"),
),
),
(Column(None, None, "column1"), ),
),
),
],
condition=binary_condition(
"and",
lhs=binary_condition(
"eq",
lhs=Column("al", None, "column3"),
rhs=Literal(None, "blabla"),
),
def _produce_output(self, expression: ColumnExpr) -> CurriedFunctionCall:
return CurriedFunctionCall(
alias=expression.alias,
internal_function=self.to_internal_function,
parameters=self.to_function_params,
)
build_query(selected_columns=[
identity(
identity(
identity(
equals(
# alias of the tuple of internal function is removed (it is not useful)
tupleElement(None, some_tuple(
alias=None), Literal(None, 1)),
Literal(None, 300),
))))
]),
id="Highly nested",
),
pytest.param(
build_query(selected_columns=[
CurriedFunctionCall(None, some_tuple(
alias="foo"), (Literal(None, "4"), ))
]),
build_query(selected_columns=[
CurriedFunctionCall(None, some_tuple(
alias=None), (Literal(None, "4"), ))
]),
id="curried function",
),
]
@pytest.mark.parametrize("input_query,expected_query", TEST_QUERIES)
def test_tuple_unaliaser(input_query, expected_query):
set_config("tuple_unaliaser_rollout", 1)
settings = HTTPQuerySettings()
TupleUnaliaser().process_query(input_query, settings)
None,
"someFunc",
(
FunctionCall(
None,
"anotherFunc",
(Column("col", None, "col"), Literal(None, 123)),
),
CurriedFunctionCall(
None,
FunctionCall(
None,
"yetAnotherOne",
(
SubscriptableReference(
"tags[release]",
Column(None, None, "tags"),
Literal(None, "release"),
),
),
),
(Column("cola", None, "cola"), Literal(None, 123)),
),
),
),
FunctionCall(
None,
"someFunc",
(
FunctionCall(
None,
function_name="f",
parameters=(FunctionCall(
alias="b",
function_name="g",
parameters=(Column(alias=None,
table_name=None,
column_name="b"), ),
), ),
),
id="Function shadows inner column",
),
pytest.param(
CurriedFunctionCall(
alias=None,
internal_function=FunctionCall(alias=None,
function_name="f",
parameters=tuple()),
parameters=(Column(alias=None, table_name=None,
column_name="a"), ),
),
{
"a":
FunctionCall(
alias="a",
function_name="f",
parameters=(Column(alias="b", table_name=None,
column_name="b"), ),
)
},
False,
CurriedFunctionCall(
alias=None,
(
FunctionCall(
None,
"f1",
(
FunctionCall("al1", "f2", (Column(None, "param1", "table1"))),
FunctionCall("al2", "f3", (Column(None, "param2", "table1"))),
),
),
"f1((f2(table1.param1) AS al1), (f3(table1.param2) AS al2))",
), # Hierarchical function call with aliases
(
CurriedFunctionCall(
None,
FunctionCall(None, "f0", (Column(None, "param1", "table1"), )),
(
FunctionCall(None, "f1", (Column(None, "param2", "table1"), )),
Column(None, "param3", "table1"),
),
),
"f0(table1.param1)(f1(table1.param2), table1.param3)",
), # Curried function call with hierarchy
(
FunctionCall(
None,
"arrayExists",
(
Lambda(
None,
("x", "y"),
FunctionCall(None, "testFunc",
(Argument(None, "x"), Argument(None, "y"))),
),
), ),
),
Literal(None, 1),
),
),
Literal(None, 0),
),
),
identity(Literal(None, None), "alias"),
id="discover craziness with replacement",
),
pytest.param(
CurriedFunctionCall(
None,
FunctionCall(None, "quantile", (Literal(None, 0.95), )),
(Column(None, None, "primary_hash"), ),
),
CurriedFunctionCall(
None,
FunctionCall(None, "quantileOrNull", (Literal(None, 0.95), )),
(Literal(None, None), ),
),
id="curried function call",
),
pytest.param(
FunctionCall(
None,
"plus",
(
Literal(None, 0),
),
"(cats.sounds AS `kittysounds`)['meow'] AS `catsound`",
),
(Column("alias", None, "c1"), "c1 AS `alias`"),
(
FunctionCall(None, "f1",
(Column(None, "t1", "c1"), Column(None, "t1", "c2"))),
"""f1(
t1.c1,
t1.c2
)""",
),
(
CurriedFunctionCall(
None,
FunctionCall(None, "f1",
(Column(None, "t1", "c1"), Column(None, "t1", "c2"))),
(Literal(None, "hello"), Literal(None, "kitty")),
),
"""f1(
t1.c1,
t1.c2
)(
'hello',
'kitty'
)""",
),
(
FunctionCall(
None,
"f1",
(
Literal(None, ""),
),
),
],
),
), # Complex query with both uniq and emptyIfNull
(
Query(
{},
TableSource("events", ColumnSet([])),
selected_columns=[
SelectedExpression(
name=None,
expression=CurriedFunctionCall(
None,
FunctionCall(None, "top", (Literal(None, 10), )),
(Column(None, None, "column1"), ),
),
)
],
),
Query(
{},
TableSource("events", ColumnSet([])),
selected_columns=[
SelectedExpression(
name=None,
expression=CurriedFunctionCall(
None,
FunctionCall(None, "topK", (Literal(None, 10), )),
(Column(None, None, "column1"), ),
"_snuba_count",
"countMerge",
(Column(None, None, "count"), ),
),
id="Test distribution count",
),
pytest.param(
"metrics_distributions",
"percentiles",
EntityKey.METRICS_DISTRIBUTIONS,
CurriedFunctionCall(
"_snuba_percentiles",
FunctionCall(
None,
"quantilesMerge",
tuple(
Literal(None, quant)
for quant in [0.5, 0.75, 0.9, 0.95, 0.99]),
),
(Column(None, None, "percentiles"), ),
),
id="Test distribution quantiles",
),
]
@pytest.mark.parametrize(
"entity_name, column_name, entity_key, translated_value", TEST_CASES)
def test_metrics_processing(
entity_name: str,
column_name: str,
"g",
(
Column("_snuba_col", None, "column"),
Literal(None, "val"),
),
)
},
"groups": {Column("_snuba_col2", None, "column2")},
},
),
id="Nested expression, part pushed down",
),
pytest.param(
CurriedFunctionCall(
"_snuba_cf",
FunctionCall("_snuba_f", "f",
(Column("_snuba_col", "events", "column"), )),
(Literal(None, "val"), ),
),
SubqueryExpression(
CurriedFunctionCall(
"_snuba_cf",
FunctionCall("_snuba_f", "f",
(Column("_snuba_col", None, "column"), )),
(Literal(None, "val"), ),
),
subquery_alias="events",
),
MainQueryExpression(
Column("_snuba_cf", "events", "_snuba_cf"),
{
"events": {
Literal("alias", 123),
Argument("alias", "arg"),
SubscriptableReference("tags[asd]", Column(None, None, "tags"),
Literal(None, "release")),
FunctionCall(
"alias",
"f",
(
Column(None, "table", "col"),
Literal(None, 123),
FunctionCall(None, "f1", (Column(None, None, "col2"), )),
),
),
CurriedFunctionCall(
None,
FunctionCall(None, "f",
(Column(None, None, "col"), Literal(None, 12))),
(Column(None, None, "col3"), ),
),
Lambda(None, ("a", "b"), FunctionCall(None, "f", (Argument(None, "a"), ))),
]
@pytest.mark.parametrize("expression", test_data)
def test_default_translation(expression: Expression) -> None:
"""
Ensures that a translation that relies on the default translation rules
produces a deep copy of the original expression.
"""
translated = expression.accept(
SnubaClickhouseMappingTranslator(TranslationMappers()))
"HAVING eq(column1, 123) "
"ORDER BY column1 ASC, table1.column2 DESC"),
),
(
# Query with complex functions
Query(
{},
TableSource("my_table", ColumnSet([])),
selected_columns=[
CurriedFunctionCall(
"my_complex_math",
FunctionCall(
None,
"doSomething",
[
Column(None, "column1", None),
Column(None, "column2", "table1"),
Column("al", "column3", None),
],
),
[Column(None, "column1", None)],
)
],
condition=binary_condition(
None,
"and",
lhs=binary_condition(
None,
"eq",
lhs=Column("al", "column3", None),
rhs=Literal(None, "blabla"),