-
Notifications
You must be signed in to change notification settings - Fork 593
feat(integrations): instrument pyreqwest tracing #5682
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: master
Are you sure you want to change the base?
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Large diffs are not rendered by default.
Large diffs are not rendered by default.
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -124,6 +124,7 @@ | |
| "Network": [ | ||
| "grpc", | ||
| "httpx", | ||
| "pyreqwest", | ||
| "requests", | ||
| ], | ||
| "Tasks": [ | ||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,186 @@ | ||
| 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 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") | ||
|
|
||
|
|
||
| class PyreqwestIntegration(Integration): | ||
| identifier = "pyreqwest" | ||
| origin = f"auto.http.{identifier}" | ||
|
|
||
| @staticmethod | ||
| def setup_once() -> None: | ||
| _patch_pyreqwest() | ||
|
|
||
|
|
||
| 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 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 | ||
|
|
||
|
|
||
| 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": | ||
| 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 | ||
| except (TypeError, AttributeError): | ||
| # In case the instance itself is immutable or doesn't allow extra attributes | ||
| pass | ||
| return original_method(self, *args, **kwargs) | ||
|
|
||
| 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) | ||
|
|
||
| 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, | ||
| ), | ||
| 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 | ||
|
|
||
| response = await next_handler.run(request) | ||
|
|
||
| span.set_http_status(response.status) | ||
|
|
||
| with capture_internal_exceptions(): | ||
| add_http_request_source(span) | ||
|
|
||
| return response | ||
|
|
||
|
|
||
| def sentry_sync_middleware(request: "Any", next_handler: "Any") -> "Any": | ||
| if sentry_sdk.get_client().get_integration(PyreqwestIntegration) is None: | ||
| return next_handler.run(request) | ||
|
|
||
| parsed_url = None | ||
| with capture_internal_exceptions(): | ||
| parsed_url = parse_url(str(request.url), sanitize=False) | ||
|
|
||
| with start_span( | ||
|
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. There is lot of same logic in both sentry_async_middleware and sentry_sync_middleware. You could extract the common stuff into a shared contextmanager for span handling |
||
| 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) | ||
|
|
||
| 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 | ||
|
|
||
| response = next_handler.run(request) | ||
|
|
||
| span.set_http_status(response.status) | ||
|
|
||
| with capture_internal_exceptions(): | ||
| add_http_request_source(span) | ||
|
|
||
| return response | ||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,3 @@ | ||
| import pytest | ||
|
|
||
| pytest.importorskip("pyreqwest") |
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,154 @@ | ||
| from http.server import BaseHTTPRequestHandler, HTTPServer | ||
| 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 | ||
|
|
||
| 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() | ||
|
|
||
|
|
||
| @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() | ||
|
|
||
| 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 | ||
| 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.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) | ||
| 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 | ||
| 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.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() | ||
|
|
||
| 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 | ||
| span = event["spans"][0] | ||
| assert span["op"] == "http.client" | ||
| 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) | ||
| 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 | ||
| span = event["spans"][0] | ||
| assert span["op"] == "http.client" | ||
| 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()], | ||
| 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"] |
Uh oh!
There was an error while loading. Please reload this page.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
You could use the types and type stubs to add types eg here ^ see eg https://github.com/MarkusSintonen/pyreqwest/blob/main/python/pyreqwest/middleware/types.py