Add support for aioredis. (#577)

* Add aioredis Instrumentation (#567) * Add aioredis instrumentation Co-authored-by: Uma Annamalai <[email protected]> Co-authored-by: Nyenty Ayuk <[email protected]> Co-authored-by: ccedacero-nr <[email protected]> * [Mega-Linter] Apply linters fixes * Bump Tests * Fix double wrapping Co-authored-by: Uma Annamalai <[email protected]> Co-authored-by: Nyenty Ayuk <[email protected]> Co-authored-by: ccedacero-nr <[email protected]> Co-authored-by: TimPansino <[email protected]> * Add aioredis test infrastructure. (#568) * Add aioredis test infra. * Fix flake8 errors. * Aredis concurrency bug reproduction. (#569) Co-authored-by: Uma Annamalai <[email protected]> Co-authored-by: Uma Annamalai <[email protected]> * Add aioredis tests (#573) * Add get and set tests. * Add more testing for aioredis. * Add aioredis testing. Co-authored-by: Tim Pansino <[email protected]> Co-authored-by: Cristian Cedacero <[email protected]> Co-authored-by: Nyenty Ayuk-Enow <[email protected]> * Patch broken tests * Final aioredis testing cleanup Co-authored-by: Nyenty Ayuk <[email protected]> Co-authored-by: ccedacero-nr <[email protected]> Co-authored-by: Uma Annamalai <[email protected]> * Parametrize multiple db tests. * Add missing arg. * Fix typo. * Add missing comma. * Add background_task decorator. * Parametrize instance info tests. * Fix formatting Co-authored-by: Tim Pansino <[email protected]> Co-authored-by: Cristian Cedacero <[email protected]> Co-authored-by: Nyenty Ayuk-Enow <[email protected]> Co-authored-by: Tim Pansino <[email protected]> Co-authored-by: ccedacero-nr <[email protected]> Co-authored-by: Uma Annamalai <[email protected]> Co-authored-by: Timothy Pansino <[email protected]> * Fix AIORedis Concurrency Bug (#574) * Add test for concurrency bug * Fix aioredis concurrency * Fix func signature * Fix ARedis Concurrency Bug (#570) * Patch aredis concurrency bug * Remove xfail marker * Format * Move fixture import Co-authored-by: Uma Annamalai <[email protected]> * Increase concurrency of redis tests (#575) * Instrument All Redis Client Methods (#576) * Initial test files * Fully instrument uninstrumented redis client methods Co-authored-by: Uma Annamalai <[email protected]> Co-authored-by: ccedacero-nr <[email protected]> Co-authored-by: Nyenty Ayuk <[email protected]> * Fix older redis client tests * Fix missing redis client method * Remove sentinel from commands list * Fix sentinels again Co-authored-by: Uma Annamalai <[email protected]> Co-authored-by: ccedacero-nr <[email protected]> Co-authored-by: Nyenty Ayuk <[email protected]> * Add aioredis v1 support (#579) * Add aioredis v1 tests * Fix aioredis v1 Co-authored-by: Uma Annamalai <[email protected]> * Adjust multiple dbs tests * Fix megalinter default base * Fix megalinter base take two * Fix aioredis version parser * Uncomment instance info tests * Fix import issues Co-authored-by: Uma Annamalai <[email protected]> Co-authored-by: Timothy Pansino <[email protected]> Co-authored-by: Uma Annamalai <[email protected]> Co-authored-by: Nyenty Ayuk <[email protected]> Co-authored-by: ccedacero-nr <[email protected]> Co-authored-by: TimPansino <[email protected]> Co-authored-by: Tim Pansino <[email protected]> Co-authored-by: Cristian Cedacero <[email protected]> Co-authored-by: Tim Pansino <[email protected]>
newrelic · Jun 28, 2022 · 77c5909 · 77c5909
1 parent 0dc8434
commit 77c5909
Show file tree

Hide file tree

Showing 21 changed files with 2,027 additions and 190 deletions.
diff --git a/.github/workflows/mega-linter.yml b/.github/workflows/mega-linter.yml
@@ -40,6 +40,7 @@ jobs:
           # All available variables are described in documentation
           # https://megalinter.github.io/configuration/
           VALIDATE_ALL_CODEBASE: ${{ github.event_name == 'push' && github.ref == 'refs/heads/main' }} # Validates all source when push on main, else just the git diff with main. Set 'true' if you always want to lint all sources
+          DEFAULT_BRANCH: ${{ github.event_name == 'pull_request' && github.base_ref || 'main' }}
           GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
           # ADD YOUR CUSTOM ENV VARIABLES HERE TO OVERRIDE VALUES OF .mega-linter.yml AT THE ROOT OF YOUR REPOSITORY
 

diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
@@ -509,11 +509,11 @@ jobs:
 
   redis:
     env:
-      TOTAL_GROUPS: 1
+      TOTAL_GROUPS: 2
 
     strategy:
       matrix:
-        group-number: [1]
+        group-number: [1, 2]
 
     runs-on: ubuntu-latest
     timeout-minutes: 30

diff --git a/newrelic/config.py b/newrelic/config.py
@@ -2603,6 +2603,14 @@ def _process_module_builtin_defaults():
         "instrument_aredis_connection",
     )
 
+    _process_module_definition("aioredis.client", "newrelic.hooks.datastore_aioredis", "instrument_aioredis_client")
+
+    _process_module_definition("aioredis.commands", "newrelic.hooks.datastore_aioredis", "instrument_aioredis_client")
+
+    _process_module_definition(
+        "aioredis.connection", "newrelic.hooks.datastore_aioredis", "instrument_aioredis_connection"
+    )
+
     _process_module_definition(
         "elasticsearch.client",
         "newrelic.hooks.datastore_elasticsearch",
@@ -2691,6 +2699,30 @@ def _process_module_builtin_defaults():
         "redis.commands.core", "newrelic.hooks.datastore_redis", "instrument_redis_commands_core"
     )
 
+    _process_module_definition(
+        "redis.commands.sentinel", "newrelic.hooks.datastore_redis", "instrument_redis_commands_sentinel"
+    )
+
+    _process_module_definition(
+        "redis.commands.json.commands", "newrelic.hooks.datastore_redis", "instrument_redis_commands_json_commands"
+    )
+
+    _process_module_definition(
+        "redis.commands.search.commands", "newrelic.hooks.datastore_redis", "instrument_redis_commands_search_commands"
+    )
+
+    _process_module_definition(
+        "redis.commands.timeseries.commands", "newrelic.hooks.datastore_redis", "instrument_redis_commands_timeseries_commands"
+    )
+
+    _process_module_definition(
+        "redis.commands.bf.commands", "newrelic.hooks.datastore_redis", "instrument_redis_commands_bf_commands"
+    )
+
+    _process_module_definition(
+        "redis.commands.graph.commands", "newrelic.hooks.datastore_redis", "instrument_redis_commands_graph_commands"
+    )
+
     _process_module_definition("motor", "newrelic.hooks.datastore_motor", "patch_motor")
 
     _process_module_definition(

diff --git a/newrelic/hooks/datastore_aioredis.py b/newrelic/hooks/datastore_aioredis.py
@@ -0,0 +1,127 @@
+# Copyright 2010 New Relic, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+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
+from newrelic.hooks.datastore_redis import (
+    _redis_client_methods,
+    _redis_multipart_commands,
+    _redis_operation_re,
+)
+
+
+def _conn_attrs_to_dict(connection):
+    host = getattr(connection, "host", None)
+    port = getattr(connection, "port", None)
+    if not host and not port and hasattr(connection, "_address"):
+        host, port = connection._address
+    return {
+        "host": host,
+        "port": port,
+        "path": getattr(connection, "path", None),
+        "db": getattr(connection, "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", 6379))
+    db = str(kwargs.get("db") or 0)
+
+    return (host, port_path_or_id, db)
+
+
+def _wrap_AioRedis_method_wrapper(module, instance_class_name, operation):
+    async def _nr_wrapper_AioRedis_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)
+
+    name = "%s.%s" % (instance_class_name, operation)
+    wrap_function_wrapper(module, name, _nr_wrapper_AioRedis_method_)
+
+
+async def wrap_Connection_send_command(wrapped, instance, args, kwargs):
+    transaction = current_transaction()
+    if not transaction:
+        return await wrapped(*args, **kwargs)
+
+    host, port_path_or_id, db = (None, None, None)
+
+    try:
+        dt = transaction.settings.datastore_tracer
+        if dt.instance_reporting.enabled or dt.database_name_reporting.enabled:
+            conn_kwargs = _conn_attrs_to_dict(instance)
+            host, port_path_or_id, db = _instance_info(conn_kwargs)
+    except Exception:
+        pass
+
+    # Older Redis clients would when sending multi part commands pass
+    # them in as separate arguments to send_command(). Need to therefore
+    # detect those and grab the next argument from the set of arguments.
+
+    operation = args[0].strip().lower()
+
+    # If it's not a multi part command, there's no need to trace it, so
+    # we can return early.
+
+    if operation.split()[0] not in _redis_multipart_commands:        # Set the datastore info on the DatastoreTrace containing this function call.
+        trace = current_trace()
+
+        # Find DatastoreTrace no matter how many other traces are inbetween
+        while trace is not None and not isinstance(trace, DatastoreTrace):
+            trace = getattr(trace, "parent", None)
+
+        if trace is not None:
+            trace.host = host
+            trace.port_path_or_id = port_path_or_id
+            trace.database_name = db
+
+        return await wrapped(*args, **kwargs)
+
+    # Convert multi args to single arg string
+
+    if operation in _redis_multipart_commands and len(args) > 1:
+        operation = "%s %s" % (operation, args[1].strip().lower())
+
+    operation = _redis_operation_re.sub("_", operation)
+
+    with DatastoreTrace(
+        product="Redis", target=None, operation=operation, host=host, port_path_or_id=port_path_or_id, database_name=db
+    ):
+        return await wrapped(*args, **kwargs)
+
+
+def instrument_aioredis_client(module):
+    # StrictRedis is just an alias of Redis, no need to wrap it as well.
+    if hasattr(module, "Redis"):
+        class_ = getattr(module, "Redis")
+        for operation in _redis_client_methods:
+            if hasattr(class_, operation):
+                _wrap_AioRedis_method_wrapper(module, "Redis", operation)
+
+
+def instrument_aioredis_connection(module):
+    if hasattr(module, "Connection"):
+        if hasattr(module.Connection, "send_command"):
+            wrap_function_wrapper(module, "Connection.send_command", wrap_Connection_send_command)
+
+    if hasattr(module, "RedisConnection"):
+        if hasattr(module.RedisConnection, "execute"):
+            wrap_function_wrapper(module, "RedisConnection.execute", wrap_Connection_send_command)
diff --git a/newrelic/hooks/datastore_aredis.py b/newrelic/hooks/datastore_aredis.py
@@ -12,11 +12,17 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-
 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
-from newrelic.hooks.datastore_redis import _conn_attrs_to_dict, _instance_info, _redis_client_methods, _redis_multipart_commands, _redis_operation_re
+from newrelic.hooks.datastore_redis import (
+    _conn_attrs_to_dict,
+    _instance_info,
+    _redis_client_methods,
+    _redis_multipart_commands,
+    _redis_operation_re,
+)
 
 
 def _wrap_Aredis_method_wrapper_(module, instance_class_name, operation):
@@ -25,35 +31,18 @@ async def _nr_wrapper_Aredis_method_(wrapped, instance, args, kwargs):
         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
+        with DatastoreTrace(product="Redis", target=None, operation=operation):
+            return await wrapped(*args, **kwargs)
 
     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):
+async def wrap_Connection_send_command(wrapped, instance, args, kwargs):
     transaction = current_transaction()
     if not transaction:
         return await wrapped(*args, **kwargs)
-        
+
     host, port_path_or_id, db = (None, None, None)
 
     try:
@@ -64,8 +53,6 @@ async def _nr_Connection_send_command_wrapper_(wrapped, instance, args, kwargs):
     except Exception:
         pass
 
-    transaction._nr_datastore_instance_info = (host, port_path_or_id, db)
-
     # Older Redis clients would when sending multi part commands pass
     # them in as separate arguments to send_command(). Need to therefore
     # detect those and grab the next argument from the set of arguments.
@@ -76,6 +63,18 @@ async def _nr_Connection_send_command_wrapper_(wrapped, instance, args, kwargs):
     # we can return early.
 
     if operation.split()[0] not in _redis_multipart_commands:
+        # Set the datastore info on the DatastoreTrace containing this function call.
+        trace = current_trace()
+
+        # Find DatastoreTrace no matter how many other traces are inbetween
+        while trace is not None and not isinstance(trace, DatastoreTrace):
+            trace = getattr(trace, "parent", None)
+
+        if trace is not None:
+            trace.host = host
+            trace.port_path_or_id = port_path_or_id
+            trace.database_name = db
+
         return await wrapped(*args, **kwargs)
 
     # Convert multi args to single arg string
@@ -89,7 +88,14 @@ async def _nr_Connection_send_command_wrapper_(wrapped, instance, args, kwargs):
         product="Redis", target=None, operation=operation, host=host, port_path_or_id=port_path_or_id, database_name=db
     ):
         return await wrapped(*args, **kwargs)
-
+
+
+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)
+
 
 def instrument_aredis_connection(module):
-    wrap_function_wrapper(module.Connection, "send_command", _nr_Connection_send_command_wrapper_)
+    wrap_function_wrapper(module.Connection, "send_command", wrap_Connection_send_command)