ref(api): parse pipeline tests, snql pipeline before treeify (#5886)

snuba/query/dsl.py

@@ -151,24 +151,12 @@ def equals(
  return binary_condition("equals", left, right)
 
 
-def and_cond(lhs: FunctionCall, rhs: FunctionCall, *args: FunctionCall) -> FunctionCall:
- """
- if only lhs and rhs are given, return and(lhs, rhs)
- otherwise (more than 2 conditions are given), returns and(lhs, and(rhs, and(...)))
- """
- if len(args) == 0:
- return binary_condition("and", lhs, rhs)
-
- sofar = args[len(args) - 1]
- for i in range(len(args) - 2, -1, -1):
- sofar = binary_condition("and", args[i], sofar)
- sofar = binary_condition("and", rhs, sofar)
- sofar = binary_condition("and", lhs, sofar)
- return sofar
+def and_cond(lhs: Expression, rhs: Expression, *args: Expression) -> FunctionCall:
+ return FunctionCall(None, "and", (lhs, rhs, *args))
 
 
-def or_cond(lhs: FunctionCall, rhs: FunctionCall) -> FunctionCall:
- return binary_condition("or", lhs, rhs)
+def or_cond(lhs: Expression, rhs: Expression, *args: Expression) -> FunctionCall:
+ return FunctionCall(None, "or", (lhs, rhs, *args))
 
 
 def in_cond(lhs: Expression, rhs: Expression) -> FunctionCall:

snuba/query/dsl_mapper.py

@@ -2,7 +2,7 @@
 
 from snuba.clickhouse.query import Query as ClickhouseQuery
 from snuba.query import LimitBy, OrderBy, SelectedExpression
-from snuba.query.conditions import get_first_level_or_conditions
+from snuba.query.composite import CompositeQuery
 from snuba.query.data_source.simple import Entity
 from snuba.query.expressions import (
  Argument,
@@ -42,8 +42,7 @@ def and_cond_repr(exp: Expression, visitor: ExpressionVisitor[str]) -> str:
 
 def or_cond_repr(exp: Expression, visitor: ExpressionVisitor[str]) -> str:
  assert isinstance(exp, FunctionCall)
- conditions = get_first_level_or_conditions(exp)
- parameters = ", ".join([arg.accept(visitor) for arg in conditions])
+ parameters = ", ".join([arg.accept(visitor) for arg in exp.parameters])
  return f"or_cond({parameters})"
 
 
@@ -225,7 +224,7 @@ def ast_repr(
  return f"[{', '.join(strings)}]"
 
 
-def query_repr(query: LogicalQuery | ClickhouseQuery) -> str:
+def query_repr(query: LogicalQuery | ClickhouseQuery | CompositeQuery[Entity]) -> str:
  visitor = DSLMapperVisitor()
  selected = ast_repr(query.get_selected_columns(), visitor)
  arrayjoin = ast_repr(query.get_arrayjoin(), visitor)

tests/pipeline/test_entity_processing_stage.py

@@ -76,8 +76,10 @@ def test_basic(
  selected_columns=[SelectedExpression("timestamp", column("timestamp"))],
  condition=and_cond(
  equals(literal(1), literal(1)),
- equals(column("org_id"), literal(1)),
- equals(column("project_id"), literal(1)),
+ and_cond(
+ equals(column("org_id"), literal(1)),
+ equals(column("project_id"), literal(1)),
+ ),
  ),
  limit=1000,
  ),

tests/query/parser/test_mql_query.py

@@ -127,27 +127,47 @@
  f.equals(
  column("granularity", None, "_snuba_granularity"), literal(60)
  ),
- in_cond(
- column("project_id", None, "_snuba_project_id"),
- f.tuple(literal(11)),
- ),
- in_cond(column("org_id", None, "_snuba_org_id"), f.tuple(literal(1))),
- f.equals(
- column("use_case_id", None, "_snuba_use_case_id"),
- literal("transactions"),
- ),
- f.greaterOrEquals(
- column("timestamp", None, "_snuba_timestamp"),
- literal(datetime(2023, 11, 23, 18, 30)),
- ),
- f.less(
- column("timestamp", None, "_snuba_timestamp"),
- literal(datetime(2023, 11, 23, 22, 30)),
- ),
- f.equals(
- column("metric_id", None, "_snuba_metric_id"), literal(123456)
+ and_cond(
+ in_cond(
+ column("project_id", None, "_snuba_project_id"),
+ f.tuple(literal(11)),
+ ),
+ and_cond(
+ in_cond(
+ column("org_id", None, "_snuba_org_id"), f.tuple(literal(1))
+ ),
+ and_cond(
+ f.equals(
+ column("use_case_id", None, "_snuba_use_case_id"),
+ literal("transactions"),
+ ),
+ and_cond(
+ f.greaterOrEquals(
+ column("timestamp", None, "_snuba_timestamp"),
+ literal(datetime(2023, 11, 23, 18, 30)),
+ ),
+ and_cond(
+ f.less(
+ column("timestamp", None, "_snuba_timestamp"),
+ literal(datetime(2023, 11, 23, 22, 30)),
+ ),
+ and_cond(
+ f.equals(
+ column(
+ "metric_id", None, "_snuba_metric_id"
+ ),
+ literal(123456),
+ ),
+ in_cond(
+ tags_raw["888"],
+ f.tuple(literal("dist1"), literal("dist2")),
+ ),
+ ),
+ ),
+ ),
+ ),
+ ),
  ),
- in_cond(tags_raw["888"], f.tuple(literal("dist1"), literal("dist2"))),
  ),
  order_by=[
  OrderBy(