Merge branch 'potel-base' into ivana/potel/custom-sampling-context-rq

Author: sentrivana

sentry_sdk/integrations/__init__.py

@@ -96,6 +96,7 @@ def iter_default_integrations(with_auto_enabling_integrations):
     "sentry_sdk.integrations.huey.HueyIntegration",
     "sentry_sdk.integrations.huggingface_hub.HuggingfaceHubIntegration",
     "sentry_sdk.integrations.langchain.LangchainIntegration",
+    "sentry_sdk.integrations.litestar.LitestarIntegration",
     "sentry_sdk.integrations.loguru.LoguruIntegration",
     "sentry_sdk.integrations.openai.OpenAIIntegration",
     "sentry_sdk.integrations.pymongo.PyMongoIntegration",

sentry_sdk/integrations/falcon.py

@@ -43,6 +43,12 @@
     FALCON3 = False
 
 
+_FALCON_UNSET = None  # type: Optional[object]
+if FALCON3:  # falcon.request._UNSET is only available in Falcon 3.0+
+    with capture_internal_exceptions():
+        from falcon.request import _UNSET as _FALCON_UNSET  # type: ignore[import-not-found, no-redef]
+
+
 class FalconRequestExtractor(RequestExtractor):
     def env(self):
         # type: () -> Dict[str, Any]
@@ -73,27 +79,23 @@ def raw_data(self):
         else:
             return None
 
-    if FALCON3:
-
-        def json(self):
-            # type: () -> Optional[Dict[str, Any]]
-            try:
-                return self.request.media
-            except falcon.errors.HTTPBadRequest:
-                return None
-
-    else:
-
-        def json(self):
-            # type: () -> Optional[Dict[str, Any]]
-            try:
-                return self.request.media
-            except falcon.errors.HTTPBadRequest:
-                # NOTE(jmagnusson): We return `falcon.Request._media` here because
-                # falcon 1.4 doesn't do proper type checking in
-                # `falcon.Request.media`. This has been fixed in 2.0.
-                # Relevant code: https://github.com/falconry/falcon/blob/1.4.1/falcon/request.py#L953
-                return self.request._media
+    def json(self):
+        # type: () -> Optional[Dict[str, Any]]
+        # fallback to cached_media = None if self.request._media is not available
+        cached_media = None
+        with capture_internal_exceptions():
+            # self.request._media is the cached self.request.media
+            # value. It is only available if self.request.media
+            # has already been accessed. Therefore, reading
+            # self.request._media will not exhaust the raw request
+            # stream (self.request.bounded_stream) because it has
+            # already been read if self.request._media is set.
+            cached_media = self.request._media
+
+        if cached_media is not _FALCON_UNSET:
+            return cached_media
+
+        return None
 
 
 class SentryFalconMiddleware:

sentry_sdk/integrations/httpx.py

@@ -2,7 +2,7 @@
 from sentry_sdk.consts import OP, SPANDATA
 from sentry_sdk.integrations import Integration, DidNotEnable
 from sentry_sdk.tracing import BAGGAGE_HEADER_NAME
-from sentry_sdk.tracing_utils import should_propagate_trace
+from sentry_sdk.tracing_utils import Baggage, should_propagate_trace
 from sentry_sdk.utils import (
     SENSITIVE_DATA_SUBSTITUTE,
     capture_internal_exceptions,
@@ -14,6 +14,7 @@
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
+    from collections.abc import MutableMapping
     from typing import Any
 
 
@@ -82,11 +83,9 @@ def send(self, request, **kwargs):
                             key=key, value=value, url=request.url
                         )
                     )
-                    if key == BAGGAGE_HEADER_NAME and request.headers.get(
-                        BAGGAGE_HEADER_NAME
-                    ):
-                        # do not overwrite any existing baggage, just append to it
-                        request.headers[key] += "," + value
+
+                    if key == BAGGAGE_HEADER_NAME:
+                        _add_sentry_baggage_to_headers(request.headers, value)
                     else:
                         request.headers[key] = value
 
@@ -178,3 +177,21 @@ async def send(self, request, **kwargs):
             return rv
 
     AsyncClient.send = send
+
+
+def _add_sentry_baggage_to_headers(headers, sentry_baggage):
+    # type: (MutableMapping[str, str], str) -> None
+    """Add the Sentry baggage to the headers.
+
+    This function directly mutates the provided headers. The provided sentry_baggage
+    is appended to the existing baggage. If the baggage already contains Sentry items,
+    they are stripped out first.
+    """
+    existing_baggage = headers.get(BAGGAGE_HEADER_NAME, "")
+    stripped_existing_baggage = Baggage.strip_sentry_baggage(existing_baggage)
+
+    separator = "," if len(stripped_existing_baggage) > 0 else ""
+
+    headers[BAGGAGE_HEADER_NAME] = (
+        stripped_existing_baggage + separator + sentry_baggage
+    )

sentry_sdk/tracing_utils.py

@@ -604,6 +604,21 @@ def serialize(self, include_third_party=False):
 
         return ",".join(items)
 
+    @staticmethod
+    def strip_sentry_baggage(header):
+        # type: (str) -> str
+        """Remove Sentry baggage from the given header.
+
+        Given a Baggage header, return a new Baggage header with all Sentry baggage items removed.
+        """
+        return ",".join(
+            (
+                item
+                for item in header.split(",")
+                if not Baggage.SENTRY_PREFIX_REGEX.match(item.strip())
+            )
+        )
+
 
 def should_propagate_trace(client, url):
     # type: (sentry_sdk.client.BaseClient, str) -> bool

tests/integrations/falcon/test_falcon.py

@@ -460,3 +460,48 @@ def test_span_origin(sentry_init, capture_events, make_client):
     (_, event) = events
 
     assert event["contexts"]["trace"]["origin"] == "auto.http.falcon"
+
+
+def test_falcon_request_media(sentry_init):
+    # test_passed stores whether the test has passed.
+    test_passed = False
+
+    # test_failure_reason stores the reason why the test failed
+    # if test_passed is False. The value is meaningless when
+    # test_passed is True.
+    test_failure_reason = "test endpoint did not get called"
+
+    class SentryCaptureMiddleware:
+        def process_request(self, _req, _resp):
+            # This capture message forces Falcon event processors to run
+            # before the request handler runs
+            sentry_sdk.capture_message("Processing request")
+
+    class RequestMediaResource:
+        def on_post(self, req, _):
+            nonlocal test_passed, test_failure_reason
+            raw_data = req.bounded_stream.read()
+
+            # If the raw_data is empty, the request body stream
+            # has been exhausted by the SDK. Test should fail in
+            # this case.
+            test_passed = raw_data != b""
+            test_failure_reason = "request body has been read"
+
+    sentry_init(integrations=[FalconIntegration()])
+
+    try:
+        app_class = falcon.App  # Falcon ≥3.0
+    except AttributeError:
+        app_class = falcon.API  # Falcon <3.0
+
+    app = app_class(middleware=[SentryCaptureMiddleware()])
+    app.add_route("/read_body", RequestMediaResource())
+
+    client = falcon.testing.TestClient(app)
+
+    client.simulate_post("/read_body", json={"foo": "bar"})
+
+    # Check that simulate_post actually calls the resource, and
+    # that the SDK does not exhaust the request body stream.
+    assert test_passed, test_failure_reason

tests/integrations/gcp/__init__.py

@@ -0,0 +1,6 @@
+import pytest
+import os
+
+
+if "gcp" not in os.environ.get("TOX_ENV_NAME", ""):
+    pytest.skip("GCP tests only run in GCP environment", allow_module_level=True)

tests/test_tracing_utils.py

@@ -1,7 +1,7 @@
 from dataclasses import asdict, dataclass
 from typing import Optional, List
 
-from sentry_sdk.tracing_utils import _should_be_included
+from sentry_sdk.tracing_utils import _should_be_included, Baggage
 import pytest
 
 
@@ -94,3 +94,24 @@ def test_should_be_included(test_case, expected):
     kwargs = asdict(test_case)
     kwargs.pop("id")
     assert _should_be_included(**kwargs) == expected
+
+
+@pytest.mark.parametrize(
+    ("header", "expected"),
+    (
+        ("", ""),
+        ("foo=bar", "foo=bar"),
+        (" foo=bar, baz =  qux ", " foo=bar, baz =  qux "),
+        ("sentry-trace_id=123", ""),
+        ("  sentry-trace_id = 123  ", ""),
+        ("sentry-trace_id=123,sentry-public_key=456", ""),
+        ("foo=bar,sentry-trace_id=123", "foo=bar"),
+        ("foo=bar,sentry-trace_id=123,baz=qux", "foo=bar,baz=qux"),
+        (
+            "foo=bar,sentry-trace_id=123,baz=qux,sentry-public_key=456",
+            "foo=bar,baz=qux",
+        ),
+    ),
+)
+def test_strip_sentry_baggage(header, expected):
+    assert Baggage.strip_sentry_baggage(header) == expected