def visit_join_clause(
self, node: JoinClause[Entity]) -> Mapping[str, SubqueryDraft]:
combined = {
**node.left_node.accept(self),
**node.right_node.accept(self)
}
for condition in node.keys:
combined[condition.left.table_alias].add_select_expression(
SelectedExpression(
# Setting a name for the selected columns in a subquery
# is not terribly useful, as this name would not
# be used anywhere.
# The external query references subquery columns by
# their aliases as this is what Clickhouse will do.
aliasify_column(condition.left.column),
Column(
aliasify_column(condition.left.column),
None,
condition.left.column,
),
))
combined[condition.right.table_alias].add_select_expression(
SelectedExpression(
aliasify_column(condition.right.column),
Column(
aliasify_column(condition.right.column),
None,
condition.right.column,
),
))
return combined
def test_query_data_source() -> None:
"""
Tests using the Query as a data source
"""
query = Query(
Table("my_table", ColumnSet([])),
selected_columns=[
SelectedExpression(
"col1",
Column(alias="col1", table_name=None, column_name="col1")),
SelectedExpression(
"some_func",
FunctionCall(
"some_func",
"f",
(Column(alias="col1", table_name=None,
column_name="col1"), ),
),
),
SelectedExpression(
None, Column(alias="col2", table_name=None,
column_name="col2")),
],
)
assert query.get_columns() == ColumnSet([("col1", Any()),
("some_func", Any()),
("_invalid_alias_2", Any())])
def process_functions(exp: Expression) -> Expression:
if isinstance(exp, FunctionCall):
if exp.function_name == "isHandled":
self.validate_parameters(exp)
return FunctionCall(
exp.alias,
"arrayExists",
(
Lambda(
None,
("x", ),
binary_condition(
BooleanFunctions.OR,
FunctionCall(None, "isNull",
(Argument(None, "x"), )),
binary_condition(
ConditionFunctions.EQ,
FunctionCall(
None,
"assumeNotNull",
(Argument(None, "x"), ),
),
Literal(None, 1),
),
),
),
Column(None, None, self.__column),
),
)
if exp.function_name == "notHandled":
self.validate_parameters(exp)
return FunctionCall(
exp.alias,
"arrayExists",
(
Lambda(
None,
("x", ),
binary_condition(
BooleanFunctions.AND,
FunctionCall(None, "isNotNull",
(Argument(None, "x"), )),
binary_condition(
ConditionFunctions.EQ,
FunctionCall(
None,
"assumeNotNull",
(Argument(None, "x"), ),
),
Literal(None, 0),
),
),
),
Column(None, None, self.__column),
),
)
return exp
def test_not_handled_processor() -> None:
columnset = ColumnSet([])
unprocessed = Query(
QueryEntity(EntityKey.EVENTS, ColumnSet([])),
selected_columns=[
SelectedExpression(name=None, expression=Column(None, None, "id")),
SelectedExpression(
"result", FunctionCall("result", "notHandled", tuple(),),
),
],
)
expected = Query(
QueryEntity(EntityKey.EVENTS, ColumnSet([])),
selected_columns=[
SelectedExpression(name=None, expression=Column(None, None, "id")),
SelectedExpression(
"result",
FunctionCall(
"result",
"arrayExists",
(
Lambda(
None,
("x",),
binary_condition(
BooleanFunctions.AND,
FunctionCall(None, "isNotNull", (Argument(None, "x"),)),
binary_condition(
ConditionFunctions.EQ,
FunctionCall(
None, "assumeNotNull", (Argument(None, "x"),)
),
Literal(None, 0),
),
),
),
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() == unprocessed.get_selected_columns()
ret = unprocessed.get_selected_columns()[1].expression.accept(
ClickhouseExpressionFormatter()
)
assert ret == (
"(arrayExists((x -> isNotNull(x) AND equals(assumeNotNull(x), 0)), exception_stacks.mechanism_handled) AS result)"
)
def test_uuid_array_column_processor(
unprocessed: Expression,
expected: Expression,
formatted_value: str,
) -> None:
unprocessed_query = Query(
Table("transactions", ColumnSet([])),
selected_columns=[
SelectedExpression("column2", Column(None, None, "column2"))
],
condition=unprocessed,
)
expected_query = Query(
Table("transactions", ColumnSet([])),
selected_columns=[
SelectedExpression("column2", Column(None, None, "column2"))
],
condition=expected,
)
UUIDArrayColumnProcessor(set(["column1", "column2"
])).process_query(unprocessed_query,
HTTPRequestSettings())
assert unprocessed_query.get_selected_columns() == [
SelectedExpression(
"column2",
FunctionCall(
None,
"arrayMap",
(
Lambda(
None,
("x", ),
FunctionCall(
None,
"replaceAll",
(
FunctionCall(None, "toString",
(Argument(None, "x"), )),
Literal(None, "-"),
Literal(None, ""),
),
),
),
Column(None, None, "column2"),
),
),
)
]
assert expected_query.get_condition() == unprocessed_query.get_condition()
condition = unprocessed_query.get_condition()
assert condition is not None
ret = condition.accept(ClickhouseExpressionFormatter())
assert ret == formatted_value
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_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",
(
combine_and_conditions(
[
binary_condition(
None,
ConditionFunctions.NEQ,
Column(None, None, "transaction_status"),
Literal(None, code),
)
for code in [0, 1, 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(notEquals(transaction_status, 0) AND notEquals(transaction_status, 1) AND notEquals(transaction_status, 2)), count()) AS perf)"
)
def add_conditions(
self,
timestamp: datetime,
offset: Optional[int],
query: Union[CompositeQuery[Entity], Query],
) -> None:
# TODO: Support composite queries with multiple entities.
from_clause = query.get_from_clause()
if not isinstance(from_clause, Entity):
raise InvalidSubscriptionError("Only simple queries are supported")
entity = get_entity(from_clause.key)
required_timestamp_column = entity.required_time_column
if required_timestamp_column is None:
raise InvalidSubscriptionError(
"Entity must have a timestamp column for subscriptions")
conditions_to_add: List[Expression] = [
binary_condition(
ConditionFunctions.EQ,
Column(None, None, "project_id"),
Literal(None, self.project_id),
),
binary_condition(
ConditionFunctions.GTE,
Column(None, None, required_timestamp_column),
Literal(None, (timestamp - self.time_window)),
),
binary_condition(
ConditionFunctions.LT,
Column(None, None, required_timestamp_column),
Literal(None, timestamp),
),
]
if offset is not None:
conditions_to_add.append(
binary_condition(
ConditionFunctions.LTE,
FunctionCall(
None,
"ifnull",
(Column(None, None, "offset"), Literal(None, 0)),
),
Literal(None, offset),
))
new_condition = combine_and_conditions(conditions_to_add)
condition = query.get_condition()
if condition:
new_condition = binary_condition(BooleanFunctions.AND, condition,
new_condition)
query.set_ast_condition(new_condition)
def test_first_level_conditions() -> None:
c1 = binary_condition(
ConditionFunctions.EQ,
Column(None, "table1", "column1"),
Literal(None, "test"),
)
c2 = binary_condition(
ConditionFunctions.EQ,
Column(None, "table2", "column2"),
Literal(None, "test"),
)
c3 = binary_condition(
ConditionFunctions.EQ,
Column(None, "table3", "column3"),
Literal(None, "test"),
)
cond = binary_condition(
BooleanFunctions.AND,
binary_condition(BooleanFunctions.AND, c1, c2),
c3,
)
assert get_first_level_and_conditions(cond) == [c1, c2, c3]
cond = binary_condition(
BooleanFunctions.AND,
FunctionCall(None, "equals",
(FunctionCall(None, "and", (c1, c2)), Literal(None, 1))),
c3,
)
assert get_first_level_and_conditions(cond) == [c1, c2, c3]
cond = binary_condition(
BooleanFunctions.OR,
binary_condition(BooleanFunctions.AND, c1, c2),
c3,
)
assert get_first_level_or_conditions(cond) == [
binary_condition(BooleanFunctions.AND, c1, c2),
c3,
]
cond = binary_condition(
ConditionFunctions.EQ,
binary_condition(BooleanFunctions.OR, c1,
binary_condition(BooleanFunctions.AND, c2, c3)),
Literal(None, 1),
)
assert get_first_level_or_conditions(cond) == [
c1,
binary_condition(BooleanFunctions.AND, c2, c3),
]
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(
QueryEntity(EntityKey.EVENTS, ColumnSet([])),
selected_columns=[
SelectedExpression(
"transaction.duration",
Column("transaction.duration", None, "duration")),
SelectedExpression("my_time", Column("my_time", None, "time")),
],
condition=condition,
groupby=[Column("my_time", None, "time")],
granularity=granularity,
)
expected = Query(
QueryEntity(EntityKey.EVENTS, 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() == unprocessed.get_selected_columns(
)
assert expected.get_condition() == unprocessed.get_condition()
ret = unprocessed.get_selected_columns()[1].expression.accept(
ClickhouseExpressionFormatter())
assert ret == formatted_column
if condition:
query_condition = unprocessed.get_condition()
assert query_condition is not None
ret = query_condition.accept(ClickhouseExpressionFormatter())
assert formatted_condition == ret
assert extract_granularity_from_query(unprocessed,
"finish_ts") == granularity
def test_aliased_cols() -> None:
"""
Test iteration whan columns have aliases. This is the expression
f2(t1.c2, f1(t1.c1, t1.c2 as a2)) as af1
"""
column1 = Column(None, "c1", "t1")
column2 = Column("a2", "c2", "t1")
function_1 = FunctionCall(None, "f1", (column1, column2))
column3 = Column(None, "c2", "t1")
function_2 = FunctionCall("af1", "f2", (column3, function_1))
expected = [column3, column1, column2, function_1, function_2]
assert list(function_2) == expected
def build_time_range(query_from: datetime, query_to: datetime) -> Expression:
return build_and(
FunctionCall(
None,
"greaterOrEquals",
(Column(None, None, "timestamp"), Literal(None, query_from)),
),
FunctionCall(
None,
"less",
(Column(None, None, "timestamp"), Literal(None, query_to)),
),
)
def test_iterate_over_query() -> None:
"""
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(ConditionFunctions.EQ, column1,
Literal(None, "1"))
prewhere = binary_condition(ConditionFunctions.EQ, column2,
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],
prewhere=prewhere,
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,
# prewhere
column2,
Literal(None, "2"),
prewhere,
]
assert list(query.get_all_expressions()) == expected_expressions
def test_column_function_translation() -> None:
assert ColumnToFunction(
None,
"ip_address",
"coalesce",
(Column(None, None, "ip_address_v4"), Column(None, None, "ip_address_v6")),
).attempt_map(
Column("ip_address", None, "ip_address"),
SnubaClickhouseMappingTranslator(TranslationMappers()),
) == FunctionCall(
"ip_address",
"coalesce",
(Column(None, None, "ip_address_v4"), Column(None, None, "ip_address_v6")),
)
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(
{},
TableSource("transactions", ColumnSet([])),
selected_columns=[
SelectedExpression(
"transaction.duration",
Column("transaction.duration", None, "duration")),
SelectedExpression("my_time", Column("my_time", None, "time")),
],
condition=condition,
granularity=granularity,
)
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_expressions_from_basic_condition() -> None:
"""
Iterates over the expressions in a basic condition
f(t1.c1) = t1.c2
"""
c = Column(None, "t1", "c1")
f1 = FunctionCall(None, "f", (c,))
c2 = Column(None, "t1", "c2")
condition = binary_condition(ConditionFunctions.EQ, f1, c2)
ret = list(condition)
expected = [c, f1, c2, condition]
assert ret == expected
def test_classify_and_replace() -> None:
condition = binary_condition(ConditionFunctions.EQ,
Column(None, "ev", "project_id"),
Literal(None, 1))
assert _classify_single_column_condition(
condition, {"ev": EntityKey.EVENTS}) == (
QualifiedCol(EntityKey.EVENTS, "project_id"),
"ev",
)
assert condition.transform(
partial(_replace_col, "ev", "project_id", "gr",
"project_id")) == binary_condition(
ConditionFunctions.EQ, Column(None, "gr", "project_id"),
Literal(None, 1))
def build_time_condition(time_columns: str, from_date: datetime,
to_date: datetime) -> Expression:
return binary_condition(
BooleanFunctions.AND,
binary_condition(
ConditionFunctions.GTE,
Column(f"_snuba_{time_columns}", None, time_columns),
Literal(None, from_date),
),
binary_condition(
ConditionFunctions.LT,
Column(f"_snuba_{time_columns}", None, time_columns),
Literal(None, to_date),
),
)
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 __replace_with_hash(self, condition: Expression) -> Expression:
match = self.__optimizable_pattern.match(condition)
if (
match is None
or match.string(KEY_COL_MAPPING_PARAM) != f"{self.__column_name}.key"
):
return condition
rhs = match.expression("right_hand_side")
assert isinstance(rhs, LiteralExpr)
key = match.string(KEY_MAPPING_PARAM).translate(ESCAPE_TRANSLATION)
return FunctionExpr(
alias=condition.alias,
function_name="has",
parameters=(
Column(
alias=None,
table_name=match.optional_string(TABLE_MAPPING_PARAM),
column_name=self.__hash_map_name,
),
FunctionExpr(
alias=None,
function_name="cityHash64",
parameters=(LiteralExpr(None, f"{key}={rhs.value}"),),
),
),
)
def function_column(col_name: str, function_name: str) -> ColumnToFunction:
return ColumnToFunction(
None,
col_name,
function_name,
(Column(None, None, col_name), ),
)
def cast_column_to_nullable(exp: Expression) -> Expression:
if isinstance(exp, Column):
if exp.column_name in self.mismatched_null_columns:
# depending on the order of the storage, this dictionary will contain
# either the nullable or non-nullable version of the column. No matter
# which one is in there, due to the mismatch on the merge table it needs to
# be cast as nullable anyways
mismatched_column = self.mismatched_null_columns[
exp.column_name]
col_is_nullable = _col_is_nullable(mismatched_column)
col_type = mismatched_column.type.for_schema()
cast_str = col_type if col_is_nullable else f"Nullable({col_type})"
return FunctionCall(
exp.alias,
"cast",
(
# move the alias up to the cast function
Column(
None,
table_name=exp.table_name,
column_name=exp.column_name,
),
Literal(None, cast_str),
),
)
return exp
def test_not_many_groups_to_exclude(query: ClickhouseQuery) -> None:
state.set_config("max_group_ids_exclude", 5)
set_project_exclude_groups(
2,
[100, 101, 102],
ReplacerState.ERRORS,
ReplacementType.
EXCLUDE_GROUPS, # Arbitrary replacement type, no impact on tests
)
PostReplacementConsistencyEnforcer("project_id",
ReplacerState.ERRORS).process_query(
query, HTTPQuerySettings())
assert query.get_condition() == build_and(
FunctionCall(
None,
"notIn",
(
FunctionCall(None, "assumeNotNull",
(Column(None, None, "group_id"), )),
FunctionCall(
None,
"tuple",
(
Literal(None, 100),
Literal(None, 101),
Literal(None, 102),
),
),
),
),
build_in("project_id", [2]),
)
assert not query.get_from_clause().final
def test_not_in_condition() -> None:
not_in_condition = binary_condition(
ConditionFunctions.NOT_IN,
Column(None, None, "tags_key"),
literals_tuple(None, [Literal(None, "t1"), Literal(None, "t2")]),
)
assert is_not_in_condition(not_in_condition)
match = is_not_in_condition_pattern(ColumnPattern(None, String("tags_key"))).match(
not_in_condition
)
assert match is not None
assert match.expression("tuple") == literals_tuple(
None, [Literal(None, "t1"), Literal(None, "t2")]
)
assert match.expression("lhs") == Column(None, None, "tags_key")
def test_functions(
default_validators: Mapping[str, FunctionCallValidator],
entity_validators: Mapping[str, FunctionCallValidator],
exception: Optional[Type[InvalidExpressionException]],
) -> None:
fn_cached = functions.default_validators
functions.default_validators = default_validators
entity_return = MagicMock()
entity_return.return_value = entity_validators
events_entity = get_entity(EntityKey.EVENTS)
cached = events_entity.get_function_call_validators
setattr(events_entity, "get_function_call_validators", entity_return)
data_source = QueryEntity(EntityKey.EVENTS, ColumnSet([]))
expression = FunctionCall(
None, "f", (Column(alias=None, table_name=None, column_name="col"), ))
if exception is None:
FunctionCallsValidator().validate(expression, data_source)
else:
with pytest.raises(exception):
FunctionCallsValidator().validate(expression, data_source)
# TODO: This should use fixture to do this
setattr(events_entity, "get_function_call_validators", cached)
functions.default_validators = fn_cached
def process_condition(exp: Expression) -> Expression:
result = CONDITION_PATTERN.match(exp)
if result is not None:
key_column = result.optional_string(KEY_COL_MAPPING_PARAM)
if key_column == "tags.key":
rhs = result.optional_string(KEY_MAPPING_PARAM)
table_name = result.optional_string(TABLE_MAPPING_PARAM)
replacement = FunctionCall(
exp.alias,
"has",
(Column(None, table_name, "tags.key"), Literal(None, rhs)),
)
assert isinstance(exp, FunctionCall)
if exp.function_name == ConditionFunctions.EQ:
replacement = FunctionCall(exp.alias, "not", (replacement,))
prev_value = query.get_experiment_value(
"empty-string-tag-condition"
)
if prev_value is not None:
return replacement if prev_value == "true" else exp
if settings.TESTING or random.random() < 0.5:
query.add_experiment("empty-string-tag-condition", "true")
return replacement
else:
query.add_experiment("empty-string-tag-condition", "false")
return exp
def test_mapping_complex_expression() -> None:
"""
Maps over an Expression container:
f0(t1.c1, fB(f())) -> f0(t1.c1, fB(f(f() as a)))
"""
f5 = FunctionCall("a", "f", ())
f4 = FunctionCall(None, "f", (f5, ))
f3 = FunctionCall(None, "f", ())
def replace_expr(e: Expression) -> Expression:
if isinstance(e, FunctionCall) and e.function_name == "f":
return f4
return e
c1 = Column(None, "c1", "t1")
f2 = FunctionCall(None, "fB", (f3, ))
f1 = FunctionCall(None, "f0", (c1, f2))
# Only the external function is going to be replaced since, when map returns a new
# column, we expect the func to have takern care of its own children.
f1 = f1.transform(replace_expr)
iterate = list(f1)
expected = [
c1,
f5,
f4,
FunctionCall(None, "fB", (f4, )),
FunctionCall(None, "f0", (c1, FunctionCall(None, "fB", (f4, )))),
]
assert iterate == expected
def test_not_many_groups_to_exclude(query: ClickhouseQuery) -> None:
state.set_config("max_group_ids_exclude", 5)
set_project_exclude_groups(2, [100, 101, 102], ReplacerState.EVENTS)
PostReplacementConsistencyEnforcer(
"project_id", ReplacerState.EVENTS
).process_query(query, HTTPRequestSettings())
assert query.get_condition_from_ast() == FunctionCall(
None,
BooleanFunctions.AND,
(
FunctionCall(
None,
"notIn",
(
FunctionCall(
None, "assumeNotNull", (Column(None, None, "group_id"),)
),
FunctionCall(
None,
"tuple",
(Literal(None, 100), Literal(None, 101), Literal(None, 102),),
),
),
),
build_in("project_id", [2]),
),
)
assert not query.get_from_clause().final
def do_post_processing(
self,
project_ids: Sequence[int],
query: Query,
request_settings: RequestSettings,
) -> None:
if not request_settings.get_turbo():
final, exclude_group_ids = get_projects_query_flags(
project_ids, self.__replacer_state_name)
if not final and exclude_group_ids:
# If the number of groups to exclude exceeds our limit, the query
# should just use final instead of the exclusion set.
max_group_ids_exclude = get_config(
"max_group_ids_exclude",
settings.REPLACER_MAX_GROUP_IDS_TO_EXCLUDE)
if len(exclude_group_ids) > max_group_ids_exclude:
query.set_final(True)
else:
query.add_conditions([(["assumeNotNull", ["group_id"]],
"NOT IN", exclude_group_ids)])
query.add_condition_to_ast(
not_in_condition(
None,
FunctionCall(None, "assumeNotNull",
(Column(None, "group_id", None), )),
[Literal(None, p) for p in exclude_group_ids],
))
else:
query.set_final(final)
def test_processing_functions() -> None:
in_condition = binary_condition(
None,
ConditionFunctions.IN,
Column(None, "tag_keys", None),
literals_tuple(
None,
[Literal(None, "t1"), Literal(None, "t2")]),
)
assert is_in_condition(in_condition)
eq_condition = binary_condition(None, ConditionFunctions.EQ,
Column(None, "test", None),
Literal(None, "1"))
assert is_binary_condition(eq_condition, ConditionFunctions.EQ)
assert not is_binary_condition(eq_condition, ConditionFunctions.NEQ)
