feat(tornado): Support span streaming

sentry_sdk/integrations/tornado.py

@@ -4,9 +4,11 @@
 
 import sentry_sdk
 from sentry_sdk.api import continue_trace
-from sentry_sdk.consts import OP
+from sentry_sdk.consts import OP, SPANDATA
 from sentry_sdk.scope import should_send_default_pii
+from sentry_sdk.traces import SegmentSource, StreamedSpan
 from sentry_sdk.tracing import TransactionSource
+from sentry_sdk.tracing_utils import has_span_streaming_enabled
 from sentry_sdk.utils import (
     HAS_REAL_CONTEXTVARS,
     CONTEXTVARS_ERROR_MESSAGE,
@@ -34,10 +36,14 @@
 
 if TYPE_CHECKING:
     from typing import Any
+    from typing import ContextManager
     from typing import Optional
     from typing import Dict
     from typing import Callable
     from typing import Generator
+    from typing import Union
+
+    from sentry_sdk.tracing import Span
 
     from sentry_sdk._types import Event, EventProcessor
 
@@ -101,6 +107,9 @@ def sentry_log_exception(
         RequestHandler.log_exception = sentry_log_exception
 
 
+_DEFAULT_TRANSACTION_NAME = "generic Tornado request"
+
+
 @contextlib.contextmanager
 def _handle_request_impl(self: "RequestHandler") -> "Generator[None, None, None]":
     integration = sentry_sdk.get_client().get_integration(TornadoIntegration)
@@ -110,6 +119,8 @@ def _handle_request_impl(self: "RequestHandler") -> "Generator[None, None, None]
         return
 
     weak_handler = weakref.ref(self)
+    client = sentry_sdk.get_client()
+    is_span_streaming_enabled = has_span_streaming_enabled(client.options)
 
     with sentry_sdk.isolation_scope() as scope:
         headers = self.request.headers
@@ -118,22 +129,91 @@ def _handle_request_impl(self: "RequestHandler") -> "Generator[None, None, None]
         processor = _make_event_processor(weak_handler)
         scope.add_event_processor(processor)
 
-        transaction = continue_trace(
-            headers,
-            op=OP.HTTP_SERVER,
-            # Like with all other integrations, this is our
-            # fallback transaction in case there is no route.
-            # sentry_urldispatcher_resolve is responsible for
-            # setting a transaction name later.
-            name="generic Tornado request",
-            source=TransactionSource.ROUTE,
-            origin=TornadoIntegration.origin,
-        )
-
-        with sentry_sdk.start_transaction(
-            transaction, custom_sampling_context={"tornado_request": self.request}
-        ):
-            yield
+        span_ctx: "ContextManager[Union[Span, StreamedSpan, None]]"
+
+        if is_span_streaming_enabled:
+            sentry_sdk.traces.continue_trace(dict(headers))
+            scope.set_custom_sampling_context({"tornado_request": self.request})
+
+            span_ctx = sentry_sdk.traces.start_span(
+                name=_DEFAULT_TRANSACTION_NAME,
+                attributes={
+                    "sentry.op": OP.HTTP_SERVER,
+                    "sentry.origin": TornadoIntegration.origin,
+                    "sentry.span.source": SegmentSource.ROUTE,
+                },
+                parent_span=None,
+            )
+        else:
+            transaction = continue_trace(
+                headers,
+                op=OP.HTTP_SERVER,
+                # Like with all other integrations, this is our
+                # fallback transaction in case there is no route.
+                # sentry_urldispatcher_resolve is responsible for
+                # setting a transaction name later.
+                name=_DEFAULT_TRANSACTION_NAME,
+                source=TransactionSource.ROUTE,
+                origin=TornadoIntegration.origin,
+            )
+            span_ctx = sentry_sdk.start_transaction(
+                transaction,
+                custom_sampling_context={"tornado_request": self.request},
+            )
+
+        with span_ctx as span:
+            if isinstance(span, StreamedSpan):
+                with capture_internal_exceptions():
+                    for attr, value in _get_request_attributes(self.request).items():
+                        span.set_attribute(attr, value)
+
+                    method = getattr(self, self.request.method.lower(), None)
+                    if method is not None:
+                        span_name = transaction_from_function(method) or ""
+                        if span_name:
+                            span.name = span_name
+                            span.set_attribute(
+                                "sentry.span.source",
+                                SegmentSource.COMPONENT,
+                            )
+
+            try:
+                yield
+            finally:
+                if isinstance(span, StreamedSpan):
+                    with capture_internal_exceptions():
+                        status_int = self.get_status()
+                        span.set_attribute(SPANDATA.HTTP_STATUS_CODE, status_int)
+                        span.status = "error" if status_int >= 400 else "ok"
+
+
+def _get_request_attributes(request: "Any") -> "Dict[str, Any]":
+    attributes = {}  # type: Dict[str, Any]
+
+    if request.method:
+        attributes[SPANDATA.HTTP_REQUEST_METHOD] = request.method.upper()
+
+    headers = _filter_headers(dict(request.headers), use_annotated_value=False)
+    for header, value in headers.items():
+        attributes[f"http.request.header.{header.lower()}"] = value
+
+    if request.query:
+        attributes[SPANDATA.HTTP_QUERY] = request.query
+
+    attributes[SPANDATA.URL_FULL] = "%s://%s%s" % (
+        request.protocol,
+        request.host,
+        request.path,
+    )
+
+    if request.protocol:
+        attributes["network.protocol.name"] = request.protocol
+
+    if should_send_default_pii() and request.remote_ip:
+        attributes["client.address"] = request.remote_ip
+        attributes["user.ip_address"] = request.remote_ip
+
+    return attributes
 
 
 @ensure_integration_enabled(TornadoIntegration)

tests/integrations/tornado/test_tornado.py

@@ -104,86 +104,143 @@ def test_basic(tornado_testcase, sentry_init, capture_events):
     assert not sentry_sdk.get_isolation_scope()._tags
 
 
+@pytest.mark.parametrize("span_streaming", [True, False])
 @pytest.mark.parametrize(
     "handler,code",
     [
         (CrashingHandler, 500),
         (HelloHandler, 200),
     ],
 )
-def test_transactions(tornado_testcase, sentry_init, capture_events, handler, code):
-    sentry_init(integrations=[TornadoIntegration()], traces_sample_rate=1.0)
-    events = capture_events()
+def test_transactions(
+    tornado_testcase,
+    sentry_init,
+    capture_events,
+    capture_items,
+    handler,
+    code,
+    span_streaming,
+):
+    sentry_init(
+        integrations=[TornadoIntegration()],
+        traces_sample_rate=1.0,
+        _experiments={"trace_lifecycle": "stream" if span_streaming else "static"},
+    )
+
+    if span_streaming:
+        items = capture_items("event", "span")
+    else:
+        events = capture_events()
+
     client = tornado_testcase(Application([(r"/hi", handler)]))
 
-    with start_transaction(name="client") as span:
-        pass
+    if span_streaming:
+        with sentry_sdk.traces.start_span(name="client") as span:
+            request_headers = dict(span._iter_headers())
+    else:
+        with start_transaction(name="client") as span:
+            pass
+        request_headers = dict(span.iter_headers())
 
     response = client.fetch(
-        "/hi", method="POST", body=b"heyoo", headers=dict(span.iter_headers())
+        "/hi", method="POST", body=b"heyoo", headers=request_headers
     )
     assert response.code == code
 
-    if code == 200:
-        client_tx, server_tx = events
-        server_error = None
-    else:
-        client_tx, server_error, server_tx = events
-
-    assert client_tx["type"] == "transaction"
-    assert client_tx["transaction"] == "client"
-    assert client_tx["transaction_info"] == {
-        "source": "custom"
-    }  # because this is just the start_transaction() above.
-
-    if server_error is not None:
-        assert server_error["exception"]["values"][0]["type"] == "ZeroDivisionError"
-        assert (
-            server_error["transaction"]
-            == "tests.integrations.tornado.test_tornado.CrashingHandler.post"
-        )
-        assert server_error["transaction_info"] == {"source": "component"}
-
-    if code == 200:
-        assert (
-            server_tx["transaction"]
-            == "tests.integrations.tornado.test_tornado.HelloHandler.post"
+    sentry_sdk.flush()
+
+    if span_streaming:
+        spans = [i.payload for i in items if i.type == "span"]
+        errors = [i.payload for i in items if i.type == "event"]
+
+        client_segment, server_segment = spans
+
+        if code == 500:
+            assert len(errors) == 1
+            server_error = errors[0]
+            assert server_error["exception"]["values"][0]["type"] == "ZeroDivisionError"
+            assert (
+                server_error["transaction"]
+                == "tests.integrations.tornado.test_tornado.CrashingHandler.post"
+            )
+            assert server_error["transaction_info"] == {"source": "component"}
+            assert (
+                server_error["contexts"]["trace"]["trace_id"]
+                == server_segment["trace_id"]
+            )
+
+        expected_handler = (
+            "tests.integrations.tornado.test_tornado.HelloHandler.post"
+            if code == 200
+            else "tests.integrations.tornado.test_tornado.CrashingHandler.post"
         )
+        assert server_segment["name"] == expected_handler
+        assert server_segment["attributes"]["sentry.span.source"] == "component"
+        assert server_segment["attributes"]["http.request.method"] == "POST"
+        assert server_segment["attributes"]["http.response.status_code"] == code
+        assert server_segment["status"] == ("ok" if code == 200 else "error")
+        assert client_segment["trace_id"] == server_segment["trace_id"]
     else:
-        assert (
-            server_tx["transaction"]
-            == "tests.integrations.tornado.test_tornado.CrashingHandler.post"
-        )
-
-    assert server_tx["transaction_info"] == {"source": "component"}
-    assert server_tx["type"] == "transaction"
-
-    request = server_tx["request"]
-    host = request["headers"]["Host"]
-    assert server_tx["request"] == {
-        "env": {"REMOTE_ADDR": "127.0.0.1"},
-        "headers": {
-            "Accept-Encoding": "gzip",
-            "Connection": "close",
-            **request["headers"],
-        },
-        "method": "POST",
-        "query_string": "",
-        "data": {"heyoo": [""]},
-        "url": "http://{host}/hi".format(host=host),
-    }
-
-    assert (
-        client_tx["contexts"]["trace"]["trace_id"]
-        == server_tx["contexts"]["trace"]["trace_id"]
-    )
+        if code == 200:
+            client_tx, server_tx = events
+            server_error = None
+        else:
+            client_tx, server_error, server_tx = events
+
+        assert client_tx["type"] == "transaction"
+        assert client_tx["transaction"] == "client"
+        assert client_tx["transaction_info"] == {
+            "source": "custom"
+        }  # because this is just the start_transaction() above.
+
+        if server_error is not None:
+            assert server_error["exception"]["values"][0]["type"] == "ZeroDivisionError"
+            assert (
+                server_error["transaction"]
+                == "tests.integrations.tornado.test_tornado.CrashingHandler.post"
+            )
+            assert server_error["transaction_info"] == {"source": "component"}
+
+        if code == 200:
+            assert (
+                server_tx["transaction"]
+                == "tests.integrations.tornado.test_tornado.HelloHandler.post"
+            )
+        else:
+            assert (
+                server_tx["transaction"]
+                == "tests.integrations.tornado.test_tornado.CrashingHandler.post"
+            )
+
+        assert server_tx["transaction_info"] == {"source": "component"}
+        assert server_tx["type"] == "transaction"
+
+        request = server_tx["request"]
+        host = request["headers"]["Host"]
+        assert server_tx["request"] == {
+            "env": {"REMOTE_ADDR": "127.0.0.1"},
+            "headers": {
+                "Accept-Encoding": "gzip",
+                "Connection": "close",
+                **request["headers"],
+            },
+            "method": "POST",
+            "query_string": "",
+            "data": {"heyoo": [""]},
+            "url": "http://{host}/hi".format(host=host),
+        }
 
-    if server_error is not None:
         assert (
-            server_error["contexts"]["trace"]["trace_id"]
+            client_tx["contexts"]["trace"]["trace_id"]
             == server_tx["contexts"]["trace"]["trace_id"]
         )
 
+        if server_error is not None:
+            assert (
+                server_error["contexts"]["trace"]["trace_id"]
+                == server_tx["contexts"]["trace"]["trace_id"]
+            )
+
 
 def test_400_not_logged(tornado_testcase, sentry_init, capture_events):
     sentry_init(integrations=[TornadoIntegration()])
@@ -438,15 +495,32 @@ def test_error_has_existing_trace_context_performance_disabled(
     )
 
 
-def test_span_origin(tornado_testcase, sentry_init, capture_events):
-    sentry_init(integrations=[TornadoIntegration()], traces_sample_rate=1.0)
-    events = capture_events()
+@pytest.mark.parametrize("span_streaming", [True, False])
+def test_span_origin(
+    tornado_testcase, sentry_init, capture_events, capture_items, span_streaming
+):
+    sentry_init(
+        integrations=[TornadoIntegration()],
+        traces_sample_rate=1.0,
+        _experiments={"trace_lifecycle": "stream" if span_streaming else "static"},
+    )
+
+    if span_streaming:
+        items = capture_items("span")
+    else:
+        events = capture_events()
+
     client = tornado_testcase(Application([(r"/hi", CrashingHandler)]))
 
     client.fetch(
         "/hi?foo=bar", headers={"Cookie": "name=value; name2=value2; name3=value3"}
     )
 
-    (_, event) = events
+    sentry_sdk.flush()
 
-    assert event["contexts"]["trace"]["origin"] == "auto.http.tornado"
+    if span_streaming:
+        (segment,) = [i.payload for i in items]
+        assert segment["attributes"]["sentry.origin"] == "auto.http.tornado"
+    else:
+        (_, event) = events
+        assert event["contexts"]["trace"]["origin"] == "auto.http.tornado"