feat(trace-items): Autocomplete for semver attributes

src/sentry/api/endpoints/organization_trace_item_attributes.py

@@ -1,3 +1,4 @@
+from collections.abc import Callable
 from datetime import datetime, timedelta
 from typing import Literal
 
@@ -25,6 +26,10 @@
 from sentry.api.serializers import serialize
 from sentry.api.utils import handle_query_errors
 from sentry.models.organization import Organization
+from sentry.models.release import Release
+from sentry.models.releaseenvironment import ReleaseEnvironment
+from sentry.models.releaseprojectenvironment import ReleaseStages
+from sentry.models.releases.release_project import ReleaseProject
 from sentry.search.eap import constants
 from sentry.search.eap.columns import ColumnDefinitions
 from sentry.search.eap.ourlogs.definitions import OURLOG_DEFINITIONS
@@ -37,6 +42,13 @@
     translate_internal_to_public_alias,
     translate_to_sentry_conventions,
 )
+from sentry.search.events.constants import (
+    RELEASE_STAGE_ALIAS,
+    SEMVER_ALIAS,
+    SEMVER_BUILD_ALIAS,
+    SEMVER_PACKAGE_ALIAS,
+)
+from sentry.search.events.filter import _flip_field_sort
 from sentry.search.events.types import SnubaParams
 from sentry.snuba.referrer import Referrer
 from sentry.tagstore.types import TagValue
@@ -346,6 +358,16 @@ def __init__(
             params=snuba_params, config=SearchResolverConfig(), definitions=definitions
         )
         self.search_type, self.attribute_key = self.resolve_attribute_key(key, snuba_params)
+        self.autocomplete_function: dict[str, Callable[[], list[TagValue]]] = (
+            {key: self.project_id_autocomplete_function for key in self.PROJECT_ID_KEYS}
+            | {key: self.project_slug_autocomplete_function for key in self.PROJECT_SLUG_KEYS}
+            | {
+                RELEASE_STAGE_ALIAS: self.release_stage_autocomplete_function,
+                SEMVER_ALIAS: self.semver_autocomplete_function,
+                SEMVER_BUILD_ALIAS: self.semver_build_autocomplete_function,
+                SEMVER_PACKAGE_ALIAS: self.semver_package_autocomplete_function,
+            }
+        )
 
     def resolve_attribute_key(
         self, key: str, snuba_params: SnubaParams
@@ -354,11 +376,10 @@ def resolve_attribute_key(
         return resolved_attr.search_type, resolved_attr.proto_definition
 
     def execute(self) -> list[TagValue]:
-        if self.key in self.PROJECT_ID_KEYS:
-            return self.project_id_autocomplete_function()
+        func = self.autocomplete_function.get(self.key)
 
-        if self.key in self.PROJECT_SLUG_KEYS:
-            return self.project_slug_autocomplete_function()
+        if func is not None:
+            return func()
 
         if self.search_type == "boolean":
             return self.boolean_autocomplete_function()
@@ -368,6 +389,146 @@ def execute(self) -> list[TagValue]:
 
         return []
 
+    def release_stage_autocomplete_function(self):
+        return [
+            TagValue(
+                key=self.key,
+                value=stage.value,
+                times_seen=None,
+                first_seen=None,
+                last_seen=None,
+            )
+            for stage in ReleaseStages
+            if not self.query or self.query in stage.value
+        ]
+
+    def semver_autocomplete_function(self):
+        versions = Release.objects.filter(version__contains="@" + self.query)
+
+        project_ids = self.snuba_params.project_ids
+        if project_ids:
+            release_projects = ReleaseProject.objects.filter(project_id__in=project_ids)
+            versions = versions.filter(id__in=release_projects.values_list("release_id", flat=True))
+
+        environment_ids = self.snuba_params.environment_ids
+        if environment_ids:
+            release_environments = ReleaseEnvironment.objects.filter(
+                environment_id__in=environment_ids
+            )
+            versions = versions.filter(
+                id__in=release_environments.values_list("release_id", flat=True)
+            )
+
+        order_by = map(_flip_field_sort, Release.SEMVER_COLS + ["package"])
+        versions = versions.filter_to_semver()  # type: ignore[attr-defined]  # mypy doesn't know about ReleaseQuerySet
+        versions = versions.annotate_prerelease_column()  # type: ignore[attr-defined]  # mypy doesn't know about ReleaseQuerySet
+        versions = versions.order_by(*order_by)
+
+        seen = set()
+        formatted_versions = []
+        # We want to format versions here in a way that makes sense for autocomplete. So we
+        # - Only include package if we think the user entered a package
+        # - Exclude build number, since it's not used as part of filtering
+        # When we don't include package, this can result in duplicate version numbers, so we
+        # also de-dupe here. This can result in less than 1000 versions returned, but we
+        # typically use very few values so this works ok.
+        for version in versions.values_list("version", flat=True)[:1000]:
+            formatted_version = version.split("@", 1)[1]
+            formatted_version = formatted_version.split("+", 1)[0]
+            if formatted_version in seen:
+                continue
+
+            seen.add(formatted_version)
+            formatted_versions.append(
+                TagValue(
+                    key=self.key,
+                    value=formatted_version,
+                    times_seen=None,
+                    first_seen=None,
+                    last_seen=None,
+                )
+            )
+
+        return formatted_versions
+
+    def semver_build_autocomplete_function(self):
+        build = self.query if self.query else ""
+        if not build.endswith("*"):
+            build += "*"
+
+        organization_id = self.snuba_params.organization_id
+        assert organization_id is not None
+
+        versions = Release.objects.filter_by_semver_build(
+            organization_id,
+            "exact",
+            build,
+            self.snuba_params.project_ids,
+        )
+
+        environment_ids = self.snuba_params.environment_ids
+        if environment_ids:
+            release_environments = ReleaseEnvironment.objects.filter(
+                environment_id__in=environment_ids
+            )
+            versions = versions.filter(
+                id__in=release_environments.values_list("release_id", flat=True)
+            )
+
+        builds = (
+            versions.values_list("build_code", flat=True).distinct().order_by("build_code")[:1000]
+        )
+
+        return [
+            TagValue(
+                key=self.key,
+                value=build,
+                times_seen=None,
+                first_seen=None,
+                last_seen=None,
+            )
+            for build in builds
+        ]
+
+    def semver_package_autocomplete_function(self):
+        packages = (
+            Release.objects.filter(
+                organization_id=self.snuba_params.organization_id, package__startswith=self.query
+            )
+            .values_list("package")
+            .distinct()
+        )
+
+        versions = Release.objects.filter(
+            organization_id=self.snuba_params.organization_id,
+            package__in=packages,
+            id__in=ReleaseProject.objects.filter(
+                project_id__in=self.snuba_params.project_ids
+            ).values_list("release_id", flat=True),
+        ).annotate_prerelease_column()  # type: ignore[attr-defined]  # mypy doesn't know about ReleaseQuerySet
+
+        environment_ids = self.snuba_params.environment_ids
+        if environment_ids:
+            release_environments = ReleaseEnvironment.objects.filter(
+                environment_id__in=environment_ids
+            )
+            versions = versions.filter(
+                id__in=release_environments.values_list("release_id", flat=True)
+            )
+
+        packages = versions.values_list("package", flat=True).distinct().order_by("package")[:1000]
+
+        return [
+            TagValue(
+                key=self.key,
+                value=package,
+                times_seen=None,
+                first_seen=None,
+                last_seen=None,
+            )
+            for package in packages
+        ]
+
     def boolean_autocomplete_function(self) -> list[TagValue]:
         return [
             TagValue(

tests/snuba/api/endpoints/test_organization_trace_item_attributes.py

@@ -7,7 +7,13 @@
 
 from sentry.exceptions import InvalidSearchQuery
 from sentry.search.eap.types import SupportedTraceItemType
-from sentry.testutils.cases import APITestCase, BaseSpansTestCase, OurLogTestCase, SnubaTestCase
+from sentry.testutils.cases import (
+    APITestCase,
+    BaseSpansTestCase,
+    OurLogTestCase,
+    SnubaTestCase,
+    SpanTestCase,
+)
 from sentry.testutils.helpers import parse_link_header
 from sentry.testutils.helpers.datetime import before_now
 from sentry.testutils.helpers.options import override_options
@@ -523,7 +529,7 @@ def test_attribute_values(self):
 
 
 class OrganizationTraceItemAttributeValuesEndpointSpansTest(
-    OrganizationTraceItemAttributeValuesEndpointBaseTest, BaseSpansTestCase
+    OrganizationTraceItemAttributeValuesEndpointBaseTest, BaseSpansTestCase, SpanTestCase
 ):
     feature_flags = {"organizations:visibility-explore-view": True}
     item_type = SupportedTraceItemType.SPANS
@@ -1271,3 +1277,145 @@ def test_pagination(self):
                 "lastSeen": mock.ANY,
             },
         ]
+
+    def test_autocomplete_release_semver_attributes(self):
+        release_1 = self.create_release(version="[email protected]+121")
+        release_2 = self.create_release(version="[email protected]+122")
+        self.store_spans(
+            [
+                self.create_span(
+                    {"sentry_tags": {"release": release_1.version}},
+                    start_ts=before_now(days=0, minutes=10),
+                ),
+                self.create_span(
+                    {"sentry_tags": {"release": release_2.version}},
+                    start_ts=before_now(days=0, minutes=10),
+                ),
+            ],
+            is_eap=True,
+        )
+
+        response = self.do_request(key="release")
+        assert response.status_code == 200
+        assert response.data == [
+            {
+                "count": mock.ANY,
+                "key": "release",
+                "value": release,
+                "name": release,
+                "firstSeen": mock.ANY,
+                "lastSeen": mock.ANY,
+            }
+            for release in ["[email protected]+121", "[email protected]+122"]
+        ]
+
+        response = self.do_request(key="release", query={"substringMatch": "121"})
+        assert response.status_code == 200
+        assert response.data == [
+            {
+                "count": mock.ANY,
+                "key": "release",
+                "value": "[email protected]+121",
+                "name": "[email protected]+121",
+                "firstSeen": mock.ANY,
+                "lastSeen": mock.ANY,
+            }
+        ]
+
+        response = self.do_request(key="release.stage")
+        assert response.status_code == 200
+        assert response.data == [
+            {
+                "count": mock.ANY,
+                "key": "release.stage",
+                "value": stage,
+                "name": stage,
+                "firstSeen": mock.ANY,
+                "lastSeen": mock.ANY,
+            }
+            for stage in ["adopted", "low_adoption", "replaced"]
+        ]
+
+        response = self.do_request(key="release.stage", query={"substringMatch": "adopt"})
+        assert response.status_code == 200
+        assert response.data == [
+            {
+                "count": mock.ANY,
+                "key": "release.stage",
+                "value": stage,
+                "name": stage,
+                "firstSeen": mock.ANY,
+                "lastSeen": mock.ANY,
+            }
+            for stage in ["adopted", "low_adoption"]
+        ]
+
+        response = self.do_request(key="release.version")
+        assert response.status_code == 200
+        assert response.data == [
+            {
+                "count": mock.ANY,
+                "key": "release.version",
+                "value": version,
+                "name": version,
+                "firstSeen": mock.ANY,
+                "lastSeen": mock.ANY,
+            }
+            for version in ["1.2.3", "2.2.4"]
+        ]
+
+        response = self.do_request(key="release.version", query={"substringMatch": "2"})
+        assert response.status_code == 200
+        assert response.data == [
+            {
+                "count": mock.ANY,
+                "key": "release.version",
+                "value": version,
+                "name": version,
+                "firstSeen": mock.ANY,
+                "lastSeen": mock.ANY,
+            }
+            for version in ["2.2.4"]
+        ]
+
+        response = self.do_request(key="release.package")
+        assert response.status_code == 200
+        assert response.data == [
+            {
+                "count": mock.ANY,
+                "key": "release.package",
+                "value": version,
+                "name": version,
+                "firstSeen": mock.ANY,
+                "lastSeen": mock.ANY,
+            }
+            for version in ["foo", "qux"]
+        ]
+
+        response = self.do_request(key="release.package", query={"substringMatch": "q"})
+        assert response.status_code == 200
+        assert response.data == [
+            {
+                "count": mock.ANY,
+                "key": "release.package",
+                "value": version,
+                "name": version,
+                "firstSeen": mock.ANY,
+                "lastSeen": mock.ANY,
+            }
+            for version in ["qux"]
+        ]
+
+        response = self.do_request(key="release.build")
+        assert response.status_code == 200
+        assert response.data == [
+            {
+                "count": mock.ANY,
+                "key": "release.build",
+                "value": version,
+                "name": version,
+                "firstSeen": mock.ANY,
+                "lastSeen": mock.ANY,
+            }
+            for version in ["121", "122"]
+        ]