feat(explorer): rpcs for getting log/metric attrs for a trace id + substring

src/sentry/seer/endpoints/seer_rpc.py

@@ -91,6 +91,8 @@
     execute_timeseries_query,
     get_issue_and_event_details,
     get_issue_details,
+    get_log_attributes_for_trace,
+    get_metric_attributes_for_trace,
     get_replay_metadata,
     get_repository_definition,
     get_trace_item_attributes,
@@ -1039,6 +1041,8 @@ def check_repository_integrations_status(*, repository_integrations: list[dict[s
     "get_trace_item_attributes": get_trace_item_attributes,
     "get_repository_definition": get_repository_definition,
     "call_custom_tool": call_custom_tool,
+    "get_log_attributes_for_trace": get_log_attributes_for_trace,
+    "get_metric_attributes_for_trace": get_metric_attributes_for_trace,
     #
     # Replays
     "get_replay_summary_logs": rpc_get_replay_summary_logs,

src/sentry/seer/explorer/tools.py

@@ -3,6 +3,8 @@
 from datetime import UTC, datetime, timedelta, timezone
 from typing import Any, cast
 
+from sentry_protos.snuba.v1.endpoint_get_trace_pb2 import GetTraceRequest
+
 from sentry import eventstore, features
 from sentry.api import client
 from sentry.api.endpoints.organization_events_timeseries import TOP_EVENTS_DATASETS
@@ -23,14 +25,21 @@
 from sentry.seer.autofix.autofix import get_all_tags_overview
 from sentry.seer.constants import SEER_SUPPORTED_SCM_PROVIDERS
 from sentry.seer.explorer.index_data import UNESCAPED_QUOTE_RE
-from sentry.seer.explorer.utils import _convert_profile_to_execution_tree, fetch_profile_data
+from sentry.seer.explorer.utils import (
+    _convert_profile_to_execution_tree,
+    fetch_profile_data,
+    parse_get_trace_rpc_response,
+)
 from sentry.seer.sentry_data_models import EAPTrace
 from sentry.services.eventstore.models import Event, GroupEvent
+from sentry.snuba.ourlogs import OurLogs
 from sentry.snuba.referrer import Referrer
 from sentry.snuba.spans_rpc import Spans
 from sentry.snuba.trace import query_trace_data
+from sentry.snuba.trace_metrics import TraceMetrics
 from sentry.snuba.utils import get_dataset
 from sentry.utils.dates import parse_stats_period
+from sentry.utils.snuba_rpc import get_trace_rpc
 
 logger = logging.getLogger(__name__)
 
@@ -51,6 +60,53 @@ def _validate_date_params(
         raise ValueError("start and end must be provided together")
 
 
+def _get_full_trace_id(
+    short_trace_id: str, organization: Organization, projects: list[Project]
+) -> str | None:
+    """
+    Get full trace id if a short id is provided. Queries EAP for a single span.
+    Use sliding 14-day windows starting from most recent, up to 90 days in the past, to avoid timeouts.
+    TODO: This query ignores the trace_id column index and can do large scans, and is a good candidate for optimization.
+    This can be done with a materialized string column for the first 8 chars and a secondary index.
+    Alternatively we can try more consistent ways of passing the full ID to Explorer.
+    """
+    now = datetime.now(timezone.utc)
+    window_days = 14
+    max_days = 90
+
+    # Slide back in time in 14-day windows
+    for days_back in range(0, max_days, window_days):
+        window_end = now - timedelta(days=days_back)
+        window_start = now - timedelta(days=min(days_back + window_days, max_days))
+
+        snuba_params = SnubaParams(
+            start=window_start,
+            end=window_end,
+            projects=projects,
+            organization=organization,
+            debug=True,
+        )
+
+        subquery_result = Spans.run_table_query(
+            params=snuba_params,
+            query_string=f"trace:{short_trace_id}",
+            selected_columns=["trace"],
+            orderby=[],
+            offset=0,
+            limit=1,
+            referrer=Referrer.SEER_RPC,
+            config=SearchResolverConfig(),
+            sampling_mode=None,
+        )
+
+        data = subquery_result.get("data")
+        full_trace_id = data[0].get("trace") if data else None
+        if full_trace_id:
+            return full_trace_id
+
+    return None
+
+
 def execute_table_query(
     *,
     org_id: int,
@@ -261,72 +317,17 @@ def get_trace_waterfall(trace_id: str, organization_id: int) -> EAPTrace | None:
 
     projects = list(Project.objects.filter(organization=organization, status=ObjectStatus.ACTIVE))
 
-    # Get full trace id if a short id is provided. Queries EAP for a single span.
-    # Use sliding 14-day windows starting from most recent, up to 90 days in the past, to avoid timeouts.
-    # TODO: This query ignores the trace_id column index and can do large scans, and is a good candidate for optimization.
-    # This can be done with a materialized string column for the first 8 chars and a secondary index.
-    # Alternatively we can try more consistent ways of passing the full ID to Explorer.
     if len(trace_id) < 32:
-        full_trace_id = None
-        now = datetime.now(timezone.utc)
-        window_days = 14
-        max_days = 90
-
-        # Slide back in time in 14-day windows
-        for days_back in range(0, max_days, window_days):
-            window_end = now - timedelta(days=days_back)
-            window_start = now - timedelta(days=min(days_back + window_days, max_days))
-
-            snuba_params = SnubaParams(
-                start=window_start,
-                end=window_end,
-                projects=projects,
-                organization=organization,
-                debug=True,
-            )
-
-            subquery_result = Spans.run_table_query(
-                params=snuba_params,
-                query_string=f"trace:{trace_id}",
-                selected_columns=["trace"],
-                orderby=[],
-                offset=0,
-                limit=1,
-                referrer=Referrer.SEER_RPC,
-                config=SearchResolverConfig(),
-                sampling_mode=None,
+        full_trace_id = _get_full_trace_id(trace_id, organization, projects)
+        if not full_trace_id:
+            logger.warning(
+                "get_trace_waterfall: No full trace id found for short trace id",
+                extra={"organization_id": organization_id, "trace_id": trace_id},
             )
-
-            data = subquery_result.get("data")
-            if not data:
-                # Temporary debug log
-                logger.warning(
-                    "get_trace_waterfall: No data returned from short id query",
-                    extra={
-                        "organization_id": organization_id,
-                        "trace_id": trace_id,
-                        "subquery_result": subquery_result,
-                        "start": window_start.isoformat(),
-                        "end": window_end.isoformat(),
-                    },
-                )
-
-            full_trace_id = data[0].get("trace") if data else None
-            if full_trace_id:
-                break
+            return None
     else:
         full_trace_id = trace_id
 
-    if not isinstance(full_trace_id, str):
-        logger.warning(
-            "get_trace_waterfall: Trace not found from short id",
-            extra={
-                "organization_id": organization_id,
-                "trace_id": trace_id,
-            },
-        )
-        return None
-
     # Get full trace data.
     start, end = default_start_end_dates()
     snuba_params = SnubaParams(
@@ -1111,3 +1112,183 @@ def get_trace_item_attributes(
     )
 
     return {"attributes": resp.data["attributes"]}
+
+
+def get_log_attributes_for_trace(
+    *,
+    org_id: int,
+    trace_id: str,
+    message_substring: str = "",
+    substring_case_sensitive: bool = True,
+    stats_period: str | None = None,
+    start: str | None = None,
+    end: str | None = None,
+    project_id: int | None = None,
+    project_slug: str | None = None,
+    sampling_mode: SAMPLING_MODES = "NORMAL",
+    limit: int | None = 50,
+) -> dict[str, Any] | None:
+
+    _validate_date_params(stats_period=stats_period, start=start, end=end)
+
+    try:
+        organization = Organization.objects.get(id=org_id)
+    except Organization.DoesNotExist:
+        logger.warning("Organization not found", extra={"org_id": org_id})
+        return None
+
+    projects = list(
+        Project.objects.filter(
+            organization=organization,
+            status=ObjectStatus.ACTIVE,
+            **({"id": project_id} if project_id else {}),
+            **({"slug": project_slug} if project_slug else {}),
+        )
+    )
+
+    # Look up full trace id if a short id is provided.
+    if len(trace_id) < 32:
+        full_trace_id = _get_full_trace_id(trace_id, organization, projects)
+        if not full_trace_id:
+            logger.warning(
+                "execute_trace_query: No full trace id found for short trace id",
+                extra={"org_id": org_id, "trace_id": trace_id},
+            )
+            return None
+    else:
+        full_trace_id = trace_id
+
+    # Build the GetTraceRequest.
+    snuba_params = SnubaParams(
+        start=datetime.fromisoformat(start) if start else None,
+        end=datetime.fromisoformat(end) if end else None,
+        stats_period=stats_period,
+        projects=projects,
+        organization=organization,
+        sampling_mode=sampling_mode,
+    )
+
+    resolver = OurLogs.get_resolver(params=snuba_params, config=SearchResolverConfig())
+    meta = resolver.resolve_meta(referrer=Referrer.SEER_RPC, sampling_mode=sampling_mode)
+    request = GetTraceRequest(
+        meta=meta,
+        trace_id=full_trace_id,
+        items=[
+            GetTraceRequest.TraceItem(
+                item_type=meta.trace_item_type,
+                attributes=None,  # Returns all attributes.
+            )
+        ],
+    )
+    if limit and not message_substring:
+        request.limit = limit
+
+    # Query EAP EndpointGetTrace
+    response = get_trace_rpc(request)
+
+    items: list[dict[str, Any]]
+    for item_group in response.item_groups:
+        items = parse_get_trace_rpc_response(item_group, resolver)
+        break  # Should only be one item group for logs
+
+    if message_substring:
+        items = (
+            [item for item in items if message_substring in item["message"]]
+            if substring_case_sensitive
+            else [item for item in items if message_substring.lower() in item["message"].lower()]
+        )
+
+    if limit:
+        items = items[:limit]
+
+    return {"data": items}
+
+
+def get_metric_attributes_for_trace(
+    *,
+    org_id: int,
+    trace_id: str,
+    metric_name: str = "",
+    metric_name_case_sensitive: bool = True,
+    stats_period: str | None = None,
+    start: str | None = None,
+    end: str | None = None,
+    project_id: int | None = None,
+    project_slug: str | None = None,
+    sampling_mode: SAMPLING_MODES = "NORMAL",
+    limit: int | None = 50,
+) -> dict[str, Any] | None:
+
+    _validate_date_params(stats_period=stats_period, start=start, end=end)
+
+    try:
+        organization = Organization.objects.get(id=org_id)
+    except Organization.DoesNotExist:
+        logger.warning("Organization not found", extra={"org_id": org_id})
+        return None
+
+    projects = list(
+        Project.objects.filter(
+            organization=organization,
+            status=ObjectStatus.ACTIVE,
+            **({"id": project_id} if project_id else {}),
+            **({"slug": project_slug} if project_slug else {}),
+        )
+    )
+
+    # Look up full trace id if a short id is provided.
+    if len(trace_id) < 32:
+        full_trace_id = _get_full_trace_id(trace_id, organization, projects)
+        if not full_trace_id:
+            logger.warning(
+                "execute_trace_query: No full trace id found for short trace id",
+                extra={"org_id": org_id, "trace_id": trace_id},
+            )
+            return None
+    else:
+        full_trace_id = trace_id
+
+    # Build the GetTraceRequest.
+    snuba_params = SnubaParams(
+        start=datetime.fromisoformat(start) if start else None,
+        end=datetime.fromisoformat(end) if end else None,
+        stats_period=stats_period,
+        projects=projects,
+        organization=organization,
+        sampling_mode=sampling_mode,
+    )
+
+    resolver = TraceMetrics.get_resolver(params=snuba_params, config=SearchResolverConfig())
+    meta = resolver.resolve_meta(referrer=Referrer.SEER_RPC, sampling_mode=sampling_mode)
+    request = GetTraceRequest(
+        meta=meta,
+        trace_id=full_trace_id,
+        items=[
+            GetTraceRequest.TraceItem(
+                item_type=meta.trace_item_type,
+                attributes=None,  # Returns all attributes.
+            )
+        ],
+    )
+    if limit and not metric_name:
+        request.limit = limit
+
+    # Query EAP EndpointGetTrace
+    response = get_trace_rpc(request)
+
+    items: list[dict[str, Any]]
+    for item_group in response.item_groups:
+        items = parse_get_trace_rpc_response(item_group, resolver)
+        break  # Should only be one item group for metrics
+
+    if metric_name:
+        items = (
+            [item for item in items if metric_name == item["metric.name"]]
+            if metric_name_case_sensitive
+            else [item for item in items if metric_name.lower() == item["metric.name"].lower()]
+        )
+
+    if limit:
+        items = items[:limit]
+
+    return {"data": items}