feat(mql): Add support for unary operator (#179)

snuba_sdk/mql/mql.py

@@ -32,8 +32,12 @@ class InvalidMQLQueryError(Exception):
 expression = term (_ expr_op _ term)*
 expr_op = "+" / "-"
 
-term = coefficient (_ term_op _ coefficient)*
+term = unary (_ term_op _ unary)*
 term_op = "*" / "/"
+
+unary = unary_op? coefficient
+unary_op = "-"
+
 coefficient = number / quoted_string / filter
 
 number = ~r"[0-9]+" ("." ~r"[0-9]+")?
@@ -109,6 +113,10 @@ class InvalidMQLQueryError(Exception):
     "/": ArithmeticOperator.DIVIDE.value,
 }
 
+UNARY_OPERATORS: Mapping[str, str] = {
+    "-": "negate",
+}
+
 
 def parse_mql(mql: str) -> Timeseries | Formula:
     """
@@ -168,21 +176,40 @@ def visit_term(
         Checks if the current node contains two term children, if so
         then merge them into a single Formula with the operator.
         """
-        coefficient_left, zero_or_more_others = children
-        assert isinstance(coefficient_left, (Formula, Timeseries, float, int, str))
+        unary_left, zero_or_more_others = children
+        assert isinstance(unary_left, (Formula, Timeseries, float, int, str))
 
         if zero_or_more_others:
             for zero_or_more in zero_or_more_others:
-                _, term_operator, _, coefficient_right, *_ = zero_or_more
-                coefficient_left = Formula(
-                    term_operator, [coefficient_left, coefficient_right]
-                )
+                _, term_operator, _, unary_right, *_ = zero_or_more
+                unary_left = Formula(term_operator, [unary_left, unary_right])
 
-        return cast(Union[Formula, Timeseries, float, int, str], coefficient_left)
+        return cast(Union[Formula, Timeseries, float, int, str], unary_left)
 
     def visit_term_op(self, node: Node, children: Sequence[Any]) -> Any:
         return TERM_OPERATORS[node.text]
 
+    def visit_unary(
+        self, node: Node, children: Sequence[Any]
+    ) -> Union[Formula, Timeseries, float, int, str]:
+        unary_op, coefficient = children
+        if unary_op:
+            if isinstance(coefficient, float) or isinstance(coefficient, int):
+                return -coefficient
+            elif isinstance(coefficient, Formula) or isinstance(
+                coefficient, Timeseries
+            ):
+                return Formula(function_name=unary_op[0], parameters=[coefficient])
+            else:
+                raise InvalidMQLQueryError(
+                    f"Unary expression not supported for type {type(coefficient)}"
+                )
+
+        return cast(Union[Formula, Timeseries, float, int, str], coefficient)
+
+    def visit_unary_op(self, node: Node, children: Sequence[Any]) -> Any:
+        return UNARY_OPERATORS[node.text]
+
     def visit_coefficient(
         self, node: Node, children: Sequence[Union[Timeseries, int, float]]
     ) -> Union[Timeseries, int, float]:

tests/test_mql.py

@@ -855,6 +855,120 @@ def test_parse_mql_base(mql_string: str, metrics_query: Formula | Timeseries) ->
         ),
         id="test expression with associativity",
     ),
+    pytest.param(
+        "-count(c:custom/page_click@none)",
+        Formula(
+            function_name="negate",
+            parameters=[
+                Timeseries(
+                    metric=Metric(mri="c:custom/page_click@none"),
+                    aggregate="count",
+                ),
+            ],
+        ),
+        id="test expression with single unary on metric",
+    ),
+    pytest.param(
+        "-(-count(c:custom/page_click@none))",
+        Formula(
+            function_name="negate",
+            parameters=[
+                Formula(
+                    function_name="negate",
+                    parameters=[
+                        Timeseries(
+                            metric=Metric(mri="c:custom/page_click@none"),
+                            aggregate="count",
+                        ),
+                    ],
+                ),
+            ],
+        ),
+        id="test expression with nested unary on metric",
+    ),
+    pytest.param(
+        "count(c:custom/page_click@none) - -1",
+        Formula(
+            function_name=ArithmeticOperator.MINUS.value,
+            parameters=[
+                Timeseries(
+                    metric=Metric(mri="c:custom/page_click@none"),
+                    aggregate="count",
+                ),
+                -1,
+            ],
+        ),
+        id="test expression with unary",
+    ),
+    pytest.param(
+        "-(count(c:custom/page_click@none) + -1)",
+        Formula(
+            function_name="negate",
+            parameters=[
+                Formula(
+                    function_name=ArithmeticOperator.PLUS.value,
+                    parameters=[
+                        Timeseries(
+                            metric=Metric(mri="c:custom/page_click@none"),
+                            aggregate="count",
+                        ),
+                        -1,
+                    ],
+                ),
+            ],
+        ),
+        id="test expression with unary on parenthesis expression",
+    ),
+    pytest.param(
+        "count(c:custom/page_click@none) + -max(d:custom/app_load@millisecond)",
+        Formula(
+            function_name=ArithmeticOperator.PLUS.value,
+            parameters=[
+                Timeseries(
+                    metric=Metric(mri="c:custom/page_click@none"),
+                    aggregate="count",
+                ),
+                Formula(
+                    function_name="negate",
+                    parameters=[
+                        Timeseries(
+                            metric=Metric(mri="d:custom/app_load@millisecond"),
+                            aggregate="max",
+                        ),
+                    ],
+                ),
+            ],
+        ),
+        id="test expression with unary on metric",
+    ),
+    pytest.param(
+        "count(c:custom/page_click@none) + (-1 + -max(d:custom/app_load@millisecond))",
+        Formula(
+            function_name=ArithmeticOperator.PLUS.value,
+            parameters=[
+                Timeseries(
+                    metric=Metric(mri="c:custom/page_click@none"),
+                    aggregate="count",
+                ),
+                Formula(
+                    function_name=ArithmeticOperator.PLUS.value,
+                    parameters=[
+                        -1,
+                        Formula(
+                            function_name="negate",
+                            parameters=[
+                                Timeseries(
+                                    metric=Metric(mri="d:custom/app_load@millisecond"),
+                                    aggregate="max",
+                                ),
+                            ],
+                        ),
+                    ],
+                ),
+            ],
+        ),
+        id="test expression with complex unary",
+    ),
 ]