filter to only the attributes available in the groupby

Author: aayush-se

src/sentry/api/endpoints/organization_trace_item_attributes_ranked.py

@@ -5,11 +5,13 @@
 
 from rest_framework.request import Request
 from rest_framework.response import Response
+from sentry_protos.snuba.v1.endpoint_trace_item_attributes_pb2 import TraceItemAttributeNamesRequest
 from sentry_protos.snuba.v1.endpoint_trace_item_stats_pb2 import (
     AttributeDistributionsRequest,
     StatsType,
     TraceItemStatsRequest,
 )
+from sentry_protos.snuba.v1.trace_item_attribute_pb2 import AttributeKey
 
 from sentry import features
 from sentry.api.api_owners import ApiOwner
@@ -27,7 +29,7 @@
 from sentry.seer.workflows.compare import keyed_rrf_score
 from sentry.snuba.referrer import Referrer
 from sentry.snuba.spans_rpc import Spans
-from sentry.utils.snuba_rpc import trace_item_stats_rpc
+from sentry.utils.snuba_rpc import attribute_names_rpc, trace_item_stats_rpc
 
 logger = logging.getLogger(__name__)
 
@@ -120,6 +122,20 @@ def get(self, request: Request, organization: Organization) -> Response:
         if query_1 == query_2:
             return Response({"rankedAttributes": []})
 
+        # First, fetch attribute names to use as an allowlist
+        attribute_names_request = TraceItemAttributeNamesRequest(
+            meta=meta,
+            type=AttributeKey.Type.TYPE_STRING,
+            limit=1000,
+        )
+        attribute_names_response = attribute_names_rpc(attribute_names_request)
+
+        allowed_attribute_names = {
+            attr.name
+            for attr in attribute_names_response.attributes
+            if attr.name and can_expose_attribute(attr.name, SupportedTraceItemType.SPANS)
+        }
+
         cohort_1, _, _ = resolver.resolve_query(query_1)
         cohort_1_request = TraceItemStatsRequest(
             filter=cohort_1,
@@ -199,7 +215,10 @@ def get(self, request: Request, organization: Organization) -> Response:
         processed_cohort_2_buckets = set()
 
         for attribute in cohort_2_data.results[0].attribute_distributions.attributes:
-            if not can_expose_attribute(attribute.attribute_name, SupportedTraceItemType.SPANS):
+            # Filter attributes to only those in the allowlist
+            if attribute.attribute_name not in allowed_attribute_names or not can_expose_attribute(
+                attribute.attribute_name, SupportedTraceItemType.SPANS
+            ):
                 continue
 
             for bucket in attribute.buckets:
@@ -208,7 +227,10 @@ def get(self, request: Request, organization: Organization) -> Response:
                 )
 
         for attribute in cohort_1_data.results[0].attribute_distributions.attributes:
-            if not can_expose_attribute(attribute.attribute_name, SupportedTraceItemType.SPANS):
+            # Filter attributes to only those in the allowlist
+            if attribute.attribute_name not in allowed_attribute_names or not can_expose_attribute(
+                attribute.attribute_name, SupportedTraceItemType.SPANS
+            ):
                 continue
             for bucket in attribute.buckets:
                 cohort_1_distribution.append((attribute.attribute_name, bucket.label, bucket.value))

tests/snuba/api/endpoints/test_organization_trace_item_attributes_ranked.py

@@ -246,8 +246,9 @@ def capture_compare(*args, **kwargs):
 
     @patch("sentry.api.endpoints.organization_trace_item_attributes_ranked.compare_distributions")
     @patch("sentry.api.endpoints.organization_trace_item_attributes_ranked.trace_item_stats_rpc")
+    @patch("sentry.api.endpoints.organization_trace_item_attributes_ranked.attribute_names_rpc")
     def test_filters_out_internal_and_private_attributes(
-        self, mock_trace_item_stats_rpc, mock_compare_distributions
+        self, mock_attribute_names_rpc, mock_trace_item_stats_rpc, mock_compare_distributions
     ) -> None:
         """Test that internal and private attributes are filtered from the response.
 
@@ -336,6 +337,23 @@ def create_mock_response():
         # Both cohorts return the same attributes
         mock_trace_item_stats_rpc.return_value = create_mock_response()
 
+        # Mock the attribute names RPC to return all the attribute names (including internal ones)
+        mock_attr_names_response = MagicMock()
+        mock_attributes = []
+        for attr_name in [
+            "sentry.device",
+            "browser",
+            "sentry.links",
+            "sentry._meta.fields.attributes.browser",
+            "sentry._internal.some_metric",
+            "__sentry_internal.debug_info",
+        ]:
+            attr = MagicMock()
+            attr.name = attr_name
+            mock_attributes.append(attr)
+        mock_attr_names_response.attributes = mock_attributes
+        mock_attribute_names_rpc.return_value = mock_attr_names_response
+
         mock_compare_distributions.return_value = {
             "results": [
                 ("sentry.device", 0.9),