def test_format_clickhouse_specific_query() -> None:
"""
Adds a few of the Clickhosue specific fields to the query.
"""
query = ClickhouseQuery(
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"),
sample=0.1,
totals=True,
limitby=(10, "environment"),
)
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 process_query(self, query: Query,
request_settings: RequestSettings) -> None:
if request_settings.get_turbo():
return
project_ids = get_project_ids_in_query_ast(query,
self.__project_column)
set_final = False
condition_to_add = None
if project_ids:
final, exclude_group_ids = get_projects_query_flags(
list(project_ids),
self.__replacer_state_name,
)
if final:
metrics.increment("final", tags={"cause": "final_flag"})
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:
metrics.increment("final", tags={"cause": "max_groups"})
set_final = True
else:
condition_to_add = (
["assumeNotNull", ["group_id"]],
"NOT IN",
exclude_group_ids,
)
query.add_condition_to_ast(
not_in_condition(
None,
FunctionCall(None, "assumeNotNull",
(Column(None, None, "group_id"), )),
[Literal(None, p) for p in exclude_group_ids],
))
else:
set_final = final
query.set_final(set_final)
if condition_to_add:
query.add_conditions([condition_to_add])
