From 3ec21fbc233a46a0483847e2a5c44688d5f08754 Mon Sep 17 00:00:00 2001 From: Terry Cain Date: Mon, 13 Nov 2017 10:24:37 +0000 Subject: [PATCH] Added Async functions for begin subsegment and capture. Also fixed typo from my last PR yield from is not valid syntax in py27 so was causing issues in the recorder module, instead recorder is subclassed to provide some py3 compatible async functions and the recorder class] will be chosen depending on if py2 or py3. Set the correct context class in all test cases as they would inherit the context of the last test ran, in some cases botocore was trying to use a async context. Makes use of async syntax consistent. Reduced duplication of boto patching functions Made the library zipsafe when loading resources Improved test for testing nested subsegments Renamed botocore_para_whitelist.json Updated README --- MANIFEST.in | 2 +- README.md | 27 ++++ aws_xray_sdk/core/__init__.py | 7 +- aws_xray_sdk/core/async_recorder.py | 69 +++++++++ aws_xray_sdk/core/patcher.py | 7 +- aws_xray_sdk/core/recorder.py | 9 +- aws_xray_sdk/core/sampling/default_sampler.py | 5 +- aws_xray_sdk/core/utils/compat.py | 1 + aws_xray_sdk/ext/aiobotocore/__init__.py | 3 + aws_xray_sdk/ext/aiobotocore/patch.py | 39 +++++ aws_xray_sdk/ext/boto_utils.py | 134 ++++++++++++++++++ aws_xray_sdk/ext/botocore/patch.py | 128 +---------------- aws_xray_sdk/ext/botocore/resources.py | 9 -- .../aws_para_whitelist.json} | 0 docs/aws_xray_sdk.ext.aiobotocore.rst | 22 +++ docs/aws_xray_sdk.ext.aiohttp.rst | 22 +++ docs/frameworks.rst | 2 +- docs/thirdparty.rst | 19 ++- tests/ext/aiobotocore/__init__.py | 0 tests/ext/aiobotocore/test_aiobotocore.py | 105 ++++++++++++++ tests/ext/botocore/test_botocore.py | 2 + tests/ext/django/test_middleware.py | 2 + tests/ext/flask/test_flask.py | 3 +- tests/ext/requests/test_requests.py | 2 + tests/ext/sqlite3/test_sqlite3.py | 2 + tests/test_async_recorder.py | 32 +++++ tests/util.py | 8 +- tox.ini | 3 +- 28 files changed, 515 insertions(+), 149 deletions(-) create mode 100644 aws_xray_sdk/core/async_recorder.py create mode 100644 aws_xray_sdk/ext/aiobotocore/__init__.py create mode 100644 aws_xray_sdk/ext/aiobotocore/patch.py create mode 100644 aws_xray_sdk/ext/boto_utils.py delete mode 100644 aws_xray_sdk/ext/botocore/resources.py rename aws_xray_sdk/ext/{botocore/para_whitelist.json => resources/aws_para_whitelist.json} (100%) create mode 100644 docs/aws_xray_sdk.ext.aiobotocore.rst create mode 100644 docs/aws_xray_sdk.ext.aiohttp.rst create mode 100644 tests/ext/aiobotocore/__init__.py create mode 100644 tests/ext/aiobotocore/test_aiobotocore.py create mode 100644 tests/test_async_recorder.py diff --git a/MANIFEST.in b/MANIFEST.in index e56307cb..1c0267e3 100644 --- a/MANIFEST.in +++ b/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 \ No newline at end of file diff --git a/README.md b/README.md index 3702cf41..cf9c3cc9 100644 --- a/README.md +++ b/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. diff --git a/aws_xray_sdk/core/__init__.py b/aws_xray_sdk/core/__init__.py index 4f25eeb1..cb94eed9 100644 --- a/aws_xray_sdk/core/__init__.py +++ b/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', diff --git a/aws_xray_sdk/core/async_recorder.py b/aws_xray_sdk/core/async_recorder.py new file mode 100644 index 00000000..36bb73d7 --- /dev/null +++ b/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) diff --git a/aws_xray_sdk/core/patcher.py b/aws_xray_sdk/core/patcher.py index 4a21a1b2..c98a4aff 100644 --- a/aws_xray_sdk/core/patcher.py +++ b/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) diff --git a/aws_xray_sdk/core/recorder.py b/aws_xray_sdk/core/recorder.py index aa41d1a8..86b03daf 100644 --- a/aws_xray_sdk/core/recorder.py +++ b/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() diff --git a/aws_xray_sdk/core/sampling/default_sampler.py b/aws_xray_sdk/core/sampling/default_sampler.py index 57a3bb28..3e029b2b 100644 --- a/aws_xray_sdk/core/sampling/default_sampler.py +++ b/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) diff --git a/aws_xray_sdk/core/utils/compat.py b/aws_xray_sdk/core/utils/compat.py index 1d60f882..494675c7 100644 --- a/aws_xray_sdk/core/utils/compat.py +++ b/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 diff --git a/aws_xray_sdk/ext/aiobotocore/__init__.py b/aws_xray_sdk/ext/aiobotocore/__init__.py new file mode 100644 index 00000000..4e8acac6 --- /dev/null +++ b/aws_xray_sdk/ext/aiobotocore/__init__.py @@ -0,0 +1,3 @@ +from .patch import patch + +__all__ = ['patch'] diff --git a/aws_xray_sdk/ext/aiobotocore/patch.py b/aws_xray_sdk/ext/aiobotocore/patch.py new file mode 100644 index 00000000..f4131414 --- /dev/null +++ b/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 diff --git a/aws_xray_sdk/ext/boto_utils.py b/aws_xray_sdk/ext/boto_utils.py new file mode 100644 index 00000000..1f1fa396 --- /dev/null +++ b/aws_xray_sdk/ext/boto_utils.py @@ -0,0 +1,134 @@ +from __future__ import absolute_import +# Need absolute import as botocore is also in the current folder for py27 +import json + +from pkg_resources import resource_filename +from botocore.exceptions import ClientError + +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, to_snake_case + + +with open(resource_filename(__name__, 'resources/aws_para_whitelist.json'), 'r') as data_file: + whitelist = json.load(data_file) + + +def inject_header(wrapped, instance, args, kwargs): + headers = args[0] + inject_trace_header(headers, xray_recorder.current_subsegment()) + return wrapped(*args, **kwargs) + + +def aws_meta_processor(wrapped, instance, args, kwargs, + return_value, exception, subsegment, stack): + region = instance.meta.region_name + + if 'operation_name' in kwargs: + operation_name = kwargs['operation_name'] + else: + operation_name = args[0] + + aws_meta = { + 'operation': operation_name, + 'region': region, + } + + if return_value: + resp_meta = return_value.get('ResponseMetadata') + if resp_meta: + aws_meta['request_id'] = resp_meta.get('RequestId') + subsegment.put_http_meta(http.STATUS, + resp_meta.get('HTTPStatusCode')) + # for service like S3 that returns special request id in response headers + if 'HTTPHeaders' in resp_meta and resp_meta['HTTPHeaders'].get('x-amz-id-2'): + aws_meta['id_2'] = resp_meta['HTTPHeaders']['x-amz-id-2'] + + elif exception: + _aws_error_handler(exception, stack, subsegment, aws_meta) + + _extract_whitelisted_params(subsegment.name, operation_name, + aws_meta, args, kwargs, return_value) + + subsegment.set_aws(aws_meta) + + +def _aws_error_handler(exception, stack, subsegment, aws_meta): + + if not exception or not isinstance(exception, ClientError): + return + + response_metadata = exception.response.get('ResponseMetadata') + + if not response_metadata: + return + + aws_meta['request_id'] = response_metadata.get('RequestId') + + status_code = response_metadata.get('HTTPStatusCode') + + subsegment.put_http_meta(http.STATUS, status_code) + if status_code == 429: + subsegment.add_throttle_flag() + if status_code / 100 == 4: + subsegment.add_error_flag() + + subsegment.add_exception(exception, stack, True) + + +def _extract_whitelisted_params(service, operation, + aws_meta, args, kwargs, response): + + # check if service is whitelisted + if service not in whitelist['services']: + return + operations = whitelist['services'][service]['operations'] + + # check if operation is whitelisted + if operation not in operations: + return + params = operations[operation] + + # record whitelisted request/response parameters + if 'request_parameters' in params: + _record_params(params['request_parameters'], args[1], aws_meta) + + if 'request_descriptors' in params: + _record_special_params(params['request_descriptors'], + args[1], aws_meta) + + if 'response_parameters' in params and response: + _record_params(params['response_parameters'], response, aws_meta) + + if 'response_descriptors' in params and response: + _record_special_params(params['response_descriptors'], + response, aws_meta) + + +def _record_params(whitelisted, actual, aws_meta): + + for key in whitelisted: + if key in actual: + snake_key = to_snake_case(key) + aws_meta[snake_key] = actual[key] + + +def _record_special_params(whitelisted, actual, aws_meta): + + for key in whitelisted: + if key in actual: + _process_descriptor(whitelisted[key], actual[key], aws_meta) + + +def _process_descriptor(descriptor, value, aws_meta): + + # "get_count" = true + if 'get_count' in descriptor and descriptor['get_count']: + value = len(value) + + # "get_keys" = true + if 'get_keys' in descriptor and descriptor['get_keys']: + value = value.keys() + + aws_meta[descriptor['rename_to']] = value \ No newline at end of file diff --git a/aws_xray_sdk/ext/botocore/patch.py b/aws_xray_sdk/ext/botocore/patch.py index eb0184e0..fe2b3bf1 100644 --- a/aws_xray_sdk/ext/botocore/patch.py +++ b/aws_xray_sdk/ext/botocore/patch.py @@ -1,12 +1,8 @@ import wrapt import botocore.client -from botocore.exceptions import ClientError 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 to_snake_case -from .resources import whitelist +from aws_xray_sdk.ext.boto_utils import inject_header, aws_meta_processor def patch(): @@ -27,7 +23,7 @@ def patch(): wrapt.wrap_function_wrapper( 'botocore.endpoint', 'Endpoint._encode_headers', - _inject_header, + inject_header, ) @@ -40,123 +36,3 @@ def _xray_traced_botocore(wrapped, instance, args, kwargs): namespace='aws', meta_processor=aws_meta_processor, ) - - -def _inject_header(wrapped, instance, args, kwargs): - headers = args[0] - inject_trace_header(headers, xray_recorder.current_subsegment()) - return wrapped(*args, **kwargs) - - -def aws_meta_processor(wrapped, instance, args, kwargs, - return_value, exception, subsegment, stack): - - region = instance.meta.region_name - - if 'operation_name' in kwargs: - operation_name = kwargs['operation_name'] - else: - operation_name = args[0] - - aws_meta = { - 'operation': operation_name, - 'region': region, - } - - if return_value: - resp_meta = return_value.get('ResponseMetadata') - if resp_meta: - aws_meta['request_id'] = resp_meta.get('RequestId') - subsegment.put_http_meta(http.STATUS, - resp_meta.get('HTTPStatusCode')) - # for service like S3 that returns special request id in response headers - if 'HTTPHeaders' in resp_meta and resp_meta['HTTPHeaders'].get('x-amz-id-2'): - aws_meta['id_2'] = resp_meta['HTTPHeaders']['x-amz-id-2'] - - elif exception: - _aws_error_handler(exception, stack, subsegment, aws_meta) - - _extract_whitelisted_params(subsegment.name, operation_name, - aws_meta, args, kwargs, return_value) - - subsegment.set_aws(aws_meta) - - -def _aws_error_handler(exception, stack, subsegment, aws_meta): - - if not exception or not isinstance(exception, ClientError): - return - - response_metadata = exception.response.get('ResponseMetadata') - - if not response_metadata: - return - - aws_meta['request_id'] = response_metadata.get('RequestId') - - status_code = response_metadata.get('HTTPStatusCode') - - subsegment.put_http_meta(http.STATUS, status_code) - if(status_code == 429): - subsegment.add_throttle_flag() - if(status_code / 100 == 4): - subsegment.add_error_flag() - - subsegment.add_exception(exception, stack, True) - - -def _extract_whitelisted_params(service, operation, - aws_meta, args, kwargs, response): - - # check if service is whitelisted - if service not in whitelist['services']: - return - operations = whitelist['services'][service]['operations'] - - # check if operation is whitelisted - if operation not in operations: - return - params = operations[operation] - - # record whitelisted request/response parameters - if 'request_parameters' in params: - _record_params(params['request_parameters'], args[1], aws_meta) - - if 'request_descriptors' in params: - _record_special_params(params['request_descriptors'], - args[1], aws_meta) - - if 'response_parameters' in params and response: - _record_params(params['response_parameters'], response, aws_meta) - - if 'response_descriptors' in params and response: - _record_special_params(params['response_descriptors'], - response, aws_meta) - - -def _record_params(whitelisted, actual, aws_meta): - - for key in whitelisted: - if key in actual: - snake_key = to_snake_case(key) - aws_meta[snake_key] = actual[key] - - -def _record_special_params(whitelisted, actual, aws_meta): - - for key in whitelisted: - if key in actual: - _process_descriptor(whitelisted[key], actual[key], aws_meta) - - -def _process_descriptor(descriptor, value, aws_meta): - - # "get_count" = true - if 'get_count' in descriptor and descriptor['get_count']: - value = len(value) - - # "get_keys" = true - if 'get_keys' in descriptor and descriptor['get_keys']: - value = value.keys() - - aws_meta[descriptor['rename_to']] = value diff --git a/aws_xray_sdk/ext/botocore/resources.py b/aws_xray_sdk/ext/botocore/resources.py deleted file mode 100644 index 12e07859..00000000 --- a/aws_xray_sdk/ext/botocore/resources.py +++ /dev/null @@ -1,9 +0,0 @@ -import os -import json - - -__location__ = os.path.realpath( - os.path.join(os.getcwd(), os.path.dirname(__file__))) - -with open(os.path.join(__location__, 'para_whitelist.json')) as data_file: - whitelist = json.load(data_file) diff --git a/aws_xray_sdk/ext/botocore/para_whitelist.json b/aws_xray_sdk/ext/resources/aws_para_whitelist.json similarity index 100% rename from aws_xray_sdk/ext/botocore/para_whitelist.json rename to aws_xray_sdk/ext/resources/aws_para_whitelist.json diff --git a/docs/aws_xray_sdk.ext.aiobotocore.rst b/docs/aws_xray_sdk.ext.aiobotocore.rst new file mode 100644 index 00000000..f753a0b5 --- /dev/null +++ b/docs/aws_xray_sdk.ext.aiobotocore.rst @@ -0,0 +1,22 @@ +aws\_xray\_sdk\.ext\.aiobotocore package +======================================== + +Submodules +---------- + +aws\_xray\_sdk\.ext\.aiobotocore\.patch module +---------------------------------------------- + +.. automodule:: aws_xray_sdk.ext.aiobotocore.patch + :members: + :undoc-members: + :show-inheritance: + + +Module contents +--------------- + +.. automodule:: aws_xray_sdk.ext.aiobotocore + :members: + :undoc-members: + :show-inheritance: diff --git a/docs/aws_xray_sdk.ext.aiohttp.rst b/docs/aws_xray_sdk.ext.aiohttp.rst new file mode 100644 index 00000000..c0cccf92 --- /dev/null +++ b/docs/aws_xray_sdk.ext.aiohttp.rst @@ -0,0 +1,22 @@ +aws\_xray\_sdk\.ext\.aiohttp package +==================================== + +Submodules +---------- + +aws\_xray\_sdk\.ext\.aiohttp\.middleware module +----------------------------------------------- + +.. automodule:: aws_xray_sdk.ext.aiohttp.middleware + :members: + :undoc-members: + :show-inheritance: + + +Module contents +--------------- + +.. automodule:: aws_xray_sdk.ext.aiohttp + :members: + :undoc-members: + :show-inheritance: diff --git a/docs/frameworks.rst b/docs/frameworks.rst index af00a01d..09bd2b71 100644 --- a/docs/frameworks.rst +++ b/docs/frameworks.rst @@ -87,7 +87,7 @@ aiohttp Server ============== For X-Ray to create a segment based on an incoming request, you need register some middleware with aiohttp. As aiohttp -is an asyncronous framework, X-Ray will also need to be configured with an ``AsyncContext`` compared to the default threadded +is an asyncronous framework, X-Ray will also need to be configured with an ``AsyncContext`` compared to the default threaded version.:: import asyncio diff --git a/docs/thirdparty.rst b/docs/thirdparty.rst index 8f3aa90c..4a3bfaa1 100644 --- a/docs/thirdparty.rst +++ b/docs/thirdparty.rst @@ -3,7 +3,7 @@ Third Party Library Support =========================== -The SDK supports boto3, botocore, requests, sqlite3 and mysql-connector. +The SDK supports aioboto3, aiobotocore, boto3, botocore, requests, sqlite3 and mysql-connector. To patch, use code like the following in the main app:: @@ -23,6 +23,8 @@ To patch specific modules:: The following modules are availble to patch:: SUPPORTED_MODULES = ( + 'aioboto3', + 'aiobotocore', 'boto3', 'botocore', 'requests', @@ -45,3 +47,18 @@ code like the following to generate a subsegment for an SQL query:: ) conn.cursor().execute('SHOW TABLES') + +Patching aioboto3 and aiobotocore +--------------------------------- + +On top of patching aioboto3 or aiobotocore, the xray_recorder also needs to be +configured to use the ``AsyncContext``. The following snippet shows how to set +up the X-Ray SDK with an Async Context, bear in mind this requires Python 3.5+:: + + from aws_xray_sdk.core.async_context import AsyncContext + from aws_xray_sdk.core import xray_recorder + # Configure X-Ray to use AsyncContext + 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 diff --git a/tests/ext/aiobotocore/__init__.py b/tests/ext/aiobotocore/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/ext/aiobotocore/test_aiobotocore.py b/tests/ext/aiobotocore/test_aiobotocore.py new file mode 100644 index 00000000..bd040aec --- /dev/null +++ b/tests/ext/aiobotocore/test_aiobotocore.py @@ -0,0 +1,105 @@ +import pytest + +import aiobotocore +from botocore.stub import Stubber, ANY + +from aws_xray_sdk.core import patch +from aws_xray_sdk.core.async_context import AsyncContext +from aws_xray_sdk.core import xray_recorder + +patch(('aiobotocore',)) + + +@pytest.fixture(scope='function') +def recorder(loop): + """ + Clean up before and after each test run + """ + xray_recorder.configure(service='test', sampling=False, context=AsyncContext(loop=loop)) + xray_recorder.clear_trace_entities() + yield xray_recorder + xray_recorder.clear_trace_entities() + + +async def test_describe_table(loop, recorder): + segment = recorder.begin_segment('name') + + req_id = '1234' + response = {'ResponseMetadata': {'RequestId': req_id, 'HTTPStatusCode': 403}} + + session = aiobotocore.get_session(loop=loop) + async with session.create_client('dynamodb', region_name='eu-west-2') as client: + with Stubber(client) as stubber: + stubber.add_response('describe_table', response, {'TableName': 'mytable'}) + await client.describe_table(TableName='mytable') + + subsegment = segment.subsegments[0] + assert subsegment.error + assert subsegment.http['response']['status'] == 403 + + aws_meta = subsegment.aws + assert aws_meta['table_name'] == 'mytable' + assert aws_meta['request_id'] == req_id + assert aws_meta['region'] == 'eu-west-2' + assert aws_meta['operation'] == 'DescribeTable' + + +async def test_list_parameter_counting(loop, recorder): + """ + Test special parameters that have shape of list are recorded + as count based on `para_whitelist.json` + """ + segment = recorder.begin_segment('name') + + queue_urls = ['url1', 'url2'] + queue_name_prefix = 'url' + response = { + 'QueueUrls': queue_urls, + 'ResponseMetadata': { + 'RequestId': '1234', + 'HTTPStatusCode': 200, + } + } + + session = aiobotocore.get_session(loop=loop) + async with session.create_client('sqs', region_name='eu-west-2') as client: + with Stubber(client) as stubber: + stubber.add_response('list_queues', response, {'QueueNamePrefix': queue_name_prefix}) + await client.list_queues(QueueNamePrefix='url') + + subsegment = segment.subsegments[0] + assert subsegment.http['response']['status'] == 200 + + aws_meta = subsegment.aws + assert aws_meta['queue_count'] == len(queue_urls) + # all whitelisted input parameters will be converted to snake case + # unless there is an explicit 'rename_to' attribute in json key + assert aws_meta['queue_name_prefix'] == queue_name_prefix + + +async def test_map_parameter_grouping(loop, recorder): + """ + Test special parameters that have shape of map are recorded + as a list of keys based on `para_whitelist.json` + """ + segment = recorder.begin_segment('name') + + response = { + 'ResponseMetadata': { + 'RequestId': '1234', + 'HTTPStatusCode': 500, + } + } + + session = aiobotocore.get_session(loop=loop) + async with session.create_client('dynamodb', region_name='eu-west-2') as client: + with Stubber(client) as stubber: + stubber.add_response('batch_write_item', response, {'RequestItems': ANY}) + await client.batch_write_item(RequestItems={'table1': [{}], 'table2': [{}]}) + + subsegment = segment.subsegments[0] + assert subsegment.fault + assert subsegment.http['response']['status'] == 500 + + aws_meta = subsegment.aws + assert sorted(aws_meta['table_names']) == ['table1', 'table2'] diff --git a/tests/ext/botocore/test_botocore.py b/tests/ext/botocore/test_botocore.py index 8db04fcc..a2d30a72 100644 --- a/tests/ext/botocore/test_botocore.py +++ b/tests/ext/botocore/test_botocore.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 patch(('botocore',)) session = botocore.session.get_session() @@ -18,6 +19,7 @@ def construct_ctx(): 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 diff --git a/tests/ext/django/test_middleware.py b/tests/ext/django/test_middleware.py index a8009193..aefa54f2 100644 --- a/tests/ext/django/test_middleware.py +++ b/tests/ext/django/test_middleware.py @@ -3,12 +3,14 @@ from django.test import TestCase from aws_xray_sdk.core import xray_recorder +from aws_xray_sdk.core.context import Context class XRayTestCase(TestCase): def setUp(self): django.setup() + xray_recorder.configure(service='test', sampling=False, context=Context()) xray_recorder.clear_trace_entities() def tearDown(self): diff --git a/tests/ext/flask/test_flask.py b/tests/ext/flask/test_flask.py index cd41d841..9fab21e7 100644 --- a/tests/ext/flask/test_flask.py +++ b/tests/ext/flask/test_flask.py @@ -2,6 +2,7 @@ from flask import Flask, render_template_string from aws_xray_sdk.ext.flask.middleware import XRayMiddleware +from aws_xray_sdk.core.context import Context from tests.util import get_new_stubbed_recorder @@ -31,7 +32,7 @@ def template(): # add X-Ray middleware to flask app recorder = get_new_stubbed_recorder() -recorder.configure(service='test', sampling=False) +recorder.configure(service='test', sampling=False, context=Context()) XRayMiddleware(app, recorder) # enable testing mode diff --git a/tests/ext/requests/test_requests.py b/tests/ext/requests/test_requests.py index 3cb62c6f..466ca0be 100644 --- a/tests/ext/requests/test_requests.py +++ b/tests/ext/requests/test_requests.py @@ -3,6 +3,7 @@ from aws_xray_sdk.core import patch from aws_xray_sdk.core import xray_recorder +from aws_xray_sdk.core.context import Context patch(('requests',)) @@ -18,6 +19,7 @@ def construct_ctx(): 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 diff --git a/tests/ext/sqlite3/test_sqlite3.py b/tests/ext/sqlite3/test_sqlite3.py index f9c131d9..86889873 100644 --- a/tests/ext/sqlite3/test_sqlite3.py +++ b/tests/ext/sqlite3/test_sqlite3.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 patch(('sqlite3',)) db = sqlite3.connect(":memory:") @@ -16,6 +17,7 @@ def construct_ctx(): 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 diff --git a/tests/test_async_recorder.py b/tests/test_async_recorder.py new file mode 100644 index 00000000..ba7964ea --- /dev/null +++ b/tests/test_async_recorder.py @@ -0,0 +1,32 @@ +from .util import get_new_stubbed_recorder +from aws_xray_sdk.core.async_context import AsyncContext + + +xray_recorder = get_new_stubbed_recorder() + + +@xray_recorder.capture_async('test_2') +async def async_method2(): + pass + + +@xray_recorder.capture_async('test_1') +async def async_method(): + await async_method2() + + +async def test_capture(loop): + xray_recorder.configure(service='test', sampling=False, context=AsyncContext(loop=loop)) + + segment = xray_recorder.begin_segment('name') + + await async_method() + + # Check subsegment is created from async_method + assert len(segment.subsegments) == 1 + assert segment.subsegments[0].name == 'test_1' + + # Check nested subsegment is created from async_method2 + subsegment = segment.subsegments[0] + assert len(subsegment.subsegments) == 1 + assert subsegment.subsegments[0].name == 'test_2' diff --git a/tests/util.py b/tests/util.py index 4ecfa737..229ecfa4 100644 --- a/tests/util.py +++ b/tests/util.py @@ -3,6 +3,7 @@ from aws_xray_sdk.core.recorder import AWSXRayRecorder from aws_xray_sdk.core.emitters.udp_emitter import UDPEmitter +from aws_xray_sdk.core.utils.compat import PY35 class StubbedEmitter(UDPEmitter): @@ -28,7 +29,12 @@ def get_new_stubbed_recorder(): """ Returns a new AWSXRayRecorder object with emitter stubbed """ - recorder = AWSXRayRecorder() + if not PY35: + recorder = AWSXRayRecorder() + else: + from aws_xray_sdk.core.async_recorder import AsyncAWSXRayRecorder + recorder = AsyncAWSXRayRecorder() + recorder.emitter = StubbedEmitter() return recorder diff --git a/tox.ini b/tox.ini index c206ceb9..83c4076a 100644 --- a/tox.ini +++ b/tox.ini @@ -17,9 +17,10 @@ deps = # Python3.5+ only deps py{35,36}: aiohttp >= 2.3.0 py{35,36}: pytest-aiohttp + py{35,36}: aiobotocore commands = - py{27,34}: coverage run --source aws_xray_sdk -m py.test tests --ignore tests/ext/aiohttp --ignore tests/test_async_local_storage.py + py{27,34}: coverage run --source aws_xray_sdk -m py.test tests --ignore tests/ext/aiohttp --ignore tests/ext/aiobotocore --ignore tests/test_async_local_storage.py --ignore tests/test_async_recorder.py py{35,36}: coverage run --source aws_xray_sdk -m py.test tests setenv =