diff --git a/src/scout_apm/async_/instruments/aioredis.py b/src/scout_apm/async_/instruments/aioredis.py new file mode 100644 index 00000000..61c4b9cc --- /dev/null +++ b/src/scout_apm/async_/instruments/aioredis.py @@ -0,0 +1,27 @@ +# 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() + with tracked_request.span(operation="Redis/{}".format(op)): + return await wrapped(*args, **kwargs) + + +@wrapt.decorator +async def wrapped_pipeline_execute(wrapped, instance, args, kwargs): + tracked_request = TrackedRequest.instance() + with tracked_request.span(operation="Redis/MULTI"): + return await wrapped(*args, **kwargs) diff --git a/src/scout_apm/instruments/__init__.py b/src/scout_apm/instruments/__init__.py index 5d0c7c9e..44aea363 100644 --- a/src/scout_apm/instruments/__init__.py +++ b/src/scout_apm/instruments/__init__.py @@ -8,7 +8,16 @@ logger = logging.getLogger(__name__) -instrument_names = ["asyncio", "elasticsearch", "jinja2", "pymongo", "redis", "urllib3"] + +instrument_names = [ + "aioredis", + "asyncio", + "elasticsearch", + "jinja2", + "pymongo", + "redis", + "urllib3", +] def ensure_all_installed(): diff --git a/src/scout_apm/instruments/aioredis.py b/src/scout_apm/instruments/aioredis.py new file mode 100644 index 00000000..0ab4f792 --- /dev/null +++ b/src/scout_apm/instruments/aioredis.py @@ -0,0 +1,75 @@ +# 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_pipeline_execute, + wrapped_redis_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 diff --git a/src/scout_apm/instruments/redis.py b/src/scout_apm/instruments/redis.py index 1262ab08..1c2fe985 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_aioredis_py36plus.py b/tests/integration/instruments/test_aioredis_py36plus.py new file mode 100644 index 00000000..f78e5392 --- /dev/null +++ b/tests/integration/instruments/test_aioredis_py36plus.py @@ -0,0 +1,180 @@ +# 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" + + +@async_test +async def test_pipeline_echo(tracked_request): + redis_conn = await get_redis_conn() + 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" + + +@async_test +async def test_execute_command_missing_argument(tracked_request): + redis_conn = await get_redis_conn() + # 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" + + +@async_test +async def test_perform_request_bad_url(tracked_request): + redis_conn = await get_redis_conn() + 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" diff --git a/tox.ini b/tox.ini index 7261c914..b88bab7d 100644 --- a/tox.ini +++ b/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