feat(chalice): Add span streaming support to Chalice integration

sentry_sdk/integrations/chalice.py

@@ -3,8 +3,15 @@
 
 import sentry_sdk
 from sentry_sdk.integrations import DidNotEnable, Integration
+from sentry_sdk.integrations._wsgi_common import _filter_headers
 from sentry_sdk.integrations.aws_lambda import _make_request_event_processor
+from sentry_sdk.traces import (
+    SpanStatus,
+    StreamedSpan,
+    get_current_span,
+)
 from sentry_sdk.tracing import TransactionSource
+from sentry_sdk.tracing_utils import has_span_streaming_enabled
 from sentry_sdk.utils import (
     capture_internal_exceptions,
     event_from_exception,
@@ -63,38 +70,88 @@
         with sentry_sdk.isolation_scope() as scope:
             with capture_internal_exceptions():
                 configured_time = app.lambda_context.get_remaining_time_in_millis()
-                scope.set_transaction_name(
-                    app.lambda_context.function_name,
-                    source=TransactionSource.COMPONENT,
-                )
-
                 scope.add_event_processor(
                     _make_request_event_processor(
                         app.current_request.to_dict(),
                         app.lambda_context,
                         configured_time,
                     )
                 )
-            try:
-                return view_function(**function_args)
-            except Exception as exc:
-                if isinstance(exc, ChaliceViewError):
+
+            if has_span_streaming_enabled(client.options):
+                current_span = get_current_span()
+                segment = None
+                if type(current_span) is StreamedSpan:
+                    # A segment already exists (created by the AWS Lambda
+                    # integration), so decorate it with Chalice attributes
+                    # The AWS Lambda integration owns the span lifecycle
+                    # (end + flush), but Chalice converts unhandled view exceptions
+                    # into 500 responses, so the error must be captured here.
+                    request_dict = app.current_request.to_dict()
+                    headers = request_dict.get("headers", {})
+
+                    header_attrs: "Dict[str, Any]" = {}
+                    for header, value in _filter_headers(
+                        headers, use_annotated_value=False
+                    ).items():
+                        header_attrs[f"http.request.header.{header.lower()}"] = value
+
+                    additional_attrs: "Dict[str, Any]" = {}
+                    if "method" in request_dict:
+                        additional_attrs["http.request.method"] = request_dict["method"]
+
+                    attributes = {
+                        "sentry.origin": ChaliceIntegration.origin,
+                        **header_attrs,
+                        **additional_attrs,
+                    }
+
+                    segment = current_span._segment
+                    segment.set_attributes(attributes)
+
+                try:
+                    return view_function(**function_args)
+                except Exception as exc:
+                    if isinstance(exc, ChaliceViewError):
+                        raise
+                    exc_info = sys.exc_info()
+                    if segment:
+                        segment.status = SpanStatus.ERROR.value
+                    sentry_event, hint = event_from_exception(
+                        exc_info,
+                        client_options=client.options,
+                        mechanism={"type": "chalice", "handled": False},
+                    )
+                    sentry_sdk.capture_event(sentry_event, hint=hint)
+                    if segment is None:
+                        client.flush()
                     raise
-                exc_info = sys.exc_info()
-                event, hint = event_from_exception(
-                    exc_info,
-                    client_options=client.options,
-                    mechanism={"type": "chalice", "handled": False},
+            else:
+                scope.set_transaction_name(
+                    app.lambda_context.function_name,
+                    source=TransactionSource.COMPONENT,
                 )
-                sentry_sdk.capture_event(event, hint=hint)
-                client.flush()
-                raise
+                try:
+                    return view_function(**function_args)
+                except Exception as exc:
+                    if isinstance(exc, ChaliceViewError):
+                        raise
+                    exc_info = sys.exc_info()
+                    sentry_event, hint = event_from_exception(
+                        exc_info,
+                        client_options=client.options,
+                        mechanism={"type": "chalice", "handled": False},
+                    )
+                    sentry_sdk.capture_event(sentry_event, hint=hint)
+                    client.flush()
+                    raise
 
     return wrapped_view_function  # type: ignore
 
 
 class ChaliceIntegration(Integration):
     identifier = "chalice"
+    origin = f"auto.function.{identifier}"
 
     @staticmethod
     def setup_once() -> None:

tests/integrations/chalice/test_chalice.py

@@ -5,11 +5,23 @@
 from chalice.local import LambdaContext, LocalGateway
 from pytest_chalice.handlers import RequestHandler
 
+import sentry_sdk
 from sentry_sdk import capture_message
 from sentry_sdk.integrations.chalice import CHALICE_VERSION, ChaliceIntegration
 from sentry_sdk.utils import parse_version
 
 
+def _populate_lambda_context(context):
+    fn = context.function_name
+    context.invoked_function_arn = (
+        f"arn:aws:lambda:us-east-1:123456789012:function:{fn}"
+    )
+    context.log_group_name = f"/aws/lambda/{fn}"
+    context.log_stream_name = "2024/01/01/[$LATEST]abcdef1234567890"
+    context.aws_request_id = "test-request-id-1234"
+    return context
+
+
 def _generate_lambda_context(self):
     # Monkeypatch of the function _generate_lambda_context
     # from the class LocalGateway
@@ -19,11 +31,12 @@ def _generate_lambda_context(self):
         timeout = 10 * 1000
     else:
         timeout = self._config.lambda_timeout * 1000
-    return LambdaContext(
+    context = LambdaContext(
         function_name=self._config.function_name,
         memory_size=self._config.lambda_memory_size,
         max_runtime_ms=timeout,
     )
+    return _populate_lambda_context(context)
 
 
 @pytest.fixture
@@ -89,8 +102,8 @@ def test_scheduled_event(app, lambda_context_args):
     def every_hour(event):
         raise Exception("schedule event!")
 
-    context = LambdaContext(
-        *lambda_context_args, max_runtime_ms=10000, time_source=time
+    context = _populate_lambda_context(
+        LambdaContext(*lambda_context_args, max_runtime_ms=10000, time_source=time)
     )
 
     lambda_event = {
@@ -160,3 +173,89 @@ def test_transaction(
     (event,) = events
     assert event["transaction"] == expected_transaction
     assert event["transaction_info"] == {"source": expected_source}
+
+
+def _make_span_streaming_app(sentry_init):
+    sentry_init(
+        integrations=[ChaliceIntegration()],
+        traces_sample_rate=1.0,
+        _experiments={"trace_lifecycle": "stream"},
+    )
+    app = Chalice(app_name="sentry_chalice")
+
+    @app.route("/message")
+    def hi():
+        capture_message("hi")
+        return {"status": "ok"}
+
+    @app.route("/boom")
+    def boom():
+        raise Exception("boom goes the dynamite!")
+
+    LocalGateway._generate_lambda_context = _generate_lambda_context
+
+    return app
+
+
+def test_span_streaming_existing_span(
+    sentry_init,
+    capture_items,
+):
+    """When a segment already exists (e.g. created by the AWS Lambda
+    integration), Chalice decorates it instead of creating a duplicate."""
+    app = _make_span_streaming_app(sentry_init)
+    client = RequestHandler(app)
+    items = capture_items("span")
+
+    with sentry_sdk.traces.start_span(
+        name="lambda_segment",
+        parent_span=None,
+        attributes={
+            "sentry.origin": "auto.function.aws_lambda",
+            "sentry.op": "function.aws",
+            "faas.name": "api_handler",
+        },
+    ):
+        response = client.get("/message")
+        assert response.status_code == 200
+
+    sentry_sdk.flush()
+
+    segment_spans = [s.payload for s in items if s.payload.get("is_segment")]
+    assert len(segment_spans) == 1
+    span = segment_spans[0]
+
+    attrs = span["attributes"]
+    assert attrs["sentry.origin"] == "auto.function.chalice"
+    assert attrs["sentry.op"] == "function.aws"
+    assert attrs["faas.name"] == "api_handler"
+    assert span["status"] == "ok"
+
+
+def test_span_streaming_existing_span_error(
+    sentry_init,
+    capture_items,
+):
+    app = _make_span_streaming_app(sentry_init)
+    client = RequestHandler(app)
+    items = capture_items("event", "span")
+
+    with sentry_sdk.traces.start_span(
+        name="lambda_segment",
+        parent_span=None,
+        attributes={"sentry.origin": "auto.function.aws_lambda"},
+    ):
+        response = client.get("/boom")
+        assert response.status_code == 500
+
+    sentry_sdk.flush()
+
+    error_items = [i for i in items if i.type == "event"]
+    assert len(error_items) == 1
+
+    segment_spans = [
+        s.payload for s in items if s.type == "span" and s.payload.get("is_segment")
+    ]
+    assert len(segment_spans) == 1
+    assert segment_spans[0]["attributes"]["sentry.origin"] == "auto.function.chalice"
+    assert segment_spans[0]["status"] == "error"