Merge pull request #6 from terrycain/aiobotocore

Add support for aiobotocore

Author: haotianw465

MANIFEST.in

@@ -1,4 +1,4 @@
-include aws_xray_sdk/ext/botocore/*.json
+include aws_xray_sdk/ext/resources/*.json
 include aws_xray_sdk/core/sampling/*.json
 include README.md
 include LICENSE

README.md

@@ -93,6 +93,17 @@ def myfunc():
 myfunc()
 ```
 
+```python
+from aws_xray_sdk.core import xray_recorder
+
+@xray_recorder.capture_async('subsegment_name')
+async def myfunc():
+    # Do something here
+
+async def main():
+    await myfunc()
+```
+
 **Trace AWS Lambda functions**
 
 ```python
@@ -149,6 +160,22 @@ xray_recorder.configure(service='fallback_name', dynamic_naming='*mysite.com*')
 XRayMiddleware(app, xray_recorder)
 ```
 
+**Add aiohttp middleware**
+```python
+from aiohttp import web
+
+from aws_xray_sdk.ext.aiohttp.middleware import middleware
+from aws_xray_sdk.core import xray_recorder
+from aws_xray_sdk.core.async_context import AsyncContext
+
+xray_recorder.configure(service='fallback_name', context=AsyncContext())
+
+app = web.Application(middlewares=[middleware])
+app.router.add_get("/", handler)
+
+web.run_app(app)
+```
+
 ## License
 
 The AWS X-Ray SDK for Python is licensed under the Apache 2.0 License. See LICENSE and NOTICE.txt for more information.

aws_xray_sdk/core/__init__.py

@@ -1,8 +1,13 @@
 from .recorder import AWSXRayRecorder
 from .patcher import patch_all, patch
+from .utils.compat import PY35
 
 
-xray_recorder = AWSXRayRecorder()
+if not PY35:
+    xray_recorder = AWSXRayRecorder()
+else:
+    from .async_recorder import AsyncAWSXRayRecorder
+    xray_recorder = AsyncAWSXRayRecorder()
 
 __all__ = [
     'patch',

aws_xray_sdk/core/async_recorder.py

@@ -0,0 +1,69 @@
+import time
+import traceback
+
+import wrapt
+
+from aws_xray_sdk.core.recorder import AWSXRayRecorder
+
+
+class AsyncAWSXRayRecorder(AWSXRayRecorder):
+    def capture_async(self, name=None):
+        """
+        A decorator that records enclosed function in a subsegment.
+        It only works with asynchronous functions.
+
+        params str name: The name of the subsegment. If not specified
+        the function name will be used.
+        """
+
+        @wrapt.decorator
+        async def wrapper(wrapped, instance, args, kwargs):
+            func_name = name
+            if not func_name:
+                func_name = wrapped.__name__
+
+            result = await self.record_subsegment_async(
+                wrapped, instance, args, kwargs,
+                name=func_name,
+                namespace='local',
+                meta_processor=None,
+            )
+
+            return result
+
+        return wrapper
+
+    async def record_subsegment_async(self, wrapped, instance, args, kwargs, name,
+                                namespace, meta_processor):
+
+        subsegment = self.begin_subsegment(name, namespace)
+
+        exception = None
+        stack = None
+        return_value = None
+
+        try:
+            return_value = await wrapped(*args, **kwargs)
+            return return_value
+        except Exception as e:
+            exception = e
+            stack = traceback.extract_stack(limit=self._max_trace_back)
+            raise
+        finally:
+            end_time = time.time()
+            if callable(meta_processor):
+                meta_processor(
+                    wrapped=wrapped,
+                    instance=instance,
+                    args=args,
+                    kwargs=kwargs,
+                    return_value=return_value,
+                    exception=exception,
+                    subsegment=subsegment,
+                    stack=stack,
+                )
+            elif exception:
+                if subsegment:
+                    subsegment.add_exception(exception, stack)
+
+            self.end_subsegment(end_time)

aws_xray_sdk/core/patcher.py

@@ -4,6 +4,7 @@
 log = logging.getLogger(__name__)
 
 SUPPORTED_MODULES = (
+    'aiobotocore',
     'botocore',
     'requests',
     'sqlite3',
@@ -23,10 +24,14 @@ def patch(modules_to_patch, raise_errors=True):
 
 
 def _patch_module(module_to_patch, raise_errors=True):
-    # boto3 depends on botocore and patch botocore is sufficient
+    # boto3 depends on botocore and patching botocore is sufficient
     if module_to_patch == 'boto3':
         module_to_patch = 'botocore'
 
+    # aioboto3 depends on aiobotocore and patching aiobotocore is sufficient
+    if module_to_patch == 'aioboto3':
+        module_to_patch = 'aiobotocore'
+
     if module_to_patch not in SUPPORTED_MODULES:
         raise Exception('module %s is currently not supported for patching'
                         % module_to_patch)

aws_xray_sdk/core/recorder.py

@@ -30,7 +30,9 @@ class AWSXRayRecorder(object):
     A global AWS X-Ray recorder that will begin/end segments/subsegments
     and send them to the X-Ray daemon. This recorder is initialized during
     loading time so you can use::
+
         from aws_xray_sdk.core import xray_recorder
+
     in your module to access it
     """
     def __init__(self):
@@ -312,15 +314,16 @@ def record_subsegment(self, wrapped, instance, args, kwargs, name,
 
         subsegment = self.begin_subsegment(name, namespace)
 
+        exception = None
+        stack = None
+        return_value = None
+
         try:
             return_value = wrapped(*args, **kwargs)
-            exception = None
-            stack = None
             return return_value
         except Exception as e:
             exception = e
             stack = traceback.extract_stack(limit=self._max_trace_back)
-            return_value = None
             raise
         finally:
             end_time = time.time()

aws_xray_sdk/core/sampling/default_sampler.py

@@ -2,13 +2,12 @@
 import json
 from random import Random
 
+from pkg_resources import resource_filename
 from .sampling_rule import SamplingRule
 from ..exceptions.exceptions import InvalidSamplingManifestError
 
-__location__ = os.path.realpath(
-    os.path.join(os.getcwd(), os.path.dirname(__file__)))
 
-with open(os.path.join(__location__, 'default_sampling_rule.json')) as f:
+with open(resource_filename(__name__, 'default_sampling_rule.json')) as f:
     default_sampling_rule = json.load(f)
 
 

aws_xray_sdk/core/utils/compat.py

@@ -2,6 +2,7 @@
 
 
 PY2 = sys.version_info < (3,)
+PY35 = sys.version_info >= (3, 5)
 
 if PY2:
     annotation_value_types = (int, long, float, bool, str)  # noqa: F821

aws_xray_sdk/ext/aiobotocore/__init__.py

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

aws_xray_sdk/ext/aiobotocore/patch.py

@@ -0,0 +1,39 @@
+import aiobotocore.client
+import wrapt
+
+from aws_xray_sdk.core import xray_recorder
+from aws_xray_sdk.ext.boto_utils import inject_header, aws_meta_processor
+
+
+def patch():
+    """
+    Patch aiobotocore client so it generates subsegments
+    when calling AWS services.
+    """
+    if hasattr(aiobotocore.client, '_xray_enabled'):
+        return
+    setattr(aiobotocore.client, '_xray_enabled', True)
+
+    wrapt.wrap_function_wrapper(
+        'aiobotocore.client',
+        'AioBaseClient._make_api_call',
+        _xray_traced_aiobotocore,
+    )
+
+    wrapt.wrap_function_wrapper(
+        'aiobotocore.endpoint',
+        'AioEndpoint._encode_headers',
+        inject_header,
+    )
+
+
+async def _xray_traced_aiobotocore(wrapped, instance, args, kwargs):
+    service = instance._service_model.metadata["endpointPrefix"]
+    result = await xray_recorder.record_subsegment_async(
+        wrapped, instance, args, kwargs,
+        name=service,
+        namespace='aws',
+        meta_processor=aws_meta_processor,
+    )
+
+    return result