def test_subscriptable() -> None:
c1 = Column(None, "t1", "tags")
l1 = Literal(None, "myTag")
s = SubscriptableReference("alias", c1, l1)
assert list(s) == [c1, l1, s]
def replace_col(e: Expression) -> Expression:
if isinstance(e, Literal):
return Literal(None, "myOtherTag")
return e
replaced = s.transform(replace_col)
l2 = Literal(None, "myOtherTag")
assert list(replaced) == [c1, l2, SubscriptableReference("alias", c1, l2)]
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 test_transformer() -> None:
query = build_query(Literal(None, "10"))
TagsTypeTransformer().process_query(query, HTTPRequestSettings())
assert query.get_selected_columns(
)[0].expression == SubscriptableReference("_snuba_tags[10]",
Column(None, None, "tags"),
Literal(None, 10))
def build_query(tag_key: Literal) -> Query:
return Query(
from_clause=None,
selected_columns=[
SelectedExpression(
"tags[10]",
SubscriptableReference("_snuba_tags[10]",
Column(None, None, "tags"), tag_key),
)
],
)
def transform_expression(exp: Expression) -> Expression:
if not isinstance(exp, SubscriptableReference):
return exp
key = exp.key
if not isinstance(key.value, str) or not key.value.isdigit():
raise InvalidExpressionException.from_args(
exp,
"Expected a string key containing an integer in subscriptable.",
)
return SubscriptableReference(exp.alias, exp.column,
Literal(None, int(key.value)))
def test_nullable_nested_translation() -> None:
translated = SubscriptableMapper(
None, "measurements", None, "measurements", nullable=True
).attempt_map(
SubscriptableReference(
"measurements[lcp]",
Column(None, None, "measurements"),
Literal(None, "lcp"),
),
SnubaClickhouseMappingTranslator(TranslationMappers()),
)
assert translated == _get_nullable_expr("measurements[lcp]")
def visit_subscriptable_reference(
self, exp: SubscriptableReference) -> SubExpression:
assert (
exp.column.table_name
), f"Invalid column expression in join: {exp}. Missing table alias"
return SubqueryExpression(
main_expression=SubscriptableReference(
exp.alias,
Column(exp.column.alias, None, exp.column.column_name),
Literal(exp.key.alias, exp.key.value),
),
subquery_alias=exp.column.table_name,
)
def transform(exp: Expression) -> Expression:
if not isinstance(exp, Column) or exp.column_name in current_aliases:
return exp
match = NESTED_COL_EXPR_RE.match(exp.column_name)
if match is None:
# This is not a tag[asd] column.
return exp
col_name = match[1]
key_name = match[2]
return SubscriptableReference(
alias=exp.column_name,
column=Column(None, None, col_name),
key=Literal(None, key_name),
)
def attempt_map(
self,
expression: SubscriptableReference,
children_translator: SnubaClickhouseStrictTranslator,
) -> Optional[SubscriptableReference]:
# TODO: remove a default for SubscriptableReference entirely.
# Since there is not SubscriptableReference in clickhouse, such
# columns have to be translated by a valid rule. They cannot have
# a default translation that makes a copy. The assertion will be
# removed as well.
column = expression.column.accept(children_translator)
assert isinstance(column, Column)
key = expression.key.accept(children_translator)
assert isinstance(key, Literal)
return SubscriptableReference(alias=expression.alias, column=column, key=key)
def test_get_all_columns_legacy() -> None:
query_body = {
"selected_columns": [
["f1", ["column1", "column2"], "f1_alias"],
["f2", [], "f2_alias"],
["formatDateTime", ["timestamp", "'%Y-%m-%d'"], "formatted_time"],
],
"aggregations": [
["count", "platform", "platforms"],
["uniq", "platform", "uniq_platforms"],
["testF", ["platform", ["anotherF", ["field2"]]], "top_platforms"],
],
"conditions": [
["tags[sentry:dist]", "IN", ["dist1", "dist2"]],
["timestamp", ">=", "2020-01-01T12:00:00"],
["timestamp", "<", "2020-01-02T12:00:00"],
["project_id", "=", 1],
],
"having": [["times_seen", ">", 1]],
"groupby": [["format_eventid", ["event_id"]]],
}
events = get_dataset("events")
snql_query = json_to_snql(query_body, "events")
query, _ = parse_snql_query(str(snql_query), events)
assert query.get_all_ast_referenced_columns() == {
Column("_snuba_column1", None, "column1"),
Column("_snuba_column2", None, "column2"),
Column("_snuba_platform", None, "platform"),
Column("_snuba_field2", None, "field2"),
Column("_snuba_tags", None, "tags"),
Column("_snuba_times_seen", None, "times_seen"),
Column("_snuba_event_id", None, "event_id"),
Column("_snuba_timestamp", None, "timestamp"),
Column("_snuba_project_id", None, "project_id"),
}
assert query.get_all_ast_referenced_subscripts() == {
SubscriptableReference(
"_snuba_tags[sentry:dist]",
Column("_snuba_tags", None, "tags"),
Literal(None, "sentry:dist"),
)
}
def test_tag_translation() -> None:
translated = SubscriptableMapper(None, "tags", None, "tags").attempt_map(
SubscriptableReference(
"tags[release]", Column(None, None, "tags"), Literal(None, "release")
),
SnubaClickhouseMappingTranslator(TranslationMappers()),
)
assert translated == FunctionCall(
"tags[release]",
"arrayElement",
(
Column(None, None, "tags.value"),
FunctionCall(
None,
"indexOf",
(Column(None, None, "tags.key"), Literal(None, "release")),
),
),
)
def test_get_all_columns() -> None:
query_body = """
MATCH (events)
SELECT f1(column1, column2) AS f1_alias,
f2() AS f2_alias,
formatDateTime(timestamp, '%Y-%m-%d') AS formatted_time,
count() AS platforms,
uniq(platform) AS uniq_platforms,
testF(platform, anotherF(field2)) AS top_platforms
BY format_eventid(event_id)
WHERE tags[sentry:dist] IN tuple('dist1', 'dist2')
AND timestamp >= toDateTime('2020-01-01 12:00:00')
AND timestamp < toDateTime('2020-01-02 12:00:00')
AND project_id = 1
HAVING times_seen > 1
"""
events = get_dataset("events")
query, _ = parse_snql_query(query_body, events)
assert query.get_all_ast_referenced_columns() == {
Column("_snuba_column1", None, "column1"),
Column("_snuba_column2", None, "column2"),
Column("_snuba_platform", None, "platform"),
Column("_snuba_field2", None, "field2"),
Column("_snuba_tags", None, "tags"),
Column("_snuba_times_seen", None, "times_seen"),
Column("_snuba_event_id", None, "event_id"),
Column("_snuba_timestamp", None, "timestamp"),
Column("_snuba_project_id", None, "project_id"),
}
assert query.get_all_ast_referenced_subscripts() == {
SubscriptableReference(
"_snuba_tags[sentry:dist]",
Column("_snuba_tags", None, "tags"),
Literal(None, "sentry:dist"),
)
}
),
pytest.param(
"MATCH (events) SELECT count() AS count BY tags[key], measurements[lcp.elementSize] WHERE a<3 AND measurements[lcp.elementSize] > 1",
LogicalQuery(
QueryEntity(
EntityKey.EVENTS, get_entity(EntityKey.EVENTS).get_data_model()
),
selected_columns=[
SelectedExpression(
"count", FunctionCall("_snuba_count", "count", tuple()),
),
SelectedExpression(
"tags[key]",
SubscriptableReference(
"_snuba_tags[key]",
Column("_snuba_tags", None, "tags"),
Literal(None, "key"),
),
),
SelectedExpression(
"measurements[lcp.elementSize]",
SubscriptableReference(
"_snuba_measurements[lcp.elementSize]",
Column("_snuba_measurements", None, "measurements",),
Literal(None, "lcp.elementSize"),
),
),
],
groupby=[
SubscriptableReference(
"_snuba_tags[key]",
TEST_CASES = [
pytest.param(
CompositeQuery(
from_clause=LogicalQuery(
from_clause=events_ent,
selected_columns=[
SelectedExpression("project_id",
Column(None, None, "project_id")),
SelectedExpression(
"count_release",
FunctionCall(
"count_release",
"uniq",
(SubscriptableReference(
None,
Column(None, None, "tags"),
Literal(None, "sentry:release"),
), ),
),
),
],
groupby=[Column(None, None, "project_id")],
condition=binary_condition(
ConditionFunctions.EQ,
Column(None, None, "project_id"),
Literal(None, 1),
),
),
selected_columns=[
SelectedExpression(
"average",
TEST_CASES = [
pytest.param(
CompositeQuery(
from_clause=LogicalQuery(
from_clause=events_ent,
selected_columns=[
SelectedExpression("project_id",
Column(None, None, "project_id")),
SelectedExpression(
"count_environment",
FunctionCall(
"count_environment",
"uniq",
(SubscriptableReference(
None,
Column(None, None, "tags"),
Literal(None, "environment"),
), ),
),
),
],
groupby=[Column(None, None, "project_id")],
condition=binary_condition(
BooleanFunctions.AND,
binary_condition(
ConditionFunctions.EQ,
Column(None, None, "project_id"),
Literal(None, 1),
),
binary_condition(
ConditionFunctions.GTE,
from snuba.query.expressions import (
Argument,
Column,
CurriedFunctionCall,
Expression,
FunctionCall,
Lambda,
Literal,
SubscriptableReference,
)
test_data = [
Column("alias", "table", "col"),
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"), ),
),
SnubaQuery(
from_clause=QueryEntity(EntityKey.EVENTS, ColumnSet([])),
selected_columns=[
SelectedExpression("alias", Column("alias", "table", "column")),
SelectedExpression(
"alias2",
FunctionCall(
"alias2",
"f1",
(Column(None, None, "column2"), Column(None, None, "column3")),
),
),
SelectedExpression(
name=None,
expression=SubscriptableReference(
None, Column(None, None, "tags"), Literal(None, "myTag")
),
),
],
),
ClickhouseQuery(
from_clause=Table("my_table", ColumnSet([])),
selected_columns=[
SelectedExpression("alias", Column("alias", "table", "column")),
SelectedExpression(
"alias2",
FunctionCall(
"alias2",
"f1",
(Column(None, None, "column2"), Column(None, None, "column3"),),
),
"f1",
(
Column("_snuba_column1", None, "column1"),
Column("_snuba_column2", None, "column2"),
),
),
),
SelectedExpression("f2_alias",
FunctionCall("_snuba_f2_alias", "f2", ())),
],
condition=with_required(
binary_condition(
"in",
SubscriptableReference(
"_snuba_tags[sentry:dist]",
Column("_snuba_tags", None, "tags"),
Literal(None, "sentry:dist"),
),
FunctionCall(
None,
"tuple",
(
Literal(None, "dist1"),
Literal(None, "dist2"),
),
),
)),
having=binary_condition(
"greater",
Column("_snuba_times_seen", None, "times_seen"),
Literal(None, 1),
(Column(None, "t1", "c1"), "t1.c1"),
(Column(None, None, "c1"), "c1"),
(Literal(None, "meowmeow"), "'meowmeow'"),
(Literal(None, 123), "123"),
(Literal(None, False), "False"),
(
Literal(None, datetime(2020, 4, 20, 16, 20)),
"datetime(2020-04-20T16:20:00)",
),
(Literal(None,
datetime(2020, 4, 20, 16, 20).date()), "date(2020-04-20)"),
(Literal(None, None), "None"),
(
SubscriptableReference(
"catsound",
column=Column(None, "cats", "sounds"),
key=Literal(None, "meow"),
),
"cats.sounds['meow'] AS `catsound`",
),
(
SubscriptableReference(
"catsound",
column=Column("kittysounds", "cats", "sounds"),
key=Literal(None, "meow"),
),
"(cats.sounds AS `kittysounds`)['meow'] AS `catsound`",
),
(Column("alias", None, "c1"), "c1 AS `alias`"),
(
FunctionCall(None, "f1",
MainQueryExpression(Literal(None, "value"), {}),
id="Basic literal",
),
pytest.param(
FunctionCall(None, "f", (Literal(None, "asd"), Literal(None, "sdv"))),
UnclassifiedExpression(
FunctionCall(None, "f",
(Literal(None, "asd"), Literal(None, "sdv")))),
MainQueryExpression(
FunctionCall(None, "f",
(Literal(None, "asd"), Literal(None, "sdv"))), {}),
id="Unclassified function",
),
pytest.param(
SubscriptableReference("_snuba_tag",
Column("_snuba_col", "events", "tags"),
Literal(None, "val")),
SubqueryExpression(
SubscriptableReference("_snuba_tag",
Column("_snuba_col", None, "tags"),
Literal(None, "val")),
"events",
),
MainQueryExpression(
Column("_snuba_tag", "events", "_snuba_tag"),
{
"events": {
SubscriptableReference(
"_snuba_tag",
Column("_snuba_col", None, "tags"),
Literal(None, "val"),
binary_condition(
"equals",
Column("_snuba_project_id", None, "project_id"),
Literal(None, 1),
),
binary_condition(
"equals",
Column("_snuba_org_id", None, "org_id"),
Literal(None, 1),
),
),
binary_condition(
"equals",
SubscriptableReference(
"_snuba_tags[asd]",
Column("_snuba_tags", None, "tags"),
Literal(None, "asd"),
),
Literal(None, 2),
),
),
),
groupby=[
Column("_snuba_project_id", None, "project_id"),
SubscriptableReference(
"_snuba_tags[asd]",
Column("_snuba_tags", None, "tags"),
Literal(None, "asd"),
),
],
),
