nit: address PR comments

Author: servusdei2018

sentry_sdk/integrations/__init__.py

@@ -94,7 +94,6 @@ def iter_default_integrations(
     "sentry_sdk.integrations.openai.OpenAIIntegration",
     "sentry_sdk.integrations.openai_agents.OpenAIAgentsIntegration",
     "sentry_sdk.integrations.pydantic_ai.PydanticAIIntegration",
-    "sentry_sdk.integrations.pyreqwest.PyreqwestIntegration",
     "sentry_sdk.integrations.pymongo.PyMongoIntegration",
     "sentry_sdk.integrations.pyramid.PyramidIntegration",
     "sentry_sdk.integrations.quart.QuartIntegration",

sentry_sdk/integrations/pyreqwest.py

@@ -15,16 +15,22 @@
     parse_url,
 )
 
-from typing import TYPE_CHECKING
-
-if TYPE_CHECKING:
-    from typing import Any
-
-
-import importlib.util
-
-if importlib.util.find_spec("pyreqwest") is None:
-    raise DidNotEnable("pyreqwest is not installed")
+from contextlib import contextmanager
+from typing import Any, Generator
+
+try:
+    from pyreqwest.client import ClientBuilder, SyncClientBuilder  # type: ignore[import-not-found]
+    from pyreqwest.request import (  # type: ignore[import-not-found]
+        Request,
+        RequestBuilder,
+        SyncRequestBuilder,
+        OneOffRequestBuilder,
+        SyncOneOffRequestBuilder,
+    )
+    from pyreqwest.middleware import Next, SyncNext  # type: ignore[import-not-found]
+    from pyreqwest.response import Response, SyncResponse  # type: ignore[import-not-found]
+except ImportError:
+    raise DidNotEnable("pyreqwest not installed or incompatible version installed")
 
 
 class PyreqwestIntegration(Integration):
@@ -38,33 +44,16 @@ def setup_once() -> None:
 
 def _patch_pyreqwest() -> None:
     # Patch Client Builders
-    try:
-        from pyreqwest.client import ClientBuilder, SyncClientBuilder  # type: ignore[import-not-found]
-
-        _patch_builder_method(ClientBuilder, "build", sentry_async_middleware)
-        _patch_builder_method(SyncClientBuilder, "build", sentry_sync_middleware)
-    except ImportError:
-        pass
+    _patch_builder_method(ClientBuilder, "build", sentry_async_middleware)
+    _patch_builder_method(SyncClientBuilder, "build", sentry_sync_middleware)
 
     # Patch Request Builders (for simple requests and manual request building)
-    try:
-        from pyreqwest.request import (  # type: ignore[import-not-found]
-            RequestBuilder,
-            SyncRequestBuilder,
-            OneOffRequestBuilder,
-            SyncOneOffRequestBuilder,
-        )
-
-        _patch_builder_method(RequestBuilder, "build", sentry_async_middleware)
-        _patch_builder_method(RequestBuilder, "build_streamed", sentry_async_middleware)
-        _patch_builder_method(SyncRequestBuilder, "build", sentry_sync_middleware)
-        _patch_builder_method(
-            SyncRequestBuilder, "build_streamed", sentry_sync_middleware
-        )
-        _patch_builder_method(OneOffRequestBuilder, "send", sentry_async_middleware)
-        _patch_builder_method(SyncOneOffRequestBuilder, "send", sentry_sync_middleware)
-    except ImportError:
-        pass
+    _patch_builder_method(RequestBuilder, "build", sentry_async_middleware)
+    _patch_builder_method(RequestBuilder, "build_streamed", sentry_async_middleware)
+    _patch_builder_method(SyncRequestBuilder, "build", sentry_sync_middleware)
+    _patch_builder_method(SyncRequestBuilder, "build_streamed", sentry_sync_middleware)
+    _patch_builder_method(OneOffRequestBuilder, "send", sentry_async_middleware)
+    _patch_builder_method(SyncOneOffRequestBuilder, "send", sentry_sync_middleware)
 
 
 def _patch_builder_method(cls: type, method_name: str, middleware: "Any") -> None:
@@ -88,21 +77,15 @@ def sentry_patched_method(self: "Any", *args: "Any", **kwargs: "Any") -> "Any":
     setattr(cls, method_name, sentry_patched_method)
 
 
-async def sentry_async_middleware(request: "Any", next_handler: "Any") -> "Any":
-    if sentry_sdk.get_client().get_integration(PyreqwestIntegration) is None:
-        return await next_handler.run(request)
-
+@contextmanager
+def _sentry_pyreqwest_span(request: "Request") -> "Generator[Any, None, None]":
     parsed_url = None
     with capture_internal_exceptions():
         parsed_url = parse_url(str(request.url), sanitize=False)
 
     with start_span(
         op=OP.HTTP_CLIENT,
-        name="%s %s"
-        % (
-            request.method,
-            parsed_url.url if parsed_url else SENSITIVE_DATA_SUBSTITUTE,
-        ),
+        name=f"{request.method} {parsed_url.url if parsed_url else SENSITIVE_DATA_SUBSTITUTE}",
         origin=PyreqwestIntegration.origin,
     ) as span:
         span.set_data(SPANDATA.HTTP_METHOD, request.method)
@@ -127,60 +110,33 @@ async def sentry_async_middleware(request: "Any", next_handler: "Any") -> "Any":
                 else:
                     request.headers[key] = value
 
-        response = await next_handler.run(request)
-
-        span.set_http_status(response.status)
+        yield span
 
     with capture_internal_exceptions():
         add_http_request_source(span)
 
-    return response
-
 
-def sentry_sync_middleware(request: "Any", next_handler: "Any") -> "Any":
+async def sentry_async_middleware(
+    request: "Request", next_handler: "Next"
+) -> "Response":
     if sentry_sdk.get_client().get_integration(PyreqwestIntegration) is None:
-        return next_handler.run(request)
+        return await next_handler.run(request)
 
-    parsed_url = None
-    with capture_internal_exceptions():
-        parsed_url = parse_url(str(request.url), sanitize=False)
+    with _sentry_pyreqwest_span(request) as span:
+        response = await next_handler.run(request)
+        span.set_http_status(response.status)
 
-    with start_span(
-        op=OP.HTTP_CLIENT,
-        name="%s %s"
-        % (
-            request.method,
-            parsed_url.url if parsed_url else SENSITIVE_DATA_SUBSTITUTE,
-        ),
-        origin=PyreqwestIntegration.origin,
-    ) as span:
-        span.set_data(SPANDATA.HTTP_METHOD, request.method)
-        if parsed_url is not None:
-            span.set_data("url", parsed_url.url)
-            span.set_data(SPANDATA.HTTP_QUERY, parsed_url.query)
-            span.set_data(SPANDATA.HTTP_FRAGMENT, parsed_url.fragment)
+    return response
 
-        if should_propagate_trace(sentry_sdk.get_client(), str(request.url)):
-            for (
-                key,
-                value,
-            ) in sentry_sdk.get_current_scope().iter_trace_propagation_headers():
-                logger.debug(
-                    "[Tracing] Adding `{key}` header {value} to outgoing request to {url}.".format(
-                        key=key, value=value, url=request.url
-                    )
-                )
 
-                if key == BAGGAGE_HEADER_NAME:
-                    add_sentry_baggage_to_headers(request.headers, value)
-                else:
-                    request.headers[key] = value
+def sentry_sync_middleware(
+    request: "Request", next_handler: "SyncNext"
+) -> "SyncResponse":
+    if sentry_sdk.get_client().get_integration(PyreqwestIntegration) is None:
+        return next_handler.run(request)
 
+    with _sentry_pyreqwest_span(request) as span:
         response = next_handler.run(request)
-
         span.set_http_status(response.status)
 
-    with capture_internal_exceptions():
-        add_http_request_source(span)
-
     return response

tests/integrations/pyreqwest/test_pyreqwest.py

@@ -2,7 +2,6 @@
 from threading import Thread
 import pytest
 
-import pyreqwest
 from pyreqwest.client import ClientBuilder, SyncClientBuilder
 from pyreqwest.simple.request import pyreqwest_get as async_pyreqwest_get
 from pyreqwest.simple.sync_request import pyreqwest_get as sync_pyreqwest_get
@@ -55,7 +54,6 @@ def clear_captured_requests():
     PyreqwestMockHandler.captured_requests.clear()
 
 
-@pytest.mark.skipif(pyreqwest is None, reason="pyreqwest not installed")
 def test_sync_client_spans(sentry_init, capture_events, server_port):
     sentry_init(integrations=[PyreqwestIntegration()], traces_sample_rate=1.0)
     events = capture_events()
@@ -75,7 +73,6 @@ def test_sync_client_spans(sentry_init, capture_events, server_port):
     assert span["origin"] == "auto.http.pyreqwest"
 
 
-@pytest.mark.skipif(pyreqwest is None, reason="pyreqwest not installed")
 @pytest.mark.asyncio
 async def test_async_client_spans(sentry_init, capture_events, server_port):
     sentry_init(integrations=[PyreqwestIntegration()], traces_sample_rate=1.0)
@@ -96,7 +93,6 @@ async def test_async_client_spans(sentry_init, capture_events, server_port):
     assert span["origin"] == "auto.http.pyreqwest"
 
 
-@pytest.mark.skipif(pyreqwest is None, reason="pyreqwest not installed")
 def test_sync_simple_request_spans(sentry_init, capture_events, server_port):
     sentry_init(integrations=[PyreqwestIntegration()], traces_sample_rate=1.0)
     events = capture_events()
@@ -112,7 +108,6 @@ def test_sync_simple_request_spans(sentry_init, capture_events, server_port):
     assert span["description"] == f"GET {url}"
 
 
-@pytest.mark.skipif(pyreqwest is None, reason="pyreqwest not installed")
 @pytest.mark.asyncio
 async def test_async_simple_request_spans(sentry_init, capture_events, server_port):
     sentry_init(integrations=[PyreqwestIntegration()], traces_sample_rate=1.0)
@@ -129,7 +124,6 @@ async def test_async_simple_request_spans(sentry_init, capture_events, server_po
     assert span["description"] == f"GET {url}"
 
 
-@pytest.mark.skipif(pyreqwest is None, reason="pyreqwest not installed")
 def test_outgoing_trace_headers(sentry_init, server_port):
     sentry_init(
         integrations=[PyreqwestIntegration()],