def all_referenced_columns(query: Query):
"""
Return the set of all columns that are used by a query.
"""
col_exprs: MutableSequence[Any] = []
if query.get_arrayjoin():
col_exprs.extend(to_list(query.get_arrayjoin()))
if query.get_groupby():
col_exprs.extend(to_list(query.get_groupby()))
if query.get_orderby():
col_exprs.extend(to_list(query.get_orderby()))
if query.get_selected_columns():
col_exprs.extend(to_list(query.get_selected_columns()))
# Conditions need flattening as they can be nested as AND/OR
if query.get_conditions():
flat_conditions = list(
chain(*[[c] if is_condition(c) else c
for c in query.get_conditions()]))
col_exprs.extend([c[0] for c in flat_conditions])
if query.get_aggregations():
col_exprs.extend([a[1] for a in query.get_aggregations()])
# Return the set of all columns referenced in any expression
return set(chain(*[columns_in_expr(ex) for ex in col_exprs]))
def conditions_expr(dataset,
conditions,
query: Query,
parsing_context: ParsingContext,
depth=0):
"""
Return a boolean expression suitable for putting in the WHERE clause of the
query. The expression is constructed by ANDing groups of OR expressions.
Expansion of columns is handled, as is replacement of columns with aliases,
if the column has already been expanded and aliased elsewhere.
"""
from snuba.clickhouse.columns import Array
if not conditions:
return ''
if depth == 0:
# dedupe conditions at top level, but keep them in order
sub = OrderedDict(
(conditions_expr(dataset, cond, query, parsing_context, depth + 1),
None) for cond in conditions)
return u' AND '.join(s for s in sub.keys() if s)
elif is_condition(conditions):
lhs, op, lit = dataset.process_condition(conditions)
# facilitate deduping IN conditions by sorting them.
if op in ('IN', 'NOT IN') and isinstance(lit, tuple):
lit = tuple(sorted(lit))
# If the LHS is a simple column name that refers to an array column
# (and we are not arrayJoining on that column, which would make it
# scalar again) and the RHS is a scalar value, we assume that the user
# actually means to check if any (or all) items in the array match the
# predicate, so we return an `any(x == value for x in array_column)`
# type expression. We assume that operators looking for a specific value
# (IN, =, LIKE) are looking for rows where any array value matches, and
# exclusionary operators (NOT IN, NOT LIKE, !=) are looking for rows
# where all elements match (eg. all NOT LIKE 'foo').
columns = dataset.get_dataset_schemas().get_read_schema().get_columns()
if (isinstance(lhs, str) and lhs in columns
and isinstance(columns[lhs].type, Array)
and columns[lhs].base_name != query.get_arrayjoin()
and not isinstance(lit, (list, tuple))):
any_or_all = 'arrayExists' if op in POSITIVE_OPERATORS else 'arrayAll'
return u'{}(x -> assumeNotNull(x {} {}), {})'.format(
any_or_all, op, escape_literal(lit),
column_expr(dataset, lhs, query, parsing_context))
else:
return u'{} {} {}'.format(
column_expr(dataset, lhs, query, parsing_context), op,
escape_literal(lit))
elif depth == 1:
sub = (conditions_expr(dataset, cond, query, parsing_context,
depth + 1) for cond in conditions)
sub = [s for s in sub if s]
res = u' OR '.join(sub)
return u'({})'.format(res) if len(sub) > 1 else res
else:
raise InvalidConditionException(str(conditions))
def test_full_query():
query = Query(
{
"selected_columns": ["c1", "c2", "c3"],
"conditions": [["c1", "=", "a"]],
"arrayjoin": "tags",
"having": [["c4", "=", "c"]],
"groupby": ["project_id"],
"aggregations": [["count()", "", "count"]],
"orderby": "event_id",
"limitby": (100, "environment"),
"sample": 10,
"limit": 100,
"offset": 50,
"totals": True,
"granularity": 60,
},
TableSource("my_table", ColumnSet([])),
)
assert query.get_selected_columns() == ["c1", "c2", "c3"]
assert query.get_aggregations() == [["count()", "", "count"]]
assert query.get_groupby() == ["project_id"]
assert query.get_conditions() == [["c1", "=", "a"]]
assert query.get_arrayjoin() == "tags"
assert query.get_having() == [["c4", "=", "c"]]
assert query.get_orderby() == "event_id"
assert query.get_limitby() == (100, "environment")
assert query.get_sample() == 10
assert query.get_limit() == 100
assert query.get_offset() == 50
assert query.has_totals() is True
assert query.get_granularity() == 60
assert query.get_data_source().format_from() == "my_table"
def conditions_expr(
dataset, conditions, query: Query, parsing_context: ParsingContext, depth=0
) -> str:
def operand_builder(val: Any) -> str:
return str(column_expr(dataset, val, query, parsing_context))
def and_builder(expressions: Sequence[str]) -> str:
return " AND ".join(e for e in expressions if e)
def or_builder(expressions: Sequence[str]) -> str:
res = " OR ".join(expressions)
return "({})".format(res) if len(expressions) > 1 else res
def unpack_array_condition_builder(lhs: str, op: str, literal: Any) -> str:
any_or_all = "arrayExists" if op in POSITIVE_OPERATORS else "arrayAll"
return "{}(x -> assumeNotNull(x {} {}), {})".format(
any_or_all, op, escape_literal(literal), lhs,
)
def simple_condition_builder(lhs: str, op: str, literal: Any) -> str:
return "{} {} {}".format(lhs, op, escape_literal(literal))
condition = parse_conditions(
operand_builder,
and_builder,
or_builder,
unpack_array_condition_builder,
simple_condition_builder,
dataset,
conditions,
query.get_arrayjoin(),
depth,
)
return condition or ""
def test_empty_query():
query = Query({})
assert query.get_selected_columns() is None
assert query.get_aggregations() is None
assert query.get_groupby() is None
assert query.get_conditions() is None
assert query.get_arrayjoin() is None
assert query.get_having() == []
assert query.get_orderby() is None
assert query.get_limitby() is None
assert query.get_sample() is None
assert query.get_limit() is None
assert query.get_offset() == 0
assert query.has_totals() is False
def test_edit_query():
query = Query(
{
"selected_columns": ["c1", "c2", "c3"],
"conditions": [["c1", "=", "a"]],
"arrayjoin": "tags",
"having": [["c4", "=", "c"]],
"groupby": ["project_id"],
"aggregations": [["count()", "", "count"]],
"orderby": "event_id",
"limitby": (100, "environment"),
"sample": 10,
"limit": 100,
"offset": 50,
"totals": True,
},
TableSource("my_table", ColumnSet([])),
)
query.set_selected_columns(["c4"])
assert query.get_selected_columns() == ["c4"]
query.set_aggregations([["different_agg()", "", "something"]])
assert query.get_aggregations() == [["different_agg()", "", "something"]]
query.add_groupby(["more", "more2"])
assert query.get_groupby() == ["project_id", "more", "more2"]
query.add_conditions([["c5", "=", "9"]])
assert query.get_conditions() == [
["c1", "=", "a"],
["c5", "=", "9"],
]
query.set_conditions([["c6", "=", "10"]])
assert query.get_conditions() == [
["c6", "=", "10"],
]
query.set_arrayjoin("not_tags")
assert query.get_arrayjoin() == "not_tags"
query.set_granularity(7200)
assert query.get_granularity() == 7200
query.set_prewhere([["pc6", "=", "10"]])
assert query.get_prewhere() == [["pc6", "=", "10"]]
def test_empty_query():
query = Query({}, TableSource("my_table", ColumnSet([])))
assert query.get_selected_columns() is None
assert query.get_aggregations() is None
assert query.get_groupby() is None
assert query.get_conditions() is None
assert query.get_arrayjoin() is None
assert query.get_having() == []
assert query.get_orderby() is None
assert query.get_limitby() is None
assert query.get_sample() is None
assert query.get_limit() is None
assert query.get_offset() == 0
assert query.has_totals() is False
assert query.get_data_source().format_from() == "my_table"
def __init__(
self,
dataset: Dataset,
query: Query,
settings: RequestSettings,
) -> None:
parsing_context = ParsingContext()
aggregate_exprs = [
column_expr(dataset, col, query, parsing_context, alias, agg)
for (agg, col, alias) in query.get_aggregations()
]
groupby = util.to_list(query.get_groupby())
group_exprs = [
column_expr(dataset, gb, query, parsing_context) for gb in groupby
]
column_names = query.get_selected_columns() or []
selected_cols = [
column_expr(dataset, util.tuplify(colname), query, parsing_context)
for colname in column_names
]
select_clause = u"SELECT {}".format(
", ".join(group_exprs + aggregate_exprs + selected_cols))
from_clause = u"FROM {}".format(query.get_data_source().format_from())
if query.get_final():
from_clause = u"{} FINAL".format(from_clause)
if not query.get_data_source().supports_sample():
sample_rate = None
else:
if query.get_sample():
sample_rate = query.get_sample()
elif settings.get_turbo():
sample_rate = snuba_settings.TURBO_SAMPLE_RATE
else:
sample_rate = None
if sample_rate:
from_clause = u"{} SAMPLE {}".format(from_clause, sample_rate)
join_clause = ""
if query.get_arrayjoin():
join_clause = u"ARRAY JOIN {}".format(query.get_arrayjoin())
where_clause = ""
if query.get_conditions():
where_clause = u"WHERE {}".format(
conditions_expr(dataset, query.get_conditions(), query,
parsing_context))
prewhere_clause = ""
if query.get_prewhere():
prewhere_clause = u"PREWHERE {}".format(
conditions_expr(dataset, query.get_prewhere(), query,
parsing_context))
group_clause = ""
if groupby:
group_clause = "GROUP BY ({})".format(", ".join(
column_expr(dataset, gb, query, parsing_context)
for gb in groupby))
if query.has_totals():
group_clause = "{} WITH TOTALS".format(group_clause)
having_clause = ""
having_conditions = query.get_having()
if having_conditions:
assert groupby, "found HAVING clause with no GROUP BY"
having_clause = u"HAVING {}".format(
conditions_expr(dataset, having_conditions, query,
parsing_context))
order_clause = ""
if query.get_orderby():
orderby = [
column_expr(dataset, util.tuplify(ob), query, parsing_context)
for ob in util.to_list(query.get_orderby())
]
orderby = [
u"{} {}".format(ob.lstrip("-"),
"DESC" if ob.startswith("-") else "ASC")
for ob in orderby
]
order_clause = u"ORDER BY {}".format(", ".join(orderby))
limitby_clause = ""
if query.get_limitby() is not None:
limitby_clause = "LIMIT {} BY {}".format(*query.get_limitby())
limit_clause = ""
if query.get_limit() is not None:
limit_clause = "LIMIT {}, {}".format(query.get_offset(),
query.get_limit())
self.__formatted_query = " ".join([
c for c in [
select_clause,
from_clause,
join_clause,
prewhere_clause,
where_clause,
group_clause,
having_clause,
order_clause,
limitby_clause,
limit_clause,
] if c
])
def __init__(
self,
dataset: Dataset,
query: Query,
settings: RequestSettings,
prewhere_conditions: Sequence[str],
) -> None:
parsing_context = ParsingContext()
aggregate_exprs = [
column_expr(dataset, col, query, parsing_context, alias, agg)
for (agg, col, alias) in query.get_aggregations()
]
groupby = util.to_list(query.get_groupby())
group_exprs = [
column_expr(dataset, gb, query, parsing_context) for gb in groupby
]
column_names = query.get_selected_columns() or []
selected_cols = [
column_expr(dataset, util.tuplify(colname), query, parsing_context)
for colname in column_names
]
select_clause = u'SELECT {}'.format(
', '.join(group_exprs + aggregate_exprs + selected_cols))
from_clause = u'FROM {}'.format(query.get_data_source().format_from())
if query.get_final():
from_clause = u'{} FINAL'.format(from_clause)
if query.get_sample():
sample_rate = query.get_sample()
elif settings.get_turbo():
sample_rate = snuba_settings.TURBO_SAMPLE_RATE
else:
sample_rate = None
if sample_rate:
from_clause = u'{} SAMPLE {}'.format(from_clause, sample_rate)
join_clause = ''
if query.get_arrayjoin():
join_clause = u'ARRAY JOIN {}'.format(query.get_arrayjoin())
where_clause = ''
if query.get_conditions():
where_clause = u'WHERE {}'.format(
conditions_expr(dataset, query.get_conditions(), query,
parsing_context))
prewhere_clause = ''
if prewhere_conditions:
prewhere_clause = u'PREWHERE {}'.format(
conditions_expr(dataset, prewhere_conditions, query,
parsing_context))
group_clause = ''
if groupby:
group_clause = 'GROUP BY ({})'.format(', '.join(
column_expr(dataset, gb, query, parsing_context)
for gb in groupby))
if query.has_totals():
group_clause = '{} WITH TOTALS'.format(group_clause)
having_clause = ''
having_conditions = query.get_having()
if having_conditions:
assert groupby, 'found HAVING clause with no GROUP BY'
having_clause = u'HAVING {}'.format(
conditions_expr(dataset, having_conditions, query,
parsing_context))
order_clause = ''
if query.get_orderby():
orderby = [
column_expr(dataset, util.tuplify(ob), query, parsing_context)
for ob in util.to_list(query.get_orderby())
]
orderby = [
u'{} {}'.format(ob.lstrip('-'),
'DESC' if ob.startswith('-') else 'ASC')
for ob in orderby
]
order_clause = u'ORDER BY {}'.format(', '.join(orderby))
limitby_clause = ''
if query.get_limitby() is not None:
limitby_clause = 'LIMIT {} BY {}'.format(*query.get_limitby())
limit_clause = ''
if query.get_limit() is not None:
limit_clause = 'LIMIT {}, {}'.format(query.get_offset(),
query.get_limit())
self.__formatted_query = ' '.join([
c for c in [
select_clause, from_clause, join_clause, prewhere_clause,
where_clause, group_clause, having_clause, order_clause,
limitby_clause, limit_clause
] if c
])