Preserve baggage for automatic Sentry events

microbootstrap/instruments/opentelemetry_instrument.py

@@ -21,6 +21,7 @@
 from opentelemetry.util._importlib_metadata import entry_points
 
 from microbootstrap.instruments.base import BaseInstrumentConfig, Instrument
+from microbootstrap.instruments.sentry_instrument import snapshot_sentry_opentelemetry_baggage
 
 
 LOGGER_OBJ: typing.Final = structlog.get_logger(__name__)
@@ -55,6 +56,10 @@ def opentelemetry_baggage_scope(
  token: typing.Final = context.attach(baggage_context)
  try:
  yield
+ except Exception as exc:
+ with contextlib.suppress(Exception):
+ snapshot_sentry_opentelemetry_baggage(exc)
+ raise
  finally:
  context.detach(token)
 

microbootstrap/instruments/sentry_instrument.py

@@ -3,13 +3,14 @@
 import functools
 import typing
 import urllib.parse
+from collections.abc import Mapping
 
 import orjson
 import pydantic
 import sentry_sdk
 from opentelemetry import baggage
 from sentry_sdk import _types as sentry_types
-from sentry_sdk.integrations import Integration # noqa: TC002
+from sentry_sdk.integrations import Integration
 
 from microbootstrap.instruments.base import BaseInstrumentConfig, Instrument
 
@@ -66,6 +67,71 @@ def enrich_sentry_event_from_structlog_log(event: sentry_types.Event, _hint: sen
 
 SENTRY_EXTRA_OTEL_TRACE_ID_KEY: typing.Final = "otelTraceID"
 SENTRY_EXTRA_OTEL_TRACE_URL_KEY: typing.Final = "otelTraceURL"
+SENTRY_OTEL_BAGGAGE_SNAPSHOT_ATTRIBUTE: typing.Final = "__microbootstrap_sentry_opentelemetry_baggage__"
+
+
+@typing.final
+class SentryOpentelemetryBaggageIntegration(Integration):
+ identifier = "microbootstrap_opentelemetry_baggage"
+
+ def __init__(self, baggage_keys: set[str], baggage_url_templates: dict[str, str]) -> None:
+ self.baggage_keys: typing.Final = frozenset(baggage_keys)
+ self.baggage_url_templates: typing.Final = dict(baggage_url_templates)
+
+ @staticmethod
+ def setup_once() -> None:
+ pass
+
+
+def snapshot_sentry_opentelemetry_baggage(exception: BaseException) -> None:
+ integration = sentry_sdk.get_client().get_integration(SentryOpentelemetryBaggageIntegration)
+ if not isinstance(integration, SentryOpentelemetryBaggageIntegration) or hasattr(
+ exception,
+ SENTRY_OTEL_BAGGAGE_SNAPSHOT_ATTRIBUTE,
+ ):
+ return
+
+ configured_keys: typing.Final = integration.baggage_keys.union(integration.baggage_url_templates)
+ snapshot: typing.Final = {key: baggage.get_baggage(key) for key in configured_keys}
+ with contextlib.suppress(AttributeError, TypeError):
+ setattr(exception, SENTRY_OTEL_BAGGAGE_SNAPSHOT_ATTRIBUTE, snapshot)
+
+
+def _find_sentry_opentelemetry_baggage_snapshot(hint: sentry_types.Hint) -> Mapping[str, object | None] | None:
+ if not isinstance(hint, Mapping):
+ return None
+
+ exc_info: typing.Final = hint.get("exc_info")
+ if not isinstance(exc_info, tuple):
+ return None
+ try:
+ _, exception_value, _ = exc_info
+ except ValueError:
+ return None
+ if not isinstance(exception_value, BaseException):
+ return None
+
+ exceptions_to_visit: list[BaseException] = [exception_value]
+ visited_exceptions: set[int] = set()
+ while exceptions_to_visit:
+ exception = exceptions_to_visit.pop()
+ if id(exception) in visited_exceptions:
+ continue
+ visited_exceptions.add(id(exception))
+ exception_snapshot = getattr(exception, SENTRY_OTEL_BAGGAGE_SNAPSHOT_ATTRIBUTE, None)
+ if isinstance(exception_snapshot, Mapping):
+ return exception_snapshot
+
+ nested_exceptions = getattr(exception, "exceptions", ())
+ if isinstance(nested_exceptions, tuple):
+ exceptions_to_visit.extend(
+ reversed([nested for nested in nested_exceptions if isinstance(nested, BaseException)])
+ )
+ if chained_exception := exception.__cause__ or (
+ None if exception.__suppress_context__ else exception.__context__
+ ):
+ exceptions_to_visit.append(chained_exception)
+ return None
 
 
 def add_trace_url_to_event(
@@ -80,27 +146,30 @@ def enrich_sentry_event_from_opentelemetry_baggage(
  baggage_keys: set[str],
  baggage_url_templates: dict[str, str],
  event: sentry_types.Event,
-_hint: sentry_types.Hint,
+hint: sentry_types.Hint,
 ) -> sentry_types.Event:
- baggage_values = {
- key: value
- for key in baggage_keys.union(baggage_url_templates)
- if (value := baggage.get_baggage(key)) is not None
- }
- if not baggage_values:
- return event
-
- if tag_values := {key: str(baggage_values[key]) for key in baggage_keys if key in baggage_values}:
- event.setdefault("tags", {}).update(tag_values)
+ configured_keys: typing.Final = baggage_keys.union(baggage_url_templates)
+ snapshot: typing.Final = _find_sentry_opentelemetry_baggage_snapshot(hint)
+ baggage_values: typing.Final = (
+ snapshot if snapshot is not None else {key: baggage.get_baggage(key) for key in configured_keys}
+ )
+
+ for key in baggage_keys:
+ if (value := baggage_values.get(key)) is not None:
+ event.setdefault("tags", {})[key] = str(value)
+ elif tags := event.get("tags"):
+ tags.pop(key, None)
 
  for key, url_template in baggage_url_templates.items():
  placeholder = f"{{{key}}}"
- if key not in baggage_values or placeholder not in url_template:
- continue
- event.setdefault("extra", {})[f"{key}_url"] = url_template.replace(
- placeholder,
- urllib.parse.quote(str(baggage_values[key]), safe=""),
- )
+ extra_key = f"{key}_url"
+ if (value := baggage_values.get(key)) is not None and placeholder in url_template:
+ event.setdefault("extra", {})[extra_key] = url_template.replace(
+ placeholder,
+ urllib.parse.quote(str(value), safe=""),
+ )
+ elif extra := event.get("extra"):
+ extra.pop(extra_key, None)
  return event
 
 
@@ -126,6 +195,17 @@ def is_ready(self) -> bool:
  return bool(self.instrument_config.sentry_dsn)
 
  def bootstrap(self) -> None:
+ baggage_integration: typing.Final = (
+ SentryOpentelemetryBaggageIntegration(
+ self.instrument_config.sentry_opentelemetry_baggage_keys,
+ self.instrument_config.sentry_opentelemetry_baggage_url_templates,
+ )
+ if (
+ self.instrument_config.sentry_opentelemetry_baggage_keys
+ or self.instrument_config.sentry_opentelemetry_baggage_url_templates
+ )
+ else None
+ )
  sentry_sdk.init(
  dsn=self.instrument_config.sentry_dsn,
  sample_rate=self.instrument_config.sentry_sample_rate,
@@ -153,7 +233,10 @@ def bootstrap(self) -> None:
  else None,
  self.instrument_config.sentry_before_send,
  ),
- integrations=self.instrument_config.sentry_integrations,
+ integrations=[
+ *self.instrument_config.sentry_integrations,
+ *([baggage_integration] if baggage_integration else []),
+ ],
  **self.instrument_config.sentry_additional_params,
  )
  if self.instrument_config.sentry_tags:

tests/bootstrappers/test_faststream.py

@@ -13,6 +13,7 @@
 from faststream.redis.opentelemetry import RedisTelemetryMiddleware
 from faststream.redis.prometheus import RedisPrometheusMiddleware
 
+from microbootstrap import opentelemetry_baggage_scope
 from microbootstrap.bootstrappers.faststream import FastStreamBootstrapper
 from microbootstrap.config.faststream import FastStreamConfig
 from microbootstrap.instruments.health_checks_instrument import HealthChecksConfig
@@ -192,6 +193,81 @@ def record_error_tag(*args: typing.Any, **kwargs: typing.Any) -> None: # noqa:
  assert conversation_id_tag not in sentry_sdk.get_isolation_scope()._tags # noqa: SLF001
 
 
+async def test_faststream_sentry_automatic_errors_use_concurrent_baggage_snapshots(
+ broker: RedisBroker,
+ minimal_sentry_config: SentryConfig,
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ channel: typing.Final = "test-channel"
+ conversation_id_tag: typing.Final = "conversation_id"
+ first_started = asyncio.Event()
+ second_started = asyncio.Event()
+ second_captured = asyncio.Event()
+ captured_tags: dict[str, str | None] = {}
+
+ init = mock.Mock()
+ monkeypatch.setattr(sentry_sdk, "init", init)
+ minimal_sentry_config.sentry_tags = None
+ minimal_sentry_config.sentry_opentelemetry_baggage_keys = {conversation_id_tag}
+
+ @broker.subscriber(channel)
+ async def handler(conversation_id: str) -> None:
+ with opentelemetry_baggage_scope({conversation_id_tag: conversation_id}):
+ if conversation_id == "first":
+ first_started.set()
+ await second_started.wait()
+ await second_captured.wait()
+ else:
+ second_started.set()
+ await first_started.wait()
+ raise ValueError(conversation_id)
+
+ FastStreamBootstrapper(FastStreamSettings()).configure_application(
+ FastStreamConfig(broker=broker)
+ ).configure_instruments(minimal_sentry_config).bootstrap()
+
+ baggage_integration: typing.Final = next(
+ integration
+ for integration in init.call_args.kwargs["integrations"]
+ if integration.identifier == "microbootstrap_opentelemetry_baggage"
+ )
+ client = mock.Mock()
+ client.get_integration.return_value = baggage_integration
+ monkeypatch.setattr(sentry_sdk, "get_client", mock.Mock(return_value=client))
+ before_send: typing.Final = init.call_args.kwargs["before_send"]
+ original_log = broker.config.logger.log
+
+ def record_automatic_error(*args: typing.Any, **kwargs: typing.Any) -> None: # noqa: ANN401
+ if kwargs.get("log_level") == logging.ERROR:
+ exception = typing.cast("ValueError", kwargs["exc_info"])
+ event = before_send(
+ {},
+ {"exc_info": (type(exception), exception, exception.__traceback__)},
+ )
+ captured_tags[str(exception)] = event.get("tags", {}).get(conversation_id_tag)
+ if str(exception) == "second":
+ second_captured.set()
+ original_log(*args, **kwargs)
+
+ monkeypatch.setattr(broker.config.logger, "log", record_automatic_error)
+
+ event_loop = asyncio.get_running_loop()
+ previous_exception_handler = event_loop.get_exception_handler()
+ event_loop.set_exception_handler(lambda *_: None)
+ try:
+ async with TestRedisBroker(broker):
+ errors: typing.Final = await asyncio.gather(
+ broker.publish("first", channel),
+ broker.publish("second", channel),
+ return_exceptions=True,
+ )
+ finally:
+ event_loop.set_exception_handler(previous_exception_handler)
+
+ assert all(isinstance(error, ValueError) for error in errors)
+ assert captured_tags == {"first": "first", "second": "second"}
+
+
 async def test_faststream_sentry_isolates_broker_configured_on_startup(
  broker: RedisBroker,
  minimal_sentry_config: SentryConfig,

tests/instruments/test_opentelemetry.py

@@ -48,6 +48,22 @@ def test_opentelemetry_baggage_scope_overrides_removes_and_restores_values() ->
  context.detach(outer_token)
 
 
+def test_opentelemetry_baggage_scope_does_not_replace_exception_when_snapshot_fails(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ monkeypatch.setattr(
+ opentelemetry_instrument,
+ "snapshot_sentry_opentelemetry_baggage",
+ Mock(side_effect=RuntimeError("snapshot failed")),
+ )
+
+ with (
+ pytest.raises(ValueError, match="application error"),
+ opentelemetry_baggage_scope({"conversation_id": "conversation-1"}),
+ ):
+ raise ValueError("application error")
+
+
 @pytest.mark.parametrize(
  ("span_kind", "expected_attribute"),
  [