Added separate instrumentation for redis.asyncio.client (#808)

* Added separate instrumentation for redis.asyncio.client

Merge main branch updates

Add tests for newrelic/config.py (#860)

Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

* Modify redis tests

* removed redis.asyncio from aioredis instrumentation

removed aioredis instrumentation in redis asyncio client

removed redis.asyncio from aioredis instrumentation

---------

Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Co-authored-by: Lalleh Rafeei <lrafeei@newrelic.com>
Co-authored-by: Lalleh Rafeei <84813886+lrafeei@users.noreply.github.com>

newrelic/config.py

@@ -2674,12 +2674,14 @@ def _process_module_builtin_defaults():
         "aioredis.connection", "newrelic.hooks.datastore_aioredis", "instrument_aioredis_connection"
     )
 
+    # Redis v4.2+
     _process_module_definition(
-        "redis.asyncio.client", "newrelic.hooks.datastore_aioredis", "instrument_aioredis_client"
+        "redis.asyncio.client", "newrelic.hooks.datastore_redis", "instrument_asyncio_redis_client"
     )
 
+    # Redis v4.2+
     _process_module_definition(
-        "redis.asyncio.commands", "newrelic.hooks.datastore_aioredis", "instrument_aioredis_client"
+        "redis.asyncio.commands", "newrelic.hooks.datastore_redis", "instrument_asyncio_redis_client"
     )
 
     _process_module_definition(

newrelic/hooks/datastore_aioredis.py

@@ -72,8 +72,6 @@ def _nr_wrapper_AioRedis_method_(wrapped, instance, args, kwargs):
             # AioRedis v2 uses a Pipeline object for a client and internally queues up pipeline commands
             if aioredis_version:
                 from aioredis.client import Pipeline
-            else:
-                from redis.asyncio.client import Pipeline
             if isinstance(instance, Pipeline):
                 return wrapped(*args, **kwargs)
 

newrelic/hooks/datastore_redis.py

@@ -16,7 +16,8 @@
 
 from newrelic.api.datastore_trace import DatastoreTrace
 from newrelic.api.transaction import current_transaction
-from newrelic.common.object_wrapper import wrap_function_wrapper
+from newrelic.common.object_wrapper import function_wrapper, wrap_function_wrapper
+
 
 _redis_client_methods = {
     "acl_cat",
@@ -504,6 +505,29 @@ def _nr_wrapper_Redis_method_(wrapped, instance, args, kwargs):
     wrap_function_wrapper(module, name, _nr_wrapper_Redis_method_)
 
 
+def _wrap_asyncio_Redis_method_wrapper(module, instance_class_name, operation):
+
+    @function_wrapper
+    async def _nr_wrapper_asyncio_Redis_async_method_(wrapped, instance, args, kwargs):
+        transaction = current_transaction()
+        if transaction is None:
+            return await wrapped(*args, **kwargs)
+
+        with DatastoreTrace(product="Redis", target=None, operation=operation):
+            return await wrapped(*args, **kwargs)
+
+    def _nr_wrapper_asyncio_Redis_method_(wrapped, instance, args, kwargs):
+        from redis.asyncio.client import Pipeline
+        if isinstance(instance, Pipeline):
+            return wrapped(*args, **kwargs)
+
+        # Method should be run when awaited, therefore we wrap in an async wrapper.
+        return _nr_wrapper_asyncio_Redis_async_method_(wrapped)(*args, **kwargs)
+
+    name = "%s.%s" % (instance_class_name, operation)
+    wrap_function_wrapper(module, name, _nr_wrapper_asyncio_Redis_method_)
+
+
 def _nr_Connection_send_command_wrapper_(wrapped, instance, args, kwargs):
     transaction = current_transaction()
 
@@ -565,6 +589,13 @@ def instrument_redis_client(module):
                 _wrap_Redis_method_wrapper_(module, "Redis", name)
 
 
+def instrument_asyncio_redis_client(module):
+    if hasattr(module, "Redis"):
+        class_ = getattr(module, "Redis")
+        for operation in _redis_client_methods:
+            if hasattr(class_, operation):
+                _wrap_asyncio_Redis_method_wrapper(module, "Redis", operation)
+
 def instrument_redis_commands_core(module):
     _instrument_redis_commands_module(module, "CoreCommands")
 

tox.ini

@@ -94,11 +94,9 @@ envlist =
     mssql-datastore_pymssql-{py37,py38,py39,py310,py311},
     mysql-datastore_pymysql-{py27,py37,py38,py39,py310,py311,pypy27,pypy38},
     solr-datastore_pysolr-{py27,py37,py38,py39,py310,py311,pypy27,pypy38},
-    redis-datastore_redis-{py27,py37,py38,pypy27,pypy38}-redis03,
     redis-datastore_redis-{py37,py38,py39,py310,py311,pypy38}-redis{0400,latest},
-    redis-datastore_aioredis-{py37,py38,py39,py310,pypy38}-aioredislatest,
-    redis-datastore_aioredis-{py37,py38,py39,py310,py311,pypy38}-redislatest,
-    redis-datastore_aredis-{py37,py38,py39,pypy38}-aredislatest,
+    redis-datastore_aioredis-{py37,py38,py39,py310,pypy38}-redis0401,
+    redis-datastore_aredis-{py37,py38,py39,pypy38}-redis0401,
     python-datastore_sqlite-{py27,py37,py38,py39,py310,py311,pypy27,pypy38},
     python-external_boto3-{py27,py37,py38,py39,py310,py311}-boto01,
     python-external_botocore-{py37,py38,py39,py310,py311}-botocorelatest,
@@ -255,11 +253,10 @@ deps =
     datastore_pysolr: pysolr<4.0
     datastore_redis-redislatest: redis
     datastore_redis-redis0400: redis<4.1
-    datastore_redis-redis03: redis<4.0
-    datastore_redis-{py27,pypy27}: rb
-    datastore_aioredis-redislatest: redis
-    datastore_aioredis-aioredislatest: aioredis
-    datastore_aredis-aredislatest: aredis
+    datastore_aioredis-redis0401: redis<4.2
+    datastore_aioredis-redis0401: aioredis
+    datastore_aredis-redis0401: redis<4.2
+    datastore_aredis-redis0401: aredis
     external_boto3-boto01: boto3<2.0
     external_boto3-boto01: moto<2.0
     external_boto3-py27: rsa<4.7.1