Repository: elastic/apm-agent-python

Issue:
Elastic support for aioredis 
**Is your feature request related to a problem? Please describe.**
I am using aioredis in my asyncio based application. aioredis is supported by a similar observability platform like newrelic. Now that I have integrated my app with the aioredis library I don't see the logs instrumentation for it.

**Describe the solution you'd like**
Similar to aiopg and asyncpg, there should be instrumentation for aioredis as well

**Describe alternatives you've considered**
N/A

**Additional context**
async libraries are getting much popularity and these should be added for a popular apm solution like elastic.



Discussion:

Resolving patch:
diff --git a/elasticapm/instrumentation/packages/asyncio/aioredis.py b/elasticapm/instrumentation/packages/asyncio/aioredis.py
new file mode 100644
--- /dev/null
+++ b/elasticapm/instrumentation/packages/asyncio/aioredis.py
@@ -0,0 +1,95 @@
+#  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 __future__ import absolute_import
+
+from elasticapm.instrumentation.packages.base import AbstractInstrumentedModule
+from elasticapm.traces import execution_context
+from elasticapm.contrib.asyncio.traces import async_capture_span
+
+
+class RedisConnectionPoolInstrumentation(AbstractInstrumentedModule):
+    name = "aioredis"
+
+    instrument_list = [("aioredis.pool", "ConnectionsPool.execute")]
+
+    def call(self, module, method, wrapped, instance, args, kwargs):
+        if len(args) > 0:
+            wrapped_name = args[0].decode()
+        else:
+            wrapped_name = self.get_wrapped_name(wrapped, instance, method)
+
+        with async_capture_span(
+            wrapped_name, span_type="db", span_subtype="redis", span_action="query", leaf=True
+        ) as span:
+            span.context["destination"] = _get_destination_info(instance)
+
+            return wrapped(*args, **kwargs)
+
+
+class RedisPipelineInstrumentation(AbstractInstrumentedModule):
+    name = "aioredis"
+
+    instrument_list = [("aioredis.commands.transaction", "Pipeline.execute")]
+
+    def call(self, module, method, wrapped, instance, args, kwargs):
+        wrapped_name = self.get_wrapped_name(wrapped, instance, method)
+
+        with async_capture_span(
+            wrapped_name, span_type="db", span_subtype="redis", span_action="query", leaf=True
+        ) as span:
+            span.context["destination"] = _get_destination_info(instance)
+
+            return wrapped(*args, **kwargs)
+
+
+class RedisConnectionInstrumentation(AbstractInstrumentedModule):
+    name = "aioredis"
+
+    instrument_list = (("aioredis.connection", "RedisConnection.execute"),)
+
+    def call(self, module, method, wrapped, instance, args, kwargs):
+        span = execution_context.get_span()
+        if span and span.subtype == "aioredis":
+            span.context["destination"] = _get_destination_info(instance)
+        return wrapped(*args, **kwargs)
+
+
+def _get_destination_info(connection):
+    destination_info = {"service": {"name": "aioredis", "resource": "redis", "type": "db"}}
+
+    if hasattr(connection, "_pool_or_conn"):
+        destination_info["port"] = connection._pool_or_conn.address[1]
+        destination_info["address"] = connection._pool_or_conn.address[0]
+    else:
+        destination_info["port"] = connection.address[1]
+        destination_info["address"] = connection.address[0]
+
+    return destination_info
diff --git a/elasticapm/instrumentation/register.py b/elasticapm/instrumentation/register.py
--- a/elasticapm/instrumentation/register.py
+++ b/elasticapm/instrumentation/register.py
@@ -79,6 +79,9 @@
             "elasticapm.instrumentation.packages.tornado.TornadoHandleRequestExceptionInstrumentation",
             "elasticapm.instrumentation.packages.tornado.TornadoRenderInstrumentation",
             "elasticapm.instrumentation.packages.asyncio.httpcore.HTTPCoreAsyncInstrumentation",
+            "elasticapm.instrumentation.packages.asyncio.aioredis.RedisConnectionPoolInstrumentation",
+            "elasticapm.instrumentation.packages.asyncio.aioredis.RedisPipelineInstrumentation",
+            "elasticapm.instrumentation.packages.asyncio.aioredis.RedisConnectionInstrumentation",
         ]
     )
 


Test patch:
diff --git a/tests/instrumentation/asyncio_tests/aioredis_tests.py b/tests/instrumentation/asyncio_tests/aioredis_tests.py
new file mode 100644
--- /dev/null
+++ b/tests/instrumentation/asyncio_tests/aioredis_tests.py
@@ -0,0 +1,125 @@
+#  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.
+
+import pytest  # isort:skip
+
+aioredis = pytest.importorskip("aioredis")  # isort:skip
+
+import os
+
+from elasticapm.conf.constants import TRANSACTION
+from elasticapm.traces import capture_span
+
+pytestmark = [pytest.mark.asyncio, pytest.mark.aioredis]
+
+if "REDIS_HOST" not in os.environ:
+    pytestmark.append(pytest.mark.skip("Skipping redis tests, no REDIS_HOST environment variable set"))
+
+
+@pytest.fixture()
+async def redis_conn():
+    _host = os.environ["REDIS_HOST"]
+    _port = os.environ.get("REDIS_PORT", 6379)
+    conn = await aioredis.create_redis_pool(f"redis://{_host}:{_port}")
+
+    yield conn
+
+    conn.close()
+    await conn.wait_closed()
+
+
+@pytest.mark.integrationtest
+async def test_pipeline(instrument, elasticapm_client, redis_conn):
+    elasticapm_client.begin_transaction("transaction.test")
+    with capture_span("test_pipeline", "test"):
+        pipeline = redis_conn.pipeline()
+        pipeline.rpush("mykey", "a", "b")
+        pipeline.expire("mykey", 1000)
+        await pipeline.execute()
+    elasticapm_client.end_transaction("MyView")
+
+    transactions = elasticapm_client.events[TRANSACTION]
+    spans = elasticapm_client.spans_for_transaction(transactions[0])
+
+    assert spans[0]["name"] in ("StrictPipeline.execute", "Pipeline.execute")
+    assert spans[0]["type"] == "db"
+    assert spans[0]["subtype"] == "redis"
+    assert spans[0]["action"] == "query"
+    assert spans[0]["context"]["destination"] == {
+        "address": os.environ.get("REDIS_HOST", "localhost"),
+        "port": int(os.environ.get("REDIS_PORT", 6379)),
+        "service": {"name": "aioredis", "resource": "redis", "type": "db"},
+    }
+
+    assert spans[1]["name"] == "test_pipeline"
+    assert spans[1]["type"] == "test"
+
+    assert len(spans) == 2
+
+
+@pytest.mark.integrationtest
+async def test_redis_client(instrument, elasticapm_client, redis_conn):
+    elasticapm_client.begin_transaction("transaction.test")
+    with capture_span("test_redis_client", "test"):
+        await redis_conn.rpush("mykey", "a", "b")
+        await redis_conn.expire("mykey", 1000)
+    elasticapm_client.end_transaction("MyView")
+
+    transactions = elasticapm_client.events[TRANSACTION]
+    spans = elasticapm_client.spans_for_transaction(transactions[0])
+
+    spans = sorted(spans, key=lambda x: x["name"])
+
+    assert {t["name"] for t in spans} == {"test_redis_client", "RPUSH", "EXPIRE"}
+
+    assert spans[0]["name"] == "EXPIRE"
+    assert spans[0]["type"] == "db"
+    assert spans[0]["subtype"] == "redis"
+    assert spans[0]["action"] == "query"
+    assert spans[0]["context"]["destination"] == {
+        "address": os.environ.get("REDIS_HOST", "localhost"),
+        "port": int(os.environ.get("REDIS_PORT", 6379)),
+        "service": {"name": "aioredis", "resource": "redis", "type": "db"},
+    }
+
+    assert spans[1]["name"] == "RPUSH"
+    assert spans[1]["type"] == "db"
+    assert spans[1]["subtype"] == "redis"
+    assert spans[1]["action"] == "query"
+    assert spans[1]["context"]["destination"] == {
+        "address": os.environ.get("REDIS_HOST", "localhost"),
+        "port": int(os.environ.get("REDIS_PORT", 6379)),
+        "service": {"name": "aioredis", "resource": "redis", "type": "db"},
+    }
+
+    assert spans[2]["name"] == "test_redis_client"
+    assert spans[2]["type"] == "test"
+
+    assert len(spans) == 3
diff --git a/tests/requirements/reqs-aioredis-2.txt b/tests/requirements/reqs-aioredis-2.txt
new file mode 100644
--- /dev/null
+++ b/tests/requirements/reqs-aioredis-2.txt
@@ -0,0 +1,2 @@
+aioredis>=2.0.0a1
+-r reqs-base.txt
diff --git a/tests/requirements/reqs-aioredis-newest.txt b/tests/requirements/reqs-aioredis-newest.txt
new file mode 100644
--- /dev/null
+++ b/tests/requirements/reqs-aioredis-newest.txt
@@ -0,0 +1,2 @@
+aioredis<2
+-r reqs-base.txt
diff --git a/tests/scripts/envs/aioredis.sh b/tests/scripts/envs/aioredis.sh
new file mode 100644
--- /dev/null
+++ b/tests/scripts/envs/aioredis.sh
@@ -0,0 +1,3 @@
+export PYTEST_MARKER="-m aioredis"
+export DOCKER_DEPS="redis"
+export REDIS_HOST="redis"