Merge 1ffd956 into 4d560d5

.travis.yml

@@ -17,11 +17,20 @@ install:
   - make install-ci
 
 services:
+  - docker
   - postgresql
   - redis
 
 before_script:
-  - psql -c 'create database travis_ci_test;' -U postgres
+  - docker-compose up -d dynamodb
+  - psql -c 'create database test;' -U postgres
+
+  # waiting for DynamoDB
+  - while [[ $(curl -so /dev/null -w '%{response_code}' localhost:8000) != '400' ]]; do sleep 1; done
+
+  # workaround for Travis CI issue #7940
+  # https://github.com/travis-ci/travis-ci/issues/7940
+  - sudo rm -f /etc/boto.cfg
 
 script:
   - make test lint

README.md

@@ -20,6 +20,7 @@ The module name is `opentracing_instrumentation`.
 ### Supported client frameworks
 
 The following libraries are instrumented for tracing in this module:
+ * [boto3](https://github.com/boto/boto3) — AWS SDK for Python
  * `urllib2`
  * `requests`
  * `SQLAlchemy`
@@ -218,6 +219,11 @@ install_all_patches()
 
 ## Development
 
+`PostgreSQL`, `Redis` and `DynamoDB` are required for certain tests.
+```bash
+docker-compose up -d
+```
+
 To prepare a development environment please execute the following commands.
 ```bash
 virtualenv env

docker-compose.yml

@@ -0,0 +1,19 @@
+version: "3"
+
+services:
+  redis:
+    image: redis
+    ports:
+      - "6379:6379"
+
+  postgres:
+    image: postgres
+    environment:
+      POSTGRES_DB: test
+    ports:
+      - "5432:5432"
+
+  dynamodb:
+    image: amazon/dynamodb-local
+    ports:
+      - "8000:8000"

opentracing_instrumentation/client_hooks/__init__.py

@@ -36,6 +36,7 @@ def install_all_patches():
 
     If a specific module is not available on the path, it is ignored.
     """
+    from . import boto3
     from . import mysqldb
     from . import psycopg2
     from . import strict_redis
@@ -45,6 +46,7 @@ def install_all_patches():
     from . import urllib2
     from . import requests
 
+    boto3.install_patches()
     mysqldb.install_patches()
     psycopg2.install_patches()
     strict_redis.install_patches()

opentracing_instrumentation/client_hooks/boto3.py

@@ -0,0 +1,83 @@
+from __future__ import absolute_import
+
+from opentracing.ext import tags
+
+from opentracing_instrumentation import utils
+from ..request_context import get_current_span
+from ._patcher import Patcher
+
+
+try:
+    from boto3.resources.action import ServiceAction
+    from botocore import xform_name
+    from botocore.exceptions import ClientError
+except ImportError:  # pragma: no cover
+    pass
+else:
+    _service_action_call = ServiceAction.__call__
+
+
+class Boto3Patcher(Patcher):
+    applicable = '_service_action_call' in globals()
+
+    def _install_patches(self):
+        ServiceAction.__call__ = self._get_call_wrapper()
+
+    def _reset_patches(self):
+        ServiceAction.__call__ = _service_action_call
+
+    @staticmethod
+    def set_request_id_tag(span, response):
+        metadata = response.get('ResponseMetadata')
+
+        # there is no ResponseMetadata for
+        # boto3:dynamodb:describe_table
+        if metadata:
+            span.set_tag('aws.request_id', metadata['RequestId'])
+
+    def _get_call_wrapper(self):
+        def call_wrapper(service, parent, *args, **kwargs):
+            """Wraps ServiceAction.__call__"""
+
+            service_name = parent.meta.service_name
+            operation_name = 'boto3:{}:{}'.format(
+                service_name,
+                xform_name(service._action_model.request.operation)
+            )
+            span = utils.start_child_span(operation_name=operation_name,
+                                          parent=get_current_span())
+
+            span.set_tag(tags.SPAN_KIND, tags.SPAN_KIND_RPC_CLIENT)
+            span.set_tag(tags.COMPONENT, 'boto3')
+            span.set_tag('boto3.service_name', service_name)
+
+            with span:
+                try:
+                    response = _service_action_call(service, parent,
+                                                    *args, **kwargs)
+                except ClientError as error:
+                    self.set_request_id_tag(span, error.response)
+                    raise
+                else:
+                    if isinstance(response, dict):
+                        self.set_request_id_tag(span, response)
+
+            return response
+
+        return call_wrapper
+
+
+patcher = Boto3Patcher()
+
+
+def set_patcher(custom_patcher):
+    global patcher
+    patcher = custom_patcher
+
+
+def install_patches():
+    patcher.install_patches()
+
+
+def reset_patches():
+    patcher.reset_patches()

setup.py

@@ -41,11 +41,14 @@
     ],
     extras_require={
         'tests': [
+            'boto3',
+            'botocore',
             'coveralls',
             'doubles',
             'flake8',
             'flake8-quotes',
             'mock',
+            'moto',
             'psycopg2-binary',
             'sqlalchemy>=1.2.0',
 

tests/opentracing_instrumentation/test_boto3.py

@@ -0,0 +1,157 @@
+import datetime
+
+import boto3
+import mock
+import pytest
+import requests
+
+from botocore.exceptions import ClientError
+from opentracing.ext import tags
+
+from opentracing_instrumentation.client_hooks import boto3 as boto3_hooks
+
+
+DYNAMODB_ENDPOINT_URL = 'http://localhost:8000'
+DYNAMODB_CONFIG = {
+    'endpoint_url': DYNAMODB_ENDPOINT_URL,
+    'aws_access_key_id': '-',
+    'aws_secret_access_key': '-',
+    'region_name': 'us-east-1',
+}
+
+
+def create_users_table(dynamodb):
+    dynamodb.create_table(
+        TableName='users',
+        KeySchema=[{
+            'AttributeName': 'username',
+            'KeyType': 'HASH'
+        }],
+        AttributeDefinitions=[{
+            'AttributeName': 'username',
+            'AttributeType': 'S'
+        }],
+        ProvisionedThroughput={
+            'ReadCapacityUnits': 9,
+            'WriteCapacityUnits': 9
+        }
+    )
+
+
+@pytest.fixture
+def dynamodb_mock():
+    import moto
+    with moto.mock_dynamodb2():
+        dynamodb = boto3.resource('dynamodb', region_name='us-east-1')
+        create_users_table(dynamodb)
+        yield dynamodb
+
+
+@pytest.fixture
+def dynamodb():
+    dynamodb = boto3.resource('dynamodb', **DYNAMODB_CONFIG)
+
+    try:
+        dynamodb.Table('users').delete()
+    except ClientError as error:
+        # you can not just use ResourceNotFoundException class
+        # to catch an error since it doesn't exist until it's raised
+        if error.__class__.__name__ != 'ResourceNotFoundException':
+            raise
+
+    create_users_table(dynamodb)
+
+    # waiting until the table exists
+    dynamodb.meta.client.get_waiter('table_exists').wait(TableName='users')
+
+    return dynamodb
+
+
+@pytest.fixture(autouse=True, scope='module')
+def patch_boto3():
+    boto3_hooks.install_patches()
+    try:
+        yield
+    finally:
+        boto3_hooks.reset_patches()
+
+
+def assert_last_span(operation, tracer, response=None):
+    span = tracer.recorder.get_spans()[-1]
+    request_id = response and response['ResponseMetadata']['RequestId']
+    assert span.operation_name == 'boto3:dynamodb:' + operation
+    assert span.tags.get(tags.SPAN_KIND) == tags.SPAN_KIND_RPC_CLIENT
+    assert span.tags.get(tags.COMPONENT) == 'boto3'
+    assert span.tags.get('boto3.service_name') == 'dynamodb'
+    assert span.tags.get('aws.request_id') == request_id
+
+
+def _test(dynamodb, tracer):
+    users = dynamodb.Table('users')
+
+    response = users.put_item(Item={
+        'username': 'janedoe',
+        'first_name': 'Jane',
+        'last_name': 'Doe',
+    })
+    assert_last_span('put_item', tracer, response)
+
+    response = users.get_item(Key={'username': 'janedoe'})
+    user = response['Item']
+    assert user['first_name'] == 'Jane'
+    assert user['last_name'] == 'Doe'
+    assert_last_span('get_item', tracer, response)
+
+    try:
+        dynamodb.Table('test').delete_item(Key={'username': 'janedoe'})
+    except ClientError as error:
+        response = error.response
+    assert_last_span('delete_item', tracer, response)
+
+    response = users.creation_date_time
+    assert isinstance(response, datetime.datetime)
+    assert_last_span('describe_table', tracer)
+
+
+def is_dynamodb_running():
+    try:
+        # feel free to suggest better solution for this check
+        response = requests.get(DYNAMODB_ENDPOINT_URL, timeout=1)
+        return response.status_code == 400
+    except requests.exceptions.ConnectionError:
+        return False
+
+
+def is_moto_presented():
+    try:
+        import moto
+        return True
+    except ImportError:
+        return False
+
+
+@pytest.mark.skipif(not is_dynamodb_running(),
+                    reason='DynamoDB is not running or cannot connect')
+def test_boto3(dynamodb, tracer):
+    _test(dynamodb, tracer)
+
+
+@pytest.mark.skipif(not is_moto_presented(),
+                    reason='moto module is not presented')
+def test_boto3_with_moto(dynamodb_mock, tracer):
+    _test(dynamodb_mock, tracer)
+
+
+@mock.patch.object(boto3_hooks, 'patcher')
+def test_set_custom_patcher(default_patcher):
+    patcher = mock.Mock()
+    boto3_hooks.set_patcher(patcher)
+
+    assert boto3_hooks.patcher is not default_patcher
+    assert boto3_hooks.patcher is patcher
+
+    boto3_hooks.install_patches()
+    boto3_hooks.reset_patches()
+
+    patcher.install_patches.assert_called_once()
+    patcher.reset_patches.assert_called_once()

tests/opentracing_instrumentation/test_postgres.py

@@ -36,7 +36,7 @@
 
 
 SKIP_REASON = 'Postgres is not running or cannot connect'
-POSTGRES_CONNECTION_STRING = 'postgresql://localhost/travis_ci_test'
+POSTGRES_CONNECTION_STRING = 'postgresql://postgres@localhost/test'
 
 
 @pytest.fixture