From e364038c74f0c03657ab1c5ef61b4f9efe264a00 Mon Sep 17 00:00:00 2001
From: Yoav Alon <yoav@orca.security>
Date: Mon, 3 Aug 2020 14:04:39 +0300
Subject: [PATCH] adding support to aioredis (#534)

---
 src/scout_apm/async_/instruments/redis.py   | 83 +++++++++++++++++++++
 src/scout_apm/instruments/redis.py          |  8 ++
 tests/integration/instruments/test_redis.py | 26 +++++--
 3 files changed, 110 insertions(+), 7 deletions(-)
 create mode 100644 src/scout_apm/async_/instruments/redis.py

diff --git a/src/scout_apm/async_/instruments/redis.py b/src/scout_apm/async_/instruments/redis.py
new file mode 100644
index 00000000..c6761b2f
--- /dev/null
+++ b/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()
diff --git a/src/scout_apm/instruments/redis.py b/src/scout_apm/instruments/redis.py
index e3baf4b2..dacf0fb8 100644
--- a/src/scout_apm/instruments/redis.py
+++ b/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
 
 
diff --git a/tests/integration/instruments/test_redis.py b/tests/integration/instruments/test_redis.py
index 30cdc9c8..fe912593 100644
--- a/tests/integration/instruments/test_redis.py
+++ b/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,9 +24,13 @@ 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):
@@ -34,7 +38,7 @@ def test_install_fail_no_redis(caplog):
     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,