From 84901044dd4d4661599f821b37d47b5b01af33e2 Mon Sep 17 00:00:00 2001 From: Michael Lavers Date: Thu, 6 Jun 2019 07:24:00 -0700 Subject: [PATCH 1/5] Add monkey patching for botocore.httpsession Closes #319 --- iopipe/contrib/trace/auto_http.py | 106 ++++++++++++++++++++++++------ 1 file changed, 85 insertions(+), 21 deletions(-) diff --git a/iopipe/contrib/trace/auto_http.py b/iopipe/contrib/trace/auto_http.py index 3931f0e6..0cb21466 100644 --- a/iopipe/contrib/trace/auto_http.py +++ b/iopipe/contrib/trace/auto_http.py @@ -3,24 +3,34 @@ 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.awsrequest import AWSResponse + from botocore.httpsession import URLLib3Session as BotocoreSession except ImportError: + AWSResponse = None 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 +70,128 @@ ) -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, "send") and hasattr( + RequestsSession.send, "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.send.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, "send") and hasattr( + BotocoreSession.send, "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.send.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, "send") and hasattr( + BotocoreVendoredSession.send, "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.send.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 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 +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 + + 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: From 9fe13e828341e6e586a64b461418126d02d7e065 Mon Sep 17 00:00:00 2001 From: Michael Lavers Date: Thu, 6 Jun 2019 07:28:35 -0700 Subject: [PATCH 2/5] Rmove AWSResponse, don't need it --- iopipe/contrib/trace/auto_http.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/iopipe/contrib/trace/auto_http.py b/iopipe/contrib/trace/auto_http.py index 0cb21466..fcd8d39a 100644 --- a/iopipe/contrib/trace/auto_http.py +++ b/iopipe/contrib/trace/auto_http.py @@ -8,10 +8,8 @@ RequestsSession = None try: - from botocore.awsrequest import AWSResponse from botocore.httpsession import URLLib3Session as BotocoreSession except ImportError: - AWSResponse = None BotocoreSession = None try: From b739ed51138e2b5ea71f00e9e0ccb2e729b44e99 Mon Sep 17 00:00:00 2001 From: Michael Lavers Date: Thu, 6 Jun 2019 07:51:34 -0700 Subject: [PATCH 3/5] Add monkey patching test --- setup.py | 1 + tests/contrib/trace/test_auto_http.py | 23 +++++++++++++++++++++++ 2 files changed, 24 insertions(+) create mode 100644 tests/contrib/trace/test_auto_http.py 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..9a03d7b2 --- /dev/null +++ b/tests/contrib/trace/test_auto_http.py @@ -0,0 +1,23 @@ +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.send, "monkey_patched") + assert not hasattr(BotocoreSession.send, "monkey_patched") + assert not hasattr(BotocoreVendoredSession.send, "monkey_patched") + + patch_requests(mock_context, None) + + assert hasattr(RequestsSession.send, "monkey_patched") + assert hasattr(BotocoreSession.send, "monkey_patched") + assert hasattr(BotocoreVendoredSession.send, "monkey_patched") + + restore_requests() + + assert not hasattr(RequestsSession.send, "monkey_patched") + assert not hasattr(BotocoreSession.send, "monkey_patched") + assert not hasattr(BotocoreVendoredSession.send, "monkey_patched") From 278c9f076ed0f31b7c3f36f87b2e6e460535256b Mon Sep 17 00:00:00 2001 From: Michael Lavers Date: Thu, 6 Jun 2019 07:57:57 -0700 Subject: [PATCH 4/5] Add mock filter to test --- tests/contrib/trace/test_auto_http.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tests/contrib/trace/test_auto_http.py b/tests/contrib/trace/test_auto_http.py index 9a03d7b2..2c92d785 100644 --- a/tests/contrib/trace/test_auto_http.py +++ b/tests/contrib/trace/test_auto_http.py @@ -10,7 +10,10 @@ def test_monkey_patching(mock_context): assert not hasattr(BotocoreSession.send, "monkey_patched") assert not hasattr(BotocoreVendoredSession.send, "monkey_patched") - patch_requests(mock_context, None) + def mock_filter(http_response): + return http_response + + patch_requests(mock_context, mock_filter) assert hasattr(RequestsSession.send, "monkey_patched") assert hasattr(BotocoreSession.send, "monkey_patched") From b207b91e03877ef6e7cf1fbaaedf5911a3b9a639 Mon Sep 17 00:00:00 2001 From: Michael Lavers Date: Thu, 6 Jun 2019 08:21:52 -0700 Subject: [PATCH 5/5] Make monkey patching python2.7 compatible --- iopipe/contrib/trace/auto_http.py | 21 +++++++++------------ tests/contrib/trace/test_auto_http.py | 18 +++++++++--------- 2 files changed, 18 insertions(+), 21 deletions(-) diff --git a/iopipe/contrib/trace/auto_http.py b/iopipe/contrib/trace/auto_http.py index fcd8d39a..1a938b20 100644 --- a/iopipe/contrib/trace/auto_http.py +++ b/iopipe/contrib/trace/auto_http.py @@ -76,9 +76,7 @@ def patch_requests_session_send(context, http_filter): if RequestsSession is None: return - if hasattr(RequestsSession, "send") and hasattr( - RequestsSession.send, "monkey_patched" - ): + if hasattr(RequestsSession, "__monkey_patched"): return def send(self, *args, **kwargs): @@ -91,7 +89,7 @@ def send(self, *args, **kwargs): return response RequestsSession.send = send - RequestsSession.send.monkey_patched = True + RequestsSession.__monkey_patched = True def patch_botocore_session_send(context, http_filter): @@ -102,9 +100,7 @@ def patch_botocore_session_send(context, http_filter): if BotocoreSession is None: return - if hasattr(BotocoreSession, "send") and hasattr( - BotocoreSession.send, "monkey_patched" - ): + if hasattr(BotocoreSession, "__monkey_patched"): return def send(self, *args, **kwargs): @@ -119,7 +115,7 @@ def send(self, *args, **kwargs): return response BotocoreSession.send = send - BotocoreSession.send.monkey_patched = True + BotocoreSession.__monkey_patched = True def patch_botocore_vendored_session_send(context, http_filter): @@ -130,9 +126,7 @@ def patch_botocore_vendored_session_send(context, http_filter): if BotocoreVendoredSession is None: return - if hasattr(BotocoreVendoredSession, "send") and hasattr( - BotocoreVendoredSession.send, "monkey_patched" - ): + if hasattr(BotocoreVendoredSession, "__monkey_patched"): return def send(self, *args, **kwargs): @@ -145,25 +139,28 @@ def send(self, *args, **kwargs): return response BotocoreVendoredSession.send = send - BotocoreVendoredSession.send.monkey_patched = True + 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""" 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): diff --git a/tests/contrib/trace/test_auto_http.py b/tests/contrib/trace/test_auto_http.py index 2c92d785..6102a268 100644 --- a/tests/contrib/trace/test_auto_http.py +++ b/tests/contrib/trace/test_auto_http.py @@ -6,21 +6,21 @@ def test_monkey_patching(mock_context): - assert not hasattr(RequestsSession.send, "monkey_patched") - assert not hasattr(BotocoreSession.send, "monkey_patched") - assert not hasattr(BotocoreVendoredSession.send, "monkey_patched") + 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.send, "monkey_patched") - assert hasattr(BotocoreSession.send, "monkey_patched") - assert hasattr(BotocoreVendoredSession.send, "monkey_patched") + assert hasattr(RequestsSession, "__monkey_patched") + assert hasattr(BotocoreSession, "__monkey_patched") + assert hasattr(BotocoreVendoredSession, "__monkey_patched") restore_requests() - assert not hasattr(RequestsSession.send, "monkey_patched") - assert not hasattr(BotocoreSession.send, "monkey_patched") - assert not hasattr(BotocoreVendoredSession.send, "monkey_patched") + assert not hasattr(RequestsSession, "__monkey_patched") + assert not hasattr(BotocoreSession, "__monkey_patched") + assert not hasattr(BotocoreVendoredSession, "__monkey_patched")