Support patching PynamoDB calls to DynamoDB

aws_xray_sdk/core/patcher.py

@@ -6,6 +6,7 @@
 SUPPORTED_MODULES = (
     'aiobotocore',
     'botocore',
+    'pynamodb',
     'requests',
     'sqlite3',
     'mysql',
@@ -19,23 +20,30 @@ def patch_all():
 
 
 def patch(modules_to_patch, raise_errors=True):
-    for m in modules_to_patch:
+    modules = set()
+    for module_to_patch in modules_to_patch:
+        # boto3 depends on botocore and patching botocore is sufficient
+        if module_to_patch == 'boto3':
+            modules.add('botocore')
+        # aioboto3 depends on aiobotocore and patching aiobotocore is sufficient
+        elif module_to_patch == 'aioboto3':
+            modules.add('aiobotocore')
+        # pynamodb requires botocore to be patched as well
+        elif module_to_patch == 'pynamodb':
+            modules.add('botocore')
+            modules.add(module_to_patch)
+        else:
+            modules.add(module_to_patch)
+    unsupported_modules = modules - set(SUPPORTED_MODULES)
+    if unsupported_modules:
+        raise Exception('modules %s are currently not supported for patching'
+                        % ', '.join(unsupported_modules))
+
+    for m in modules:
         _patch_module(m, raise_errors)
 
 
 def _patch_module(module_to_patch, raise_errors=True):
-    # 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)
-
     try:
         _patch(module_to_patch)
     except Exception:

aws_xray_sdk/ext/pynamodb/__init__.py

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

aws_xray_sdk/ext/pynamodb/patch.py

@@ -0,0 +1,63 @@
+import botocore.vendored.requests.sessions
+import json
+import wrapt
+
+from aws_xray_sdk.core import xray_recorder
+from aws_xray_sdk.core.models import http
+from aws_xray_sdk.ext.boto_utils import _extract_whitelisted_params
+
+
+def patch():
+    """Patch PynamoDB so it generates subsegements when calling DynamoDB."""
+    if hasattr(botocore.vendored.requests.sessions, '_xray_enabled'):
+        return
+    setattr(botocore.vendored.requests.sessions, '_xray_enabled', True)
+
+    wrapt.wrap_function_wrapper(
+        'botocore.vendored.requests.sessions',
+        'Session.send',
+        _xray_traced_pynamodb,
+    )
+
+
+def _xray_traced_pynamodb(wrapped, instance, args, kwargs):
+
+    # Check if it's a request to DynamoDB and return otherwise.
+    try:
+        service = args[0].headers['X-Amz-Target'].decode('utf-8').split('_')[0]
+    except KeyError:
+        return wrapped(*args, **kwargs)
+    if service.lower() != 'dynamodb':
+        return wrapped(*args, **kwargs)
+
+    return xray_recorder.record_subsegment(
+        wrapped, instance, args, kwargs,
+        name='dynamodb',
+        namespace='aws',
+        meta_processor=pynamodb_meta_processor,
+    )
+
+
+def pynamodb_meta_processor(wrapped, instance, args, kwargs, return_value,
+                            exception, subsegment, stack):
+    operation_name = args[0].headers['X-Amz-Target'].decode('utf-8').split('.')[1]
+    region = args[0].url.split('.')[1]
+    request_id = return_value.headers.get('x-amzn-RequestId')
+
+    aws_meta = {
+        'operation': operation_name,
+        'request_id': request_id,
+        'region': region
+    }
+
+    if exception:
+        subsegment.add_error_flag()
+        subsegment.add_exception(exception, stack, True)
+
+    subsegment.put_http_meta(http.STATUS, return_value.status_code)
+
+    _extract_whitelisted_params(subsegment.name, operation_name,
+                                aws_meta, [None, json.loads(args[0].body)],
+                                None, return_value.json())
+
+    subsegment.set_aws(aws_meta)

docs/aws_xray_sdk.ext.pynamodb.rst

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

docs/index.rst

@@ -23,8 +23,9 @@ Currently supported web frameworks and libraries:
 * flask
 * boto3
 * botocore
+* pynamodb
 * requests
-* sqlite3 
+* sqlite3
 * mysql-connector
 
 You must have the X-Ray daemon running to use the SDK.

docs/thirdparty.rst

@@ -6,10 +6,11 @@ Third Party Library Support
 Patching Supported Libraries
 ----------------------------
 
-The SDK supports aioboto3, aiobotocore, boto3, botocore, requests, sqlite3 and mysql-connector.
+The SDK supports aioboto3, aiobotocore, boto3, botocore, pynamodb, requests, sqlite3 and
+mysql-connector.
 
 To patch, use code like the following in the main app::
-    
+
     from aws_xray_sdk.core import patch_all
 
     patch_all()
@@ -30,12 +31,16 @@ The following modules are availble to patch::
         'aiobotocore',
         'boto3',
         'botocore',
+        'pynamodb',
         'requests',
         'sqlite3',
         'mysql',
     )
 
-Patching boto3 and botocore are equivalent since boto3 depends on botocore
+Patching boto3 and botocore are equivalent since boto3 depends on botocore.
+
+Patching pynamodb applies the botocore patch as well, as it uses the logic from the botocore
+patch to apply the trace header.
 
 Patching mysql
 ----------------------------
@@ -67,4 +72,4 @@ up the X-Ray SDK with an Async Context, bear in mind this requires Python 3.5+::
     xray_recorder.configure(service='service_name', context=AsyncContext())
 
 See :ref:`Configure Global Recorder <configurations>` for more information about
-configuring the ``xray_recorder``.
+configuring the ``xray_recorder``.

tests/ext/pynamodb/test_pynamodb.py

@@ -0,0 +1,74 @@
+import botocore.session
+import pytest
+from botocore.exceptions import ClientError
+from pynamodb.attributes import UnicodeAttribute
+from pynamodb.exceptions import VerboseClientError
+from pynamodb.models import Model
+
+from aws_xray_sdk.core import patch
+from aws_xray_sdk.core import xray_recorder
+from aws_xray_sdk.core.context import Context
+
+patch(('pynamodb',))
+
+
+@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_exception():
+    class SampleModel(Model):
+        class Meta:
+            region = 'us-west-2'
+            table_name = 'mytable'
+
+        sample_attribute = UnicodeAttribute(hash_key=True)
+
+    try:
+        SampleModel.describe_table()
+    except VerboseClientError:
+        pass
+
+    subsegments = xray_recorder.current_segment().subsegments
+    assert len(subsegments) == 1
+    subsegment = subsegments[0]
+    assert subsegment.name == 'dynamodb'
+    assert len(subsegment.subsegments) == 0
+    assert subsegment.error
+
+    aws_meta = subsegment.aws
+    assert aws_meta['region'] == 'us-west-2'
+    assert aws_meta['operation'] == 'DescribeTable'
+    assert aws_meta['table_name'] == 'mytable'
+
+
+def test_only_dynamodb_calls_are_traced():
+    """Test only a single subsegment is created for other AWS services.
+
+    As the pynamodb patch applies the botocore patch as well, we need
+    to ensure that only one subsegment is created for all calls not
+    made by PynamoDB. As PynamoDB calls botocore differently than the
+    botocore patch expects we also just get a single subsegment per
+    PynamoDB call.
+    """
+    session = botocore.session.get_session()
+    s3 = session.create_client('s3', region_name='us-west-2')
+    try:
+        s3.get_bucket_location(Bucket='mybucket')
+    except ClientError:
+        pass
+
+    subsegments = xray_recorder.current_segment().subsegments
+    assert len(subsegments) == 1
+    assert subsegments[0].name == 's3'
+    assert len(subsegments[0].subsegments) == 0

tox.ini

@@ -14,6 +14,7 @@ deps =
     flask >= 0.10
     # the sdk doesn't support earlier version of django
     django >= 1.10, <2.0
+    pynamodb
     # Python3.5+ only deps
     py{35,36}: aiohttp >= 2.3.0
     py{35,36}: pytest-aiohttp