diff --git a/lite_bootstrap/bootstrappers/faststream_bootstrapper.py b/lite_bootstrap/bootstrappers/faststream_bootstrapper.py index b84a1bb..18323e2 100644 --- a/lite_bootstrap/bootstrappers/faststream_bootstrapper.py +++ b/lite_bootstrap/bootstrappers/faststream_bootstrapper.py @@ -4,6 +4,7 @@ import logging import typing from collections.abc import Mapping, Sequence +from types import TracebackType from lite_bootstrap import import_checker from lite_bootstrap.bootstrappers.base import BaseBootstrapper @@ -20,6 +21,24 @@ from faststream._internal.logger.params_storage import ManualLoggerStorage from faststream.asgi import AsgiFastStream, AsgiResponse from faststream.asgi import get as handle_get + from faststream.middlewares import BaseMiddleware + + class _SentryIsolationMiddleware(BaseMiddleware[typing.Any, typing.Any]): + async def on_receive(self) -> None: + self._scope_context = sentry_sdk.isolation_scope() + self._scope_context.__enter__() + + async def after_processed( + self, + exc_type: type[BaseException] | None = None, + exc_val: BaseException | None = None, + exc_tb: TracebackType | None = None, + ) -> None: + self._scope_context.__exit__(exc_type, exc_val, exc_tb) + + +if import_checker.is_sentry_installed: + import sentry_sdk if import_checker.is_structlog_installed: import structlog @@ -227,13 +246,20 @@ def bootstrap(self) -> None: broker.add_middleware(middleware_cls(registry=self.collector_registry)) +class FastStreamSentryInstrument(SentryInstrument["FastStreamConfig"]): + def bootstrap(self) -> None: + super().bootstrap() + for broker in self.bootstrap_config.application.brokers: + broker.insert_middleware(_SentryIsolationMiddleware) + + class FastStreamBootstrapper(BaseBootstrapper["AsgiFastStream", FastStreamConfig]): __slots__ = "bootstrap_config", "instruments" instruments_types: typing.ClassVar = [ FastStreamOpenTelemetryInstrument, PyroscopeInstrument, - SentryInstrument, + FastStreamSentryInstrument, FastStreamHealthChecksInstrument, FastStreamLoggingInstrument, FastStreamPrometheusInstrument, diff --git a/tests/test_faststream_bootstrap.py b/tests/test_faststream_bootstrap.py index 5d41db5..a4858e7 100644 --- a/tests/test_faststream_bootstrap.py +++ b/tests/test_faststream_bootstrap.py @@ -5,14 +5,17 @@ import typing import uuid import warnings +from types import TracebackType from unittest.mock import AsyncMock, patch import faststream.asgi import prometheus_client import pytest +import sentry_sdk import structlog from faststream._internal.broker import BrokerUsecase from faststream._internal.logger.params_storage import ManualLoggerStorage +from faststream.middlewares import BaseMiddleware from faststream.redis import RedisBroker, TestRedisBroker from faststream.redis.opentelemetry import RedisTelemetryMiddleware from faststream.redis.prometheus import RedisPrometheusMiddleware @@ -124,6 +127,58 @@ async def test_faststream_bootstrap(broker: RedisBroker) -> None: assert not bootstrapper.is_bootstrapped +async def test_faststream_sentry_context_is_isolated_per_message(broker: RedisBroker) -> None: + transport = SentryTestTransport() + config = FastStreamConfig( + application=faststream.asgi.AsgiFastStream(broker), + sentry_dsn="https://testdsn@localhost/1", + sentry_additional_params={"transport": transport}, + ) + bootstrapper = FastStreamBootstrapper(bootstrap_config=config) + + class CaptureFailure(BaseMiddleware[typing.Any, typing.Any]): + async def after_processed( + self, + exc_type: type[BaseException] | None = None, + exc_val: BaseException | None = None, + exc_tb: TracebackType | None = None, + ) -> None: + if exc_type is not None and exc_val is not None: + sentry_sdk.capture_exception((exc_type, exc_val, exc_tb)) + + broker.add_middleware(CaptureFailure) + + @broker.subscriber("messages") + async def consume(message: str) -> None: + if message == "first": + sentry_sdk.set_tag("customer", "alice") + sentry_sdk.add_breadcrumb(message="first message breadcrumb") + return + sentry_sdk.set_tag("message", "second") + sentry_sdk.add_breadcrumb(message="second message breadcrumb") + msg = "second message failed" + raise ValueError(msg) + + with sentry_sdk.isolation_scope(): + bootstrapper.bootstrap() + try: + async with TestRedisBroker(broker): + await broker.publish("first", "messages") + with pytest.raises(ValueError, match="second message failed"): + await broker.publish("second", "messages") + finally: + bootstrapper.teardown() + + (event,) = [event for envelope in transport.mock_envelopes if (event := envelope.get_event())] + assert event["exception"]["values"][-1]["value"] == "second message failed" + assert event["tags"] == {"message": "second"} + breadcrumbs = event.get("breadcrumbs") + assert isinstance(breadcrumbs, dict) + messages = [crumb.get("message") for crumb in breadcrumbs["values"]] + assert "second message breadcrumb" in messages + assert "first message breadcrumb" not in messages + + async def test_faststream_bootstrap_health_check_wo_broker() -> None: bootstrap_config = build_faststream_config() bootstrapper = FastStreamBootstrapper(bootstrap_config=bootstrap_config) @@ -624,7 +679,7 @@ def _middleware_types(broker: BrokerUsecase[typing.Any, typing.Any, typing.Any]) def _bootstrap_with_brokers( brokers: list[BrokerUsecase[typing.Any, typing.Any, typing.Any]], **overrides: object ) -> FastStreamBootstrapper: - config = dataclasses.replace(build_faststream_config(broker=brokers[0]), **overrides) + config = dataclasses.replace(build_faststream_config(broker=brokers[0]), sentry_dsn=None, **overrides) for extra in brokers[1:]: config.application.add_broker(extra) bootstrapper = FastStreamBootstrapper(bootstrap_config=config)