Merge branch 'master' into improve-stacktraces

Author: haotianw465

CHANGELOG.rst

@@ -2,6 +2,10 @@
 CHANGELOG
 =========
 
+1.1.2
+=====
+* bugfix: Fixed an issue on PynamoDB patcher where the capture didn't handle client timeout.
+
 1.1.1
 =====
 * bugfix: Handle Aiohttp Exceptions as valid responses `PR59 <https://github.com/aws/aws-xray-sdk-python/pull/59>`_.

aws_xray_sdk/core/patcher.py

@@ -11,6 +11,7 @@
     'sqlite3',
     'mysql',
     'httplib',
+    'pymongo',
 )
 
 _PATCHED_MODULES = set()

aws_xray_sdk/ext/pymongo/__init__.py

@@ -0,0 +1,4 @@
+# Copyright © 2018 Clarity Movement Co. All rights reserved.
+from .patch import patch
+
+__all__ = ['patch']

aws_xray_sdk/ext/pymongo/patch.py

@@ -0,0 +1,55 @@
+# Copyright © 2018 Clarity Movement Co. All rights reserved.
+from pymongo import monitoring
+from aws_xray_sdk.core import xray_recorder
+
+
+class XrayCommandListener(monitoring.CommandListener):
+    """
+    A listener that traces all pymongo db commands to AWS Xray.
+    Creates a subsegment for each mongo db conmmand.
+
+    name: '[email protected]:27017'
+    records all available information provided by pymongo,
+    except for `command` and `reply`. They may contain business secrets.
+    If you insist to record them, specify `record_full_documents=True`.
+    """
+
+    def __init__(self, record_full_documents):
+        super(XrayCommandListener, self).__init__()
+        self.record_full_documents = record_full_documents
+
+    def started(self, event):
+        host, port = event.connection_id
+        host_and_port_str = f'{host}:{port}'
+
+        subsegment = xray_recorder.begin_subsegment(
+            f'{event.database_name}@{host_and_port_str}', 'remote')
+        subsegment.put_annotation('mongodb_command_name', event.command_name)
+        subsegment.put_annotation('mongodb_connection_id', host_and_port_str)
+        subsegment.put_annotation('mongodb_database_name', event.database_name)
+        subsegment.put_annotation('mongodb_operation_id', event.operation_id)
+        subsegment.put_annotation('mongodb_request_id', event.request_id)
+        if self.record_full_documents:
+            subsegment.put_metadata('mongodb_command', event.command)
+
+    def succeeded(self, event):
+        subsegment = xray_recorder.current_subsegment()
+        subsegment.put_annotation('mongodb_duration_micros', event.duration_micros)
+        if self.record_full_documents:
+            subsegment.put_metadata('mongodb_reply', event.reply)
+        xray_recorder.end_subsegment()
+
+    def failed(self, event):
+        subsegment = xray_recorder.current_subsegment()
+        subsegment.add_fault_flag()
+        subsegment.put_annotation('mongodb_duration_micros', event.duration_micros)
+        subsegment.put_metadata('failure', event.failure)
+        xray_recorder.end_subsegment()
+
+
+def patch(record_full_documents=False):
+    # ensure `patch()` is idempotent
+    if hasattr(monitoring, '_xray_enabled'):
+        return
+    setattr(monitoring, '_xray_enabled', True)
+    monitoring.register(XrayCommandListener(record_full_documents))

aws_xray_sdk/ext/pynamodb/patch.py

@@ -42,22 +42,24 @@ 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
     }
 
+    # in case of client timeout the return value will be empty
+    if return_value is not None:
+        aws_meta['request_id'] = return_value.headers.get('x-amzn-RequestId')
+        subsegment.put_http_meta(http.STATUS, return_value.status_code)
+
     if exception:
         subsegment.add_error_flag()
         subsegment.add_exception(exception, stack, True)
 
-    subsegment.put_http_meta(http.STATUS, return_value.status_code)
-
+    resp = return_value.json() if return_value else None
     _extract_whitelisted_params(subsegment.name, operation_name, aws_meta,
                                 [None, json.loads(args[0].body.decode('utf-8'))],
-                                None, return_value.json())
+                                None, resp)
 
     subsegment.set_aws(aws_meta)

aws_xray_sdk/version.py

@@ -1 +1 @@
-VERSION = '1.1.1'
+VERSION = '1.1.2'

tests/ext/pynamodb/test_pynamodb.py

@@ -54,6 +54,29 @@ class Meta:
     assert aws_meta['table_name'] == 'mytable'
 
 
+def test_empty_response():
+    from aws_xray_sdk.ext.pynamodb.patch import pynamodb_meta_processor
+    subsegment = xray_recorder.begin_subsegment('test')
+
+    class TempReq(object):
+        def __init__(self):
+            self.headers = {'X-Amz-Target': 'ddb.ListTables'.encode('utf-8')}
+            self.url = 'ddb.us-west-2'
+            self.body = '{}'.encode('utf-8')
+
+    prepared_request = TempReq()
+    args = [prepared_request]
+
+    pynamodb_meta_processor(wrapped=None, instance=None, args=args,
+                            kwargs=None, return_value=None,
+                            exception=None, subsegment=subsegment,
+                            stack=None)
+
+    aws_meta = subsegment.aws
+    assert aws_meta['region'] == 'us-west-2'
+    assert aws_meta['operation'] == 'ListTables'
+
+
 def test_only_dynamodb_calls_are_traced():
     """Test only a single subsegment is created for other AWS services.
 

tox.ini

@@ -1,8 +1,8 @@
 [tox]
 envlist =
     py{27,34,35,36}
-    py36-aiohttp3
-    py35-aiohttp3
+    py36-aiohttp2
+    py35-aiohttp2
     coverage-report
 
 skip_missing_interpreters = True
@@ -21,38 +21,38 @@ deps =
     django >= 1.10, <2.0
     pynamodb
     # Python3.5+ only deps
-    py{35,36}: aiohttp >= 2.3.0,<3.0.0
+    py{35,36}: aiohttp >= 3.0.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/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 --ignore tests/ext/aiohttp/test_client.py
+    py{35,36}: coverage run --source aws_xray_sdk -m py.test tests
 
 setenv =
     DJANGO_SETTINGS_MODULE = tests.ext.django.app.settings
     AWS_SECRET_ACCESS_KEY = fake_key
     AWS_ACCESS_KEY_ID=fake_id
 
-[testenv:py35-aiohttp3]
+[testenv:py35-aiohttp2]
 deps =
     pytest > 3.0.0
-    aiohttp >= 3.0.0
+    aiohttp >= 2.3.0,<3.0.0
     pytest-aiohttp
     coverage
 
 commands =
-    py{35}: coverage run --source aws_xray_sdk -m py.test tests/ext/aiohttp
+    py{35}: coverage run --source aws_xray_sdk -m py.test tests/ext/aiohttp --ignore tests/ext/aiohttp/test_client.py
 
-[testenv:py36-aiohttp3]
+[testenv:py36-aiohttp2]
 deps =
     pytest > 3.0.0
-    aiohttp >= 3.0.0
+    aiohttp >= 2.3.0,<3.0.0
     pytest-aiohttp
     coverage
 
 commands =
-    py{36}: coverage run --source aws_xray_sdk -m py.test tests/ext/aiohttp
+    py{36}: coverage run --source aws_xray_sdk -m py.test tests/ext/aiohttp --ignore tests/ext/aiohttp/test_client.py
 
 [testenv:coverage-report]
 deps = coverage