Instrument async Redis clients similarly to sync clients

open-telemetry · Apr 29, 2022 · a7a39b4 · a7a39b4
1 parent f2628b4
commit a7a39b4
Showing 1 changed file with 25 additions and 0 deletions.
diff --git a/...n/opentelemetry-instrumentation-redis/src/opentelemetry/instrumentation/redis/__init__.py b/...n/opentelemetry-instrumentation-redis/src/opentelemetry/instrumentation/redis/__init__.py
@@ -102,6 +102,10 @@ def response_hook(span, instance, response):
     typing.Callable[[Span, redis.connection.Connection, Any], None]
 ]
 
+_REDIS_ASYNCIO_VERSION = (4, 2, 0)
+if redis.VERSION > _REDIS_ASYNCIO_VERSION:
+    import redis.asyncio
+
 
 def _set_connection_attributes(span, conn):
     if not span.is_recording():
@@ -176,6 +180,22 @@ def _traced_execute_pipeline(func, instance, args, kwargs):
         f"{pipeline_class}.immediate_execute_command",
         _traced_execute_command,
     )
+    if redis.VERSION >= _REDIS_ASYNCIO_VERSION:
+        wrap_function_wrapper(
+            "redis.asyncio",
+            f"{redis_class}.execute_command",
+            _traced_execute_command,
+        )
+        wrap_function_wrapper(
+            "redis.asyncio.client",
+            f"{pipeline_class}.execute",
+            _traced_execute_pipeline,
+        )
+        wrap_function_wrapper(
+            "redis.asyncio.client",
+            f"{pipeline_class}.immediate_execute_command",
+            _traced_execute_command,
+        )
 
 
 class RedisInstrumentor(BaseInstrumentor):
@@ -222,3 +242,8 @@ def _uninstrument(self, **kwargs):
             unwrap(redis.Redis, "pipeline")
             unwrap(redis.client.Pipeline, "execute")
             unwrap(redis.client.Pipeline, "immediate_execute_command")
+        if redis.VERSION >= _REDIS_ASYNCIO_VERSION:
+            unwrap(redis.asyncio.Redis, "execute_command")
+            unwrap(redis.asyncio.Redis, "pipeline")
+            unwrap(redis.asyncio.client.Pipeline, "execute")
+            unwrap(redis.asyncio.client.Pipeline, "immediate_execute_command")