feat(snql): Support array_join in snql (#29019)

This adds support for the `array_join` function in snql. Also keeps it private
so callers have to explicitly opt in to use it.

src/sentry/search/events/base.py

@@ -1,4 +1,4 @@
-from typing import List, Mapping, Set
+from typing import List, Mapping, Optional, Set
 
 from django.utils.functional import cached_property
 from snuba_sdk.aliased_expression import AliasedExpression
@@ -12,9 +12,12 @@
 
 
 class QueryBase:
-    def __init__(self, dataset: Dataset, params: ParamsType):
-        self.params = params
+    def __init__(
+        self, dataset: Dataset, params: ParamsType, functions_acl: Optional[List[str]] = None
+    ):
         self.dataset = dataset
+        self.params = params
+        self.functions_acl = set() if functions_acl is None else functions_acl
 
         # Function is a subclass of CurriedFunction
         self.where: List[WhereType] = []

src/sentry/search/events/builder.py

@@ -24,11 +24,12 @@ def __init__(
         orderby: Optional[List[str]] = None,
         auto_aggregations: bool = False,
         use_aggregate_conditions: bool = False,
+        functions_acl: Optional[List[str]] = None,
         limit: Optional[int] = 50,
         offset: Optional[int] = 0,
         limitby: Optional[Tuple[str, int]] = None,
     ):
-        super().__init__(dataset, params)
+        super().__init__(dataset, params, functions_acl)
 
         # TODO: implement this in `resolve_select`
         self.auto_aggregations = auto_aggregations

src/sentry/search/events/fields.py

@@ -2166,8 +2166,10 @@ def validate(self):
 class QueryFields(QueryBase):
     """Field logic for a snql query"""
 
-    def __init__(self, dataset: Dataset, params: ParamsType):
-        super().__init__(dataset, params)
+    def __init__(
+        self, dataset: Dataset, params: ParamsType, functions_acl: Optional[List[str]] = None
+    ):
+        super().__init__(dataset, params, functions_acl)
 
         self.field_alias_converter: Mapping[str, Callable[[str], SelectType]] = {
             # NOTE: `ISSUE_ALIAS` simply maps to the id, meaning that post processing
@@ -2620,8 +2622,16 @@ def __init__(self, dataset: Dataset, params: ParamsType):
                     ),
                     default_result_type="number",
                 ),
+                SnQLFunction(
+                    "array_join",
+                    required_args=[StringArrayColumn("column")],
+                    snql_aggregate=lambda args, alias: Function(
+                        "arrayJoin", [args["column"]], alias
+                    ),
+                    default_result_type="string",
+                    private=True,
+                ),
                 # TODO: implement these
-                SnQLFunction("array_join", snql_aggregate=self._resolve_unimplemented_function),
                 SnQLFunction("histogram", snql_aggregate=self._resolve_unimplemented_function),
                 SnQLFunction("percentage", snql_aggregate=self._resolve_unimplemented_function),
                 SnQLFunction("t_test", snql_aggregate=self._resolve_unimplemented_function),
@@ -2768,7 +2778,9 @@ def resolve_function(self, function: str, match: Optional[Match[str]] = None) ->
             raise NotImplementedError("Aggregate aliases not implemented in snql field parsing yet")
 
         name, arguments, alias = self.parse_function(match)
-        snql_function = self.function_converter.get(name)
+        snql_function = self.function_converter[name]
+        if not snql_function.is_accessible(self.functions_acl):
+            raise InvalidSearchQuery(f"{snql_function.name}: no access to private function")
 
         arguments = snql_function.format_as_arguments(name, arguments, self.params)
         for arg in snql_function.args:

src/sentry/search/events/filter.py

@@ -1037,8 +1037,10 @@ def format_search_filter(term, params):
 class QueryFilter(QueryFields):
     """Filter logic for a snql query"""
 
-    def __init__(self, dataset: Dataset, params: ParamsType):
-        super().__init__(dataset, params)
+    def __init__(
+        self, dataset: Dataset, params: ParamsType, functions_acl: Optional[List[str]] = None
+    ):
+        super().__init__(dataset, params, functions_acl)
 
         self.search_filter_converter: Mapping[
             str, Callable[[SearchFilter], Optional[WhereType]]

src/sentry/snuba/discover.py

@@ -240,6 +240,7 @@ def query(
             orderby=orderby,
             auto_aggregations=auto_aggregations,
             use_aggregate_conditions=use_aggregate_conditions,
+            functions_acl=functions_acl,
             limit=limit,
             offset=offset,
         )

tests/sentry/snuba/test_discover.py

@@ -2937,14 +2937,49 @@ def test_count_with_or(self):
         assert data[0]["transaction"] == "a" * 32
         assert data[0]["count"] == 1
 
-    def test_access_to_private_functions(self):
-        # using private functions directly without access should error
-        with pytest.raises(InvalidSearchQuery, match="array_join: no access to private function"):
-            discover.query(
-                selected_columns=["array_join(tags.key)"],
+    def test_array_join(self):
+        data = load_data("transaction", timestamp=before_now(seconds=3))
+        data["measurements"] = {
+            "fp": {"value": 1000},
+            "fcp": {"value": 1000},
+            "lcp": {"value": 1000},
+        }
+        self.store_event(data=data, project_id=self.project.id)
+
+        for use_snql in [False, True]:
+            results = discover.query(
+                selected_columns=["array_join(measurements_key)"],
                 query="",
-                params={"project_id": [self.project.id]},
+                params={
+                    "project_id": [self.project.id],
+                    "start": self.two_min_ago,
+                    "end": self.now,
+                },
+                functions_acl=["array_join"],
+                use_snql=use_snql,
             )
+            assert {"fcp", "fp", "lcp"} == {
+                row["array_join_measurements_key"] for row in results["data"]
+            }
+
+    def test_access_to_private_functions(self):
+        for use_snql in [False, True]:
+            # using private functions directly without access should error
+            with pytest.raises(
+                InvalidSearchQuery, match="array_join: no access to private function"
+            ):
+                discover.query(
+                    selected_columns=["array_join(tags.key)"],
+                    query="",
+                    params={
+                        "project_id": [self.project.id],
+                        "start": self.two_min_ago,
+                        "end": self.now,
+                    },
+                    use_snql=use_snql,
+                )
+
+        # TODO: test the following with `use_snql=True` once histogram is using snql
 
         # using private functions in an aggregation without access should error
         with pytest.raises(InvalidSearchQuery, match="histogram: no access to private function"):