diff --git a/aws_xray_sdk/core/patcher.py b/aws_xray_sdk/core/patcher.py index c98a4aff..f4386264 100644 --- a/aws_xray_sdk/core/patcher.py +++ b/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: diff --git a/aws_xray_sdk/ext/pynamodb/__init__.py b/aws_xray_sdk/ext/pynamodb/__init__.py new file mode 100644 index 00000000..4e8acac6 --- /dev/null +++ b/aws_xray_sdk/ext/pynamodb/__init__.py @@ -0,0 +1,3 @@ +from .patch import patch + +__all__ = ['patch'] diff --git a/aws_xray_sdk/ext/pynamodb/patch.py b/aws_xray_sdk/ext/pynamodb/patch.py new file mode 100644 index 00000000..86a2396c --- /dev/null +++ b/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) diff --git a/docs/aws_xray_sdk.ext.pynamodb.rst b/docs/aws_xray_sdk.ext.pynamodb.rst new file mode 100644 index 00000000..13447d6f --- /dev/null +++ b/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: diff --git a/docs/index.rst b/docs/index.rst index b08ec16f..02b033c4 100644 --- a/docs/index.rst +++ b/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. diff --git a/docs/thirdparty.rst b/docs/thirdparty.rst index bd23d38c..fedea80e 100644 --- a/docs/thirdparty.rst +++ b/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 ` for more information about -configuring the ``xray_recorder``. \ No newline at end of file +configuring the ``xray_recorder``. diff --git a/tests/ext/pynamodb/__init__.py b/tests/ext/pynamodb/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/ext/pynamodb/test_pynamodb.py b/tests/ext/pynamodb/test_pynamodb.py new file mode 100644 index 00000000..ee292f92 --- /dev/null +++ b/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 diff --git a/tox.ini b/tox.ini index b428245d..c1539742 100644 --- a/tox.ini +++ b/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