feat(starlette): Support new failed_request_status_codes (getsentry#3563)

szokeasaurusrex · web-flow · commit 6489fa0e9dc2 · 2024-09-25T11:33:16.000+02:00
Add support for passing `failed_request_status_codes` to the `StarletteIntegration` and `FastApiIntegration` constructors as a `Set[int]`, while maintaining backwards-compatibility with the old format.
diff --git a/sentry_sdk/integrations/__init__.py b/sentry_sdk/integrations/__init__.py
@@ -16,6 +16,9 @@
     from typing import Type
 
 
+_DEFAULT_FAILED_REQUEST_STATUS_CODES = frozenset(range(500, 600))
+
+
 _installer_lock = Lock()
 
 # Set of all integration identifiers we have attempted to install
diff --git a/sentry_sdk/integrations/_wsgi_common.py b/sentry_sdk/integrations/_wsgi_common.py
@@ -210,7 +210,7 @@ def _filter_headers(headers):
 
 
 def _in_http_status_code_range(code, code_ranges):
-    # type: (int, list[HttpStatusCodeRange]) -> bool
+    # type: (object, list[HttpStatusCodeRange]) -> bool
     for target in code_ranges:
         if isinstance(target, int):
             if code == target:
@@ -226,3 +226,18 @@ def _in_http_status_code_range(code, code_ranges):
             )
 
     return False
+
+
+class HttpCodeRangeContainer:
+    """
+    Wrapper to make it possible to use list[HttpStatusCodeRange] as a Container[int].
+    Used for backwards compatibility with the old `failed_request_status_codes` option.
+    """
+
+    def __init__(self, code_ranges):
+        # type: (list[HttpStatusCodeRange]) -> None
+        self._code_ranges = code_ranges
+
+    def __contains__(self, item):
+        # type: (object) -> bool
+        return _in_http_status_code_range(item, self._code_ranges)
diff --git a/sentry_sdk/integrations/aiohttp.py b/sentry_sdk/integrations/aiohttp.py
@@ -5,7 +5,11 @@
 import sentry_sdk
 from sentry_sdk.api import continue_trace
 from sentry_sdk.consts import OP, SPANSTATUS, SPANDATA
-from sentry_sdk.integrations import Integration, DidNotEnable
+from sentry_sdk.integrations import (
+    _DEFAULT_FAILED_REQUEST_STATUS_CODES,
+    Integration,
+    DidNotEnable,
+)
 from sentry_sdk.integrations.logging import ignore_logger
 from sentry_sdk.sessions import track_session
 from sentry_sdk.integrations._wsgi_common import (
@@ -61,7 +65,6 @@
 
 
 TRANSACTION_STYLE_VALUES = ("handler_name", "method_and_path_pattern")
-_DEFAULT_FAILED_REQUEST_STATUS_CODES = frozenset(range(500, 600))
 
 
 class AioHttpIntegration(Integration):
diff --git a/sentry_sdk/integrations/starlette.py b/sentry_sdk/integrations/starlette.py
@@ -1,12 +1,18 @@
 import asyncio
 import functools
+import warnings
+from collections.abc import Set
 from copy import deepcopy
 
 import sentry_sdk
 from sentry_sdk.consts import OP
-from sentry_sdk.integrations import DidNotEnable, Integration
+from sentry_sdk.integrations import (
+    DidNotEnable,
+    Integration,
+    _DEFAULT_FAILED_REQUEST_STATUS_CODES,
+)
 from sentry_sdk.integrations._wsgi_common import (
-    _in_http_status_code_range,
+    HttpCodeRangeContainer,
     _is_json_content_type,
     request_body_within_bounds,
 )
@@ -30,7 +36,7 @@
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
-    from typing import Any, Awaitable, Callable, Dict, Optional, Tuple
+    from typing import Any, Awaitable, Callable, Container, Dict, Optional, Tuple, Union
 
     from sentry_sdk._types import Event, HttpStatusCodeRange
 
@@ -76,23 +82,37 @@ class StarletteIntegration(Integration):
 
     def __init__(
         self,
-        transaction_style="url",
-        failed_request_status_codes=None,
-        middleware_spans=True,
+        transaction_style="url",  # type: str
+        failed_request_status_codes=_DEFAULT_FAILED_REQUEST_STATUS_CODES,  # type: Union[Set[int], list[HttpStatusCodeRange], None]
+        middleware_spans=True,  # type: bool
     ):
-        # type: (str, Optional[list[HttpStatusCodeRange]], bool) -> None
+        # type: (...) -> None
         if transaction_style not in TRANSACTION_STYLE_VALUES:
             raise ValueError(
                 "Invalid value for transaction_style: %s (must be in %s)"
                 % (transaction_style, TRANSACTION_STYLE_VALUES)
             )
         self.transaction_style = transaction_style
         self.middleware_spans = middleware_spans
-        self.failed_request_status_codes = (
-            [range(500, 599)]
-            if failed_request_status_codes is None
-            else failed_request_status_codes
-        )  # type: list[HttpStatusCodeRange]
+
+        if isinstance(failed_request_status_codes, Set):
+            self.failed_request_status_codes = (
+                failed_request_status_codes
+            )  # type: Container[int]
+        else:
+            warnings.warn(
+                "Passing a list or None for failed_request_status_codes is deprecated. "
+                "Please pass a set of int instead.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+
+            if failed_request_status_codes is None:
+                self.failed_request_status_codes = _DEFAULT_FAILED_REQUEST_STATUS_CODES
+            else:
+                self.failed_request_status_codes = HttpCodeRangeContainer(
+                    failed_request_status_codes
+                )
 
     @staticmethod
     def setup_once():
@@ -226,9 +246,7 @@ async def _sentry_patched_exception_handler(self, *args, **kwargs):
                     is_http_server_error = (
                         hasattr(exp, "status_code")
                         and isinstance(exp.status_code, int)
-                        and _in_http_status_code_range(
-                            exp.status_code, integration.failed_request_status_codes
-                        )
+                        and exp.status_code in integration.failed_request_status_codes
                     )
                     if is_http_server_error:
                         _capture_exception(exp, handled=True)
diff --git a/tests/integrations/fastapi/test_fastapi.py b/tests/integrations/fastapi/test_fastapi.py
@@ -1,6 +1,7 @@
 import json
 import logging
 import threading
+import warnings
 from unittest import mock
 
 import pytest
@@ -505,20 +506,28 @@ def test_transaction_name_in_middleware(
     )
 
 
-@test_starlette.parametrize_test_configurable_status_codes
-def test_configurable_status_codes(
+@test_starlette.parametrize_test_configurable_status_codes_deprecated
+def test_configurable_status_codes_deprecated(
     sentry_init,
     capture_events,
     failed_request_status_codes,
     status_code,
     expected_error,
 ):
+    with pytest.warns(DeprecationWarning):
+        starlette_integration = StarletteIntegration(
+            failed_request_status_codes=failed_request_status_codes
+        )
+
+    with pytest.warns(DeprecationWarning):
+        fast_api_integration = FastApiIntegration(
+            failed_request_status_codes=failed_request_status_codes
+        )
+
     sentry_init(
         integrations=[
-            StarletteIntegration(
-                failed_request_status_codes=failed_request_status_codes
-            ),
-            FastApiIntegration(failed_request_status_codes=failed_request_status_codes),
+            starlette_integration,
+            fast_api_integration,
         ]
     )
 
@@ -537,3 +546,36 @@ async def _error():
         assert len(events) == 1
     else:
         assert not events
+
+
+@test_starlette.parametrize_test_configurable_status_codes
+def test_configurable_status_codes(
+    sentry_init,
+    capture_events,
+    failed_request_status_codes,
+    status_code,
+    expected_error,
+):
+    integration_kwargs = {}
+    if failed_request_status_codes is not None:
+        integration_kwargs["failed_request_status_codes"] = failed_request_status_codes
+
+    with warnings.catch_warnings():
+        warnings.simplefilter("error", DeprecationWarning)
+        starlette_integration = StarletteIntegration(**integration_kwargs)
+        fastapi_integration = FastApiIntegration(**integration_kwargs)
+
+    sentry_init(integrations=[starlette_integration, fastapi_integration])
+
+    events = capture_events()
+
+    app = FastAPI()
+
+    @app.get("/error")
+    async def _error():
+        raise HTTPException(status_code)
+
+    client = TestClient(app)
+    client.get("/error")
+
+    assert len(events) == int(expected_error)
diff --git a/tests/integrations/starlette/test_starlette.py b/tests/integrations/starlette/test_starlette.py
@@ -6,6 +6,7 @@
 import os
 import re
 import threading
+import warnings
 from unittest import mock
 
 import pytest
@@ -1133,7 +1134,22 @@ def test_span_origin(sentry_init, capture_events):
         assert span["origin"] == "auto.http.starlette"
 
 
-parametrize_test_configurable_status_codes = pytest.mark.parametrize(
+class NonIterableContainer:
+    """Wraps any container and makes it non-iterable.
+
+    Used to test backwards compatibility with our old way of defining failed_request_status_codes, which allowed
+    passing in a list of (possibly non-iterable) containers. The Python standard library does not provide any built-in
+    non-iterable containers, so we have to define our own.
+    """
+
+    def __init__(self, inner):
+        self.inner = inner
+
+    def __contains__(self, item):
+        return item in self.inner
+
+
+parametrize_test_configurable_status_codes_deprecated = pytest.mark.parametrize(
     "failed_request_status_codes,status_code,expected_error",
     [
         (None, 500, True),
@@ -1150,28 +1166,29 @@ def test_span_origin(sentry_init, capture_events):
         ([range(400, 403), 500, 501], 501, True),
         ([range(400, 403), 500, 501], 503, False),
         ([], 500, False),
+        ([NonIterableContainer(range(500, 600))], 500, True),
+        ([NonIterableContainer(range(500, 600))], 404, False),
     ],
 )
-"""Test cases for configurable status codes.
+"""Test cases for configurable status codes (deprecated API).
 Also used by the FastAPI tests.
 """
 
 
-@parametrize_test_configurable_status_codes
-def test_configurable_status_codes(
+@parametrize_test_configurable_status_codes_deprecated
+def test_configurable_status_codes_deprecated(
     sentry_init,
     capture_events,
     failed_request_status_codes,
     status_code,
     expected_error,
 ):
-    sentry_init(
-        integrations=[
-            StarletteIntegration(
-                failed_request_status_codes=failed_request_status_codes
-            )
-        ]
-    )
+    with pytest.warns(DeprecationWarning):
+        starlette_integration = StarletteIntegration(
+            failed_request_status_codes=failed_request_status_codes
+        )
+
+    sentry_init(integrations=[starlette_integration])
 
     events = capture_events()
 
@@ -1191,3 +1208,59 @@ async def _error(request):
         assert len(events) == 1
     else:
         assert not events
+
+
+parametrize_test_configurable_status_codes = pytest.mark.parametrize(
+    ("failed_request_status_codes", "status_code", "expected_error"),
+    (
+        (None, 500, True),
+        (None, 400, False),
+        ({500, 501}, 500, True),
+        ({500, 501}, 401, False),
+        ({*range(400, 500)}, 401, True),
+        ({*range(400, 500)}, 500, False),
+        ({*range(400, 600)}, 300, False),
+        ({*range(400, 600)}, 403, True),
+        ({*range(400, 600)}, 503, True),
+        ({*range(400, 403), 500, 501}, 401, True),
+        ({*range(400, 403), 500, 501}, 405, False),
+        ({*range(400, 403), 500, 501}, 501, True),
+        ({*range(400, 403), 500, 501}, 503, False),
+        (set(), 500, False),
+    ),
+)
+
+
+@parametrize_test_configurable_status_codes
+def test_configurable_status_codes(
+    sentry_init,
+    capture_events,
+    failed_request_status_codes,
+    status_code,
+    expected_error,
+):
+    integration_kwargs = {}
+    if failed_request_status_codes is not None:
+        integration_kwargs["failed_request_status_codes"] = failed_request_status_codes
+
+    with warnings.catch_warnings():
+        warnings.simplefilter("error", DeprecationWarning)
+        starlette_integration = StarletteIntegration(**integration_kwargs)
+
+    sentry_init(integrations=[starlette_integration])
+
+    events = capture_events()
+
+    async def _error(_):
+        raise HTTPException(status_code)
+
+    app = starlette.applications.Starlette(
+        routes=[
+            starlette.routing.Route("/error", _error, methods=["GET"]),
+        ],
+    )
+
+    client = TestClient(app)
+    client.get("/error")
+
+    assert len(events) == int(expected_error)
