Fix redis.asyncio sync operations wrapper

[bc291]

Overview

redis.asyncio: Fix synchronous methods wrapped in coroutines

Related Github Issue

#623

Testing

Few tests added - I'm not quite sure if put in the right place

Bug reproduction

To reproduce this behaviour sync method of async redis may be used.
As example redis.pubsub() could be invoked, followed by pattern subscription:

from redis import asyncio as aioredis

async_redis = aioredis.from_url(REDIS_URL, encoding="utf8", decode_responses=True)

async def redis_test():
    pubsub = async_redis.pubsub()
    await pubsub.psubscribe("channel:*")

This results in AttributeError: 'coroutine' object has no attribute 'psubscribe':

Traceback (most recent call last):
  File "/opt/poetry-venv/lib/python3.11/site-packages/newrelic/api/asgi_application.py", line 363, in nr_async_asgi
    return await coro
           ^^^^^^^^^^
  File "/opt/poetry-venv/lib/python3.11/site-packages/newrelic/common/async_proxy.py", line 148, in __next__
    return self.send(None)
           ^^^^^^^^^^^^^^^
  File "/opt/poetry-venv/lib/python3.11/site-packages/newrelic/common/async_proxy.py", line 120, in send
    return self.__wrapped__.send(value)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/poetry-venv/lib/python3.11/site-packages/fastapi/applications.py", line 270, in __call__
    await super().__call__(scope, receive, send)
  File "/opt/poetry-venv/lib/python3.11/site-packages/newrelic/api/asgi_application.py", line 363, in nr_async_asgi
    return await coro
           ^^^^^^^^^^
  File "/opt/poetry-venv/lib/python3.11/site-packages/newrelic/common/async_proxy.py", line 148, in __next__
    return self.send(None)
           ^^^^^^^^^^^^^^^
  File "/opt/poetry-venv/lib/python3.11/site-packages/newrelic/common/async_proxy.py", line 120, in send
    return self.__wrapped__.send(value)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  ...
  File "/opt/poetry-venv/lib/python3.11/site-packages/fastapi/routing.py", line 161, in run_endpoint_function
    return await dependant.call(**values)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "test.py", line 14, in redis_test
    await pubsub.psubscribe("channel:*")
          ^^^^^^^^^^^^^^^^^
AttributeError: 'coroutine' object has no attribute 'psubscribe'

Culprit

Bug perhaps lays here

Logic assumes that every redis.* interaction is an asynchronous operation and it breaks usage of e.g. pubsub(), monitor(), client() and many more (see solution no 3 below).

Solution

inspect.iscoroutinefunction - rejected

inspect.iscoroutinefunction may be problematic here because -> coroutine happens after invocation:

from redis.asyncio import client
import inspect

redis = client.Redis()

redis.get('test')  # <coroutine object Redis.execute_command at ...>

redis.get  # <bound method BasicKeyCommands.get of Redis<ConnectionPool<Connection<host=localhost,port=6379,db=0>>>>

inspect.iscoroutinefunction(redis.get)  # False
inspect.isawaitable(redis.get('test'))  # True

So we can invoke method prematurely, but I think that's going to break DatastoreTrace

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.
    aioredis_version = get_package_version_tuple("aioredis")
    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.
        return wrapped(*args, **kwargs)
    else:
         # 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)

    wrapped_instance = wrapped(*args, **kwargs)  # pre invocation, isawaitable may be used

    if not inspect.isawaitable(wrapped_instance):
        return wrapped_instance

    # Method should be run when awaited, therefore we wrap in an async wrapper.
    return _nr_wrapper_AioRedis_async_method_(wrapped_instance)(*args, **kwargs)

inspect.signature - rejected

Another solution would be to inspect.signature() of wrapped function, e.g.:

inspect.signature(redis.get).return_annotation  # typing.Union[typing.Awaitable, typing.Any]
inspect.signature(redis.set).return_annotation  # typing.Union[typing.Awaitable, typing.Any]

# may raise TypeError when arg 1 is not a class
if not issubclass(typing.Awaitable, inspect.signature(wrapped).return_annotation):
    pass

but after some investigation this is not reliable across python versions.

predefined sync methods list - imlemented in this PR

Using _redis_client_methods and some inspect.signature(wrapped).return_annotation we can compile sync method list:

_redis_client_sync_methods = {'renamenx', 'watch', 'sort_ro', 'command_getkeysandflags', 'bzmpop', 'pubsub',
 'script_debug', 'sentinel_reset', 'memory_help', 'srandmember', 'lpop', 'unwatch',
 'command', 'blmpop', 'command_docs', 'hello', 'lpos', 'debug_segfault', 'zrevrangebyscore',
 'memory_doctor', 'psetex', 'sentinel_get_master_addr_by_name', 'sort', 'zlexcount',
 'latency_graph', 'spop', 'command_info', 'sentinel_master', 'bitfield', 'psync', 'monitor',
 'acl_dryrun', 'sentinel_failover', 'sentinel_set', 'sentinel_slaves', 'sentinel_remove',
 'pexpiretime', 'shutdown', 'bgrewriteaof', 'auth', 'sentinel_sentinels', 'client',
 'latency_histogram', 'sentinel_monitor', 'sentinel_flushconfig', 'latency_doctor',
 'rpop', 'failover', 'expiretime', 'sentinel_ckquorum', 'lcs', 'sentinel_masters'}

and simply use that in wrapping function decision.

[CLAassistant]

Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you all sign our Contributor License Agreement before we can accept your contribution.
3 out of 4 committers have signed the CLA.

✅ bc291
✅ ahmedhr
✅ lrafeei
❌ mergify[bot]
_{You have signed the CLA already but the status is still pending? Let us recheck it.}

[codecov-commenter]

Codecov Report

❗ No coverage uploaded for pull request base (develop-swap-redis-asyncio-commits@d7d0968). Click here to learn what that means.
The diff coverage is n/a.

@@                          Coverage Diff                          @@
##             develop-swap-redis-asyncio-commits     #782   +/-   ##
=====================================================================
  Coverage                                      ?   81.69%           
=====================================================================
  Files                                         ?      185           
  Lines                                         ?    18638           
  Branches                                      ?     3265           
=====================================================================
  Hits                                          ?    15226           
  Misses                                        ?     2496           
  Partials                                      ?      916           

📣 We’re building smart automated test selection to slash your CI/CD build times. Learn more

[lrafeei]

Hi @bc291 ! If you could sign the CLA we'll get this PR merged. I've moved some tests around and tweaked the existing instrumentation, but overall, this looks great--thanks for the contribution!

[bc291]

@lrafeei
sorry, did not get notification about this one 🤔

[lrafeei]

@bc291 No worries--I'll reopen this and if you sign the CLA we can get this one merged