Repository: getsentry/sentry

Issue:
Support measurement metrics in new metrics/traces sample endpoint



Discussion:

Resolving patch:
diff --git a/src/sentry/sentry_metrics/querying/samples_list.py b/src/sentry/sentry_metrics/querying/samples_list.py
--- a/src/sentry/sentry_metrics/querying/samples_list.py
+++ b/src/sentry/sentry_metrics/querying/samples_list.py
@@ -2,13 +2,13 @@
 from datetime import datetime
 from typing import Any
 
-from snuba_sdk import And, Condition, Op, Or
+from snuba_sdk import And, Column, Condition, Function, Op, Or
 
 from sentry import options
 from sentry.search.events.builder import QueryBuilder, SpansIndexedQueryBuilder
 from sentry.search.events.types import QueryBuilderConfig, SnubaParams
 from sentry.snuba.dataset import Dataset
-from sentry.snuba.metrics.naming_layer.mri import SpanMRI, TransactionMRI
+from sentry.snuba.metrics.naming_layer.mri import SpanMRI, TransactionMRI, is_measurement, parse_mri
 from sentry.snuba.referrer import Referrer
 
 
@@ -86,10 +86,9 @@ def get_spans_by_key(self, span_ids: list[tuple[str, str, str]]):
 
 class SegmentsSamplesListExecutor(SamplesListExecutor):
     @classmethod
+    @abstractmethod
     def mri_to_column(cls, mri) -> str | None:
-        if mri == TransactionMRI.DURATION.value:
-            return "duration"
-        return None
+        raise NotImplementedError
 
     @classmethod
     def supports(cls, mri: str) -> bool:
@@ -102,6 +101,16 @@ def execute(self, offset, limit):
     def get_span_keys(self, offset: int, limit: int) -> list[tuple[str, str, str]]:
         rounded_timestamp = f"rounded_timestamp({self.rollup})"
 
+        """
+        When getting examples for a segment, it's actually much faster to read it
+        from the transactions dataset compared to the spans dataset as it's a much
+        smaller dataset.
+
+        One consideration here is that there is an one to one mapping between a
+        transaction to a segment today. If this relationship changes, we'll have to
+        rethink how to fetch segment samples a little as the transactions dataset
+        may not contain all the necessary data.
+        """
         builder = QueryBuilder(
             Dataset.Transactions,
             self.params,
@@ -114,6 +123,8 @@ def get_span_keys(self, offset: int, limit: int) -> list[tuple[str, str, str]]:
             config=QueryBuilderConfig(functions_acl=["rounded_timestamp", "example"]),
         )
 
+        builder.add_conditions(self.get_additional_conditions())
+
         query_results = builder.run_query(self.referrer.value)
         result = builder.process_results(query_results)
 
@@ -126,6 +137,42 @@ def get_span_keys(self, offset: int, limit: int) -> list[tuple[str, str, str]]:
             for row in result["data"]
         ]
 
+    @abstractmethod
+    def get_additional_conditions(self) -> list[Condition]:
+        raise NotImplementedError
+
+
+class TransactionDurationSamplesListExecutor(SegmentsSamplesListExecutor):
+    @classmethod
+    def mri_to_column(cls, mri) -> str | None:
+        if mri == TransactionMRI.DURATION.value:
+            return "duration"
+        return None
+
+    def get_additional_conditions(self) -> list[Condition]:
+        return []
+
+
+class MeasurementsSamplesListExecutor(SegmentsSamplesListExecutor):
+    @classmethod
+    def mri_to_column(cls, mri) -> str | None:
+        name = cls.measurement_name(mri)
+        if name is not None:
+            return f"measurements[{name}]"
+
+        return None
+
+    @classmethod
+    def measurement_name(cls, mri) -> str | None:
+        parsed_mri = parse_mri(mri)
+        if parsed_mri is not None and is_measurement(parsed_mri):
+            return parsed_mri.name[len("measurements:") :]
+        return None
+
+    def get_additional_conditions(self) -> list[Condition]:
+        name = self.measurement_name(self.mri)
+        return [Condition(Function("has", [Column("measurements.key"), name]), Op.EQ, 1)]
+
 
 class SpansSamplesListExecutor(SamplesListExecutor):
     MRI_MAPPING = {
@@ -187,7 +234,8 @@ def get_span_keys(self, offset: int, limit: int) -> list[tuple[str, str, str]]:
 
 SAMPLE_LIST_EXECUTORS = [
     SpansSamplesListExecutor,
-    SegmentsSamplesListExecutor,
+    TransactionDurationSamplesListExecutor,
+    MeasurementsSamplesListExecutor,
 ]
 
 


Test patch:
diff --git a/tests/sentry/api/endpoints/test_organization_metrics.py b/tests/sentry/api/endpoints/test_organization_metrics.py
--- a/tests/sentry/api/endpoints/test_organization_metrics.py
+++ b/tests/sentry/api/endpoints/test_organization_metrics.py
@@ -217,3 +217,59 @@ def test_transaction_duration_samples(self):
         expected = {int(span_id, 16) for span_id in span_ids}
         actual = {int(row["id"], 16) for row in response.data["data"]}
         assert actual == expected
+
+    def test_measurement_samples(self):
+        good_span_ids = [uuid4().hex[:16] for _ in range(1)]
+        bad_span_ids = [uuid4().hex[:16] for _ in range(1)]
+        for i, (good_span_id, bad_span_id) in enumerate(zip(good_span_ids, bad_span_ids)):
+            ts = before_now(days=i, minutes=10)
+
+            # first write to the transactions dataset
+            data = load_data("transaction", timestamp=ts)
+            # bad span ids will not have the measurement
+            data["measurements"] = {}
+            data["contexts"]["trace"]["span_id"] = bad_span_id
+            self.store_event(
+                data=data,
+                project_id=self.project.id,
+            )
+
+            # next write to the spans dataset
+            self.store_segment(
+                self.project.id,
+                uuid4().hex,
+                uuid4().hex,
+                span_id=bad_span_id,
+                timestamp=ts,
+            )
+
+            # first write to the transactions dataset
+            data = load_data("transaction", timestamp=ts)
+            # good span ids will have the measurement
+            data["measurements"] = {"lcp": {"value": 10}}
+            data["contexts"]["trace"]["span_id"] = good_span_id
+            self.store_event(
+                data=data,
+                project_id=self.project.id,
+            )
+
+            # next write to the spans dataset
+            self.store_segment(
+                self.project.id,
+                uuid4().hex,
+                uuid4().hex,
+                span_id=good_span_id,
+                timestamp=ts,
+            )
+
+        query = {
+            "mri": "d:transactions/measurements.lcp@millisecond",
+            "field": ["id"],
+            "project": [self.project.id],
+            "statsPeriod": "14d",
+        }
+        response = self.do_request(query)
+        assert response.status_code == 200, response.data
+        expected = {int(span_id, 16) for span_id in good_span_ids}
+        actual = {int(row["id"], 16) for row in response.data["data"]}
+        assert actual == expected