def extractor(condition: Expression) -> Set[Tuple[str, ...]]:
match = is_in_condition_pattern(key_pattern).match(condition)
if match is None:
return set()
function = match.expression("sequence")
if not isinstance(function, FunctionCallExpr) or function.function_name not in (
"tuple",
"array",
):
return set()
parameters: Set[Tuple[str, ...]] = set()
for tuple_param in function.parameters:
if (
not isinstance(tuple_param, FunctionCallExpr)
or tuple_param.function_name != "tuple"
):
return set()
parameters.add(
tuple(
param.value
for param in tuple_param.parameters
if isinstance(param, LiteralExpr) and isinstance(param.value, str)
)
)
return parameters
def get_project_ids_in_condition(
condition: Expression) -> Optional[Set[int]]:
"""
Extract project ids from an expression. Returns None if no project
if condition is found. It returns an empty set of conflicting project_id
conditions are found.
"""
match = FunctionCall(
None,
String(ConditionFunctions.EQ),
(
Column(column_name=String(project_column)),
Literal(value=Param("project_id", Any(int))),
),
).match(condition)
if match is not None:
return {match.integer("project_id")}
match = is_in_condition_pattern(
Column(column_name=String(project_column))).match(condition)
if match is not None:
projects = match.expression("tuple")
assert isinstance(projects, FunctionCallExpr)
return {
l.value
for l in projects.parameters
if isinstance(l, LiteralExpr) and isinstance(l.value, int)
}
match = FunctionCall(
None,
Param(
"operator",
Or([String(BooleanFunctions.AND),
String(BooleanFunctions.OR)]),
),
(Param("lhs", AnyExpression()), Param("rhs", AnyExpression())),
).match(condition)
if match is not None:
lhs_projects = get_project_ids_in_condition(
match.expression("lhs"))
rhs_projects = get_project_ids_in_condition(
match.expression("rhs"))
if lhs_projects is None:
return rhs_projects
elif rhs_projects is None:
return lhs_projects
else:
return (lhs_projects & rhs_projects if match.string("operator")
== BooleanFunctions.AND else lhs_projects
| rhs_projects)
return None
def extractor(condition: Expression) -> Set[str]:
match = is_in_condition_pattern(key_pattern).match(condition)
if match is None:
return set()
function = match.expression("sequence")
assert isinstance(function, FunctionCallExpr)
return {
param.value
for param in function.parameters
if isinstance(param, LiteralExpr) and isinstance(param.value, str)
}
def test_in_condition() -> None:
in_condition = binary_condition(
ConditionFunctions.IN,
Column(None, None, "tags_key"),
literals_tuple(None, [Literal(None, "t1"), Literal(None, "t2")]),
)
assert is_in_condition(in_condition)
match = is_in_condition_pattern(ColumnPattern(None, String("tags_key"))).match(
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 _get_mapping_keys_in_condition(
condition: Expression, column_name: str
) -> Optional[Set[str]]:
"""
Finds the top level conditions that include filter based on the arrayJoin.
This is meant to be used to find the keys the query is filtering the arrayJoin
on.
We can only apply the arrayFilter optimization to arrayJoin conditions
that are not in OR with other columns. To simplify the problem, we only
consider those conditions that are included in the first level of the query:
[['tagskey' '=' 'a'],['col' '=' 'b'],['col2' '=' 'c']] works
[[['tagskey' '=' 'a'], ['col2' '=' 'b']], ['tagskey' '=' 'c']] does not
If we encounter an OR condition we return None, which means we cannot
safely apply the optimization. Empty set means we did not find any
suitable arrayJoin for optimization in this condition but that does
not disqualify the whole query in the way the OR condition does.
"""
keys_found = set()
conditions = get_first_level_and_conditions(condition)
for c in conditions:
if is_binary_condition(c, BooleanFunctions.OR):
return None
match = FunctionCall(
None,
String(ConditionFunctions.EQ),
(array_join_pattern(column_name), Literal(None, Param("key", Any(str)))),
).match(c)
if match is not None:
keys_found.add(match.string("key"))
match = is_in_condition_pattern(array_join_pattern(column_name)).match(c)
if match is not None:
function = match.expression("tuple")
assert isinstance(function, FunctionCallExpr)
keys_found |= {
lit.value
for lit in function.parameters
if isinstance(lit, LiteralExpr) and isinstance(lit.value, str)
}
return keys_found