feat(profiling): support nodestore as a possible backend for span-scoped flamegraph (#52166)

This is meant to be used temporarily to build a poc since the
`indexed_span` dataset is currently only available for internal
projects.

Author: viglia

src/sentry/api/endpoints/organization_profiling_profiles.py

@@ -57,15 +57,18 @@ def get(self, request: Request, organization: Organization) -> HttpResponse:
 
         span_group = request.query_params.get("spans.group", None)
         if span_group is not None:
+            backend = request.query_params.get("backend", "indexed_spans")
             profile_ids = get_profile_ids_with_spans(
                 organization.id,
                 project_ids[0],
                 params,
                 span_group,
+                backend,
                 request.query_params.get("query", None),
             )
         else:
             profile_ids = get_profile_ids(params, request.query_params.get("query", None))
+
         kwargs: Dict[str, Any] = {
             "method": "POST",
             "path": f"/organizations/{organization.id}/projects/{project_ids[0]}/flamegraph",

src/sentry/profiles/flamegraph.py

@@ -1,8 +1,10 @@
 from datetime import datetime
 from typing import Any, Dict, List, Optional, Tuple
 
+from rest_framework.exceptions import ParseError
 from snuba_sdk import Column, Condition, Entity, Function, Op, Query, Request
 
+from sentry import eventstore
 from sentry.search.events.builder import QueryBuilder
 from sentry.search.events.types import ParamsType
 from sentry.snuba.dataset import Dataset, EntityKey
@@ -78,6 +80,7 @@ def get_span_intervals(
             Condition(Column("timestamp"), Op.LT, params["end"]),
         ],
     )
+
     request = Request(
         dataset=Dataset.SpansIndexed.value,
         app_id="default",
@@ -87,17 +90,59 @@ def get_span_intervals(
             "organization_id": organization_id,
         },
     )
-    return raw_snql_query(
+    data = raw_snql_query(
         request,
         referrer=Referrer.API_STARFISH_PROFILE_FLAMEGRAPH.value,
     )["data"]
+    spans_interval = []
+    for row in data:
+        start_ns = (int(datetime.fromisoformat(row["start_timestamp"]).timestamp()) * 10**9) + (
+            row["start_ms"] * 10**6
+        )
+        end_ns = (int(datetime.fromisoformat(row["end_timestamp"]).timestamp()) * 10**9) + (
+            row["end_ms"] * 10**6
+        )
+        interval = {}
+        interval["transaction_id"] = row["transaction_id"]
+        interval["start_ns"] = str(start_ns)
+        interval["end_ns"] = str(end_ns)
+        spans_interval.append(interval)
+    return spans_interval
+
+
+def get_span_intervals_from_nodestore(
+    project_id: str,
+    span_group: str,
+    transaction_ids: List[str],
+) -> List[Dict[str, Any]]:
+    spans_interval = []
+    for id in transaction_ids:
+        nodestore_event = eventstore.backend.get_event_by_id(project_id, id)
+        data = nodestore_event.data
+        for span in data.get("spans", []):
+            if span["hash"] == span_group:
+
+                start_sec, start_us = map(int, str(span["start_timestamp"]).split("."))
+                end_sec, end_us = map(int, str(span["timestamp"]).split("."))
+
+                start_ns = (start_sec * 10**9) + (start_us * 10**3)
+                end_ns = (end_sec * 10**9) + (end_us * 10**3)
+
+                interval = {}
+                interval["transaction_id"] = id
+                interval["start_ns"] = str(start_ns)
+                interval["end_ns"] = str(end_ns)
+
+                spans_interval.append(interval)
+    return spans_interval
 
 
 def get_profile_ids_with_spans(
     organization_id: str,
     project_id: str,
     params: ParamsType,
     span_group: str,
+    backend: str,
     query: Optional[str] = None,
 ):
     data = query_profiles_data(
@@ -115,24 +160,29 @@ def get_profile_ids_with_spans(
         row["id"]: (row["profile.id"], []) for row in data
     }
 
-    data = get_span_intervals(
-        project_id,
-        span_group,
-        list(transaction_to_prof.keys()),
-        organization_id,
-        params,
-    )
-
-    for row in data:
-        start_ns = (int(datetime.fromisoformat(row["start_timestamp"]).timestamp()) * 10**9) + (
-            row["start_ms"] * 10**6
+    if backend == "nodestore":
+        data = get_span_intervals_from_nodestore(
+            project_id,
+            span_group,
+            list(transaction_to_prof.keys()),
         )
-        end_ns = (int(datetime.fromisoformat(row["end_timestamp"]).timestamp()) * 10**9) + (
-            row["end_ms"] * 10**6
+    elif backend == "indexed_spans":
+        data = get_span_intervals(
+            project_id,
+            span_group,
+            list(transaction_to_prof.keys()),
+            organization_id,
+            params,
+        )
+    else:
+        raise ParseError(
+            detail="Backend not supported: choose between 'indexed_spans' or 'nodestore'."
         )
-        transaction_id = row["transaction_id"]
 
-        transaction_to_prof[transaction_id][1].append({"start": str(start_ns), "end": str(end_ns)})
+    for row in data:
+        transaction_to_prof[row["transaction_id"]][1].append(
+            {"start": row["start_ns"], "end": row["end_ns"]}
+        )
 
     profile_ids = [tup[0] for tup in transaction_to_prof.values()]
     spans = [tup[1] for tup in transaction_to_prof.values()]