Instrument aioredis

Rebuild of #552 with tracking in an aioredis-specific instrument submodule.

Author: yoav-orca

src/scout_apm/async_/instruments/aioredis.py

@@ -0,0 +1,35 @@
+# coding=utf-8
+from __future__ import absolute_import, division, print_function, unicode_literals
+
+import wrapt
+
+from scout_apm.core.tracked_request import TrackedRequest
+
+
+@wrapt.decorator
+async def wrapped_redis_execute(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_pipeline_execute(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/aioredis.py

@@ -0,0 +1,73 @@
+# coding=utf-8
+from __future__ import absolute_import, division, print_function, unicode_literals
+
+import logging
+
+try:
+    from aioredis import Redis
+except ImportError:  # pragma: no cover
+    Redis = None
+
+try:
+    from aioredis.commands import Pipeline
+except ImportError:
+    Pipeline = None
+
+# The async_ module can only be shipped on Python 3.6+
+try:
+    from scout_apm.async_.instruments.aioredis import (
+        wrapped_redis_execute,
+        wrapped_pipeline_execute,
+    )
+except ImportError:
+    wrapped_redis_execute = None
+    wrapped_pipeline_execute = None
+
+logger = logging.getLogger(__name__)
+
+have_patched_redis_execute = False
+have_patched_pipeline_execute = False
+
+
+def ensure_installed():
+    global have_patched_redis_execute, have_patched_pipeline_execute
+
+    logger.debug("Instrumenting aioredis.")
+
+    if Redis is None:
+        logger.debug("Couldn't import aioredis.Redis - probably not installed.")
+    elif wrapped_redis_execute is None:
+        logger.debug(
+            "Couldn't import scout_apm.async_.instruments.aioredis - probably using Python < 3.6."
+        )
+    elif not have_patched_redis_execute:
+        try:
+            Redis.execute = wrapped_redis_execute(Redis.execute)
+        except Exception as exc:
+            logger.warning(
+                "Failed to instrument aioredis.Redis.execute: %r",
+                exc,
+                exc_info=exc,
+            )
+        else:
+            have_patched_redis_execute = True
+
+    if Pipeline is None:
+        logger.debug(
+            "Couldn't import aioredis.commands.Pipeline - probably not installed."
+        )
+    elif wrapped_pipeline_execute is None:
+        logger.debug(
+            "Couldn't import scout_apm.async_.instruments.aioredis - probably using Python < 3.6."
+        )
+    elif not have_patched_pipeline_execute:
+        try:
+            Pipeline.execute = wrapped_pipeline_execute(Pipeline.execute)
+        except Exception as exc:
+            logger.warning(
+                "Failed to instrument aioredis.commands.Pipeline.execute: %r",
+                exc,
+                exc_info=exc,
+            )
+        else:
+            have_patched_pipeline_execute = True

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_aioredis_py36plus.py

@@ -0,0 +1,168 @@
+# coding=utf-8
+from __future__ import absolute_import, division, print_function, unicode_literals
+
+import logging
+import os
+
+import aioredis
+import pytest
+
+from scout_apm.instruments.aioredis import ensure_installed
+from tests.compat import mock
+from tests.tools import async_test
+
+
+async def get_redis_conn():
+    ensure_installed()
+    # e.g. export REDIS_URL="redis://localhost:6379/0"
+    if "REDIS_URL" not in os.environ:
+        raise pytest.skip("Redis isn't available")
+    conn = await aioredis.create_connection(os.environ["REDIS_URL"])
+    return aioredis.Redis(conn)
+
+
+def test_ensure_installed_twice(caplog):
+    ensure_installed()
+    ensure_installed()
+
+    assert caplog.record_tuples == 2 * [
+        (
+            "scout_apm.instruments.aioredis",
+            logging.DEBUG,
+            "Instrumenting aioredis.",
+        )
+    ]
+
+
+def test_ensure_installed_fail_no_redis_execute(caplog):
+    mock_not_patched = mock.patch(
+        "scout_apm.instruments.aioredis.have_patched_redis_execute", new=False
+    )
+    mock_redis = mock.patch("scout_apm.instruments.aioredis.Redis")
+    with mock_not_patched, mock_redis as mocked_redis:
+        del mocked_redis.execute
+
+        ensure_installed()
+
+    assert len(caplog.record_tuples) == 2
+    assert caplog.record_tuples[0] == (
+        "scout_apm.instruments.aioredis",
+        logging.DEBUG,
+        "Instrumenting aioredis.",
+    )
+    logger, level, message = caplog.record_tuples[1]
+    assert logger == "scout_apm.instruments.aioredis"
+    assert level == logging.WARNING
+    assert message.startswith(
+        "Failed to instrument aioredis.Redis.execute: AttributeError"
+    )
+
+
+def test_ensure_installed_fail_no_wrapped_redis_execute(caplog):
+    mock_not_patched = mock.patch(
+        "scout_apm.instruments.aioredis.have_patched_redis_execute", new=False
+    )
+    mock_wrapped_redis_execute = mock.patch(
+        "scout_apm.instruments.aioredis.wrapped_redis_execute", new=None
+    )
+    with mock_not_patched, mock_wrapped_redis_execute:
+        ensure_installed()
+
+    assert len(caplog.record_tuples) == 2
+    assert caplog.record_tuples[0] == (
+        "scout_apm.instruments.aioredis",
+        logging.DEBUG,
+        "Instrumenting aioredis.",
+    )
+    assert caplog.record_tuples[1] == (
+        "scout_apm.instruments.aioredis",
+        logging.DEBUG,
+        "Couldn't import scout_apm.async_.instruments.aioredis - probably using Python < 3.6.",
+    )
+
+
+def test_ensure_installed_fail_no_pipeline_execute(caplog):
+    mock_not_patched = mock.patch(
+        "scout_apm.instruments.aioredis.have_patched_pipeline_execute", new=False
+    )
+    mock_pipeline = mock.patch("scout_apm.instruments.aioredis.Pipeline")
+    with mock_not_patched, mock_pipeline as mocked_pipeline:
+        del mocked_pipeline.execute
+
+        ensure_installed()
+
+    assert len(caplog.record_tuples) == 2
+    assert caplog.record_tuples[0] == (
+        "scout_apm.instruments.aioredis",
+        logging.DEBUG,
+        "Instrumenting aioredis.",
+    )
+    logger, level, message = caplog.record_tuples[1]
+    assert logger == "scout_apm.instruments.aioredis"
+    assert level == logging.WARNING
+    assert message.startswith(
+        "Failed to instrument aioredis.commands.Pipeline.execute: AttributeError"
+    )
+
+
+def test_ensure_installed_fail_no_wrapped_pipeline_execute(caplog):
+    mock_not_patched = mock.patch(
+        "scout_apm.instruments.aioredis.have_patched_pipeline_execute", new=False
+    )
+    mock_wrapped_pipeline_execute = mock.patch(
+        "scout_apm.instruments.aioredis.wrapped_pipeline_execute", new=None
+    )
+    with mock_not_patched, mock_wrapped_pipeline_execute:
+        ensure_installed()
+
+    assert len(caplog.record_tuples) == 2
+    assert caplog.record_tuples[0] == (
+        "scout_apm.instruments.aioredis",
+        logging.DEBUG,
+        "Instrumenting aioredis.",
+    )
+    assert caplog.record_tuples[1] == (
+        "scout_apm.instruments.aioredis",
+        logging.DEBUG,
+        "Couldn't import scout_apm.async_.instruments.aioredis - probably using Python < 3.6.",
+    )
+
+
+@async_test
+async def test_echo(tracked_request):
+    redis_conn = await get_redis_conn()
+
+    await redis_conn.echo("Hello World!")
+
+    assert len(tracked_request.complete_spans) == 1
+    assert tracked_request.complete_spans[0].operation == "Redis/ECHO"
+
+
+# def test_pipeline_echo(redis_conn, tracked_request):
+#     with redis_conn.pipeline() as p:
+#         p.echo("Hello World!")
+#         p.execute()
+
+#     assert len(tracked_request.complete_spans) == 1
+#     assert tracked_request.complete_spans[0].operation == "Redis/MULTI"
+
+
+# def test_execute_command_missing_argument(redis_conn, tracked_request):
+#     # Redis instrumentation doesn't crash if op is missing.
+#     # This raises a TypeError (Python 3) or IndexError (Python 2)
+#     # when calling the original method.
+#     with pytest.raises(IndexError):
+#         redis_conn.execute_command()
+
+#     assert len(tracked_request.complete_spans) == 1
+#     assert tracked_request.complete_spans[0].operation == "Redis/Unknown"
+
+
+# def test_perform_request_bad_url(redis_conn, tracked_request):
+#     with pytest.raises(TypeError):
+#         # Redis instrumentation doesn't crash if op has the wrong type.
+#         # This raises a TypeError when calling the original method.
+#         redis_conn.execute_command(None)
+
+#     assert len(tracked_request.complete_spans) == 1
+#     assert tracked_request.complete_spans[0].operation == "Redis/None"

tox.ini

@@ -16,6 +16,7 @@ passenv =
     MONGODB_URL
     REDIS_URL
 deps =
+    aioredis ; python_version >= "3.6"
     bottle
     cherrypy
     celery!=4.4.4  # https://github.com/celery/celery/issues/6153