Skip to content

Commit

Permalink
Merge branch 'main' into fix-local-scoped-package-reporting
Browse files Browse the repository at this point in the history
  • Loading branch information
hmstepanek authored Jun 26, 2023
2 parents c64b670 + a7dfe33 commit 8dcef38
Show file tree
Hide file tree
Showing 2 changed files with 12 additions and 8 deletions.
19 changes: 11 additions & 8 deletions newrelic/hooks/datastore_aioredis.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,13 +15,13 @@
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
from newrelic.common.object_wrapper import function_wrapper, wrap_function_wrapper
from newrelic.common.package_version_utils import get_package_version_tuple
from newrelic.hooks.datastore_redis import (
_redis_client_methods,
_redis_multipart_commands,
_redis_operation_re,
)
from newrelic.common.package_version_utils import get_package_version_tuple


def _conn_attrs_to_dict(connection):
Expand All @@ -39,14 +39,13 @@ def _conn_attrs_to_dict(connection):

def _instance_info(kwargs):
host = kwargs.get("host") or "localhost"
port_path_or_id = str(kwargs.get("port") or kwargs.get("path", 6379))
port_path_or_id = str(kwargs.get("path") or kwargs.get("port", 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):

@function_wrapper
async def _nr_wrapper_AioRedis_async_method_(wrapped, instance, args, kwargs):
transaction = current_transaction()
Expand All @@ -55,7 +54,7 @@ async def _nr_wrapper_AioRedis_async_method_(wrapped, instance, args, kwargs):

with DatastoreTrace(product="Redis", target=None, operation=operation):
return await wrapped(*args, **kwargs)

def _nr_wrapper_AioRedis_method_(wrapped, instance, args, kwargs):
# Check for transaction and return early if found.
# Method will return synchronously without executing,
Expand All @@ -64,6 +63,7 @@ def _nr_wrapper_AioRedis_method_(wrapped, instance, args, kwargs):
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):
# Method will return synchronously without executing,
# it will be added to the command stack and run later.
Expand All @@ -80,7 +80,6 @@ def _nr_wrapper_AioRedis_method_(wrapped, instance, args, kwargs):
# Method should be run when awaited, therefore we wrap in an async wrapper.
return _nr_wrapper_AioRedis_async_method_(wrapped)(*args, **kwargs)


name = "%s.%s" % (instance_class_name, operation)
wrap_function_wrapper(module, name, _nr_wrapper_AioRedis_method_)

Expand Down Expand Up @@ -109,7 +108,9 @@ async def wrap_Connection_send_command(wrapped, instance, args, kwargs):
# 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.
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
Expand Down Expand Up @@ -161,7 +162,9 @@ def wrap_RedisConnection_execute(wrapped, instance, args, kwargs):
# 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.
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
Expand Down
1 change: 1 addition & 0 deletions newrelic/hooks/datastore_redis.py
Original file line number Diff line number Diff line change
Expand Up @@ -394,6 +394,7 @@
"unsubscribe",
"unwatch",
"wait",
"waitaof",
"watch",
"xack",
"xadd",
Expand Down

0 comments on commit 8dcef38

Please sign in to comment.