def do_query(
query: ClickhouseQuery,
request_settings: RequestSettings,
) -> QueryResult:
from_date_ast, to_date_ast = get_time_range(query, "timestamp")
assert from_date_ast is not None and isinstance(
from_date_ast, datetime)
assert to_date_ast is not None and isinstance(to_date_ast, datetime)
conditions = query.get_conditions() or []
from_date_str = next(
(condition[2] for condition in conditions
if condition[0] == "timestamp" and condition[1] == ">="),
None,
)
to_date_str = next(
(condition[2] for condition in conditions
if condition[0] == "timestamp" and condition[1] == "<"),
None,
)
assert from_date_str == from_date_ast.isoformat()
assert to_date_str == to_date_ast.isoformat()
found_timestamps.append(
(from_date_ast.isoformat(), to_date_ast.isoformat()))
return QueryResult({"data": []}, {})
def _replace_condition(query: Query, field: str, operator: str,
new_literal: Union[str, List[AnyType]]) -> None:
query.set_conditions([
cond if not _identify_condition(cond, field, operator) else
[field, operator, new_literal]
for cond in query.get_conditions() or []
])
def test_without_turbo_without_projects_needing_final(
query: ClickhouseQuery) -> None:
PostReplacementConsistencyEnforcer("project_id", None).process_query(
query, HTTPRequestSettings())
assert query.get_conditions() == [("project_id", "IN", [2])]
assert query.get_condition_from_ast() == build_in("project_id", [2])
assert not query.get_final()
def _update_conditions(self, query: Query,
prewhere_conditions: Sequence[Condition]) -> None:
conditions = query.get_conditions()
# This should never ne None at this point, but for mypy this can be None.
assert conditions is not None
query.set_conditions(
[cond for cond in conditions if cond not in prewhere_conditions])
query.set_prewhere(prewhere_conditions)
def test_too_many_groups_to_exclude(query: ClickhouseQuery) -> None:
state.set_config("max_group_ids_exclude", 2)
set_project_exclude_groups(2, [100, 101, 102], ReplacerState.EVENTS)
PostReplacementConsistencyEnforcer("project_id",
ReplacerState.EVENTS).process_query(
query, HTTPRequestSettings())
assert query.get_conditions() == [("project_id", "IN", [2])]
assert query.get_condition_from_ast() == build_in("project_id", [2])
assert query.get_final()
def _get_prewhere_candidates(
self, query: Query, prewhere_keys: Sequence[str]
) -> Sequence[Tuple[Iterable[str], Condition]]:
# Add any condition to PREWHERE if:
# - It is a single top-level condition (not OR-nested), and
# - Any of its referenced columns are in prewhere_keys
conditions = query.get_conditions()
if not conditions:
return []
return [(util.columns_in_expr(cond[0]), cond) for cond in conditions
if util.is_condition(cond) and cond[1] in ALLOWED_OPERATORS
and any(col in prewhere_keys
for col in util.columns_in_expr(cond[0]))]
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())
expected = [
("project_id", "IN", [2]),
(["assumeNotNull", ["group_id"]], "NOT IN", [100, 101, 102]),
]
assert query.get_conditions() == expected
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_final()
def process_query(self, query: Query,
request_settings: RequestSettings) -> None:
conditions = query.get_conditions()
if not conditions:
return
# Enable the processor only if we have enough data in the flattened
# columns. Which have been deployed at BEGINNING_OF_TIME. If the query
# starts earlier than that we do not apply the optimization.
if self.__beginning_of_time:
apply_optimization = False
for condition in conditions:
if (is_condition(condition) and isinstance(condition[0], str)
and condition[0] in self.__timestamp_cols
and condition[1] in (">=", ">")
and isinstance(condition[2], str)):
try:
start_ts = parse_datetime(condition[2])
if (start_ts -
self.__beginning_of_time).total_seconds() > 0:
apply_optimization = True
except Exception:
# We should not get here, it means the from timestamp is malformed
# Returning here is just for safety
logger.error(
"Cannot parse start date for NestedFieldOptimizer: %r",
condition,
)
return
if not apply_optimization:
return
# Do not use flattened tags if tags are being unpacked anyway. In that case
# using flattened tags only implies loading an additional column thus making
# the query heavier and slower
if self.__has_tags(query.get_arrayjoin_from_ast()):
return
if query.get_groupby_from_ast():
for expression in query.get_groupby_from_ast():
if self.__has_tags(expression):
return
if self.__has_tags(query.get_having_from_ast()):
return
if query.get_orderby_from_ast():
for orderby in query.get_orderby_from_ast():
if self.__has_tags(orderby.expression):
return
new_conditions = []
positive_like_expression: List[str] = []
negative_like_expression: List[str] = []
for c in conditions:
keyvalue = self.__is_optimizable(c, self.__nested_col)
if not keyvalue:
new_conditions.append(c)
else:
expression = f"{escape_field(keyvalue.nested_col_key)}={escape_field(keyvalue.value)}"
if keyvalue.operand == Operand.EQ:
positive_like_expression.append(expression)
else:
negative_like_expression.append(expression)
if positive_like_expression:
# Positive conditions "=" are all merged together in one LIKE expression
positive_like_expression = sorted(positive_like_expression)
like_formatted = f"%|{'|%|'.join(positive_like_expression)}|%"
new_conditions.append(
[self.__flattened_col, "LIKE", like_formatted])
for expression in negative_like_expression:
# Negative conditions "!=" cannot be merged together. We can still transform
# them into NOT LIKE statements, but each condition has to be one
# statement.
not_like_formatted = f"%|{expression}|%"
new_conditions.append(
[self.__flattened_col, "NOT LIKE", not_like_formatted])
query.set_conditions(new_conditions)
def get_project_ids_in_query(query: Query,
project_column: str) -> Optional[Set[int]]:
"""
Finds the project ids this query is filtering according to the legacy query
representation.
It returns a set of referenced project_ids if relevant conditions are found or
None if no project_id condition is found in the query. Empty set means that multiple
conflicting project id conditions were found.
This function looks into first level AND conditions and second level OR conditions
but it cannot support project ids used as function parameters.
Specific limitations:
- If a project_id is a parameter of a function that returns the project_id itself
it is not supported. It would be very hard to support every function without a
whitelist/blacklist of allowed functions in Snuba queries.
- boolean functions are not supported. So we do not unpack and/or/not conditions
expressed as functions. We will be able to do that with the AST.
- does not exclude projects referenced in NOT conditions.
We are going to try to lift as many of these limitations as possible. So, please,
do not rely on them for the correctness of your code.
"""
def find_project_id_sets(
conditions: Sequence[Condition], ) -> Sequence[Set[int]]:
"""
Scans a potentially nested sequence of conditions.
For each simple condition adds to the output the set of project ids referenced
by the condition.
For each nested condition, it assumes it is a union of simple conditions
(which is the only supported valid case by the Query object) and adds the union of the
referenced project ids to the output.
"""
project_id_sets: List[Set[int]] = list()
for c in conditions:
if is_condition(c):
# This is a simple condition. Can extract the project ids directly.
# Supports these kinds of conditions
# ["col", "=", 1]
# ["col", "IN", [1,2,3]]
# ["col", "IN", (1,2,3)]
if c[0] == project_column:
if c[1] == "=" and isinstance(c[2], int):
project_id_sets.append({c[2]})
elif c[1] == "IN" and all(
isinstance(project, int) for project in c[2]):
project_id_sets.append(set(c[2]))
elif all(is_condition(second_level) for second_level in c):
# This is supposed to be a union of simple conditions. Need to union
# the sets of project ids.
sets_to_unite = find_project_id_sets(c)
if sets_to_unite:
project_id_sets.append(
reduce(lambda x, y: x | y, sets_to_unite))
else:
raise ValueError(f"Invalid condition {conditions}")
return project_id_sets
all_project_id_sets = find_project_id_sets(query.get_conditions() or [])
if not all_project_id_sets:
return None
return reduce(lambda x, y: x & y, all_project_id_sets)
def execute(
self,
query: Query,
request_settings: RequestSettings,
runner: SplitQueryRunner,
) -> Optional[QueryResult]:
"""
If a query is:
- ORDER BY timestamp DESC
- has no grouping
- has an offset/limit
- has a large time range
We know we have to reverse-sort the entire set of rows to return the small
chunk at the end of the time range, so optimistically split the time range
into smaller increments, and start with the last one, so that we can potentially
avoid querying the entire range.
"""
limit = query.get_limit()
if limit is None or query.get_groupby():
return None
if query.get_offset() >= 1000:
return None
orderby = query.get_orderby()
if not orderby or orderby[0] != f"-{self.__timestamp_col}":
return None
conditions = query.get_conditions() or []
from_date_str = next(
(condition[2] for condition in conditions
if _identify_condition(condition, self.__timestamp_col, ">=")),
None,
)
to_date_str = next(
(condition[2] for condition in conditions
if _identify_condition(condition, self.__timestamp_col, "<")),
None,
)
from_date_ast, to_date_ast = get_time_range(query,
self.__timestamp_col)
if not from_date_str or not to_date_str:
return None
date_align, split_step = state.get_configs([("date_align_seconds", 1),
("split_step", 3600)
] # default 1 hour
)
to_date = util.parse_datetime(to_date_str, date_align)
from_date = util.parse_datetime(from_date_str, date_align)
if from_date != from_date_ast:
logger.warning(
"Mismatch in start date on time splitter.",
extra={
"ast": str(from_date_ast),
"legacy": str(from_date)
},
exc_info=True,
)
metrics.increment("mismatch.ast_from_date")
remaining_offset = query.get_offset()
overall_result = None
split_end = to_date
split_start = max(split_end - timedelta(seconds=split_step), from_date)
total_results = 0
while split_start < split_end and total_results < limit:
# We need to make a copy to use during the query execution because we replace
# the start-end conditions on the query at each iteration of this loop.
split_query = copy.deepcopy(query)
_replace_condition(split_query, self.__timestamp_col, ">=",
split_start.isoformat())
_replace_ast_condition(split_query, self.__timestamp_col, ">=",
LiteralExpr(None, split_start))
_replace_condition(split_query, self.__timestamp_col, "<",
split_end.isoformat())
_replace_ast_condition(split_query, self.__timestamp_col, "<",
LiteralExpr(None, split_end))
# Because its paged, we have to ask for (limit+offset) results
# and set offset=0 so we can then trim them ourselves.
split_query.set_offset(0)
split_query.set_limit(limit - total_results + remaining_offset)
# At every iteration we only append the "data" key from the results returned by
# the runner. The "extra" key is only populated at the first iteration of the
# loop and never changed.
result = runner(split_query, request_settings)
if overall_result is None:
overall_result = result
else:
overall_result.result["data"].extend(result.result["data"])
if remaining_offset > 0 and len(overall_result.result["data"]) > 0:
to_trim = min(remaining_offset,
len(overall_result.result["data"]))
overall_result.result["data"] = overall_result.result["data"][
to_trim:]
remaining_offset -= to_trim
total_results = len(overall_result.result["data"])
if total_results < limit:
if len(result.result["data"]) == 0:
# If we got nothing from the last query, expand the range by a static factor
split_step = split_step * STEP_GROWTH
else:
# If we got some results but not all of them, estimate how big the time
# range should be for the next query based on how many results we got for
# our last query and its time range, and how many we have left to fetch.
remaining = limit - total_results
split_step = split_step * math.ceil(
remaining / float(len(result.result["data"])))
# Set the start and end of the next query based on the new range.
split_end = split_start
try:
split_start = max(
split_end - timedelta(seconds=split_step), from_date)
except OverflowError:
split_start = from_date
return overall_result
