address potential data loss issue - https://github.com/getsentry/sentry-python/pull/6503\#discussion_r3357864384

Author: ericapisani

sentry_sdk/integrations/chalice.py

@@ -11,7 +11,7 @@
     CLOUD_PLATFORM,
     CLOUD_PROVIDER,
 )
-from sentry_sdk.traces import SegmentSource
+from sentry_sdk.traces import SegmentSource, SpanStatus
 from sentry_sdk.tracing import TransactionSource
 from sentry_sdk.tracing_utils import has_span_streaming_enabled
 from sentry_sdk.utils import (
@@ -51,23 +51,27 @@ def __call__(self, event: "Any", context: "Any") -> "Any":
                 )
 
             if has_span_streaming_enabled(client.options):
-                with sentry_sdk.traces.start_span(
+                span = sentry_sdk.traces.start_span(
                     name=context.function_name,
                     parent_span=None,
                     attributes=_get_lambda_span_attributes(context),
-                ):
-                    try:
-                        return ChaliceEventSourceHandler.__call__(self, event, context)
-                    except Exception:
-                        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()
-                        reraise(*exc_info)
+                )
+                try:
+                    return ChaliceEventSourceHandler.__call__(self, event, context)
+                except Exception:
+                    exc_info = sys.exc_info()
+                    span.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)
+                    reraise(*exc_info)
+                finally:
+                    span.end()
+                    client.flush()
+
             else:
                 try:
                     return ChaliceEventSourceHandler.__call__(self, event, context)
@@ -115,29 +119,32 @@ def wrapped_view_function(**function_args: "Any") -> "Any":
                 if "method" in request_dict:
                     additional_attrs["http.request.method"] = request_dict["method"]
 
-                with sentry_sdk.traces.start_span(
+                span = sentry_sdk.traces.start_span(
                     name=aws_context.function_name,
                     parent_span=None,
                     attributes={
                         **_get_lambda_span_attributes(aws_context),
                         **header_attrs,
                         **additional_attrs,
                     },
-                ):
-                    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()
+                )
+                try:
+                    return view_function(**function_args)
+                except Exception as exc:
+                    if isinstance(exc, ChaliceViewError):
                         raise
+                    exc_info = sys.exc_info()
+                    span.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)
+                    raise
+                finally:
+                    span.end()
+                    client.flush()
             else:
                 scope.set_transaction_name(
                     app.lambda_context.function_name,

tests/integrations/chalice/test_chalice.py

@@ -236,6 +236,7 @@ def test_span_streaming_basic(
     assert attrs["aws.log.stream.names"] == ["2024/01/01/[$LATEST]abcdef1234567890"]
     assert attrs["http.request.method"] == "GET"
     assert attrs["http.request.header.x-custom-header"] == "test-value"
+    assert span["status"] == "ok"
 
 
 def test_span_streaming_error(
@@ -263,6 +264,29 @@ def test_span_streaming_error(
     attrs = segment_spans[0]["attributes"]
     assert attrs["sentry.op"] == "function.aws"
     assert attrs["sentry.origin"] == "auto.function.chalice"
+    assert segment_spans[0]["status"] == "error"
+
+
+def test_span_streaming_error_flush_ordering(
+    sentry_init,
+    capture_items,
+):
+    """The handler's own client.flush() must send the segment span.
+
+    If flush runs before the span ends, the segment won't be in the
+    transport without an extra sentry_sdk.flush() call after the request.
+    On Lambda, the worker can freeze right after the response, so there's
+    no second chance.
+    """
+    app = _make_span_streaming_app(sentry_init)
+    client = RequestHandler(app)
+    items = capture_items("span")
+
+    response = client.get("/boom")
+    assert response.status_code == 500
+
+    segment_spans = [s.payload for s in items if s.payload.get("is_segment")]
+    assert len(segment_spans) == 1
 
 
 def test_span_streaming_scheduled_event(
@@ -328,3 +352,4 @@ def every_hour(event):
     assert attrs["faas.version"] == "$LATEST"
     assert attrs["aws.log.group.names"] == ["/aws/lambda/lambda_name"]
     assert attrs["aws.log.stream.names"] == ["2024/01/01/[$LATEST]abcdef1234567890"]
+    assert segment_spans[0]["status"] == "error"