8000 Fix rust tracing · getsentry/sentry-python@5870c74 · GitHub
[go: up one dir, main page]
Skip to content

Commit 5870c74

Browse files
sl0thentr0pysl0thentr0py
committed
Fix rust tracing
1 parent 4c07934 commit 5870c74

File tree

3 files changed

+76
-84
lines changed

3 files changed

+76
-84
lines changed

‎sentry_sdk/integrations/rust_tracing.py

Lines changed: 27 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -37,11 +37,9 @@
3737
import sentry_sdk
3838
from sentry_sdk.integrations import Integration
3939
from sentry_sdk.scope import should_send_default_pii
40-
from sentry_sdk.tracing import Span as SentrySpan
40+
from sentry_sdk.tracing import POTelSpan as SentrySpan
4141
from sentry_sdk.utils import SENSITIVE_DATA_SUBSTITUTE
4242

43-
TraceState = Optional[Tuple[Optional[SentrySpan], SentrySpan]]
44-
4543

4644
class RustTracingLevel(Enum):
4745
Trace: str = "TRACE"
@@ -171,7 +169,7 @@ def _include_tracing_fields(self) -> bool:
171169
else self.include_tracing_fields
172170
)
173171

174-
def on_event(self, event: str, _span_state: TraceState) -> None:
172+
def on_event(self, event: str, _span_state: Optional[SentrySpan]) -> None:
175173
deserialized_event = json.loads(event)
176174
metadata = deserialized_event.get("metadata", {})
177175

@@ -185,7 +183,7 @@ def on_event(self, event: str, _span_state: TraceState) -> None:
185183
elif event_type == EventTypeMapping.Event:
186184
process_event(deserialized_event)
187185

188-
def on_new_span(self, attrs: str, span_id: str) -> TraceState:
186+
def on_new_span(self, attrs: str, span_id: str) -> Optional[SentrySpan]:
189187
attrs = json.loads(attrs)
190188
metadata = attrs.get("metadata", {})
191189

@@ -205,48 +203,35 @@ def on_new_span(self, attrs: str, span_id: str) -> TraceState:
205203
else:
206204
sentry_span_name = "<unknown>"
207205

208-
kwargs = {
209-
"op": "function",
210-
"name": sentry_span_name,
211-
"origin": self.origin,
212-
}
213-
214-
scope = sentry_sdk.get_current_scope()
215-
parent_sentry_span = scope.span
216-
if parent_sentry_span:
217-
sentry_span = parent_sentry_span.start_child(**kwargs)
218-
else:
219-
sentry_span = scope.start_span(**kwargs)
206+
span = sentry_sdk.start_span(
207+
op="function",
208+
name=sentry_span_name,
209+
origin=self.origin,
210+
only_if_parent=True,
211+
)
212+
span.__enter__()
220213

221214
fields = metadata.get("fields", [])
222215
for field in fields:
223216
if self._include_tracing_fields():
224-
sentry_span.set_data(field, attrs.get(field))
225-
else:
226-
sentry_span.set_data(field, SENSITIVE_DATA_SUBSTITUTE)
227-
228-
scope.span = sentry_span
229-
return (parent_sentry_span, sentry_span)
230-
231-
def on_close(self, span_id: str, span_state: TraceState) -> None:
232-
if span_state is None:
233-
return
234-
235-
parent_sentry_span, sentry_span = span_state
236-
sentry_span.finish()
237-
sentry_sdk.get_current_scope().span = parent_sentry_span
238-
239-
def on_record(self, span_id: str, values: str, span_state: TraceState) -> None:
240-
if span_state is None:
241-
return
242-
_parent_sentry_span, sentry_span = span_state
243-
244-
deserialized_values = json.loads(values)
245-
for key, value in deserialized_values.items():
246-
if self._include_tracing_fields():
247-
sentry_span.set_data(key, value)
217+
span.set_data(field, attrs.get(field))
248218
else:
249-
sentry_span.set_data(key, SENSITIVE_DATA_SUBSTITUTE)
219+
span.set_data(field, SENSITIVE_DATA_SUBSTITUTE)
220+
221+
return span
222+
223+
def on_close(self, span_id: str, span: Optional[SentrySpan]) -> None:
224+
if span is not None:
225+
span.__exit__(None, None, None)
226+
227+
def on_record(self, span_id: str, values: str, span: Optional[SentrySpan]) -> None:
228+
if span is not None:
229+
deserialized_values = json.loads(values)
230+
for key, value in deserialized_values.items():
231+
if self._include_tracing_fields():
232+
span.set_data(key, value)
233+
else:
234+
span.set_data(key, SENSITIVE_DATA_SUBSTITUTE)
250235

251236

252237
class RustTracingIntegration(Integration):

‎sentry_sdk/tracing.py

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import json
12
import uuid
23
import random
34
import time
@@ -1631,8 +1632,12 @@ def finish(self, end_timestamp=None):
16311632

16321633
def to_json(self):
16331634
# type: () -> dict[str, Any]
1634-
# TODO-neel-potel for sampling context
1635-
pass
1635+
"""
1636+
Only meant for testing. Not used internally anymore.
1637+
"""
1638+
if not isinstance(self._otel_span, ReadableSpan):
1639+
return {}
1640+
return json.loads(self._otel_span.to_json())
16361641

16371642
def get_trace_context(self):
16381643
# type: () -> dict[str, Any]

‎tests/integrations/rust_tracing/test_rust_tracing.py

Lines changed: 42 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,8 @@
10000
1111
RustTracingLevel,
1212
EventTypeMapping,
1313
)
14-
from sentry_sdk import start_transaction, capture_message
14+
from sentry_sdk import start_span, capture_message
15+
from tests.conftest import ApproxDict
1516

1617

1718
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):
7475
sentry_init(integrations=[integration], traces_sample_rate=1.0)
7576

7677
events = capture_events()
77-
with start_transaction():
78+
with start_span():
7879
rust_tracing.new_span(RustTracingLevel.Info, 3)
7980

8081
sentry_first_rust_span = sentry_sdk.get_current_span()
81-
_, rust_first_rust_span = rust_tracing.spans[3]
82+
rust_first_rust_span = rust_tracing.spans[3]
8283

8384
assert sentry_first_rust_span == rust_first_rust_span
8485

@@ -102,7 +103,7 @@ def test_on_new_span_on_close(sentry_init, capture_events):
102103
data = span["data"]
103104
assert data["use_memoized"]
104105
assert data["index"] == 10
105-
assert data["version"] is None
106+
assert "version" not in data
106107

107108

108109
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):
115116
sentry_init(integrations=[integration], traces_sample_rate=1.0)
116117

117118
events = capture_events()
118-
with start_transaction():
119+
with start_span():
119120
original_sentry_span = sentry_sdk.get_current_span()
120121

121122
rust_tracing.new_span(RustTracingLevel.Info, 3, index_arg=10)
122123
sentry_first_rust_span = sentry_sdk.get_current_span()
123-
_, rust_first_rust_span = rust_tracing.spans[3]
124+
rust_first_rust_span = rust_tracing.spans[3]
124125

125126
# Use a different `index_arg` value for the inner span to help
126127
# distinguish the two at the end of the test
127128
rust_tracing.new_span(RustTracingLevel.Info, 5, index_arg=9)
128129
sentry_second_rust_span = sentry_sdk.get_current_span()
129-
rust_parent_span, rust_second_rust_span = rust_tracing.spans[5]
130+
rust_second_rust_span = rust_tracing.spans[5]
130131

131132
assert rust_second_rust_span == sentry_second_rust_span
132-
assert rust_parent_span == sentry_first_rust_span
133-
assert rust_parent_span == rust_first_rust_span
134-
assert rust_parent_span != rust_second_rust_span
135133

136134
rust_tracing.close_span(5)
137135

@@ -171,12 +169,12 @@ def test_nested_on_new_span_on_close(sentry_init, capture_events):
171169
first_span_data = first_span["data"]
172170
assert first_span_data["use_memoized"]
173171
assert first_span_data["index"] == 10
174-
assert first_span_data["version"] is None
172+
assert "version" not in first_span_data
175173

176174
second_span_data = second_span["data"]
177175
assert second_span_data["use_memoized"]
178176
assert second_span_data["index"] == 9
179-
assert second_span_data["version"] is None
177+
assert "version" not in second_span_data
180178

181179

182180
def test_on_new_span_without_transaction(sentry_init):
@@ -207,7 +205,7 @@ def test_on_event_exception(sentry_init, capture_events):
207205
events = capture_events()
208206
sentry_sdk.get_isolation_scope().clear_breadcrumbs()
209207

210-
with start_transaction():
208+
with start_span():
211209
rust_tracing.new_span(RustTracingLevel.Info, 3)
212210

213211
# Mapped to Exception
@@ -243,7 +241,7 @@ def test_on_event_breadcrumb(sentry_init, capture_events):
243241
events = capture_events()
244242
sentry_sdk.get_isolation_scope().clear_breadcrumbs()
245243

246- 1241
with start_transaction():
244+
with start_span():
247245
rust_tracing.new_span(RustTracingLevel.Info, 3)
248246

249247
# Mapped to Breadcrumb
@@ -274,7 +272,7 @@ def test_on_event_event(sentry_init, capture_events):
274272
events = capture_events()
275273
sentry_sdk.get_isolation_scope().clear_breadcrumbs()
276274

277-
with start_transaction():
275+
with start_span():
278276
rust_tracing.new_span(RustTracingLevel.Info, 3)
279277

280278
# Mapped to Event
@@ -311,7 +309,7 @@ def test_on_event_ignored(sentry_init, capture_events):
311309
events = capture_events()
312310
sentry_sdk.get_isolation_scope().clear_breadcrumbs()
313311

314-
with start_transaction():
312+
with start_span():
315313
rust_tracing.new_span(RustTracingLevel.Info, 3)
316314

317315
# Ignored
@@ -344,7 +342,7 @@ def span_filter(metadata: Dict[str, object]) -> bool:
344342
sentry_init(integrations=[integration], traces_sample_rate=1.0)
345343

346344
events = capture_events()
347-
with start_transaction():
345+
with start_span():
348346
original_sentry_span = sentry_sdk.get_current_span()
349347

350348
# Span is not ignored
@@ -377,16 +375,16 @@ def test_record(sentry_init):
377375
)
378376
sentry_init(integrations=[integration], traces_sample_rate=1.0)
379377

380-
with start_transaction():
378+
with start_span():
381379
rust_tracing.new_span(RustTracingLevel.Info, 3)
382380

383381
span_before_record = sentry_sdk.get_current_span().to_json()
384-
assert span_before_record["data"]["version"] is None
382+
assert "version" not in span_before_record["attributes"]
385383

386384
rust_tracing.record(3)
387385

388386
span_after_record = sentry_sdk.get_current_span().to_json()
389-
assert span_after_record["data"]["version"] == "memoized"
387+
assert span_after_record["attributes"]["version"] == "memoized"
390388

391389

392390
def test_record_in_ignored_span(sentry_init):
@@ -403,18 +401,18 @@ def span_filter(metadata: Dict[str, object]) -> bool:
403401
)
404402
sentry_init(integrations=[integration], traces_sample_rate=1.0)
405403

406-
with start_transaction():
404+
with start_span():
407405
rust_tracing.new_span(RustTracingLevel.Info, 3)
408406

409407
span_before_record = sentry_sdk.get_current_span().to_json()
410-
assert span_before_record["data"]["version"] is None
408+
assert "version" not in span_before_record["attributes"]
411409

412410
rust_tracing.new_span(RustTracingLevel.Trace, 5)
413411
rust_tracing.record(5)
414412

415413
# `on_record()` should not do anything to the current Sentry span if the associated Rust span was ignored
416414
span_after_record = sentry_sdk.get_current_span().to_json()
417-
assert span_after_record["data"]["version"] is None
415+
assert "version" not in span_after_record["attributes"]
418416

419417

420418
@pytest.mark.parametrize(
@@ -443,33 +441,37 @@ def test_include_tracing_fields(
443441
traces_sample_rate=1.0,
444442
send_default_pii=send_default_pii,
445443
)
446-
with start_transaction():
444+
with start_span():
447445
rust_tracing.new_span(RustTracingLevel.Info, 3)
448446

449447
span_before_record = sentry_sdk.get_current_span().to_json()
450448
if tracing_fields_expected:
451-
assert span_before_record["data"]["version"] is None
449+
assert "version" not in span_before_record["attributes"]
452450
else:
453-
assert span_before_record["data"]["version"] == "[Filtered]"
451+
assert span_before_record["attributes"]["version"] == "[Filtered]"
454452

455453
rust_tracing.record(3)
456454

457455
span_after_record = sentry_sdk.get_current_span().to_json()
458456

459457
if tracing_fields_expected:
460-
assert span_after_record["data"] == {
461-
"thread.id": mock.ANY,
462-
"thread.name": mock.ANY,
463-
"use_memoized": True,
464-
"version": "memoized",
465-
"index": 10,
466-
}
458+
assert span_after_record["attributes"] == ApproxDict(
459+
{
460+
"thread.id": mock.ANY,
461+
"thread.name": mock.ANY,
462+
"use_memoized": True,
463+
"version": "memoized",
464+
"index": 10,
465+
}
466+
)
467467

468468
else:
469-
assert span_after_record["data"] == {
470-
"thread.id": mock.ANY,
471-
"thread.name": mock.ANY,
472-
"use_memoized": "[Filtered]",
473-
"version": "[Filtered]",
474-
"index": "[Filtered]",
475-
}
469+
assert span_after_record["attributes"] == ApproxDict(
470+
{
471+
"thread.id": mock.ANY,
472+
"thread.name": mock.ANY,
473+
"use_memoized": "[Filtered]",
474+
"version": "[Filtered]",
475+
"index": "[Filtered]",
476+
}
477+
)

0 commit comments

Comments
 (0)
0