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 =