def test_timeseries_column_format_expressions(granularity: int,
ast_value: FunctionCall,
formatted_value: str) -> None:
unprocessed = Query(
{"granularity": granularity},
TableSource("transactions", ColumnSet([])),
selected_columns=[
SelectedExpression(
"transaction.duration",
Column("transaction.duration", None, "duration")),
SelectedExpression("my_time", Column("my_time", None, "time")),
],
)
expected = Query(
{"granularity": granularity},
TableSource("transactions", ColumnSet([])),
selected_columns=[
SelectedExpression(
"transaction.duration",
Column("transaction.duration", None, "duration")),
SelectedExpression(ast_value.alias, ast_value),
],
)
dataset = TransactionsDataset()
for processor in dataset.get_query_processors():
if isinstance(processor, TimeSeriesColumnProcessor):
processor.process_query(unprocessed, HTTPRequestSettings())
assert (expected.get_selected_columns_from_ast() ==
unprocessed.get_selected_columns_from_ast())
ret = unprocessed.get_selected_columns_from_ast()[1].expression.accept(
ClickhouseExpressionFormatter())
assert ret == formatted_value
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 query_runner(query: Query, settings: RequestSettings,
reader: Reader[SqlQuery]) -> QueryResult:
assert query.get_selected_columns_from_ast() == [
SelectedExpression(
"tags[transaction]",
Column("tags[transaction]", None, "transaction")),
SelectedExpression(
"contexts[browser.name]",
FunctionCall(
"contexts[browser.name]",
"arrayElement",
(
Column(None, None, "contexts.value"),
FunctionCall(
None,
"indexOf",
(
Column(None, None, "contexts.key"),
Literal(None, "browser.name"),
),
),
),
),
),
]
return QueryResult({}, {})
def test_format_clickhouse_specific_query() -> None:
"""
Adds a few of the Clickhosue specific fields to the query.
"""
query = Query(
{
"sample": 0.1,
"totals": True,
"limitby": (10, "environment")
},
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"),
)
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_failure_rate_format_expressions() -> None:
unprocessed = Query(
{},
TableSource("events", ColumnSet([])),
selected_columns=[
SelectedExpression(name=None,
expression=Column(None, None, "column2")),
SelectedExpression("perf", FunctionCall("perf", "failure_rate",
())),
],
)
expected = Query(
{},
TableSource("events", ColumnSet([])),
selected_columns=[
SelectedExpression(name=None,
expression=Column(None, None, "column2")),
SelectedExpression(
"perf",
divide(
FunctionCall(
None,
"countIf",
(binary_condition(
None,
ConditionFunctions.NOT_IN,
Column(None, None, "transaction_status"),
FunctionCall(
None,
"tuple",
(
Literal(alias=None, value=0),
Literal(alias=None, value=1),
Literal(alias=None, value=2),
),
),
), ),
),
count(),
"perf",
),
),
],
)
failure_rate_processor(ColumnSet([])).process_query(
unprocessed, HTTPRequestSettings())
assert (expected.get_selected_columns_from_ast() ==
unprocessed.get_selected_columns_from_ast())
ret = unprocessed.get_selected_columns_from_ast()[1].expression.accept(
ClickhouseExpressionFormatter())
assert ret == (
"(divide(countIf(notIn(transaction_status, tuple(0, 1, 2))), count()) AS perf)"
)
def test_timeseries_format_expressions(
granularity: int,
condition: Optional[FunctionCall],
exp_column: FunctionCall,
exp_condition: Optional[FunctionCall],
formatted_column: str,
formatted_condition: str,
) -> None:
unprocessed = Query(
{"granularity": granularity},
TableSource("transactions", ColumnSet([])),
selected_columns=[
SelectedExpression(
"transaction.duration",
Column("transaction.duration", None, "duration")),
SelectedExpression("my_time", Column("my_time", None, "time")),
],
condition=condition,
)
expected = Query(
{"granularity": granularity},
TableSource("transactions", ColumnSet([])),
selected_columns=[
SelectedExpression(
"transaction.duration",
Column("transaction.duration", None, "duration")),
SelectedExpression(exp_column.alias, exp_column),
],
condition=exp_condition,
)
entity = TransactionsEntity()
processors = entity.get_query_processors()
for processor in processors:
if isinstance(processor, TimeSeriesProcessor):
processor.process_query(unprocessed, HTTPRequestSettings())
assert (expected.get_selected_columns_from_ast() ==
unprocessed.get_selected_columns_from_ast())
assert expected.get_condition_from_ast(
) == unprocessed.get_condition_from_ast()
ret = unprocessed.get_selected_columns_from_ast()[1].expression.accept(
ClickhouseExpressionFormatter())
assert ret == formatted_column
if condition:
ret = unprocessed.get_condition_from_ast().accept(
ClickhouseExpressionFormatter())
assert formatted_condition == ret
def test_transaction_column_format_expressions() -> None:
unprocessed = Query(
{},
TableSource("events", ColumnSet([])),
selected_columns=[
SelectedExpression(
"transaction.duration",
Column("transaction.duration", None, "duration")),
SelectedExpression("the_event_id",
Column("the_event_id", None, "event_id")),
],
)
expected = Query(
{},
TableSource("events", ColumnSet([])),
selected_columns=[
SelectedExpression(
"transaction.duration",
Column("transaction.duration", None, "duration")),
SelectedExpression(
"the_event_id",
FunctionCall(
"the_event_id",
"replaceAll",
(
FunctionCall(
None,
"toString",
(Column(None, None, "event_id"), ),
),
Literal(None, "-"),
Literal(None, ""),
),
),
),
],
)
TransactionColumnProcessor().process_query(unprocessed,
HTTPRequestSettings())
assert (expected.get_selected_columns_from_ast() ==
unprocessed.get_selected_columns_from_ast())
formatted = unprocessed.get_selected_columns_from_ast(
)[1].expression.accept(ClickhouseExpressionFormatter())
assert formatted == "(replaceAll(toString(event_id), '-', '') AS the_event_id)"
def test_col_split(
dataset_name: str,
id_column: str,
project_column: str,
timestamp_column: str,
first_query_data: Sequence[MutableMapping[str, Any]],
second_query_data: Sequence[MutableMapping[str, Any]],
) -> None:
def do_query(
query: ClickhouseQuery,
request_settings: RequestSettings,
reader: Reader[SqlQuery],
) -> QueryResult:
selected_cols = query.get_selected_columns()
assert selected_cols == [
c.expression.column_name
for c in query.get_selected_columns_from_ast() or []
if isinstance(c.expression, Column)
]
if selected_cols == list(first_query_data[0].keys()):
return QueryResult({"data": first_query_data}, {})
elif selected_cols == list(second_query_data[0].keys()):
return QueryResult({"data": second_query_data}, {})
else:
raise ValueError(f"Unexpected selected columns: {selected_cols}")
events = get_dataset(dataset_name)
query = ClickhouseQuery(
LogicalQuery(
{
"selected_columns": list(second_query_data[0].keys()),
"conditions": [""],
"orderby": "events.event_id",
"sample": 10,
"limit": 100,
"offset": 50,
},
events.get_all_storages()[0].get_schema().get_data_source(),
selected_columns=[
SelectedExpression(name=col_name,
expression=Column(None, None, col_name))
for col_name in second_query_data[0].keys()
],
))
strategy = SimpleQueryPlanExecutionStrategy(
ClickhouseCluster("localhost", 1024, "default", "", "default", 80,
set(), True),
[],
[
ColumnSplitQueryStrategy(id_column, project_column,
timestamp_column),
TimeSplitQueryStrategy(timestamp_col=timestamp_column),
],
)
strategy.execute(query, HTTPRequestSettings(), do_query)
def build_query(
selected_columns: Optional[Sequence[Expression]] = None,
condition: Optional[Expression] = None,
having: Optional[Expression] = None,
) -> ClickhouseQuery:
return ClickhouseQuery(
SnubaQuery(
{},
None,
selected_columns=[
SelectedExpression(name=s.alias, expression=s)
for s in selected_columns or []
],
condition=condition,
having=having,
))
def test_events_column_format_expressions() -> None:
unprocessed = Query(
{},
TableSource("events", ColumnSet([])),
selected_columns=[
SelectedExpression("dr_claw", Column("dr_claw", None, "culprit")),
SelectedExpression("the_group_id",
Column("the_group_id", None, "group_id")),
SelectedExpression("the_message",
Column("the_message", None, "message")),
],
)
expected = Query(
{},
TableSource("events", ColumnSet([])),
selected_columns=[
SelectedExpression("dr_claw", Column("dr_claw", None, "culprit")),
SelectedExpression(
"the_group_id",
FunctionCall(
"the_group_id",
"nullIf",
(
Column(None, None, "group_id"),
Literal(None, 0),
),
),
),
SelectedExpression(
"the_message",
FunctionCall(
"the_message",
"coalesce",
(
Column(None, None, "search_message"),
Column(None, None, "message"),
),
),
),
],
)
EventsColumnProcessor().process_query(unprocessed, HTTPRequestSettings())
assert (expected.get_selected_columns_from_ast() ==
unprocessed.get_selected_columns_from_ast())
expected = (
"(nullIf(group_id, 0) AS the_group_id)",
"(coalesce(search_message, message) AS the_message)",
)
for idx, column in enumerate(
unprocessed.get_selected_columns_from_ast()[1:]):
formatted = column.expression.accept(ClickhouseExpressionFormatter())
assert expected[idx] == formatted
def build_selected_expressions(
raw_expressions: Sequence[Any], ) -> List[SelectedExpression]:
output = []
for raw_expression in raw_expressions:
exp = parse_expression(tuplify(raw_expression),
entity.get_data_model(), set())
output.append(
SelectedExpression(
# An expression in the query can be a string or a
# complex list with an alias. In the second case
# we trust the parser to find the alias.
name=raw_expression
if isinstance(raw_expression, str) else exp.alias,
expression=exp,
))
return output
def test_iterate_over_query():
"""
Creates a query with the new AST and iterate over all expressions.
"""
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(None, ConditionFunctions.EQ, column1,
Literal(None, "1"))
orderby = OrderBy(OrderByDirection.ASC, function_2)
query = Query(
{},
TableSource("my_table", ColumnSet([])),
selected_columns=[SelectedExpression("alias", function_1)],
array_join=None,
condition=condition,
groupby=[function_1],
having=None,
order_by=[orderby],
)
expected_expressions = [
# selected columns
column1,
column2,
function_1,
# condition
column1,
Literal(None, "1"),
condition,
# groupby
column1,
column2,
function_1,
# order by
column2,
function_2,
]
assert list(query.get_all_expressions()) == expected_expressions
def test_handled_processor_invalid() -> None:
columnset = ColumnSet([])
unprocessed = Query(
{},
TableSource("events", columnset),
selected_columns=[
SelectedExpression(
"result",
FunctionCall(
"result",
"isHandled",
(Column(None, None, "type"), ),
),
),
],
)
processor = handled_functions.HandledFunctionsProcessor(
"exception_stacks.mechanism_handled", columnset)
with pytest.raises(InvalidExpressionException):
processor.process_query(unprocessed, HTTPRequestSettings())
def test_invalid_datetime() -> None:
unprocessed = Query(
{},
TableSource("transactions", ColumnSet([])),
selected_columns=[
SelectedExpression(
"transaction.duration",
Column("transaction.duration", None, "duration")),
],
condition=binary_condition(
None,
ConditionFunctions.EQ,
Column("my_time", None, "time"),
Literal(None, ""),
),
)
entity = TransactionsEntity()
processors = entity.get_query_processors()
for processor in processors:
if isinstance(processor, TimeSeriesProcessor):
with pytest.raises(InvalidQueryException):
processor.process_query(unprocessed, HTTPRequestSettings())
def _parse_query_impl(body: MutableMapping[str, Any], entity: Entity) -> Query:
def build_selected_expressions(
raw_expressions: Sequence[Any], ) -> List[SelectedExpression]:
output = []
for raw_expression in raw_expressions:
exp = parse_expression(tuplify(raw_expression),
entity.get_data_model(), set())
output.append(
SelectedExpression(
# An expression in the query can be a string or a
# complex list with an alias. In the second case
# we trust the parser to find the alias.
name=raw_expression
if isinstance(raw_expression, str) else exp.alias,
expression=exp,
))
return output
aggregations = []
for aggregation in body.get("aggregations", []):
if not isinstance(aggregation, Sequence):
raise ParsingException((
f"Invalid aggregation structure {aggregation}. "
"It must be a sequence containing expression, column and alias."
))
aggregation_function = aggregation[0]
column_expr = aggregation[1]
column_expr = column_expr if column_expr else []
alias = aggregation[2]
alias = alias if alias else None
aggregations.append(
SelectedExpression(
name=alias,
expression=parse_aggregation(
aggregation_function,
column_expr,
alias,
entity.get_data_model(),
set(),
),
))
groupby_clause = build_selected_expressions(
to_list(body.get("groupby", [])))
select_clause = (
groupby_clause + aggregations +
build_selected_expressions(body.get("selected_columns", [])))
array_join_cols = set()
arrayjoin = body.get("arrayjoin")
# TODO: Properly detect all array join columns in all clauses of the query.
# This is missing an arrayJoin in condition with an alias that is then
# used in the select.
if arrayjoin:
array_join_cols.add(arrayjoin)
array_join_expr: Optional[Expression] = parse_expression(
body["arrayjoin"], entity.get_data_model(), {arrayjoin})
else:
array_join_expr = None
for select_expr in select_clause:
if isinstance(select_expr.expression, FunctionCall):
if select_expr.expression.function_name == "arrayJoin":
parameters = select_expr.expression.parameters
if len(parameters) != 1:
raise ParsingException(
"arrayJoin(...) only accepts a single parameter.")
if isinstance(parameters[0], Column):
array_join_cols.add(parameters[0].column_name)
else:
# We only accepts columns or functions that do not
# reference columns. We could not say whether we are
# actually arrayjoining on the values of the column
# if it is nested in an arbitrary function. But
# functions of literals are fine.
for e in parameters[0]:
if isinstance(e, Column):
raise ParsingException(
"arrayJoin(...) cannot contain columns nested in functions."
)
where_expr = parse_conditions_to_expr(body.get("conditions", []), entity,
array_join_cols)
having_expr = parse_conditions_to_expr(body.get("having", []), entity,
array_join_cols)
orderby_exprs = []
for orderby in to_list(body.get("orderby", [])):
if isinstance(orderby, str):
match = NEGATE_RE.match(orderby)
if match is None:
raise ParsingException((
f"Invalid Order By clause {orderby}. If the Order By is a string, "
"it must respect the format `[-]column`"))
direction, col = match.groups()
orderby = col
elif is_function(orderby):
match = NEGATE_RE.match(orderby[0])
if match is None:
raise ParsingException((
f"Invalid Order By clause {orderby}. If the Order By is an expression, "
"the function name must respect the format `[-]func_name`"
))
direction, col = match.groups()
orderby = [col] + orderby[1:]
else:
raise ParsingException(
(f"Invalid Order By clause {orderby}. The Clause was neither "
"a string nor a function call."))
orderby_parsed = parse_expression(tuplify(orderby),
entity.get_data_model(), set())
orderby_exprs.append(
OrderBy(
OrderByDirection.DESC
if direction == "-" else OrderByDirection.ASC,
orderby_parsed,
))
return Query(
body,
None,
selected_columns=select_clause,
array_join=array_join_expr,
condition=where_expr,
groupby=[g.expression for g in groupby_clause],
having=having_expr,
order_by=orderby_exprs,
)
BooleanFunctions,
ConditionFunctions,
binary_condition,
)
from snuba.query.expressions import Column, CurriedFunctionCall, FunctionCall, Literal
from snuba.query.logical import OrderBy, OrderByDirection, Query, SelectedExpression
from snuba.request.request_settings import HTTPRequestSettings
test_cases = [
pytest.param(
# Simple query with aliases and multiple tables
Query(
{},
TableSource("my_table", ColumnSet([])),
selected_columns=[
SelectedExpression("column1", Column(None, None, "column1")),
SelectedExpression("column2", Column(None, "table1",
"column2")),
SelectedExpression("column3", Column("al", None, "column3")),
],
condition=binary_condition(
None,
"eq",
lhs=Column("al", None, "column3"),
rhs=Literal(None, "blabla"),
),
groupby=[
Column(None, None, "column1"),
Column(None, "table1", "column2"),
Column("al", None, "column3"),
Column(None, None, "column4"),
from snuba.query.matchers import (
Column as ColumnMatch,
String as StringMatch,
MatchResult,
Param,
)
from snuba.query.processors.pattern_replacer import PatternReplacer
from snuba.request.request_settings import HTTPRequestSettings
test_data = [
pytest.param(
Query(
{},
TableSource("events", ColumnSet([])),
selected_columns=[
SelectedExpression(name=None,
expression=Column(None, None, "column1")),
SelectedExpression(name=None,
expression=Column(None, None, "column2")),
],
),
Query(
{},
TableSource("events", ColumnSet([])),
selected_columns=[
SelectedExpression(
name=None,
expression=FunctionCall(
None,
"nullIf",
(Column(None, None, "column1"), Literal(None, "")),
),
from snuba.query.processors.custom_function import (
CustomFunction,
InvalidCustomFunctionCall,
partial_function,
simple_function,
)
from snuba.query.validation.signature import Column as ColType
from snuba.request.request_settings import HTTPRequestSettings
TEST_CASES = [
pytest.param(
Query(
{},
None,
selected_columns=[
SelectedExpression("column1", Column("column1", None,
"column1")),
],
groupby=[Column("column1", None, "column1")],
condition=binary_condition(
None,
"equals",
FunctionCall("group_id", "f",
(Column("something", None, "something"), )),
Literal(None, 1),
),
),
Query(
{},
None,
selected_columns=[
SelectedExpression("column1", Column("column1", None,
Column,
FunctionCall,
Literal,
SubscriptableReference,
)
from snuba.query.logical import Query as SnubaQuery
from snuba.query.logical import SelectedExpression
test_cases = [
pytest.param(
TranslationMappers(),
SnubaQuery(
body={},
data_source=TableSource("my_table", 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")
),
),
],
def test_tags_expander() -> None:
query_body = {
"selected_columns": [
["f1", ["tags_key", "column2"], "f1_alias"],
["f2", [], "f2_alias"],
],
"aggregations": [
["count", "platform", "platforms"],
["testF", ["platform", "tags_value"], "top_platforms"],
],
"conditions": [["tags_key", "=", "tags_key"]],
"having": [["tags_value", "IN", ["tag"]]],
}
events = get_dataset("events")
query = parse_query(query_body, events)
processor = TagsExpanderProcessor()
request_settings = HTTPRequestSettings()
processor.process_query(query, request_settings)
assert query.get_selected_columns_from_ast() == [
SelectedExpression(
"platforms",
FunctionCall("platforms", "count",
(Column("platform", None, "platform"), )),
),
SelectedExpression(
"top_platforms",
FunctionCall(
"top_platforms",
"testF",
(
Column("platform", None, "platform"),
FunctionCall("tags_value", "arrayJoin",
(Column(None, None, "tags.value"), )),
),
),
),
SelectedExpression(
"f1_alias",
FunctionCall(
"f1_alias",
"f1",
(
FunctionCall("tags_key", "arrayJoin",
(Column(None, None, "tags.key"), )),
Column("column2", None, "column2"),
),
),
),
SelectedExpression("f2_alias", FunctionCall("f2_alias", "f2",
tuple())),
]
assert query.get_condition_from_ast() == binary_condition(
None,
OPERATOR_TO_FUNCTION["="],
FunctionCall("tags_key", "arrayJoin",
(Column(None, None, "tags.key"), )),
Literal(None, "tags_key"),
)
assert query.get_having_from_ast() == in_condition(
None,
FunctionCall("tags_value", "arrayJoin",
(Column(None, None, "tags.value"), )),
[Literal(None, "tag")],
)
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)
from snuba.query.logical import OrderBy, OrderByDirection, Query, SelectedExpression
from snuba.query.snql.parser import parse_snql_query
test_cases = [
pytest.param(
"MATCH(blah)WHEREa<3COLLECT4-5,3*g(c),c",
Query(
{},
None,
selected_columns=[
SelectedExpression(
"4-5",
FunctionCall(
None,
"minus",
(
Literal(None, 4),
Literal(None, 5),
),
),
),
SelectedExpression(
"3*g(c)",
FunctionCall(
None,
"multiply",
(
Literal(None, 3),
FunctionCall(
None,
"g",
def test_events_boolean_context() -> None:
columns = ColumnSet([
("device_charging", Nullable(UInt(8))),
("contexts", Nested([("key", String()), ("value", String())])),
])
query = ClickhouseQuery(
LogicalQuery(
{},
TableSource("events", columns),
selected_columns=[
SelectedExpression(
"contexts[device.charging]",
FunctionCall(
"contexts[device.charging]",
"arrayElement",
(
Column(None, None, "contexts.value"),
FunctionCall(
None,
"indexOf",
(
Column(None, None, "contexts.key"),
Literal(None, "device.charging"),
),
),
),
),
)
],
))
expected = ClickhouseQuery(
LogicalQuery(
{},
TableSource("events", columns),
selected_columns=[
SelectedExpression(
"contexts[device.charging]",
FunctionCall(
"contexts[device.charging]",
"multiIf",
(
binary_condition(
None,
ConditionFunctions.EQ,
FunctionCall(
None,
"toString",
(Column(None, None, "device_charging"), ),
),
Literal(None, ""),
),
Literal(None, ""),
binary_condition(
None,
ConditionFunctions.IN,
FunctionCall(
None,
"toString",
(Column(None, None, "device_charging"), ),
),
literals_tuple(None, [
Literal(None, "1"),
Literal(None, "True")
]),
),
Literal(None, "True"),
Literal(None, "False"),
),
),
)
],
))
settings = HTTPRequestSettings()
MappingColumnPromoter({
"contexts": {
"device.charging": "device_charging"
}
}).process_query(query, settings)
EventsBooleanContextsProcessor().process_query(query, settings)
assert (query.get_selected_columns_from_ast() ==
expected.get_selected_columns_from_ast())
from snuba.query.logical import OrderBy, OrderByDirection, Query, SelectedExpression
from snuba.query.parser import parse_query
from snuba.query.parser.exceptions import AliasShadowingException, CyclicAliasException
test_cases = [
pytest.param(
{
"selected_columns": ["column1"],
"groupby": ["column2", "column3"],
"aggregations": [["test_func", "column4", "test_func_alias"]],
},
Query(
{},
TableSource("events", ColumnSet([])),
selected_columns=[
SelectedExpression("column2", Column("column2", None,
"column2")),
SelectedExpression("column3", Column("column3", None,
"column3")),
SelectedExpression(
"test_func_alias",
FunctionCall(
"test_func_alias",
"test_func",
(Column("column4", None, "column4"), ),
),
),
SelectedExpression("column1", Column("column1", None,
"column1")),
],
groupby=[
Column("column2", None, "column2"),
def test_apdex_format_expressions() -> None:
unprocessed = Query(
{},
TableSource("events", ColumnSet([])),
selected_columns=[
SelectedExpression(name=None, expression=Column(None, None, "column2")),
SelectedExpression(
"perf",
FunctionCall(
"perf", "apdex", (Column(None, None, "column1"), Literal(None, 300))
),
),
],
)
expected = Query(
{},
TableSource("events", ColumnSet([])),
selected_columns=[
SelectedExpression(name=None, expression=Column(None, None, "column2")),
SelectedExpression(
"perf",
divide(
plus(
FunctionCall(
None,
"countIf",
(
binary_condition(
None,
ConditionFunctions.LTE,
Column(None, None, "column1"),
Literal(None, 300),
),
),
),
divide(
FunctionCall(
None,
"countIf",
(
binary_condition(
None,
BooleanFunctions.AND,
binary_condition(
None,
ConditionFunctions.GT,
Column(None, None, "column1"),
Literal(None, 300),
),
binary_condition(
None,
ConditionFunctions.LTE,
Column(None, None, "column1"),
multiply(
Literal(None, 300), Literal(None, 4)
),
),
),
),
),
Literal(None, 2),
),
),
FunctionCall(None, "count", (),),
"perf",
),
),
],
)
apdex_processor(ColumnSet([])).process_query(unprocessed, HTTPRequestSettings())
assert (
expected.get_selected_columns_from_ast()
== unprocessed.get_selected_columns_from_ast()
)
ret = unprocessed.get_selected_columns_from_ast()[1].expression.accept(
ClickhouseExpressionFormatter()
)
assert ret == (
"(divide(plus(countIf(lessOrEquals(column1, 300)), "
"divide(countIf(greater(column1, 300) AND "
"lessOrEquals(column1, multiply(300, 4))), 2)), count()) AS perf)"
)
def test_replace_expression():
"""
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(None, ConditionFunctions.EQ, function_1,
Literal(None, "1"))
orderby = OrderBy(OrderByDirection.ASC, function_2)
query = Query(
{},
TableSource("my_table", ColumnSet([])),
selected_columns=[SelectedExpression("alias", function_1)],
array_join=None,
condition=condition,
groupby=[function_1],
having=None,
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(
{},
TableSource("my_table", ColumnSet([])),
selected_columns=[
SelectedExpression(
"alias", FunctionCall("alias", "tag", (Literal(None, "f1"), )))
],
array_join=None,
condition=binary_condition(
None,
ConditionFunctions.EQ,
FunctionCall("alias", "tag", (Literal(None, "f1"), )),
Literal(None, "1"),
),
groupby=[FunctionCall("alias", "tag", (Literal(None, "f1"), ))],
having=None,
order_by=[orderby],
)
assert (query.get_selected_columns_from_ast() ==
expected_query.get_selected_columns_from_ast())
assert query.get_condition_from_ast(
) == expected_query.get_condition_from_ast()
assert query.get_groupby_from_ast() == expected_query.get_groupby_from_ast(
)
assert query.get_having_from_ast() == expected_query.get_having_from_ast()
assert query.get_orderby_from_ast() == expected_query.get_orderby_from_ast(
)
assert list(query.get_all_expressions()) == list(
expected_query.get_all_expressions())
ConditionFunctions,
binary_condition,
)
from snuba.query.expressions import Column, FunctionCall, Literal
from snuba.query.logical import Query, SelectedExpression
from snuba.querylog.query_metadata import ClickhouseQueryProfile, FilterProfile
from snuba.state import safe_dumps
test_cases = [
pytest.param(
ClickhouseQuery(
Query(
{},
TableSource("events", ColumnSet([])),
selected_columns=[
SelectedExpression("column2",
Column("column2", None, "column2")),
SelectedExpression(
"something",
FunctionCall(
"something",
"arrayJoin",
(Column(None, None, "contexts.key"), ),
),
),
],
condition=binary_condition(
None,
BooleanFunctions.AND,
binary_condition(
None,
ConditionFunctions.GTE,
def visit_selected_expression(
self, node: Node, visited_children: Tuple[Expression, Any]
) -> SelectedExpression:
exp, _ = visited_children
return SelectedExpression(node.text.strip(), exp)
from snuba.clickhouse.columns import ColumnSet
from snuba.datasets.schemas.tables import TableSource
from snuba.query.expressions import Column, CurriedFunctionCall, FunctionCall, Literal
from snuba.query.logical import Query, SelectedExpression
from snuba.query.processors.basic_functions import BasicFunctionsProcessor
from snuba.request.request_settings import HTTPRequestSettings
test_data = [
(
Query(
{},
TableSource("events", ColumnSet([])),
selected_columns=[
SelectedExpression(
"alias",
FunctionCall("alias", "uniq",
(Column(None, None, "column1"), )),
),
SelectedExpression(
"alias2",
FunctionCall("alias2", "emptyIfNull",
(Column(None, None, "column2"), )),
),
],
),
Query(
{},
TableSource("events", ColumnSet([])),
selected_columns=[
SelectedExpression(
"alias",
def test_handled_processor() -> None:
columnset = ColumnSet([])
unprocessed = Query(
{},
TableSource("events", columnset),
selected_columns=[
SelectedExpression(name=None, expression=Column(None, None, "id")),
SelectedExpression(
"result",
FunctionCall(
"result",
"isHandled",
tuple(),
),
),
],
)
expected = Query(
{},
TableSource("events", columnset),
selected_columns=[
SelectedExpression(name=None, expression=Column(None, None, "id")),
SelectedExpression(
"result",
FunctionCall(
"result",
"arrayExists",
(
Lambda(
None,
("x", ),
binary_condition(
None,
BooleanFunctions.OR,
FunctionCall(None, "isNull",
(Argument(None, "x"), )),
binary_condition(
None,
ConditionFunctions.EQ,
FunctionCall(None, "assumeNotNull",
(Argument(None, "x"), )),
Literal(None, 1),
),
),
),
Column(None, None,
"exception_stacks.mechanism_handled"),
),
),
),
],
)
processor = handled_functions.HandledFunctionsProcessor(
"exception_stacks.mechanism_handled", columnset)
processor.process_query(unprocessed, HTTPRequestSettings())
assert (expected.get_selected_columns_from_ast() ==
unprocessed.get_selected_columns_from_ast())
ret = unprocessed.get_selected_columns_from_ast()[1].expression.accept(
ClickhouseExpressionFormatter())
assert ret == (
"(arrayExists((x -> (isNull(x) OR equals(assumeNotNull(x), 1))), exception_stacks.mechanism_handled) AS result)"
)