def test_mock_consumer() -> None:
storage = get_writable_storage(StorageKey.ERRORS)
strategy = KafkaConsumerStrategyFactory(
None,
lambda message: None,
build_mock_batch_writer(storage, True, TestingMetricsBackend(), 100,
50),
max_batch_size=1,
max_batch_time=1,
processes=None,
input_block_size=None,
output_block_size=None,
initialize_parallel_transform=None,
).create(lambda message: None)
strategy.submit(
Message(
Partition(Topic("events"), 0),
1,
KafkaPayload(None, b"INVALID MESSAGE", []),
datetime.now(),
))
strategy.close()
strategy.join()
# If the mock was not applied correctly we would have data in Clickhouse
reader = storage.get_cluster().get_reader()
result = reader.execute(
FormattedQuery([StringNode("SELECT count() as c from errors_local")]))
assert result["data"] == [{"c": 0}]
def _build_optional_string_node(
name: str,
expression: Optional[Expression],
formatter: ExpressionVisitor[str],
) -> Optional[StringNode]:
return (StringNode(f"{name} {expression.accept(formatter)}")
if expression is not None else None)
def _build_optional_string_node(
name: str,
expression: Optional[Expression],
formatter: ClickhouseExpressionFormatterBase,
) -> Optional[StringNode]:
return (StringNode(f"{name} {expression.accept(formatter)}")
if expression is not None else None)
def _format_select(query: AbstractQuery,
formatter: ClickhouseExpressionFormatter) -> StringNode:
selected_cols = [
e.expression.accept(formatter)
for e in query.get_selected_columns_from_ast()
]
return StringNode(f"SELECT {', '.join(selected_cols)}")
def visit_join_clause(self, node: JoinClause[Table]) -> FormattedNode:
join_type = f"{node.join_type.value} " if node.join_type else ""
modifier = f"{node.join_modifier.value} " if node.join_modifier else ""
return SequenceNode([
node.left_node.accept(self),
StringNode(f"{modifier}{join_type}JOIN"),
node.right_node.accept(self),
StringNode("ON"),
SequenceNode(
[
StringNode(
f"{k.left.table_alias}.{k.left.column}={k.right.table_alias}.{k.right.column}"
) for k in node.keys
],
" AND ",
),
])
def _format_granularity(
query: AbstractQuery, formatter: ExpressionVisitor[str]
) -> Optional[StringNode]:
ast_granularity = query.get_granularity()
return (
StringNode(f"GRANULARITY {ast_granularity}")
if ast_granularity is not None
else None
)
def _visit_simple_source(self, data_source: Table) -> StringNode:
"""
Formats a simple table data source.
"""
final = " FINAL" if data_source.final else ""
sample = (f" SAMPLE {data_source.sampling_rate}"
if data_source.sampling_rate else "")
return StringNode(f"{data_source.table_name}{final}{sample}")
def _format_arrayjoin(
query: AbstractQuery,
formatter: ExpressionVisitor[str]) -> Optional[StringNode]:
array_join = query.get_arrayjoin()
if array_join is not None:
column_likes_joined = [el.accept(formatter) for el in array_join]
return StringNode("ARRAY JOIN {}".format(
",".join(column_likes_joined)))
return None
def _format_limitby(
query: AbstractQuery,
formatter: ClickhouseExpressionFormatterBase) -> Optional[StringNode]:
ast_limitby = query.get_limitby()
if ast_limitby is not None:
return StringNode("LIMIT {} BY {}".format(
ast_limitby.limit, ast_limitby.expression.accept(formatter)))
return None
def _format_groupby(query: AbstractQuery,
formatter: ExpressionVisitor[str]) -> Optional[StringNode]:
group_clause: Optional[StringNode] = None
ast_groupby = query.get_groupby()
if ast_groupby:
groupby_expressions = [e.accept(formatter) for e in ast_groupby]
group_clause_str = f"{', '.join(groupby_expressions)}"
if query.has_totals():
group_clause_str = f"{group_clause_str} WITH TOTALS"
group_clause = StringNode(f"GROUP BY {group_clause_str}")
return group_clause
def _format_orderby(query: AbstractQuery,
formatter: ExpressionVisitor[str]) -> Optional[StringNode]:
ast_orderby = query.get_orderby()
if ast_orderby:
orderby = [
f"{e.expression.accept(formatter)} {e.direction.value}"
for e in ast_orderby
]
return StringNode(f"ORDER BY {', '.join(orderby)}")
else:
return None
def visit_join_clause(self, node: JoinClause[Entity]) -> StringNode:
left = f"LEFT {node.left_node.accept(self)}"
type = f"TYPE {node.join_type}"
right = f"RIGHT {node.right_node.accept(self)}\n"
on = "".join(
[
f"{c.left.table_alias}.{c.left.column} {c.right.table_alias}.{c.right.column}"
for c in node.keys
]
)
return StringNode(f"{left} {type} {right} ON {on}")
def _format_limitby(query: AbstractQuery,
formatter: ExpressionVisitor[str]) -> Optional[StringNode]:
ast_limitby = query.get_limitby()
if ast_limitby is not None:
columns_accepted = [
column.accept(formatter) for column in ast_limitby.columns
]
return StringNode("LIMIT {} BY {}".format(ast_limitby.limit,
",".join(columns_accepted)))
return None
def _visit_simple_source(self, data_source: Entity) -> StringNode:
sample_val = data_source.sample
sample_str = f" SAMPLE {sample_val}" if sample_val is not None else ""
return StringNode(f"{data_source.human_readable_id}{sample_str}")
def _format_limit(
query: AbstractQuery,
formatter: ClickhouseExpressionFormatterBase) -> Optional[StringNode]:
ast_limit = query.get_limit()
return (StringNode(f"LIMIT {ast_limit} OFFSET {query.get_offset()}")
if ast_limit is not None else None)
def visit_individual_node(self, node: IndividualNode[Entity]) -> StringNode:
return StringNode(f"{node.alias}, {self.visit(node.data_source)}")
def _format_limit(query: AbstractQuery,
formatter: ExpressionVisitor[str]) -> Optional[StringNode]:
ast_limit = query.get_limit()
return (StringNode(f"LIMIT {ast_limit} OFFSET {query.get_offset()}")
if ast_limit is not None else None)
def _format_select(query: AbstractQuery,
formatter: ExpressionVisitor[str]) -> StringNode:
selected_cols = [
e.expression.accept(formatter) for e in query.get_selected_columns()
]
return StringNode(f"SELECT {', '.join(selected_cols)}")
right_node=node_group,
keys=[
JoinCondition(
left=JoinConditionExpression("err", "group_id"),
right=JoinConditionExpression("groups", "id"),
),
JoinCondition(
left=JoinConditionExpression("err", "project_id"),
right=JoinConditionExpression("groups", "project_id"),
),
],
join_type=JoinType.INNER,
join_modifier=JoinModifier.SEMI,
),
SequenceNode([
PaddingNode(None, StringNode("errors_local"), "err"),
StringNode("SEMI INNER JOIN"),
PaddingNode(None, StringNode("groupedmessage_local"), "groups"),
StringNode("ON"),
SequenceNode(
[
StringNode("err.group_id=groups.id"),
StringNode("err.project_id=groups.project_id"),
],
" AND ",
),
]),
("errors_local err SEMI INNER JOIN groupedmessage_local groups "
"ON err.group_id=groups.id AND err.project_id=groups.project_id"),
id="Simple join",
),
def test_composite_query() -> None:
query = FormattedQuery(
[
StringNode("SELECT avg(a)"),
PaddingNode(
"FROM",
FormattedSubQuery(
[
StringNode("SELECT t_a.a, t_b.b"),
PaddingNode(
"FROM",
SequenceNode(
[
PaddingNode(
None,
FormattedSubQuery(
[
StringNode("SELECT a, b"),
StringNode("FROM somewhere"),
]
),
"t_a",
),
StringNode("INNER SEMI JOIN"),
PaddingNode(
None,
FormattedSubQuery(
[
StringNode("SELECT a, b"),
StringNode("FROM somewhere_else"),
]
),
"t_b",
),
StringNode("ON t_a.a = t_b.b"),
],
),
),
],
),
),
StringNode("WHERE something something"),
],
)
assert query.get_sql(format="JSON") == (
"SELECT avg(a) FROM "
"(SELECT t_a.a, t_b.b FROM "
"(SELECT a, b FROM somewhere) t_a "
"INNER SEMI JOIN "
"(SELECT a, b FROM somewhere_else) t_b "
"ON t_a.a = t_b.b) "
"WHERE something something "
"FORMAT JSON"
)
assert query.structured() == [
"SELECT avg(a)",
[
"FROM",
[
"SELECT t_a.a, t_b.b",
[
"FROM",
[
[["SELECT a, b", "FROM somewhere"], "t_a"],
"INNER SEMI JOIN",
[["SELECT a, b", "FROM somewhere_else"], "t_b"],
"ON t_a.a = t_b.b",
],
],
],
],
"WHERE something something",
]