httplib support (#19)

* initial httplib patch

* fix unittests

* fix warnings

* strip url in requests as well

* updates based on review

* add some docs

* refactor based on review

* rename to getresponse to be clearer

Author: thehesiod

.gitignore

@@ -14,8 +14,9 @@ dist
 *.egg-info
 .tox
 .python-version
+.pytest_cache
 
 pip-selfcheck.json
 
 .coverage*
-htmlcov
+htmlcov

aws_xray_sdk/core/patcher.py

@@ -10,6 +10,7 @@
     'requests',
     'sqlite3',
     'mysql',
+    'httplib',
 )
 
 _PATCHED_MODULES = set()

aws_xray_sdk/ext/httplib/__init__.py

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

aws_xray_sdk/ext/httplib/patch.py

@@ -0,0 +1,129 @@
+from collections import namedtuple
+import sys
+import wrapt
+
+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
+
+import ssl
+
+if sys.version_info >= (3, 0, 0):
+    httplib_client_module = 'http.client'
+    import http.client as httplib
+else:
+    httplib_client_module = 'httplib'
+    import httplib
+
+
+_XRAY_PROP = '_xray_prop'
+_XRay_Data = namedtuple('xray_data', ['method', 'host', 'url'])
+
+
+def http_response_processor(wrapped, instance, args, kwargs, return_value,
+                            exception, subsegment, stack):
+    xray_data = getattr(instance, _XRAY_PROP)
+
+    subsegment.put_http_meta(http.METHOD, xray_data.method)
+    subsegment.put_http_meta(http.URL, xray_data.url)
+
+    if return_value:
+        subsegment.put_http_meta(http.STATUS, return_value.code)
+
+        # propagate to response object
+        xray_data = _XRay_Data('READ', xray_data.host, xray_data.url)
+        setattr(return_value, _XRAY_PROP, xray_data)
+
+    if exception:
+        subsegment.add_exception(exception, stack)
+
+
+def _xray_traced_http_getresponse(wrapped, instance, args, kwargs):
+    if 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)
+
+    return xray_recorder.record_subsegment(
+        wrapped, instance, args, kwargs,
+        name=strip_url(xray_data.url),
+        namespace='remote',
+        meta_processor=http_response_processor,
+    )
+
+
+def http_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
+    subsegment.put_http_meta(http.METHOD, xray_data.method)
+    subsegment.put_http_meta(http.URL, xray_data.url)
+
+    if exception:
+        subsegment.add_exception(exception, stack)
+
+
+def _prep_request(wrapped, instance, args, kwargs):
+    def decompose_args(method, url, body, headers, encode_chunked):
+        inject_trace_header(headers, xray_recorder.current_subsegment())
+
+        # we have to check against sock because urllib3's HTTPSConnection inherit's from http.client.HTTPConnection
+        scheme = 'https' if isinstance(instance.sock, ssl.SSLSocket) else 'http'
+        xray_url = '{}://{}{}'.format(scheme, instance.host, url)
+        xray_data = _XRay_Data(method, instance.host, xray_url)
+        setattr(instance, _XRAY_PROP, xray_data)
+
+        # we add a segment here in case connect fails
+        return xray_recorder.record_subsegment(
+            wrapped, instance, args, kwargs,
+            name=strip_url(xray_data.url),
+            namespace='remote',
+            meta_processor=http_request_processor
+        )
+
+    return decompose_args(*args, **kwargs)
+
+
+def http_read_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
+    subsegment.put_http_meta(http.METHOD, xray_data.method)
+    subsegment.put_http_meta(http.URL, xray_data.url)
+    subsegment.put_http_meta(http.STATUS, instance.status)
+
+    if exception:
+        subsegment.add_exception(exception, stack)
+
+
+def _xray_traced_http_client_read(wrapped, instance, args, kwargs):
+    xray_data = getattr(instance, _XRAY_PROP)
+
+    return xray_recorder.record_subsegment(
+        wrapped, instance, args, kwargs,
+        name=strip_url(xray_data.url),
+        namespace='remote',
+        meta_processor=http_read_processor
+    )
+
+
+def patch():
+    wrapt.wrap_function_wrapper(
+        httplib_client_module,
+        'HTTPConnection._send_request',
+        _prep_request
+    )
+
+    wrapt.wrap_function_wrapper(
+        httplib_client_module,
+        'HTTPConnection.getresponse',
+        _xray_traced_http_getresponse
+    )
+
+    wrapt.wrap_function_wrapper(
+        httplib_client_module,
+        'HTTPResponse.read',
+        _xray_traced_http_client_read
+    )

aws_xray_sdk/ext/requests/patch.py

@@ -2,7 +2,7 @@
 
 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
+from aws_xray_sdk.ext.util import inject_trace_header, strip_url
 
 
 def patch():
@@ -26,7 +26,7 @@ def _xray_traced_requests(wrapped, instance, args, kwargs):
 
     return xray_recorder.record_subsegment(
         wrapped, instance, args, kwargs,
-        name=url,
+        name=strip_url(url),
         namespace='remote',
         meta_processor=requests_processor,
     )

aws_xray_sdk/ext/util.py

@@ -89,3 +89,13 @@ def to_snake_case(name):
     s1 = first_cap_re.sub(r'\1_\2', name)
     # handle acronym words
     return all_cap_re.sub(r'\1_\2', s1).lower()
+
+
+# ? is not a valid entity, and we don't want things after the ? for the segment name
+def strip_url(url: str):
+    """
+    Will generate a valid url string for use as a segment name
+    :param url: url to strip
+    :return: validated url string
+    """
+    return url.partition('?')[0] if url else url

docs/aws_xray_sdk.ext.httplib.rst

@@ -0,0 +1,22 @@
+aws\_xray\_sdk\.ext\.httplib package
+=====================================
+
+Submodules
+----------
+
+aws\_xray\_sdk\.ext\.httplib\.patch module
+-------------------------------------------
+
+.. automodule:: aws_xray_sdk.ext.httplib.patch
+    :members:
+    :undoc-members:
+    :show-inheritance:
+
+
+Module contents
+---------------
+
+.. automodule:: aws_xray_sdk.ext.httplib
+    :members:
+    :undoc-members:
+    :show-inheritance:

docs/aws_xray_sdk.ext.rst

@@ -14,6 +14,7 @@ Subpackages
     aws_xray_sdk.ext.mysql
     aws_xray_sdk.ext.requests
     aws_xray_sdk.ext.sqlite3
+    aws_xray_sdk.ext.httplib
 
 Submodules
 ----------

docs/thirdparty.rst

@@ -6,7 +6,7 @@ Third Party Library Support
 Patching Supported Libraries
 ----------------------------
 
-The SDK supports aioboto3, aiobotocore, boto3, botocore, pynamodb, requests, sqlite3 and
+The SDK supports aioboto3, aiobotocore, boto3, botocore, pynamodb, requests, sqlite3, httplib and
 mysql-connector.
 
 To patch, use code like the following in the main app::
@@ -35,6 +35,7 @@ The following modules are availble to patch::
         'requests',
         'sqlite3',
         'mysql',
+        'httplib',
     )
 
 Patching boto3 and botocore are equivalent since boto3 depends on botocore.
@@ -73,3 +74,10 @@ up the X-Ray SDK with an Async Context, bear in mind this requires Python 3.5+::
 
 See :ref:`Configure Global Recorder <configurations>` for more information about
 configuring the ``xray_recorder``.
+
+Patching httplib
+----------------
+
+httplib is a low-level python module which is used by several third party modules, so
+by enabling patching to this module you can gain patching of many modules "for free."
+Some examples of modules that depend on httplib: requests and httplib2

tests/ext/httplib/__init__.py

@@ -0,0 +1,115 @@
+import pytest
+import sys
+
+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.util import strip_url
+
+if sys.version_info >= (3, 0, 0):
+    import http.client as httplib
+    from urllib.parse import urlparse
+else:
+    import httplib
+    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'
+
+
+@pytest.fixture(autouse=True)
+def construct_ctx():
+    """
+    Clean up context storage on each test run and begin a segment
+    so that later subsegment can be attached. After each test run
+    it cleans up context storage again.
+    """
+    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()
+
+
+def _do_req(url, method='GET'):
+    parts = urlparse(url)
+    host, _, port = parts.netloc.partition(':')
+    if port == '':
+        port = None
+    conn = httplib.HTTPConnection(parts.netloc, port)
+
+    path = '{}?{}'.format(parts.path, parts.query) if parts.query else parts.path
+    conn.request(method, path)
+    resp = conn.getresponse()
+
+
+def test_ok():
+    status_code = 200
+    url = 'http://{}/status/{}?foo=bar&baz=foo'.format(BASE_URL, status_code)
+    _do_req(url)
+    subsegment = xray_recorder.current_segment().subsegments[1]
+    assert subsegment.name == strip_url(url)
+
+    http_meta = subsegment.http
+    assert http_meta['request']['url'] == url
+    assert http_meta['request']['method'].upper() == 'GET'
+    assert http_meta['response']['status'] == status_code
+
+
+def test_error():
+    status_code = 400
+    url = 'http://{}/status/{}'.format(BASE_URL, status_code)
+    _do_req(url, 'POST')
+    subsegment = xray_recorder.current_segment().subsegments[1]
+    assert subsegment.name == url
+    assert subsegment.error
+
+    http_meta = subsegment.http
+    assert http_meta['request']['url'] == url
+    assert http_meta['request']['method'].upper() == 'POST'
+    assert http_meta['response']['status'] == status_code
+
+
+def test_throttle():
+    status_code = 429
+    url = 'http://{}/status/{}'.format(BASE_URL, status_code)
+    _do_req(url, 'HEAD')
+    subsegment = xray_recorder.current_segment().subsegments[1]
+    assert subsegment.name == url
+    assert subsegment.error
+    assert subsegment.throttle
+
+    http_meta = subsegment.http
+    assert http_meta['request']['url'] == url
+    assert http_meta['request']['method'].upper() == 'HEAD'
+    assert http_meta['response']['status'] == status_code
+
+
+def test_fault():
+    status_code = 500
+    url = 'http://{}/status/{}'.format(BASE_URL, status_code)
+    _do_req(url, 'PUT')
+    subsegment = xray_recorder.current_segment().subsegments[1]
+    assert subsegment.name == url
+    assert subsegment.fault
+
+    http_meta = subsegment.http
+    assert http_meta['request']['url'] == url
+    assert http_meta['request']['method'].upper() == 'PUT'
+    assert http_meta['response']['status'] == status_code
+
+
+def test_invalid_url():
+    try:
+        _do_req('http://doesnt.exist')
+    except Exception:
+        # prevent uncatch exception from breaking test run
+        pass
+    subsegment = xray_recorder.current_segment().subsegments[0]
+    assert subsegment.fault
+
+    exception = subsegment.cause['exceptions'][0]
+    assert exception.type == 'gaierror'

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.util import strip_url
 
 
 patch(('requests',))
@@ -28,10 +29,10 @@ def construct_ctx():
 
 def test_ok():
     status_code = 200
-    url = 'http://{}/status/{}'.format(BASE_URL, status_code)
+    url = 'http://{}/status/{}?foo=bar'.format(BASE_URL, status_code)
     requests.get(url)
     subsegment = xray_recorder.current_segment().subsegments[0]
-    assert subsegment.name == url
+    assert subsegment.name == strip_url(url)
 
     http_meta = subsegment.http
     assert http_meta['request']['url'] == url