httplib fixes

CHANGELOG.rst

@@ -6,6 +6,7 @@ Unreleased
 ==========
 * feature: Add support for SQLAlchemy and Flask-SQLAlcemy. `PR14 <https://github.com/aws/aws-xray-sdk-python/pull/14>`_.
 * feature: Add support for PynamoDB calls to DynamoDB. `PR13 <https://github.com/aws/aws-xray-sdk-python/pull/13>`_.
+* feature: Add support for httplib calls `PR19 <https://github.com/aws/aws-xray-sdk-python/pull/19>`_.
 * bugfix:  Drop invalid annotation keys and log a warning. `PR22 <https://github.com/aws/aws-xray-sdk-python/pull/22>`_.
 * bugfix:  Respect `with` statement on cursor objects in dbapi2 patcher. `PR17 <https://github.com/aws/aws-xray-sdk-python/pull/17>`_.
 

aws_xray_sdk/ext/httplib/__init__.py

@@ -1,3 +1,3 @@
-from .patch import patch
+from .patch import patch, unpatch
 
-__all__ = ['patch']
+__all__ = ['patch', 'unpatch']

aws_xray_sdk/ext/httplib/patch.py

@@ -4,14 +4,16 @@
 
 from aws_xray_sdk.core import xray_recorder
 from aws_xray_sdk.core.models import http
-from aws_xray_sdk.ext.util import inject_trace_header, strip_url
+from aws_xray_sdk.ext.util import inject_trace_header, strip_url, unwrap
 
 import ssl
 
 if sys.version_info >= (3, 0, 0):
+    PY2 = False
     httplib_client_module = 'http.client'
     import http.client as httplib
 else:
+    PY2 = True
     httplib_client_module = 'httplib'
     import httplib
 
@@ -39,7 +41,7 @@ def http_response_processor(wrapped, instance, args, kwargs, return_value,
 
 
 def _xray_traced_http_getresponse(wrapped, instance, args, kwargs):
-    if kwargs.get('buffering', False):
+    if not PY2 and kwargs.get('buffering', False):
         return wrapped(*args, **kwargs)  # ignore py2 calls that fail as 'buffering` only exists in py2.
 
     xray_data = getattr(instance, _XRAY_PROP)
@@ -52,8 +54,8 @@ def _xray_traced_http_getresponse(wrapped, instance, args, kwargs):
     )
 
 
-def http_request_processor(wrapped, instance, args, kwargs, return_value,
-                        exception, subsegment, stack):
+def http_send_request_processor(wrapped, instance, args, kwargs, return_value,
+                                exception, subsegment, stack):
     xray_data = getattr(instance, _XRAY_PROP)
 
     # we don't delete the attr as we can have multiple reads
@@ -64,7 +66,7 @@ def http_request_processor(wrapped, instance, args, kwargs, return_value,
         subsegment.add_exception(exception, stack)
 
 
-def _prep_request(wrapped, instance, args, kwargs):
+def _send_request(wrapped, instance, args, kwargs):
     def decompose_args(method, url, body, headers, encode_chunked=False):
         inject_trace_header(headers, xray_recorder.current_subsegment())
 
@@ -79,7 +81,7 @@ def decompose_args(method, url, body, headers, encode_chunked=False):
             wrapped, instance, args, kwargs,
             name=strip_url(xray_data.url),
             namespace='remote',
-            meta_processor=http_request_processor
+            meta_processor=http_send_request_processor
         )
 
     return decompose_args(*args, **kwargs)
@@ -110,10 +112,17 @@ def _xray_traced_http_client_read(wrapped, instance, args, kwargs):
 
 
 def patch():
+    """ patch the built-in urllib/httplib/httplib.client methods for tracing"""
+
+    # we set an attribute to avoid double unwrapping
+    if getattr(httplib, '__xray_patch', False):
+        return
+    setattr(httplib, '__xray_patch', True)
+
     wrapt.wrap_function_wrapper(
         httplib_client_module,
         'HTTPConnection._send_request',
-        _prep_request
+        _send_request
     )
 
     wrapt.wrap_function_wrapper(
@@ -127,3 +136,16 @@ def patch():
         'HTTPResponse.read',
         _xray_traced_http_client_read
     )
+
+
+def unpatch():
+    """ unpatch any previously patched modules """
+    if not getattr(httplib, '__xray_patch', False):
+        return
+    setattr(httplib, '__xray_patch', False)
+
+    # send_request encapsulates putrequest, putheader[s], and endheaders
+    # NOTE that requests
+    unwrap(httplib.HTTPConnection, '_send_request')
+    unwrap(httplib.HTTPConnection, 'getresponse')
+    unwrap(httplib.HTTPConnection, 'read')

aws_xray_sdk/ext/util.py

@@ -3,6 +3,8 @@
 from aws_xray_sdk.core.models.trace_header import TraceHeader
 from aws_xray_sdk.core.models import http
 
+import wrapt
+
 
 first_cap_re = re.compile('(.)([A-Z][a-z]+)')
 all_cap_re = re.compile('([a-z0-9])([A-Z])')
@@ -99,3 +101,14 @@ def strip_url(url):
     :return: validated url string
     """
     return url.partition('?')[0] if url else url
+
+
+def unwrap(obj, attr):
+    """
+    Will unwrap a `wrapt` attribute
+    :param obj: base object
+    :param attr: attribute on `obj` to unwrap
+    """
+    f = getattr(obj, attr, None)
+    if f and isinstance(f, wrapt.ObjectProxy) and hasattr(f, '__wrapped__'):
+        setattr(obj, attr, f.__wrapped__)

tests/ext/httplib/test_httplib.py

@@ -4,6 +4,7 @@
 from aws_xray_sdk.core import patch
 from aws_xray_sdk.core import xray_recorder
 from aws_xray_sdk.core.context import Context
+from aws_xray_sdk.ext.httplib import unpatch
 from aws_xray_sdk.ext.util import strip_url
 
 if sys.version_info >= (3, 0, 0):
@@ -14,8 +15,6 @@
     from urlparse import urlparse
 
 
-patch(('httplib',))
-
 # httpbin.org is created by the same author of requests to make testing http easy.
 BASE_URL = 'httpbin.org'
 
@@ -27,11 +26,14 @@ def construct_ctx():
     so that later subsegment can be attached. After each test run
     it cleans up context storage again.
     """
+    patch(('httplib',))
     xray_recorder.configure(service='test', sampling=False, context=Context())
     xray_recorder.clear_trace_entities()
     xray_recorder.begin_segment('name')
+
     yield
     xray_recorder.clear_trace_entities()
+    unpatch()
 
 
 def _do_req(url, method='GET'):