def _format_query_body(query: Query) -> Mapping[str, Any]:
expression_formatter = TracingExpressionFormatter()
formatted = {
"SELECT": [[e.name, e.expression.accept(expression_formatter)]
for e in query.get_selected_columns_from_ast()],
"GROUPBY":
[e.accept(expression_formatter) for e in query.get_groupby_from_ast()],
"ORDERBY": [[e.expression.accept(expression_formatter), e.direction]
for e in query.get_orderby_from_ast()],
}
array_join = query.get_arrayjoin_from_ast()
if array_join:
formatted["ARRAYJOIN"] = array_join.accept(expression_formatter)
condition = query.get_condition_from_ast()
if condition:
formatted["WHERE"] = condition.accept(expression_formatter)
having = query.get_having_from_ast()
if having:
formatted["HAVING"] = having.accept(expression_formatter)
limitby = query.get_limitby()
if limitby:
formatted["LIMITBY"] = {
"LIMIT": limitby.limit,
"BY": limitby.expression.accept(expression_formatter),
}
limit = query.get_limit()
if limit:
formatted["LIMIT"] = limit
offset = query.get_offset()
if offset:
formatted["OFFSET"] = offset
return formatted
def _replace_time_condition(
self,
query: Query,
from_date: datetime,
from_exp: FunctionCall,
to_date: datetime,
to_exp: FunctionCall,
) -> None:
max_days, date_align = state.get_configs(
[("max_days", None), ("date_align_seconds", 1)]
)
def align_fn(dt: datetime) -> datetime:
assert isinstance(date_align, int)
return dt - timedelta(seconds=(dt - dt.min).seconds % date_align)
from_date, to_date = align_fn(from_date), align_fn(to_date)
assert from_date <= to_date
if max_days is not None and (to_date - from_date).days > max_days:
from_date = to_date - timedelta(days=max_days)
def replace_cond(exp: Expression) -> Expression:
if not isinstance(exp, FunctionCall):
return exp
elif exp == from_exp:
return replace(
exp, parameters=(from_exp.parameters[0], Literal(None, from_date)),
)
elif exp == to_exp:
return replace(
exp, parameters=(to_exp.parameters[0], Literal(None, to_date))
)
return exp
condition = query.get_condition_from_ast()
top_level = get_first_level_and_conditions(condition) if condition else []
new_top_level = list(map(replace_cond, top_level))
query.set_ast_condition(combine_and_conditions(new_top_level))
def validate_required_conditions(
self, query: Query, alias: Optional[str] = None
) -> bool:
if not self._required_filter_columns and not self._required_time_column:
return True
condition = query.get_condition_from_ast()
top_level = get_first_level_and_conditions(condition) if condition else []
if not top_level:
return False
alias_match = AnyOptionalString() if alias is None else StringMatch(alias)
def build_match(
col: str, ops: Sequence[str], param_type: Any
) -> Or[Expression]:
# The IN condition has to be checked separately since each parameter
# has to be checked individually.
column_match = ColumnMatch(alias_match, StringMatch(col))
return Or(
[
FunctionCallMatch(
Or([StringMatch(op) for op in ops]),
(column_match, LiteralMatch(AnyMatch(param_type))),
),
FunctionCallMatch(
StringMatch(ConditionFunctions.IN),
(
column_match,
FunctionCallMatch(
Or([StringMatch("array"), StringMatch("tuple")]),
all_parameters=LiteralMatch(AnyMatch(param_type)),
),
),
),
]
)
if self._required_filter_columns:
for col in self._required_filter_columns:
match = build_match(col, [ConditionFunctions.EQ], int)
found = any(match.match(cond) for cond in top_level)
if not found:
return False
if self._required_time_column:
match = build_match(
self._required_time_column, [ConditionFunctions.EQ], datetime,
)
found = any(match.match(cond) for cond in top_level)
if found:
return True
lower, upper = get_time_range_expressions(
top_level, self._required_time_column, alias
)
if not lower or not upper:
return False
# At this point we have valid conditions. However we need to align them and
# make sure they don't exceed the max_days. Replace the conditions.
self._replace_time_condition(query, *lower, *upper)
return True