feat(measurements): Add experimental set_measurement api on transaction (getsentry#1359)

sl0thentr0py · web-flow · commit 7a3b0e5b6bed · 2022-05-06T09:16:39.000Z
diff --git a/sentry_sdk/_types.py b/sentry_sdk/_types.py
@@ -48,3 +48,34 @@
     ]
     SessionStatus = Literal["ok", "exited", "crashed", "abnormal"]
     EndpointType = Literal["store", "envelope"]
+
+    DurationUnit = Literal[
+        "nanosecond",
+        "microsecond",
+        "millisecond",
+        "second",
+        "minute",
+        "hour",
+        "day",
+        "week",
+    ]
+
+    InformationUnit = Literal[
+        "bit",
+        "byte",
+        "kilobyte",
+        "kibibyte",
+        "megabyte",
+        "mebibyte",
+        "gigabyte",
+        "gibibyte",
+        "terabyte",
+        "tebibyte",
+        "petabyte",
+        "pebibyte",
+        "exabyte",
+        "exbibyte",
+    ]
+
+    FractionUnit = Literal["ratio", "percent"]
+    MeasurementUnit = Union[DurationUnit, InformationUnit, FractionUnit, str]
diff --git a/sentry_sdk/consts.py b/sentry_sdk/consts.py
@@ -33,6 +33,7 @@
             "record_sql_params": Optional[bool],
             "smart_transaction_trimming": Optional[bool],
             "propagate_tracestate": Optional[bool],
+            "custom_measurements": Optional[bool],
         },
         total=False,
     )
diff --git a/sentry_sdk/tracing.py b/sentry_sdk/tracing.py
@@ -20,7 +20,7 @@
     from typing import Tuple
     from typing import Iterator
 
-    from sentry_sdk._types import SamplingContext
+    from sentry_sdk._types import SamplingContext, MeasurementUnit
 
 
 class _SpanRecorder(object):
@@ -487,6 +487,7 @@ class Transaction(Span):
         "_sentry_tracestate",
         # tracestate data from other vendors, of the form `dogs=yes,cats=maybe`
         "_third_party_tracestate",
+        "_measurements",
     )
 
     def __init__(
@@ -515,6 +516,7 @@ def __init__(
         # first time an event needs it for inclusion in the captured data
         self._sentry_tracestate = sentry_tracestate
         self._third_party_tracestate = third_party_tracestate
+        self._measurements = {}  # type: Dict[str, Any]
 
     def __repr__(self):
         # type: () -> str
@@ -594,17 +596,30 @@ def finish(self, hub=None):
         # to be garbage collected
         self._span_recorder = None
 
-        return hub.capture_event(
-            {
-                "type": "transaction",
-                "transaction": self.name,
-                "contexts": {"trace": self.get_trace_context()},
-                "tags": self._tags,
-                "timestamp": self.timestamp,
-                "start_timestamp": self.start_timestamp,
-                "spans": finished_spans,
-            }
-        )
+        event = {
+            "type": "transaction",
+            "transaction": self.name,
+            "contexts": {"trace": self.get_trace_context()},
+            "tags": self._tags,
+            "timestamp": self.timestamp,
+            "start_timestamp": self.start_timestamp,
+            "spans": finished_spans,
+        }
+
+        if has_custom_measurements_enabled():
+            event["measurements"] = self._measurements
+
+        return hub.capture_event(event)
+
+    def set_measurement(self, name, value, unit=""):
+        # type: (str, float, MeasurementUnit) -> None
+        if not has_custom_measurements_enabled():
+            logger.debug(
+                "[Tracing] Experimental custom_measurements feature is disabled"
+            )
+            return
+
+        self._measurements[name] = {"value": value, "unit": unit}
 
     def to_json(self):
         # type: () -> Dict[str, Any]
@@ -727,4 +742,5 @@ def _set_initial_sampling_decision(self, sampling_context):
     has_tracing_enabled,
     is_valid_sample_rate,
     maybe_create_breadcrumbs_from_span,
+    has_custom_measurements_enabled,
 )
diff --git a/sentry_sdk/tracing_utils.py b/sentry_sdk/tracing_utils.py
@@ -406,6 +406,13 @@ def has_tracestate_enabled(span=None):
     return bool(options and options["_experiments"].get("propagate_tracestate"))
 
 
+def has_custom_measurements_enabled():
+    # type: () -> bool
+    client = sentry_sdk.Hub.current.client
+    options = client and client.options
+    return bool(options and options["_experiments"].get("custom_measurements"))
+
+
 # Circular imports
 
 if MYPY:
diff --git a/tests/tracing/test_misc.py b/tests/tracing/test_misc.py
@@ -246,3 +246,31 @@ def test_has_tracestate_enabled(sentry_init, tracestate_enabled):
         assert has_tracestate_enabled() is True
     else:
         assert has_tracestate_enabled() is False
+
+
+def test_set_meaurement(sentry_init, capture_events):
+    sentry_init(traces_sample_rate=1.0, _experiments={"custom_measurements": True})
+
+    events = capture_events()
+
+    transaction = start_transaction(name="measuring stuff")
+
+    with pytest.raises(TypeError):
+        transaction.set_measurement()
+
+    with pytest.raises(TypeError):
+        transaction.set_measurement("metric.foo")
+
+    transaction.set_measurement("metric.foo", 123)
+    transaction.set_measurement("metric.bar", 456, unit="second")
+    transaction.set_measurement("metric.baz", 420.69, unit="custom")
+    transaction.set_measurement("metric.foobar", 12, unit="percent")
+    transaction.set_measurement("metric.foobar", 17.99, unit="percent")
+
+    transaction.finish()
+
+    (event,) = events
+    assert event["measurements"]["metric.foo"] == {"value": 123, "unit": ""}
+    assert event["measurements"]["metric.bar"] == {"value": 456, "unit": "second"}
+    assert event["measurements"]["metric.baz"] == {"value": 420.69, "unit": "custom"}
+    assert event["measurements"]["metric.foobar"] == {"value": 17.99, "unit": "percent"}

Original file line number	Diff line number	Diff line change
`@@ -33,6 +33,7 @@`
`33`	`33`	`"record_sql_params": Optional[bool],`
`34`	`34`	`"smart_transaction_trimming": Optional[bool],`
`35`	`35`	`"propagate_tracestate": Optional[bool],`
	`36`	`+ "custom_measurements": Optional[bool],`
`36`	`37`	`},`
`37`	`38`	`total=False,`
`38`	`39`	`)`