Add aredis method wrapper. (#488)

* Add aredis method wrapper. Co-authored-by: Lalleh Rafeei <[email protected]> Co-authored-by: Jessica Shortz <[email protected]> Co-authored-by: Charly DeFreitas <[email protected]> * Drop py310 from aredis test matrix. Co-authored-by: Lalleh Rafeei <[email protected]> Co-authored-by: Jessica Shortz <[email protected]> Co-authored-by: Charly DeFreitas <[email protected]>
newrelic · Feb 25, 2022 · ec02b06 · ec02b06
1 parent 6295f25
commit ec02b06
Show file tree

Hide file tree

Showing 4 changed files with 46 additions and 196 deletions.
diff --git a/newrelic/config.py b/newrelic/config.py
@@ -2547,6 +2547,8 @@ def _process_module_builtin_defaults():
 
     _process_module_definition("solr", "newrelic.hooks.datastore_solrpy", "instrument_solrpy")
 
+    _process_module_definition("aredis.client", "newrelic.hooks.datastore_aredis", "instrument_aredis_client")
+
     _process_module_definition(
         "aredis.connection",
         "newrelic.hooks.datastore_aredis",

diff --git a/newrelic/hooks/datastore_aredis.py b/newrelic/hooks/datastore_aredis.py
@@ -12,196 +12,41 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import re
 
 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.hooks.datastore_redis import _conn_attrs_to_dict, _instance_info, _redis_client_methods, _redis_multipart_commands, _redis_operation_re
 
-_redis_client_methods = (
-    "bgrewriteaof",
-    "bgsave",
-    "client_kill",
-    "client_list",
-    "client_getname",
-    "client_setname",
-    "config_get",
-    "config_set",
-    "config_resetstat",
-    "config_rewrite",
-    "dbsize",
-    "debug_object",
-    "echo",
-    "flushall",
-    "flushdb",
-    "info",
-    "lastsave",
-    "object",
-    "ping",
-    "save",
-    "sentinel",
-    "sentinel_get_master_addr_by_name",
-    "sentinel_master",
-    "sentinel_masters",
-    "sentinel_monitor",
-    "sentinel_remove",
-    "sentinel_sentinels",
-    "sentinel_set",
-    "sentinel_slaves",
-    "shutdown",
-    "slaveof",
-    "slowlog_get",
-    "slowlog_reset",
-    "time",
-    "append",
-    "bitcount",
-    "bitop",
-    "bitpos",
-    "decr",
-    "delete",
-    "dump",
-    "exists",
-    "expire",
-    "expireat",
-    "get",
-    "getbit",
-    "getrange",
-    "getset",
-    "incr",
-    "incrby",
-    "incrbyfloat",
-    "keys",
-    "mget",
-    "mset",
-    "msetnx",
-    "move",
-    "persist",
-    "pexpire",
-    "pexpireat",
-    "psetex",
-    "pttl",
-    "randomkey",
-    "rename",
-    "renamenx",
-    "restore",
-    "set",
-    "setbit",
-    "setex",
-    "setnx",
-    "setrange",
-    "strlen",
-    "substr",
-    "ttl",
-    "type",
-    "watch",
-    "unwatch",
-    "blpop",
-    "brpop",
-    "brpoplpush",
-    "lindex",
-    "linsert",
-    "llen",
-    "lpop",
-    "lpush",
-    "lpushx",
-    "lrange",
-    "lrem",
-    "lset",
-    "ltrim",
-    "rpop",
-    "rpoplpush",
-    "rpush",
-    "rpushx",
-    "sort",
-    "scan",
-    "scan_iter",
-    "sscan",
-    "sscan_iter",
-    "hscan",
-    "hscan_inter",
-    "zscan",
-    "zscan_iter",
-    "sadd",
-    "scard",
-    "sdiff",
-    "sdiffstore",
-    "sinter",
-    "sinterstore",
-    "sismember",
-    "smembers",
-    "smove",
-    "spop",
-    "srandmember",
-    "srem",
-    "sunion",
-    "sunionstore",
-    "zadd",
-    "zcard",
-    "zcount",
-    "zincrby",
-    "zinterstore",
-    "zlexcount",
-    "zrange",
-    "zrangebylex",
-    "zrangebyscore",
-    "zrank",
-    "zrem",
-    "zremrangebylex",
-    "zremrangebyrank",
-    "zremrangebyscore",
-    "zrevrange",
-    "zrevrangebyscore",
-    "zrevrank",
-    "zscore",
-    "zunionstore",
-    "pfadd",
-    "pfcount",
-    "pfmerge",
-    "hdel",
-    "hexists",
-    "hget",
-    "hgetall",
-    "hincrby",
-    "hincrbyfloat",
-    "hkeys",
-    "hlen",
-    "hset",
-    "hsetnx",
-    "hmset",
-    "hmget",
-    "hvals",
-    "publish",
-    "eval",
-    "evalsha",
-    "script_exists",
-    "script_flush",
-    "script_kill",
-    "script_load",
-    "setex",
-    "lrem",
-    "zadd",
-)
-
-_redis_multipart_commands = set(["client", "cluster", "command", "config", "debug", "sentinel", "slowlog", "script"])
-
-_redis_operation_re = re.compile(r"[-\s]+")
-
-
-def _conn_attrs_to_dict(connection):
-    return {
-        "host": getattr(connection, "host", None),
-        "port": getattr(connection, "port", None),
-        "path": getattr(connection, "path", None),
-        "db": getattr(connection, "db", None),
-    }
-
-
-def _instance_info(kwargs):
-    host = kwargs.get("host") or "localhost"
-    port_path_or_id = str(kwargs.get("port") or kwargs.get("path", "unknown"))
-    db = str(kwargs.get("db") or 0)
-
-    return (host, port_path_or_id, db)
+
+def _wrap_Aredis_method_wrapper_(module, instance_class_name, operation):
+    async def _nr_wrapper_Aredis_method_(wrapped, instance, args, kwargs):
+        transaction = current_transaction()
+        if transaction is None:
+            return await wrapped(*args, **kwargs)
+
+        dt = DatastoreTrace(product="Redis", target=None, operation=operation)
+
+        transaction._nr_datastore_instance_info = (None, None, None)
+
+        with dt:
+            result = await wrapped(*args, **kwargs)
+
+            host, port_path_or_id, db = transaction._nr_datastore_instance_info
+            dt.host = host
+            dt.port_path_or_id = port_path_or_id
+            dt.database_name = db
+            return result
+
+    name = "%s.%s" % (instance_class_name, operation)
+    wrap_function_wrapper(module, name, _nr_wrapper_Aredis_method_)
+
+
+def instrument_aredis_client(module):
+    if hasattr(module, "StrictRedis"):
+        for name in _redis_client_methods:
+            if hasattr(module.StrictRedis, name):
+                _wrap_Aredis_method_wrapper_(module, "StrictRedis", name)
 
 
 async def _nr_Connection_send_command_wrapper_(wrapped, instance, args, kwargs):

diff --git a/tests/datastore_aredis/test_execute_command.py b/tests/datastore_aredis/test_execute_command.py
@@ -19,6 +19,7 @@
 
 from testing_support.fixtures import (validate_transaction_metrics,
     override_application_settings)
+from testing_support.fixture.event_loop import event_loop as loop
 from testing_support.db_settings import redis_settings
 from testing_support.util import instance_hostname
 
@@ -67,20 +68,22 @@
         (_instance_metric_name, None)
 )
 
-def exercise_redis_multi_args(client):
-    client.execute_command('CLIENT', 'LIST', parse='LIST')
+async def exercise_redis_multi_args(client):
+    await client.execute_command('CLIENT', 'LIST', parse='LIST')
+
+async def exercise_redis_single_arg(client):
+    await client.execute_command('CLIENT LIST')
 
-def exercise_redis_single_arg(client):
-    client.execute_command('CLIENT LIST')
 
 @override_application_settings(_enable_instance_settings)
 @validate_transaction_metrics(
-    'test_execute_command:test_strict_redis_execute_command_two_args_enable',
-    scoped_metrics=_enable_scoped_metrics,
-    rollup_metrics=_enable_rollup_metrics,
-    background_task=True)
+        'test_execute_command:test_strict_redis_execute_command_two_args_enable',
+        scoped_metrics=_enable_scoped_metrics,
+        rollup_metrics=_enable_rollup_metrics,
+        background_task=True)
 @background_task()
-def test_strict_redis_execute_command_two_args_enable():
+def test_strict_redis_execute_command_two_args_enable(loop):
     r = aredis.StrictRedis(host=DB_SETTINGS['host'],
             port=DB_SETTINGS['port'], db=0)
-    exercise_redis_multi_args(r)
+    loop.run_until_complete(exercise_redis_multi_args(r))
+
diff --git a/tox.ini b/tox.ini
@@ -84,7 +84,7 @@ envlist =
     solr-datastore_pysolr-{py27,py36,py37,py38,py39,py310,pypy,pypy3},
     redis-datastore_redis-{py27,py36,py37,py38,pypy,pypy3}-redis03,
     redis-datastore_redis-{py36,py37,py38,py39,py310,pypy3}-redis{0400,latest},
-    redis-datastore_aredis-{py36,py37,py38,py39,py310,pypy3}-aredislatest,
+    redis-datastore_aredis-{py36,py37,py38,py39,pypy3}-aredislatest,
     solr-datastore_solrpy-{py27,pypy}-solrpy{00,01},
     python-datastore_sqlite-{py27,py36,py37,py38,py39,py310,pypy,pypy3},
     memcached-datastore_umemcache-{py27,pypy},