From 5870c744b6b428b09d6eb3fe198be111a1de06bf Mon Sep 17 00:00:00 2001 From: Neel Shah Date: Fri, 22 Nov 2024 14:54:22 +0100 Subject: [PATCH] Fix rust tracing --- sentry_sdk/integrations/rust_tracing.py | 69 ++++++---------- sentry_sdk/tracing.py | 9 +- .../rust_tracing/test_rust_tracing.py | 82 ++++++++++--------- 3 files changed, 76 insertions(+), 84 deletions(-) diff --git a/sentry_sdk/integrations/rust_tracing.py b/sentry_sdk/integrations/rust_tracing.py index ae52c850c3..d394ba5712 100644 --- a/sentry_sdk/integrations/rust_tracing.py +++ b/sentry_sdk/integrations/rust_tracing.py @@ -37,11 +37,9 @@ import sentry_sdk from sentry_sdk.integrations import Integration from sentry_sdk.scope import should_send_default_pii -from sentry_sdk.tracing import Span as SentrySpan +from sentry_sdk.tracing import POTelSpan as SentrySpan from sentry_sdk.utils import SENSITIVE_DATA_SUBSTITUTE -TraceState = Optional[Tuple[Optional[SentrySpan], SentrySpan]] - class RustTracingLevel(Enum): Trace: str = "TRACE" @@ -171,7 +169,7 @@ def _include_tracing_fields(self) -> bool: else self.include_tracing_fields ) - def on_event(self, event: str, _span_state: TraceState) -> None: + def on_event(self, event: str, _span_state: Optional[SentrySpan]) -> None: deserialized_event = json.loads(event) metadata = deserialized_event.get("metadata", {}) @@ -185,7 +183,7 @@ def on_event(self, event: str, _span_state: TraceState) -> None: elif event_type == EventTypeMapping.Event: process_event(deserialized_event) - def on_new_span(self, attrs: str, span_id: str) -> TraceState: + def on_new_span(self, attrs: str, span_id: str) -> Optional[SentrySpan]: attrs = json.loads(attrs) metadata = attrs.get("metadata", {}) @@ -205,48 +203,35 @@ def on_new_span(self, attrs: str, span_id: str) -> TraceState: else: sentry_span_name = "" - kwargs = { - "op": "function", - "name": sentry_span_name, - "origin": self.origin, - } - - scope = sentry_sdk.get_current_scope() - parent_sentry_span = scope.span - if parent_sentry_span: - sentry_span = parent_sentry_span.start_child(**kwargs) - else: - sentry_span = scope.start_span(**kwargs) + span = sentry_sdk.start_span( + op="function", + name=sentry_span_name, + origin=self.origin, + only_if_parent=True, + ) + span.__enter__() fields = metadata.get("fields", []) for field in fields: if self._include_tracing_fields(): - sentry_span.set_data(field, attrs.get(field)) - else: - sentry_span.set_data(field, SENSITIVE_DATA_SUBSTITUTE) - - scope.span = sentry_span - return (parent_sentry_span, sentry_span) - - def on_close(self, span_id: str, span_state: TraceState) -> None: - if span_state is None: - return - - parent_sentry_span, sentry_span = span_state - sentry_span.finish() - sentry_sdk.get_current_scope().span = parent_sentry_span - - def on_record(self, span_id: str, values: str, span_state: TraceState) -> None: - if span_state is None: - return - _parent_sentry_span, sentry_span = span_state - - deserialized_values = json.loads(values) - for key, value in deserialized_values.items(): - if self._include_tracing_fields(): - sentry_span.set_data(key, value) + span.set_data(field, attrs.get(field)) else: - sentry_span.set_data(key, SENSITIVE_DATA_SUBSTITUTE) + span.set_data(field, SENSITIVE_DATA_SUBSTITUTE) + + return span + + def on_close(self, span_id: str, span: Optional[SentrySpan]) -> None: + if span is not None: + span.__exit__(None, None, None) + + def on_record(self, span_id: str, values: str, span: Optional[SentrySpan]) -> None: + if span is not None: + deserialized_values = json.loads(values) + for key, value in deserialized_values.items(): + if self._include_tracing_fields(): + span.set_data(key, value) + else: + span.set_data(key, SENSITIVE_DATA_SUBSTITUTE) class RustTracingIntegration(Integration): diff --git a/sentry_sdk/tracing.py b/sentry_sdk/tracing.py index 59533106b8..4a4f28552b 100644 --- a/sentry_sdk/tracing.py +++ b/sentry_sdk/tracing.py @@ -1,3 +1,4 @@ +import json import uuid import random import time @@ -1631,8 +1632,12 @@ def finish(self, end_timestamp=None): def to_json(self): # type: () -> dict[str, Any] - # TODO-neel-potel for sampling context - pass + """ + Only meant for testing. Not used internally anymore. + """ + if not isinstance(self._otel_span, ReadableSpan): + return {} + return json.loads(self._otel_span.to_json()) def get_trace_context(self): # type: () -> dict[str, Any] diff --git a/tests/integrations/rust_tracing/test_rust_tracing.py b/tests/integrations/rust_tracing/test_rust_tracing.py index 893fc86966..77f07649b2 100644 --- a/tests/integrations/rust_tracing/test_rust_tracing.py +++ b/tests/integrations/rust_tracing/test_rust_tracing.py @@ -11,7 +11,8 @@ RustTracingLevel, EventTypeMapping, ) -from sentry_sdk import start_transaction, capture_message +from sentry_sdk import start_span, capture_message +from tests.conftest import ApproxDict def _test_event_type_mapping(metadata: Dict[str, object]) -> EventTypeMapping: @@ -74,11 +75,11 @@ def test_on_new_span_on_close(sentry_init, capture_events): sentry_init(integrations=[integration], traces_sample_rate=1.0) events = capture_events() - with start_transaction(): + with start_span(): rust_tracing.new_span(RustTracingLevel.Info, 3) sentry_first_rust_span = sentry_sdk.get_current_span() - _, rust_first_rust_span = rust_tracing.spans[3] + rust_first_rust_span = rust_tracing.spans[3] assert sentry_first_rust_span == rust_first_rust_span @@ -102,7 +103,7 @@ def test_on_new_span_on_close(sentry_init, capture_events): data = span["data"] assert data["use_memoized"] assert data["index"] == 10 - assert data["version"] is None + assert "version" not in data def test_nested_on_new_span_on_close(sentry_init, capture_events): @@ -115,23 +116,20 @@ def test_nested_on_new_span_on_close(sentry_init, capture_events): sentry_init(integrations=[integration], traces_sample_rate=1.0) events = capture_events() - with start_transaction(): + with start_span(): original_sentry_span = sentry_sdk.get_current_span() rust_tracing.new_span(RustTracingLevel.Info, 3, index_arg=10) sentry_first_rust_span = sentry_sdk.get_current_span() - _, rust_first_rust_span = rust_tracing.spans[3] + rust_first_rust_span = rust_tracing.spans[3] # Use a different `index_arg` value for the inner span to help # distinguish the two at the end of the test rust_tracing.new_span(RustTracingLevel.Info, 5, index_arg=9) sentry_second_rust_span = sentry_sdk.get_current_span() - rust_parent_span, rust_second_rust_span = rust_tracing.spans[5] + rust_second_rust_span = rust_tracing.spans[5] assert rust_second_rust_span == sentry_second_rust_span - assert rust_parent_span == sentry_first_rust_span - assert rust_parent_span == rust_first_rust_span - assert rust_parent_span != rust_second_rust_span rust_tracing.close_span(5) @@ -171,12 +169,12 @@ def test_nested_on_new_span_on_close(sentry_init, capture_events): first_span_data = first_span["data"] assert first_span_data["use_memoized"] assert first_span_data["index"] == 10 - assert first_span_data["version"] is None + assert "version" not in first_span_data second_span_data = second_span["data"] assert second_span_data["use_memoized"] assert second_span_data["index"] == 9 - assert second_span_data["version"] is None + assert "version" not in second_span_data def test_on_new_span_without_transaction(sentry_init): @@ -207,7 +205,7 @@ def test_on_event_exception(sentry_init, capture_events): events = capture_events() sentry_sdk.get_isolation_scope().clear_breadcrumbs() - with start_transaction(): + with start_span(): rust_tracing.new_span(RustTracingLevel.Info, 3) # Mapped to Exception @@ -243,7 +241,7 @@ def test_on_event_breadcrumb(sentry_init, capture_events): events = capture_events() sentry_sdk.get_isolation_scope().clear_breadcrumbs() - with start_transaction(): + with start_span(): rust_tracing.new_span(RustTracingLevel.Info, 3) # Mapped to Breadcrumb @@ -274,7 +272,7 @@ def test_on_event_event(sentry_init, capture_events): events = capture_events() sentry_sdk.get_isolation_scope().clear_breadcrumbs() - with start_transaction(): + with start_span(): rust_tracing.new_span(RustTracingLevel.Info, 3) # Mapped to Event @@ -311,7 +309,7 @@ def test_on_event_ignored(sentry_init, capture_events): events = capture_events() sentry_sdk.get_isolation_scope().clear_breadcrumbs() - with start_transaction(): + with start_span(): rust_tracing.new_span(RustTracingLevel.Info, 3) # Ignored @@ -344,7 +342,7 @@ def span_filter(metadata: Dict[str, object]) -> bool: sentry_init(integrations=[integration], traces_sample_rate=1.0) events = capture_events() - with start_transaction(): + with start_span(): original_sentry_span = sentry_sdk.get_current_span() # Span is not ignored @@ -377,16 +375,16 @@ def test_record(sentry_init): ) sentry_init(integrations=[integration], traces_sample_rate=1.0) - with start_transaction(): + with start_span(): rust_tracing.new_span(RustTracingLevel.Info, 3) span_before_record = sentry_sdk.get_current_span().to_json() - assert span_before_record["data"]["version"] is None + assert "version" not in span_before_record["attributes"] rust_tracing.record(3) span_after_record = sentry_sdk.get_current_span().to_json() - assert span_after_record["data"]["version"] == "memoized" + assert span_after_record["attributes"]["version"] == "memoized" def test_record_in_ignored_span(sentry_init): @@ -403,18 +401,18 @@ def span_filter(metadata: Dict[str, object]) -> bool: ) sentry_init(integrations=[integration], traces_sample_rate=1.0) - with start_transaction(): + with start_span(): rust_tracing.new_span(RustTracingLevel.Info, 3) span_before_record = sentry_sdk.get_current_span().to_json() - assert span_before_record["data"]["version"] is None + assert "version" not in span_before_record["attributes"] rust_tracing.new_span(RustTracingLevel.Trace, 5) rust_tracing.record(5) # `on_record()` should not do anything to the current Sentry span if the associated Rust span was ignored span_after_record = sentry_sdk.get_current_span().to_json() - assert span_after_record["data"]["version"] is None + assert "version" not in span_after_record["attributes"] @pytest.mark.parametrize( @@ -443,33 +441,37 @@ def test_include_tracing_fields( traces_sample_rate=1.0, send_default_pii=send_default_pii, ) - with start_transaction(): + with start_span(): rust_tracing.new_span(RustTracingLevel.Info, 3) span_before_record = sentry_sdk.get_current_span().to_json() if tracing_fields_expected: - assert span_before_record["data"]["version"] is None + assert "version" not in span_before_record["attributes"] else: - assert span_before_record["data"]["version"] == "[Filtered]" + assert span_before_record["attributes"]["version"] == "[Filtered]" rust_tracing.record(3) span_after_record = sentry_sdk.get_current_span().to_json() if tracing_fields_expected: - assert span_after_record["data"] == { - "thread.id": mock.ANY, - "thread.name": mock.ANY, - "use_memoized": True, - "version": "memoized", - "index": 10, - } + assert span_after_record["attributes"] == ApproxDict( + { + "thread.id": mock.ANY, + "thread.name": mock.ANY, + "use_memoized": True, + "version": "memoized", + "index": 10, + } + ) else: - assert span_after_record["data"] == { - "thread.id": mock.ANY, - "thread.name": mock.ANY, - "use_memoized": "[Filtered]", - "version": "[Filtered]", - "index": "[Filtered]", - } + assert span_after_record["attributes"] == ApproxDict( + { + "thread.id": mock.ANY, + "thread.name": mock.ANY, + "use_memoized": "[Filtered]", + "version": "[Filtered]", + "index": "[Filtered]", + } + )