Add instrumentation for aiobotocore (#1520)

* Add instrumentation for aiobotocore

* CHANGELOG fix

* `call` needs to be async to not close the
context manager before the coro gets awaited.
refactor common code to limit duplication in
sync and async code.

* Exclude aiobotocore from py3.6

* Add a note about why _call and call are separate

Co-authored-by: Colton Myers <colton.myers@gmail.com>

.ci/.jenkins_exclude.yml

@@ -237,3 +237,8 @@ exclude:
     FRAMEWORK: aiomysql-newest
   - PYTHON_VERSION: python-3.10  # getting "loop argument must agree with lock" error
     FRAMEWORK: aiomysql-newest
+  # aiobotocore
+  - PYTHON_VERSION: pypy-3
+    FRAMEWORK: aiobotocore-newest
+  - PYTHON_VERSION: python-3.6
+    FRAMEWORK: aiobotocore-newest

.ci/.jenkins_framework.yml

@@ -52,3 +52,4 @@ FRAMEWORK:
   - prometheus_client-newest
   - sanic-newest
   - aiomysql-newest
+  - aiobotocore-newest

.ci/.jenkins_framework_full.yml

@@ -84,3 +84,4 @@ FRAMEWORK:
   - httplib2-newest
   - sanic-20.12
   - sanic-newest
+  - aiobotocore-newest

CHANGELOG.asciidoc

@@ -29,6 +29,19 @@ endif::[]
 //===== Bug fixes
 //
 
+=== Unreleased
+
+// Unreleased changes go here
+// When the next release happens, nest these changes under the "Python Agent version 6.x" heading
+[float]
+===== Features
+
+* Add instrumentation for https://github.com/aio-libs/aiobotocore[`aiobotocore`] {pull}xxx[#xxx]
+
+//[float]
+//===== Bug fixes
+//
+
 
 [[release-notes-6.x]]
 === Python Agent version 6.x

docs/supported-technologies.asciidoc

@@ -570,6 +570,24 @@ Collected trace data for all services:
 
 Additionally, some services collect more specific data
 
+[float]
+[[automatic-instrumentation-aiobotocore]]
+==== AWS Aiobotocore
+
+Library: `aiobotocore` (`>=2.2.0`)
+
+Instrumented methods:
+
+ * `aiobotocore.client.BaseClient._make_api_call`
+
+Collected trace data for all services:
+
+ * AWS region (e.g. `eu-central-1`)
+ * AWS service name (e.g. `s3`)
+ * operation name (e.g. `ListBuckets`)
+
+Additionally, some services collect more specific data
+
 [float]
 [[automatic-instrumentation-s3]]
 ===== S3

elasticapm/instrumentation/packages/asyncio/aiobotocore.py

@@ -0,0 +1,49 @@
+#  BSD 3-Clause License
+#
+#  Copyright (c) 2019, Elasticsearch BV
+#  All rights reserved.
+#
+#  Redistribution and use in source and binary forms, with or without
+#  modification, are permitted provided that the following conditions are met:
+#
+#  * Redistributions of source code must retain the above copyright notice, this
+#    list of conditions and the following disclaimer.
+#
+#  * Redistributions in binary form must reproduce the above copyright notice,
+#    this list of conditions and the following disclaimer in the documentation
+#    and/or other materials provided with the distribution.
+#
+#  * Neither the name of the copyright holder nor the names of its
+#    contributors may be used to endorse or promote products derived from
+#    this software without specific prior written permission.
+#
+#  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+#  AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+#  IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+#  DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+#  FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+#  DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+#  SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+#  CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+#  OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+#  OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+from elasticapm.contrib.asyncio.traces import async_capture_span
+from elasticapm.instrumentation.packages.botocore import BotocoreInstrumentation, span_modifiers
+
+
+class AioBotocoreInstrumentation(BotocoreInstrumentation):
+    name = "aiobotocore"
+
+    instrument_list = [("aiobotocore.client", "AioBaseClient._make_api_call")]
+
+    capture_span_ctx = async_capture_span
+
+    async def call(self, module, method, wrapped, instance, args, kwargs):
+        service = self._get_service(instance)
+
+        ctx = self._call(service, instance, args, kwargs)
+        async with ctx as span:
+            if service in span_modifiers:
+                span_modifiers[service](span, args, kwargs)
+            return await wrapped(*args, **kwargs)

elasticapm/instrumentation/packages/botocore.py

@@ -52,19 +52,18 @@ class BotocoreInstrumentation(AbstractInstrumentedModule):
 
     instrument_list = [("botocore.client", "BaseClient._make_api_call")]
 
-    def call(self, module, method, wrapped, instance, args, kwargs):
+    capture_span_ctx = capture_span
+
+    def _call(self, service, instance, args, kwargs):
+        """
+        This is split out from `call()` so that it can be re-used by the
+        aiobotocore instrumentation without duplicating all of this code.
+        """
         if "operation_name" in kwargs:
             operation_name = kwargs["operation_name"]
         else:
             operation_name = args[0]
 
-        service_model = instance.meta.service_model
-        if hasattr(service_model, "service_id"):  # added in boto3 1.7
-            service = service_model.service_id
-        else:
-            service = service_model.service_name.upper()
-            service = endpoint_to_service_id.get(service, service)
-
         parsed_url = urllib.parse.urlparse(instance.meta.endpoint_url)
         context = {
             "destination": {
@@ -81,14 +80,29 @@ def call(self, module, method, wrapped, instance, args, kwargs):
         if not handler_info:
             handler_info = handle_default(operation_name, service, instance, args, kwargs, context)
 
-        with capture_span(
+        return self.capture_span_ctx(
             handler_info.signature,
             span_type=handler_info.span_type,
             leaf=True,
             span_subtype=handler_info.span_subtype,
             span_action=handler_info.span_action,
             extra=handler_info.context,
-        ) as span:
+        )
+
+    def _get_service(self, instance):
+        service_model = instance.meta.service_model
+        if hasattr(service_model, "service_id"):  # added in boto3 1.7
+            service = service_model.service_id
+        else:
+            service = service_model.service_name.upper()
+            service = endpoint_to_service_id.get(service, service)
+        return service
+
+    def call(self, module, method, wrapped, instance, args, kwargs):
+        service = self._get_service(instance)
+
+        ctx = self._call(service, instance, args, kwargs)
+        with ctx as span:
             if service in span_modifiers:
                 span_modifiers[service](span, args, kwargs)
             return wrapped(*args, **kwargs)

elasticapm/instrumentation/register.py

@@ -86,6 +86,7 @@
             "elasticapm.instrumentation.packages.asyncio.aioredis.RedisPipelineInstrumentation",
             "elasticapm.instrumentation.packages.asyncio.aioredis.RedisConnectionInstrumentation",
             "elasticapm.instrumentation.packages.asyncio.aiomysql.AioMySQLInstrumentation",
+            "elasticapm.instrumentation.packages.asyncio.aiobotocore.AioBotocoreInstrumentation",
         ]
     )
 

setup.cfg

@@ -162,6 +162,7 @@ markers =
     prometheus_client
     sanic
     jinja2
+    aiobotocore
 
 [isort]
 line_length=120