feat(processor): Handle casted columns in UUID column processor (#7552)

The UUID column processor handles simple conditions like `equals(column,
'value')` but fails when the column is casted like `equals(cast(column,
'String'), 'value')`. This adds support for the casted case

Author: Zylphrex

snuba/query/processors/physical/type_converters.py

@@ -69,27 +69,50 @@ def __init__(self, columns: Set[str], optimize_ordering: bool = False):
         )
 
         col = Param("col", ColumnMatch(None, column_match))
+        casted_col = FunctionCallMatch(String("cast"), (col, LiteralMatch(String("String"))))
 
         self.__condition_matcher = Or(
             [
                 FunctionCallMatch(operator, (literal, col)),
                 FunctionCallMatch(operator, (col, literal)),
                 FunctionCallMatch(Param("operator", String("has")), (col, literal)),
+                # These are the same as above but the column is wrapped in `cast(..., 'String')`
+                FunctionCallMatch(operator, (casted_col, literal)),
+                FunctionCallMatch(operator, (literal, casted_col)),
+                FunctionCallMatch(Param("operator", String("has")), (casted_col, literal)),
             ]
         )
 
-        self.__in_condition_matcher = FunctionCallMatch(
-            in_operators,
-            (
-                col,
-                Param(
-                    "literals",
-                    FunctionCallMatch(
-                        Or([String("array"), String("tuple")]),
-                        all_parameters=LiteralMatch(),
+        self.__in_condition_matcher = Or(
+            [
+                FunctionCallMatch(
+                    in_operators,
+                    (
+                        col,
+                        Param(
+                            "literals",
+                            FunctionCallMatch(
+                                Or([String("array"), String("tuple")]),
+                                all_parameters=LiteralMatch(),
+                            ),
+                        ),
                     ),
                 ),
-            ),
+                # These are the same as above but the column is wrapped in `cast(..., 'String')`
+                FunctionCallMatch(
+                    in_operators,
+                    (
+                        casted_col,
+                        Param(
+                            "literals",
+                            FunctionCallMatch(
+                                Or([String("array"), String("tuple")]),
+                                all_parameters=LiteralMatch(),
+                            ),
+                        ),
+                    ),
+                ),
+            ]
         )
 
         self.__unoptimizable_condition_matcher = Or(
@@ -100,9 +123,7 @@ def __init__(self, columns: Set[str], optimize_ordering: bool = False):
         )
 
     def process_query(self, query: Query, query_settings: QuerySettings) -> None:
-        query.transform_expressions(
-            self._process_expressions, skip_transform_condition=True
-        )
+        query.transform_expressions(self._process_expressions, skip_transform_condition=True)
 
         condition = query.get_condition()
         if condition is not None:
@@ -117,9 +138,7 @@ def process_query(self, query: Query, query_settings: QuerySettings) -> None:
 
     def __strip_column_alias(self, exp: Expression) -> Expression:
         assert isinstance(exp, Column)
-        return Column(
-            alias=None, table_name=exp.table_name, column_name=exp.column_name
-        )
+        return Column(alias=None, table_name=exp.table_name, column_name=exp.column_name)
 
     def __contains_unoptimizable_condition(self, exp: Expression) -> bool:
         """
@@ -144,9 +163,7 @@ def assert_literal(lit: Expression) -> Literal:
                 match.string("operator"),
                 (
                     self.__strip_column_alias(match.expression("col")),
-                    self._translate_literal(
-                        assert_literal(match.expression("literal"))
-                    ),
+                    self._translate_literal(assert_literal(match.expression("literal"))),
                 ),
             )
 

tests/query/processors/test_uuid_column_processor.py

@@ -81,12 +81,8 @@
                 None,
                 "tuple",
                 (
-                    Literal(
-                        None, str(uuid.UUID("a7d67cf7-9677-4551-a95b-e6543cacd459"))
-                    ),
-                    Literal(
-                        None, str(uuid.UUID("a7d67cf7-9677-4551-a95b-e6543cacd45a"))
-                    ),
+                    Literal(None, str(uuid.UUID("a7d67cf7-9677-4551-a95b-e6543cacd459"))),
+                    Literal(None, str(uuid.UUID("a7d67cf7-9677-4551-a95b-e6543cacd45a"))),
                 ),
             ),
         ),
@@ -231,6 +227,34 @@
         "equals(column2, 'a7d67cf7-9677-4551-a95b-e6543cacd460') AND equals(column1, 'a7d67cf7-9677-4551-a95b-e6543cacd459')",
         id="equals(column2, 'a7d67cf7-9677-4551-a95b-e6543cacd460') AND equals(column1, 'a7d67cf7-9677-4551-a95b-e6543cacd459')",
     ),
+    pytest.param(
+        binary_condition(
+            ConditionFunctions.EQ,
+            FunctionCall(None, "cast", (Column(None, None, "column1"), Literal(None, "String"))),
+            Literal(None, "a7d67cf796774551a95be6543cacd459"),
+        ),
+        binary_condition(
+            ConditionFunctions.EQ,
+            Column(None, None, "column1"),
+            Literal(None, str(uuid.UUID("a7d67cf7-9677-4551-a95b-e6543cacd459"))),
+        ),
+        "equals(column1, 'a7d67cf7-9677-4551-a95b-e6543cacd459')",
+        id="equals(column1, 'a7d67cf7-9677-4551-a95b-e6543cacd459')",
+    ),
+    pytest.param(
+        binary_condition(
+            ConditionFunctions.IN,
+            FunctionCall(None, "cast", (Column(None, None, "column1"), Literal(None, "String"))),
+            FunctionCall(None, "tuple", (Literal(None, "a7d67cf796774551a95be6543cacd459"),)),
+        ),
+        binary_condition(
+            ConditionFunctions.IN,
+            Column(None, None, "column1"),
+            FunctionCall(None, "tuple", (Literal(None, "a7d67cf7-9677-4551-a95b-e6543cacd459"),)),
+        ),
+        "in(column1, tuple('a7d67cf7-9677-4551-a95b-e6543cacd459'))",
+        id="in(column1, tuple('a7d67cf7-9677-4551-a95b-e6543cacd459'))",
+    ),
 ]