diff --git a/iopipe/contrib/trace/auto_http.py b/iopipe/contrib/trace/auto_http.py index 3931f0e6..1a938b20 100644 --- a/iopipe/contrib/trace/auto_http.py +++ b/iopipe/contrib/trace/auto_http.py @@ -3,24 +3,32 @@ import uuid try: - from requests.sessions import Session + from requests.sessions import Session as RequestsSession except ImportError: - Session = None + RequestsSession = None try: - from botocore.vendored.requests.sessions import Session as BotocoreSession + from botocore.httpsession import URLLib3Session as BotocoreSession except ImportError: BotocoreSession = None +try: + from botocore.vendored.requests.sessions import Session as BotocoreVendoredSession +except ImportError: + BotocoreVendoredSession = None + from iopipe.compat import urlparse from .util import ensure_utf8 -if Session is not None: - original_session_send = Session.send +if RequestsSession is not None: + original_requests_session_send = RequestsSession.send if BotocoreSession is not None: original_botocore_session_send = BotocoreSession.send +if BotocoreVendoredSession is not None: + original_botocore_vendored_session_send = BotocoreVendoredSession.send + INCLUDE_HEADERS = [ "accept", "accept-encoding", @@ -60,74 +68,125 @@ ) -def patch_session_send(context, http_filter): +def patch_requests_session_send(context, http_filter): """ - Monkey patches requests' Session class, if available. Overloads the + Monkey patches requests' session class, if available. Overloads the send method to add tracing and metrics collection. """ - if Session is None: + if RequestsSession is None: + return + + if hasattr(RequestsSession, "__monkey_patched"): return def send(self, *args, **kwargs): id = ensure_utf8(str(uuid.uuid4())) with context.iopipe.mark(id): - response = original_session_send(self, *args, **kwargs) + response = original_requests_session_send(self, *args, **kwargs) trace = context.iopipe.mark.measure(id) context.iopipe.mark.delete(id) collect_metrics_for_response(response, context, trace, http_filter) return response - Session.send = send + RequestsSession.send = send + RequestsSession.__monkey_patched = True def patch_botocore_session_send(context, http_filter): """ - Monkey patches botocore's vendored requests, if available. Overloads the - Session class' send method to add tracing and metric collection. + Monkey patches botocore's session, if available. Overloads the + session class' send method to add tracing and metric collection. """ if BotocoreSession is None: return + if hasattr(BotocoreSession, "__monkey_patched"): + return + def send(self, *args, **kwargs): id = str(uuid.uuid4()) with context.iopipe.mark(id): response = original_botocore_session_send(self, *args, **kwargs) trace = context.iopipe.mark.measure(id) context.iopipe.mark.delete(id) - collect_metrics_for_response(response, context, trace, http_filter) + collect_metrics_for_response( + response, context, trace, http_filter, http_request=args[0] + ) return response BotocoreSession.send = send + BotocoreSession.__monkey_patched = True -def restore_session_send(): - """Restores the original Session send method""" - if Session is not None: - Session.send = original_session_send +def patch_botocore_vendored_session_send(context, http_filter): + """ + Monkey patches botocore's vendored requests, if available. Overloads the + session class' send method to add tracing and metric collection. + """ + if BotocoreVendoredSession is None: + return + + if hasattr(BotocoreVendoredSession, "__monkey_patched"): + return + + def send(self, *args, **kwargs): + id = str(uuid.uuid4()) + with context.iopipe.mark(id): + response = original_botocore_vendored_session_send(self, *args, **kwargs) + trace = context.iopipe.mark.measure(id) + context.iopipe.mark.delete(id) + collect_metrics_for_response(response, context, trace, http_filter) + return response + + BotocoreVendoredSession.send = send + BotocoreVendoredSession.__monkey_patched = True + + +def restore_requests_session_send(): + """Restores the original requests session send method""" + if RequestsSession is not None: + RequestsSession.send = original_requests_session_send + delattr(RequestsSession, "__monkey_patched") def restore_botocore_session_send(): - """Restores the original botocore Session send method""" + """Restores the original botocore session send method""" if BotocoreSession is not None: BotocoreSession.send = original_botocore_session_send + delattr(BotocoreSession, "__monkey_patched") + + +def restore_botocore_vendored_session_send(): + """Restores the original botocore vendored session send method""" + if BotocoreVendoredSession is not None: + BotocoreVendoredSession.send = original_botocore_vendored_session_send + delattr(BotocoreVendoredSession, "__monkey_patched") def patch_requests(context, http_filter): - patch_session_send(context, http_filter) + patch_requests_session_send(context, http_filter) patch_botocore_session_send(context, http_filter) + patch_botocore_vendored_session_send(context, http_filter) def restore_requests(): - restore_session_send() + restore_requests_session_send() restore_botocore_session_send() + restore_botocore_vendored_session_send() -def collect_metrics_for_response(http_response, context, trace, http_filter): +def collect_metrics_for_response( + http_response, context, trace, http_filter, http_request=None +): """ Collects relevant metrics from a requests Response object and adds them to the IOpipe context. """ + # We make copies to let the user mutate these objects via http_filter http_response = copy.deepcopy(http_response) + if http_request is not None: + http_response.request = copy.deepcopy(http_request) + if http_filter is not None and callable(http_filter): http_response = http_filter(http_response) if http_response is False: diff --git a/setup.py b/setup.py index 298fc372..788691c1 100644 --- a/setup.py +++ b/setup.py @@ -15,6 +15,7 @@ install_requires.append("futures") tests_require = [ + "botocore==1.12.162", "jmespath>=0.7.1,<1.0.0", "mock", "more-itertools<6.0.0", diff --git a/tests/contrib/trace/test_auto_http.py b/tests/contrib/trace/test_auto_http.py new file mode 100644 index 00000000..6102a268 --- /dev/null +++ b/tests/contrib/trace/test_auto_http.py @@ -0,0 +1,26 @@ +from botocore.httpsession import URLLib3Session as BotocoreSession +from botocore.vendored.requests.sessions import Session as BotocoreVendoredSession +from requests.sessions import Session as RequestsSession + +from iopipe.contrib.trace.auto_http import patch_requests, restore_requests + + +def test_monkey_patching(mock_context): + assert not hasattr(RequestsSession, "__monkey_patched") + assert not hasattr(BotocoreSession, "__monkey_patched") + assert not hasattr(BotocoreVendoredSession, "__monkey_patched") + + def mock_filter(http_response): + return http_response + + patch_requests(mock_context, mock_filter) + + assert hasattr(RequestsSession, "__monkey_patched") + assert hasattr(BotocoreSession, "__monkey_patched") + assert hasattr(BotocoreVendoredSession, "__monkey_patched") + + restore_requests() + + assert not hasattr(RequestsSession, "__monkey_patched") + assert not hasattr(BotocoreSession, "__monkey_patched") + assert not hasattr(BotocoreVendoredSession, "__monkey_patched")