adding support for aioredis (#534)

src/scout_apm/async_/instruments/redis.py

@@ -0,0 +1,83 @@
+# coding=utf-8
+from __future__ import absolute_import, division, print_function, unicode_literals
+
+import logging
+
+import wrapt
+
+from scout_apm.core.tracked_request import TrackedRequest
+
+try:
+    import aioredis
+except ImportError:
+    aioredis = None
+else:
+    from aioredis import Redis as AioRedis
+    from aioredis.commands import Pipeline as AioPipeline
+
+logger = logging.getLogger(__name__)
+
+
+have_patched_aioredis_execute = False
+have_patched_aiopipeline_execute = False
+
+
+def ensure_async_installed():
+    global have_patched_aioredis_execute, have_patched_aiopipeline_execute
+
+    if aioredis is None:
+        logger.debug("Couldn't import aioredis - probably not installed")
+        return
+
+    if wrapped_execute_async is None or wrapped_execute_command_async is None:
+        logger.debug("Couldn't import async wrapper - probably async not supported")
+        return
+
+    if not have_patched_aioredis_execute:
+        try:
+            AioRedis.execute = wrapped_execute_command_async(AioRedis.execute)
+        except Exception as exc:
+            logger.warning(
+                "Failed to instrument aioredis.Redis.execute: %r", exc, exc_info=exc
+            )
+        else:
+            have_patched_aioredis_execute = True
+
+    if not have_patched_aiopipeline_execute:
+        try:
+            AioPipeline.execute = wrapped_execute_command_async(AioPipeline.execute)
+        except Exception as exc:
+            logger.warning(
+                "Failed to instrument aioredis.Redis.execute: %r", exc, exc_info=exc
+            )
+        else:
+            have_patched_aiopipeline_execute = True
+
+
+@wrapt.decorator
+async def wrapped_execute_command_async(wrapped, instance, args, kwargs):
+    try:
+        op = args[0]
+        if isinstance(op, bytes):
+            op = op.decode()
+    except (IndexError, TypeError):
+        op = "Unknown"
+
+    tracked_request = TrackedRequest.instance()
+    tracked_request.start_span(operation="Redis/{}".format(op))
+
+    try:
+        return await wrapped(*args, **kwargs)
+    finally:
+        tracked_request.stop_span()
+
+
+@wrapt.decorator
+async def wrapped_execute_async(wrapped, instance, args, kwargs):
+    tracked_request = TrackedRequest.instance()
+    tracked_request.start_span(operation="Redis/MULTI")
+
+    try:
+        return await wrapped(*args, **kwargs)
+    finally:
+        tracked_request.stop_span()

src/scout_apm/instruments/redis.py

@@ -19,6 +19,11 @@
         from redis import StrictRedis as Redis
         from redis.client import BasePipeline as Pipeline
 
+try:
+    from scout_apm.async_.instruments.redis import ensure_async_installed
+except ImportError:
+    ensure_async_installed = None
+
 logger = logging.getLogger(__name__)
 
 
@@ -56,6 +61,9 @@ def ensure_installed():
             else:
                 have_patched_pipeline_execute = True
 
+    if ensure_async_installed is not None:
+        ensure_async_installed()
+
     return True
 
 

tests/integration/instruments/test_redis.py

@@ -7,7 +7,7 @@
 import pytest
 import redis
 
-from scout_apm.instruments.redis import ensure_installed
+from scout_apm.instruments.redis import ensure_async_installed, ensure_installed
 from tests.compat import mock
 
 
@@ -24,17 +24,21 @@ def test_ensure_installed_twice(caplog):
     ensure_installed()
     ensure_installed()
 
-    assert caplog.record_tuples == 2 * [
-        ("scout_apm.instruments.redis", logging.DEBUG, "Instrumenting redis.",)
-    ]
+    log_lines = (
+        "scout_apm.instruments.redis",
+        logging.DEBUG,
+        "Instrumenting redis.",
+    )
+
+    assert caplog.record_tuples.count(log_lines) == 2
 
 
 def test_install_fail_no_redis(caplog):
     mock_no_redis = mock.patch("scout_apm.instruments.redis.redis", new=None)
     with mock_no_redis:
         ensure_installed()
 
-    assert caplog.record_tuples == [
+    assert caplog.record_tuples[:2] == [
         ("scout_apm.instruments.redis", logging.DEBUG, "Instrumenting redis.",),
         (
             "scout_apm.instruments.redis",
@@ -54,7 +58,11 @@ def test_ensure_installed_fail_no_redis_execute_command(caplog):
 
         ensure_installed()
 
-    assert len(caplog.record_tuples) == 2
+    if ensure_async_installed is None:
+        assert len(caplog.record_tuples) == 2
+    else:
+        assert len(caplog.record_tuples) == 3
+
     assert caplog.record_tuples[0] == (
         "scout_apm.instruments.redis",
         logging.DEBUG,
@@ -78,7 +86,11 @@ def test_ensure_installed_fail_no_pipeline_execute(caplog):
 
         ensure_installed()
 
-    assert len(caplog.record_tuples) == 2
+    if ensure_async_installed is None:
+        assert len(caplog.record_tuples) == 2
+    else:
+        assert len(caplog.record_tuples) == 3
+
     assert caplog.record_tuples[0] == (
         "scout_apm.instruments.redis",
         logging.DEBUG,