feat(seer): Return custom attribute context from get_attribute_names

src/sentry/seer/assisted_query/traces_tools.py

@@ -6,9 +6,9 @@
 from sentry.models.apikey import ApiKey
 from sentry.models.organization import Organization
 from sentry.seer.sentry_data_models import (
+    AttributeMeta,
     AttributeNamesResponse,
     AttributeValuesResponse,
-    BuiltInField,
 )
 
 logger = logging.getLogger(__name__)
@@ -124,8 +124,9 @@ def get_attribute_names(
         item_type: Type of trace item (default: "spans", can be "spans", "logs", etc.)
         include_context: When True, include the context metadata (brief,
             examples, deprecation, etc.) for each attribute and attach it to the
-            matching built-in fields in the response. Today the metadata comes
-            from the sentry conventions; custom attribute context is planned.
+            matching field in the response. Sentry-owned attributes get their
+            metadata from the sentry conventions or a column definition; custom
+            attributes get user-authored metadata.
 
     Returns:
         Dictionary with attributes:
@@ -138,25 +139,28 @@ def get_attribute_names(
                 {"key": "span.op", "type": "string", "context": {...}},
                 {"key": "span.duration", "type": "number", "context": None},
                 ...
+            ],
+            "custom_fields": [
+                {"key": "custom_field_1", "type": "string", "context": {...}},
+                ...
             ]
         }
 
-        Each built-in field's "context" is only populated when expand="context"
-        is requested (and the attribute maps to a known convention); otherwise it
-        is None. Convention-backed attributes that aren't hardcoded built-ins
-        (e.g. http.route) are also appended to "built_in_fields" so their context
-        isn't lost.
+        A field's "context" is only populated when expand="context" is requested
+        and metadata exists for the attribute; otherwise it is None.
+        Convention-backed attributes that aren't hardcoded built-ins (e.g.
+        http.route) are appended to "built_in_fields" so their context isn't lost.
+        Custom attributes go to "custom_fields" when they have authored context,
+        and appear only in "fields" otherwise.
     """
     organization = Organization.objects.get(id=org_id)
 
     api_key = ApiKey(organization_id=org_id, scope_list=API_KEY_SCOPES)
 
     fields: dict[str, list[str]] = {"string": [], "number": []}
-    # Maps an attribute name to its context, populated only when the caller
-    # passes include_context=True. Used below to attach context to the built-in
-    # fields (which is where Seer reads attribute context from). Today the
-    # context comes from the sentry conventions, but custom attribute context is
-    # planned, at which point user-defined attributes will be populated too.
+    # Maps an attribute name to its context, populated only when the caller passes
+    # include_context=True. Attached below to the built-in and custom fields, which
+    # is where Seer reads attribute context from.
     context_by_name: dict[str, dict[str, Any]] = {}
     # Maps an attribute name to its type ("string"/"number"), so context-bearing
     # attributes that aren't hardcoded built-ins can still be emitted as fields.
@@ -203,24 +207,27 @@ def get_attribute_names(
 
     hardcoded_fields = _get_built_in_fields(item_type)
     built_in_fields = [
-        BuiltInField(**f, context=context_by_name.get(f["key"])) for f in hardcoded_fields
+        AttributeMeta(**f, context=context_by_name.get(f["key"])) for f in hardcoded_fields
     ]
 
     # Context-bearing attributes that aren't hardcoded built-ins (e.g. the
-    # convention http.route, or Sentry-defined fields like span.description) would
-    # otherwise be dropped. Surface them through built_in_fields too, since that's
-    # where Seer reads attribute context from. We promote conventions
-    # (isConvention=True) and Sentry-defined attributes (source=sentry);
-    # user-authored custom context is intentionally left out.
+    # convention http.route) would otherwise be dropped. Conventions and
+    # Sentry-defined fields join built_in_fields; user-authored context goes to
+    # custom_fields, being org-specific rather than part of Sentry's schema.
     hardcoded_keys = {f["key"] for f in hardcoded_fields}
+    custom_fields = []
     for name, context in context_by_name.items():
         if name in hardcoded_keys:
             continue
-        if not (context.get("isConvention") or source_by_name.get(name) == "sentry"):
-            continue
-        built_in_fields.append(BuiltInField(key=name, type=type_by_name[name], context=context))
-
-    return AttributeNamesResponse(fields=fields, built_in_fields=built_in_fields)
+        field = AttributeMeta(key=name, type=type_by_name[name], context=context)
+        if context.get("isConvention") or source_by_name.get(name) == "sentry":
+            built_in_fields.append(field)
+        elif context.get("isCustom"):
+            custom_fields.append(field)
+
+    return AttributeNamesResponse(
+        fields=fields, built_in_fields=built_in_fields, custom_fields=custom_fields
+    )
 
 
 def get_attribute_values_with_substring(

src/sentry/seer/sentry_data_models.py

@@ -205,21 +205,28 @@ class SpanAttributesResponse(BaseModel):
     attributes: list[SpanAttribute]
 
 
-class BuiltInField(BaseModel):
+class AttributeMeta(BaseModel):
+    """An attribute's name, type, and optional context, as listed in
+    ``AttributeNamesResponse``."""
+
     key: str
     type: str
-    # Attribute metadata (brief, examples, isDeprecated, replacementAttribute,
-    # ...) for the attribute, populated when the caller requests
-    # `expand="context"`; otherwise None. Today the metadata comes from the
-    # sentry conventions, so only attributes that map to a known convention
-    # carry it, but custom attribute context is planned and will populate this
-    # for user-defined attributes too.
+    # Metadata (brief, examples, isDeprecated, replacementAttribute, ...),
+    # populated when the caller requests `expand="context"`; otherwise None.
+    # Sourced from the sentry conventions or a column definition for Sentry-owned
+    # attributes, and authored by the org for custom ones.
     context: dict[str, Any] | None = None
 
 
 class AttributeNamesResponse(BaseModel):
     fields: dict[str, list[str]]
-    built_in_fields: list[BuiltInField]
+    # Sentry-owned attributes (conventions and Sentry-defined fields), whose
+    # context comes from the sentry conventions or a column definition.
+    built_in_fields: list[AttributeMeta]
+    # Custom (user-defined) attributes that carry user-authored context. Only
+    # populated when the caller requests context and the organization has
+    # authored some; attributes without context are still listed in `fields`.
+    custom_fields: list[AttributeMeta] = []
 
 
 class AttributeBucket(BaseModel):
@@ -238,7 +245,7 @@ class MetricMetadataRow(BaseModel):
     count: int
     # Authored context (brief, details) for the metric, populated only when the
     # caller passes include_context=True (and the metric has context); otherwise
-    # None. Mirrors the attributes context shape (see BuiltInField.context).
+    # None. Mirrors the attributes context shape (see AttributeMeta.context).
     context: dict[str, Any] | None = None
 
 

tests/sentry/seer/assisted_query/test_traces_tools.py

@@ -0,0 +1,145 @@
+from unittest.mock import MagicMock, patch
+
+from sentry.seer.assisted_query.traces_tools import get_attribute_names
+from sentry.testutils.cases import TestCase
+
+
+def _attribute(name, attr_type="string", source="user", context=None):
+    attribute = {
+        "key": name,
+        "name": name,
+        "attributeType": attr_type,
+        "attributeSource": {"source_type": source},
+    }
+    if context is not None:
+        attribute["context"] = context
+    return attribute
+
+
+class TestGetAttributeNames(TestCase):
+    def setUp(self) -> None:
+        super().setUp()
+        self.org = self.create_organization()
+        self.project = self.create_project(organization=self.org)
+
+    def _mock_responses(self, mock_client_cls, string_attrs, number_attrs=None):
+        """Stub the two per-type calls get_attribute_names makes (string, then number)."""
+        mock_client = mock_client_cls.return_value
+        string_response, number_response = MagicMock(), MagicMock()
+        string_response.data = string_attrs
+        number_response.data = number_attrs or []
+        mock_client.get.side_effect = [string_response, number_response]
+        return mock_client
+
+    @patch("sentry.seer.assisted_query.traces_tools.ApiClient")
+    def test_custom_context_goes_to_custom_fields(self, mock_client_cls: MagicMock) -> None:
+        self._mock_responses(
+            mock_client_cls,
+            [
+                # A convention, which belongs in built_in_fields.
+                _attribute(
+                    "device.class",
+                    context={"isConvention": True, "brief": "Device class"},
+                ),
+                # User-authored context, which belongs in custom_fields.
+                _attribute(
+                    "my_custom_attr",
+                    context={
+                        "isCustom": True,
+                        "brief": "Set by the checkout service",
+                        "details": ["Only on payment spans."],
+                        "examples": ["visa"],
+                    },
+                ),
+                # No context at all, so it appears only in `fields`.
+                _attribute("undescribed_attr", context={}),
+            ],
+        )
+
+        result = get_attribute_names(
+            org_id=self.org.id,
+            project_ids=[self.project.id],
+            stats_period="7d",
+            include_context=True,
+        )
+
+        assert [field.key for field in result.custom_fields] == ["my_custom_attr"]
+        assert result.custom_fields[0].type == "string"
+        assert result.custom_fields[0].context == {
+            "isCustom": True,
+            "brief": "Set by the checkout service",
+            "details": ["Only on payment spans."],
+            "examples": ["visa"],
+        }
+
+        built_in_keys = {field.key for field in result.built_in_fields}
+        assert "device.class" in built_in_keys
+        # Custom context must never leak into the Sentry-owned fields.
+        assert "my_custom_attr" not in built_in_keys
+        assert "undescribed_attr" not in built_in_keys
+
+        # Every attribute is still listed in `fields`, with or without context.
+        assert result.fields["string"] == [
+            "device.class",
+            "my_custom_attr",
+            "undescribed_attr",
+        ]
+
+    @patch("sentry.seer.assisted_query.traces_tools.ApiClient")
+    def test_number_typed_custom_context(self, mock_client_cls: MagicMock) -> None:
+        self._mock_responses(
+            mock_client_cls,
+            [],
+            [_attribute("cart_total", "number", context={"isCustom": True, "brief": "In cents"})],
+        )
+
+        result = get_attribute_names(
+            org_id=self.org.id,
+            project_ids=[self.project.id],
+            stats_period="7d",
+            include_context=True,
+        )
+
+        assert len(result.custom_fields) == 1
+        assert result.custom_fields[0].key == "cart_total"
+        assert result.custom_fields[0].type == "number"
+
+    @patch("sentry.seer.assisted_query.traces_tools.ApiClient")
+    def test_sentry_source_attribute_stays_built_in(self, mock_client_cls: MagicMock) -> None:
+        # A Sentry-defined attribute carries isConvention=False, so it is routed by
+        # its `sentry` source rather than being mistaken for custom context.
+        self._mock_responses(
+            mock_client_cls,
+            [
+                _attribute(
+                    "span.description",
+                    source="sentry",
+                    context={"isConvention": False, "brief": "Span description"},
+                )
+            ],
+        )
+
+        result = get_attribute_names(
+            org_id=self.org.id,
+            project_ids=[self.project.id],
+            stats_period="7d",
+            include_context=True,
+        )
+
+        assert "span.description" in {field.key for field in result.built_in_fields}
+        assert result.custom_fields == []
+
+    @patch("sentry.seer.assisted_query.traces_tools.ApiClient")
+    def test_no_context_requested(self, mock_client_cls: MagicMock) -> None:
+        mock_client = self._mock_responses(mock_client_cls, [_attribute("my_custom_attr")])
+
+        result = get_attribute_names(
+            org_id=self.org.id,
+            project_ids=[self.project.id],
+            stats_period="7d",
+        )
+
+        assert result.custom_fields == []
+        # Context is only requested from the endpoint when asked for.
+        _args, kwargs = mock_client.get.call_args
+        assert "expand" not in kwargs["params"]