feat(uptime): Optionally include uptime stats in organization_trace_meta

This adds `include_uptime` to the meta endpoint as well. When it's passed, we'll include the number of uptime spans in the meta as well

Relies on #96821

Author: wedamija

src/sentry/api/endpoints/organization_trace_meta.py

@@ -1,9 +1,11 @@
+import logging
 from concurrent.futures import ThreadPoolExecutor
 from typing import TypedDict
 
 from django.http import HttpRequest, HttpResponse
 from rest_framework.request import Request
 from rest_framework.response import Response
+from sentry_protos.snuba.v1.endpoint_trace_item_table_pb2 import TraceItemTableResponse
 
 from sentry.api.api_publish_status import ApiPublishStatus
 from sentry.api.base import region_silo_endpoint
@@ -21,15 +23,32 @@
 from sentry.snuba.referrer import Referrer
 from sentry.snuba.rpc_dataset_common import RPCBase, TableQuery
 from sentry.snuba.spans_rpc import Spans
+from sentry.snuba.trace import _run_uptime_results_query, _uptime_results_query
 
+logger = logging.getLogger(__name__)
 
-class SerializedResponse(TypedDict):
+
+class SerializedResponse(TypedDict, total=False):
     logs: int
     errors: int
     performance_issues: int
     span_count: int
     transaction_child_count_map: SnubaData
     span_count_map: dict[str, int]
+    uptime_checks: int  # Only present when include_uptime is True
+
+
+def extract_uptime_count(uptime_result: list[TraceItemTableResponse]) -> int:
+    """Safely extract uptime count from query result."""
+    if not uptime_result:
+        return 0
+
+    first_result = uptime_result[0]
+    if not first_result.column_values:
+        return 0
+
+    first_column = first_result.column_values[0]
+    return len(first_column.results) if first_column.results else 0
 
 
 @region_silo_endpoint
@@ -151,9 +170,11 @@ def get(self, request: Request, organization: Organization, trace_id: str) -> Ht
                 query=f"trace:{trace_id}",
                 limit=1,
             )
+            include_uptime = request.GET.get("include_uptime", "0") == "1"
+            max_workers = 3 + (1 if include_uptime else 0)
             with ThreadPoolExecutor(
                 thread_name_prefix=__name__,
-                max_workers=3,
+                max_workers=max_workers,
             ) as query_thread_pool:
                 spans_future = query_thread_pool.submit(
                     self.query_span_data, trace_id, snuba_params
@@ -165,15 +186,31 @@ def get(self, request: Request, organization: Organization, trace_id: str) -> Ht
                     errors_query.run_query, Referrer.API_TRACE_VIEW_GET_EVENTS.value
                 )
 
+                uptime_future = None
+                if include_uptime:
+                    uptime_query = _uptime_results_query(snuba_params, trace_id)
+                    uptime_future = query_thread_pool.submit(
+                        _run_uptime_results_query, uptime_query
+                    )
+
             results = spans_future.result()
             perf_issues = perf_issues_future.result()
             errors = errors_future.result()
             results["errors"] = errors
 
-        return Response(self.serialize(results, perf_issues))
-
-    def serialize(self, results: dict[str, EAPResponse], perf_issues: int) -> SerializedResponse:
-        return {
+            uptime_count = None
+            if uptime_future:
+                try:
+                    uptime_result = uptime_future.result()
+                    uptime_count = extract_uptime_count(uptime_result)
+                except Exception:
+                    logger.exception("Failed to fetch uptime results")
+        return Response(self.serialize(results, perf_issues, uptime_count))
+
+    def serialize(
+        self, results: dict[str, EAPResponse], perf_issues: int, uptime_count: int | None = None
+    ) -> SerializedResponse:
+        response: SerializedResponse = {
             # Values can be null if there's no result
             "logs": results["logs_meta"]["data"][0].get("count()") or 0,
             "errors": results["errors"]["data"][0].get("errors") or 0,
@@ -184,3 +221,6 @@ def serialize(self, results: dict[str, EAPResponse], perf_issues: int) -> Serial
                 row["span.op"]: row["count()"] for row in results["spans_op_count"]["data"]
             },
         }
+        if uptime_count is not None:
+            response["uptime_checks"] = uptime_count
+        return response

tests/snuba/api/endpoints/test_organization_trace_meta.py

@@ -4,6 +4,7 @@
 from django.urls import NoReverseMatch, reverse
 
 from sentry.testutils.helpers.datetime import before_now
+from tests.sentry.uptime.endpoints.test_base import UptimeResultEAPTestCase
 from tests.snuba.api.endpoints.test_organization_events_trace import (
     OrganizationEventsTraceEndpointBase,
 )
@@ -61,6 +62,7 @@ def test_bad_ids(self) -> None:
         assert data["performance_issues"] == 0
         assert data["span_count"] == 0
         assert data["span_count_map"] == {}
+        assert "uptime_checks" not in data  # Should not be present without include_uptime param
 
         # Invalid trace id
         with pytest.raises(NoReverseMatch):
@@ -146,3 +148,98 @@ def test_with_invalid_date(self) -> None:
                     format="json",
                 )
         assert response.status_code == 400, response.content
+
+
+class OrganizationTraceMetaUptimeTest(OrganizationEventsTraceEndpointBase, UptimeResultEAPTestCase):
+    url_name = "sentry-api-0-organization-trace-meta"
+    FEATURES = ["organizations:trace-spans-format"]
+
+    def create_uptime_check(self, trace_id=None, **kwargs):
+        defaults = {
+            "trace_id": trace_id or self.trace_id,
+            "scheduled_check_time": self.day_ago,
+        }
+        defaults.update(kwargs)
+        return self.create_eap_uptime_result(**defaults)
+
+    def test_trace_meta_without_uptime_param(self) -> None:
+        """Test that uptime_checks field is NOT present when include_uptime is not set"""
+        self.load_trace(is_eap=True)
+        uptime_result = self.create_uptime_check()
+        self.store_uptime_results([uptime_result])
+        with self.feature(self.FEATURES):
+            response = self.client.get(
+                self.url,
+                data={"project": -1},
+                format="json",
+            )
+
+        assert response.status_code == 200
+        data = response.data
+        assert "uptime_checks" not in data
+        assert data["errors"] == 0
+        assert data["performance_issues"] == 2
+        assert data["span_count"] == 19
+
+    def test_trace_meta_with_uptime_param(self) -> None:
+        """Test that uptime_checks shows correct count when include_uptime=1"""
+        self.load_trace(is_eap=True)
+
+        uptime_results = [
+            self.create_uptime_check(check_status="success"),
+            self.create_uptime_check(check_status="failure"),
+            self.create_uptime_check(check_status="success"),
+        ]
+        self.store_uptime_results(uptime_results)
+
+        with self.feature(self.FEATURES):
+            response = self.client.get(
+                self.url,
+                data={"project": "-1", "include_uptime": "1"},
+                format="json",
+            )
+
+        assert response.status_code == 200
+        data = response.data
+        assert "uptime_checks" in data
+        assert data["uptime_checks"] == 3
+        assert data["errors"] == 0
+        assert data["performance_issues"] == 2
+        assert data["span_count"] == 19
+
+    def test_trace_meta_no_uptime_results(self) -> None:
+        """Test that uptime_checks is 0 when there are no uptime results"""
+        self.load_trace(is_eap=True)
+
+        with self.feature(self.FEATURES):
+            response = self.client.get(
+                self.url,
+                data={"project": "-1", "include_uptime": "1"},
+                format="json",
+            )
+
+        assert response.status_code == 200
+        data = response.data
+        assert "uptime_checks" in data
+        assert data["uptime_checks"] == 0
+        assert data["errors"] == 0
+        assert data["performance_issues"] == 2
+        assert data["span_count"] == 19
+
+    def test_trace_meta_different_trace_id(self) -> None:
+        """Test that uptime results from different traces are not counted"""
+        self.load_trace(is_eap=True)
+        other_trace_id = uuid4().hex
+        uptime_result = self.create_uptime_check(trace_id=other_trace_id)
+        self.store_uptime_results([uptime_result])
+
+        with self.feature(self.FEATURES):
+            response = self.client.get(
+                self.url,
+                data={"project": "-1", "include_uptime": "1"},
+                format="json",
+            )
+        assert response.status_code == 200
+        data = response.data
+        assert "uptime_checks" in data
+        assert data["uptime_checks"] == 0