Add monkey patching for botocore.httpsession

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:

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",

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")