add instrumentation for aiohttp.client

Signed-off-by: Benjamin Wohlwend <[email protected]>

Author: beniwohli

elasticapm/__init__.py

@@ -26,5 +26,4 @@
 from elasticapm.traces import set_transaction_result  # noqa: F401
 
 if sys.version_info >= (3, 5):
-    from elasticapm.contrib.asyncio.traces import async_capture_span #  noqa: F401
-
+    from elasticapm.contrib.asyncio.traces import async_capture_span  # noqa: F401

elasticapm/contrib/asyncio/traces.py

@@ -1,4 +1,5 @@
 import functools
+
 from elasticapm.traces import capture_span, error_logger, get_transaction
 from elasticapm.utils import get_name_from_func
 

elasticapm/instrumentation/packages/aiohttp.py

@@ -0,0 +1,35 @@
+from elasticapm import async_capture_span
+from elasticapm.instrumentation.packages.asyncio_base import \
+    AsyncAbstractInstrumentedModule
+from elasticapm.utils import default_ports
+from elasticapm.utils.compat import urlparse
+
+
+def get_host_from_url(url):
+    parsed_url = urlparse.urlparse(url)
+    host = parsed_url.hostname or " "
+
+    if (
+        parsed_url.port and
+        default_ports.get(parsed_url.scheme) != parsed_url.port
+    ):
+        host += ":" + str(parsed_url.port)
+
+    return host
+
+
+class AioHttpClientInstrumentation(AsyncAbstractInstrumentedModule):
+    name = 'aiohttp_client'
+
+    instrument_list = [
+        ("aiohttp.client", "ClientSession._request"),
+    ]
+
+    async def call(self, module, method, wrapped, instance, args, kwargs):
+        method = kwargs['method'] if 'method' in kwargs else args[0]
+        url = kwargs['method'] if 'method' in kwargs else args[1]
+
+        signature = " ".join([method.upper(), get_host_from_url(url)])
+
+        async with async_capture_span(signature, "ext.http.aiohttp", {'url': url}, leaf=True):
+            return await wrapped(*args, **kwargs)

elasticapm/instrumentation/packages/asyncio_base.py

@@ -0,0 +1,6 @@
+from elasticapm.instrumentation.packages.base import AbstractInstrumentedModule
+
+
+class AsyncAbstractInstrumentedModule(AbstractInstrumentedModule):
+    async def call(self, module, method, wrapped, instance, args, kwargs):
+        raise NotImplementedError()

elasticapm/instrumentation/register.py

@@ -1,3 +1,5 @@
+import sys
+
 from elasticapm.utils.module_import import import_string
 
 _cls_register = {
@@ -25,6 +27,11 @@
     'elasticapm.instrumentation.packages.django.template.DjangoTemplateSourceInstrumentation',
 }
 
+if sys.version_info >= (3, 5):
+    _cls_register.update([
+        'elasticapm.instrumentation.packages.aiohttp.AioHttpClientInstrumentation',
+    ])
+
 
 def register(cls):
     _cls_register.add(cls)