httplib support

aws_xray_sdk/core/patcher.py

@@ -10,6 +10,7 @@
     'requests',
     'sqlite3',
     'mysql',
+    'httplib',  # TODO: perhaps requests should map to this below
 )
 
 _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,126 @@
+import wrapt
+from collections import namedtuple
+
+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
+
+import ssl
+
+
+_XRAY_PROP = '_xray_prop'
+_XRay_Data = namedtuple('xray_data', ['method', 'host', 'url'])
+
+
+# ? is not a valid entity, and we don't want things after the ? for the segment name
+def _strip_url(url: str):
+    return url.partition('?')[0]
+
+
+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_client(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)
+
+        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)
+    subsegment.apply_status_code(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(
+        'http.client',
+        'HTTPConnection._send_request',
+        _prep_request
+    )
+
+    wrapt.wrap_function_wrapper(
+        'http.client',
+        'HTTPConnection.getresponse',
+        _xray_traced_http_client
+    )
+
+    wrapt.wrap_function_wrapper(
+        'http.client',
+        'HTTPResponse.read',
+        _xray_traced_http_client_read
+    )

aws_xray_sdk/ext/requests/patch.py

@@ -3,6 +3,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.httplib.patch import _strip_url
 
 
 def patch():
@@ -26,7 +27,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,
     )

tests/ext/httplib/test_httplib.py

@@ -0,0 +1,95 @@
+import pytest
+import requests
+
+from aws_xray_sdk.core import patch
+from aws_xray_sdk.core import xray_recorder
+from aws_xray_sdk.core.context import Context
+
+
+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 test_ok():
+    status_code = 200
+    url = 'http://{}/status/{}?foo=bar&baz=foo'.format(BASE_URL, status_code)
+    requests.get(url)
+    subsegment = xray_recorder.current_segment().subsegments[1]
+    assert subsegment.name == 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)
+    requests.post(url)
+    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)
+    requests.head(url)
+    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)
+    requests.put(url)
+    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:
+        requests.get('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 == 'NewConnectionError'

tests/ext/requests/test_requests.py

@@ -28,7 +28,7 @@ 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