feat(integrations): instrument pyreqwest tracing #5682
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -123,6 +123,7 @@ | |
| "Network": [ | ||
| "grpc", | ||
| "httpx", | ||
| "pyreqwest", | ||
| "requests", | ||
| ], | ||
| "Tasks": [ | ||
|
|
||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,136 @@ | ||
| import sentry_sdk | ||
| from sentry_sdk import start_span | ||
| 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, | ||
| add_http_request_source, | ||
| add_sentry_baggage_to_headers, | ||
| ) | ||
| from sentry_sdk.utils import ( | ||
| SENSITIVE_DATA_SUBSTITUTE, | ||
| capture_internal_exceptions, | ||
| logger, | ||
| parse_url, | ||
| ) | ||
|
|
||
| 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, | ||
| 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): | ||
| identifier = "pyreqwest" | ||
| origin = f"auto.http.{identifier}" | ||
|
|
||
| @staticmethod | ||
| def setup_once() -> None: | ||
| _patch_pyreqwest() | ||
|
|
||
|
|
||
| def _patch_pyreqwest() -> None: | ||
| # Patch Client Builders | ||
| _patch_builder_method(ClientBuilder, "build", sentry_async_middleware) | ||
| _patch_builder_method(SyncClientBuilder, "build", sentry_sync_middleware) | ||
|
|
||
| # Patch Request Builders | ||
| _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: | ||
| if not hasattr(cls, method_name): | ||
| return | ||
|
|
||
| original_method = getattr(cls, method_name) | ||
|
|
||
| def sentry_patched_method(self: "Any", *args: "Any", **kwargs: "Any") -> "Any": | ||
servusdei2018 marked this conversation as resolved.
Show resolved
Hide resolved
|
||
| if not getattr(self, "_sentry_instrumented", False): | ||
| integration = sentry_sdk.get_client().get_integration(PyreqwestIntegration) | ||
| if integration is not None: | ||
| self.with_middleware(middleware) | ||
| try: | ||
| self._sentry_instrumented = True | ||
|
There was a problem hiding this comment. Did we have some test to check if this works and doesn't always give the exception? (as these are native classes) |
||
| except (TypeError, AttributeError): | ||
| # In case the instance itself is immutable or doesn't allow extra attributes | ||
| pass | ||
servusdei2018 marked this conversation as resolved.
Show resolved
Hide resolved
|
||
| return original_method(self, *args, **kwargs) | ||
|
|
||
| setattr(cls, method_name, sentry_patched_method) | ||
|
|
||
|
|
||
| @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=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) | ||
| 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) | ||
|
|
||
| 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 | ||
|
|
||
| yield span | ||
|
|
||
| with capture_internal_exceptions(): | ||
| add_http_request_source(span) | ||
|
|
||
|
|
||
| async def sentry_async_middleware( | ||
| request: "Request", next_handler: "Next" | ||
| ) -> "Response": | ||
| if sentry_sdk.get_client().get_integration(PyreqwestIntegration) is None: | ||
| return await next_handler.run(request) | ||
|
|
||
| with _sentry_pyreqwest_span(request) as span: | ||
| response = await next_handler.run(request) | ||
| span.set_http_status(response.status) | ||
|
|
||
| return response | ||
|
|
||
|
|
||
| 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) | ||
|
|
||
| return response | ||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,3 @@ | ||
| import pytest | ||
|
|
||
| pytest.importorskip("pyreqwest") |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,152 @@ | ||
| from http.server import BaseHTTPRequestHandler, HTTPServer | ||
| from threading import Thread | ||
| import pytest | ||
|
|
||
| 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 | ||
|
|
||
| from sentry_sdk import start_transaction | ||
| from sentry_sdk.consts import SPANDATA | ||
| from sentry_sdk.integrations.pyreqwest import PyreqwestIntegration | ||
| from tests.conftest import get_free_port | ||
|
|
||
|
|
||
| class PyreqwestMockHandler(BaseHTTPRequestHandler): | ||
| captured_requests = [] | ||
|
|
||
| def do_GET(self) -> None: | ||
| self.captured_requests.append( | ||
| { | ||
| "path": self.path, | ||
| "headers": {k.lower(): v for k, v in self.headers.items()}, | ||
| } | ||
| ) | ||
|
|
||
| code = 200 | ||
| if "/status/" in self.path: | ||
| try: | ||
| code = int(self.path.split("/")[-1]) | ||
| except (ValueError, IndexError): | ||
| code = 200 | ||
|
|
||
| self.send_response(code) | ||
| self.end_headers() | ||
| self.wfile.write(b"OK") | ||
|
|
||
| def log_message(self, format: str, *args: object) -> None: | ||
| pass | ||
|
|
||
|
|
||
| @pytest.fixture(scope="module") | ||
| def server_port(): | ||
| port = get_free_port() | ||
| server = HTTPServer(("localhost", port), PyreqwestMockHandler) | ||
| thread = Thread(target=server.serve_forever) | ||
| thread.daemon = True | ||
| thread.start() | ||
| yield port | ||
| server.shutdown() | ||
|
|
||
|
|
||
| @pytest.fixture(autouse=True) | ||
| def clear_captured_requests(): | ||
| PyreqwestMockHandler.captured_requests.clear() | ||
|
|
||
|
|
||
| def test_sync_client_spans(sentry_init, capture_events, server_port): | ||
| sentry_init(integrations=[PyreqwestIntegration()], traces_sample_rate=1.0) | ||
| events = capture_events() | ||
|
|
||
| url = f"http://localhost:{server_port}/hello" | ||
| with start_transaction(name="test_transaction"): | ||
| client = SyncClientBuilder().build() | ||
| response = client.get(url).build().send() | ||
| assert response.status == 200 | ||
|
|
||
| (event,) = events | ||
| assert len(event["spans"]) == 1 | ||
| span = event["spans"][0] | ||
| assert span["op"] == "http.client" | ||
| assert span["description"] == f"GET {url}" | ||
| assert span["data"]["url"] == url | ||
| assert span["data"][SPANDATA.HTTP_STATUS_CODE] == 200 | ||
| assert span["origin"] == "auto.http.pyreqwest" | ||
|
|
||
|
|
||
| @pytest.mark.asyncio | ||
| async def test_async_client_spans(sentry_init, capture_events, server_port): | ||
| sentry_init(integrations=[PyreqwestIntegration()], traces_sample_rate=1.0) | ||
| events = capture_events() | ||
|
|
||
| url = f"http://localhost:{server_port}/hello" | ||
| async with ClientBuilder().build() as client: | ||
| with start_transaction(name="test_transaction"): | ||
| response = await client.get(url).build().send() | ||
| assert response.status == 200 | ||
|
|
||
| (event,) = events | ||
| assert len(event["spans"]) == 1 | ||
| span = event["spans"][0] | ||
| assert span["op"] == "http.client" | ||
| assert span["description"] == f"GET {url}" | ||
| assert span["data"]["url"] == url | ||
| assert span["data"][SPANDATA.HTTP_STATUS_CODE] == 200 | ||
| assert span["origin"] == "auto.http.pyreqwest" | ||
|
|
||
|
|
||
| def test_sync_simple_request_spans(sentry_init, capture_events, server_port): | ||
| sentry_init(integrations=[PyreqwestIntegration()], traces_sample_rate=1.0) | ||
| events = capture_events() | ||
|
|
||
| url = f"http://localhost:{server_port}/hello-simple" | ||
| with start_transaction(name="test_transaction"): | ||
| response = sync_pyreqwest_get(url).send() | ||
| assert response.status == 200 | ||
|
|
||
| (event,) = events | ||
| assert len(event["spans"]) == 1 | ||
| span = event["spans"][0] | ||
| assert span["op"] == "http.client" | ||
| assert span["description"] == f"GET {url}" | ||
|
|
||
|
|
||
| @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) | ||
| events = capture_events() | ||
|
|
||
| url = f"http://localhost:{server_port}/hello-simple-async" | ||
| with start_transaction(name="test_transaction"): | ||
| response = await async_pyreqwest_get(url).send() | ||
| assert response.status == 200 | ||
|
|
||
| (event,) = events | ||
| assert len(event["spans"]) == 1 | ||
| span = event["spans"][0] | ||
| assert span["op"] == "http.client" | ||
| assert span["description"] == f"GET {url}" | ||
|
|
||
|
|
||
| def test_outgoing_trace_headers(sentry_init, server_port): | ||
| sentry_init( | ||
| integrations=[PyreqwestIntegration()], | ||
| traces_sample_rate=1.0, | ||
| trace_propagation_targets=["localhost"], | ||
| ) | ||
|
|
||
| url = f"http://localhost:{server_port}/trace" | ||
| with start_transaction( | ||
| name="test_transaction", trace_id="01234567890123456789012345678901" | ||
| ): | ||
| client = SyncClientBuilder().build() | ||
| response = client.get(url).build().send() | ||
| assert response.status == 200 | ||
|
|
||
| assert len(PyreqwestMockHandler.captured_requests) == 1 | ||
| headers = PyreqwestMockHandler.captured_requests[0]["headers"] | ||
|
|
||
| assert "sentry-trace" in headers | ||
| assert headers["sentry-trace"].startswith("01234567890123456789012345678901") | ||
| assert "baggage" in headers | ||
| assert "sentry-trace_id=01234567890123456789012345678901" in headers["baggage"] |
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Uh oh!
There was an error while loading. Please reload this page.