Support redis.asyncio (#744)

* Support `redis.asyncio` * Fix `flake8` issues Co-authored-by: Timothy Pansino <[email protected]>
newrelic · Jan 26, 2023 · ca2c0ee · ca2c0ee
1 parent 01ceca7
commit ca2c0ee
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 11 deletions.
diff --git a/newrelic/config.py b/newrelic/config.py
@@ -2665,6 +2665,14 @@ def _process_module_builtin_defaults():
         "aioredis.connection", "newrelic.hooks.datastore_aioredis", "instrument_aioredis_connection"
     )
 
+    _process_module_definition("redis.asyncio.client", "newrelic.hooks.datastore_aioredis", "instrument_aioredis_client")
+
+    _process_module_definition("redis.asyncio.commands", "newrelic.hooks.datastore_aioredis", "instrument_aioredis_client")
+
+    _process_module_definition(
+        "redis.asyncio.connection", "newrelic.hooks.datastore_aioredis", "instrument_aioredis_connection"
+    )
+
     _process_module_definition(
         "elasticsearch.client",
         "newrelic.hooks.datastore_elasticsearch",

diff --git a/newrelic/hooks/datastore_aioredis.py b/newrelic/hooks/datastore_aioredis.py
@@ -12,24 +12,26 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from newrelic.api.datastore_trace import DatastoreTrace, DatastoreTraceWrapper
+from newrelic.api.datastore_trace import DatastoreTrace
 from newrelic.api.time_trace import current_trace
 from newrelic.api.transaction import current_transaction
-from newrelic.common.object_wrapper import wrap_function_wrapper, function_wrapper, FunctionWrapper
+from newrelic.common.object_wrapper import wrap_function_wrapper, function_wrapper
 from newrelic.hooks.datastore_redis import (
     _redis_client_methods,
     _redis_multipart_commands,
     _redis_operation_re,
 )
 
-from newrelic.common.async_wrapper import async_wrapper
 
-import aioredis
-
-try:
-    AIOREDIS_VERSION = lambda: tuple(int(x) for x in getattr(aioredis, "__version__").split("."))
-except Exception:
-    AIOREDIS_VERSION = lambda: (0, 0, 0)
+def get_aioredis_version():
+    try:
+        import aioredis as aioredis_legacy
+    except ModuleNotFoundError:
+        return None
+    try:
+        return tuple(int(x) for x in getattr(aioredis_legacy, "__version__").split("."))
+    except Exception:
+        return 0, 0, 0
 
 
 def _conn_attrs_to_dict(connection):
@@ -68,7 +70,8 @@ def _nr_wrapper_AioRedis_method_(wrapped, instance, args, kwargs):
         # Check for transaction and return early if found.
         # Method will return synchronously without executing,
         # it will be added to the command stack and run later.
-        if AIOREDIS_VERSION() < (2,):
+        aioredis_version = get_aioredis_version()
+        if aioredis_version and aioredis_version < (2,):
             # AioRedis v1 uses a RedisBuffer instead of a real connection for queueing up pipeline commands
             from aioredis.commands.transaction import _RedisBuffer
             if isinstance(instance._pool_or_conn, _RedisBuffer):
@@ -77,7 +80,10 @@ def _nr_wrapper_AioRedis_method_(wrapped, instance, args, kwargs):
                 return wrapped(*args, **kwargs)
         else:
             # AioRedis v2 uses a Pipeline object for a client and internally queues up pipeline commands
-            from aioredis.client import Pipeline
+            if aioredis_version:
+                from aioredis.client import Pipeline
+            else:
+                from redis.asyncio.client import Pipeline
             if isinstance(instance, Pipeline):
                 return wrapped(*args, **kwargs)