aioredis instrumentation fix with extra context and statement (#1406)

* aioredis instrumentation fix with extra context and statement

* CHANGELOG

* Add db context to existing span, rather than a new span

Also add to redis, not just aioredis

* Simplify args checking

Co-authored-by: Colton Myers <colton.myers@gmail.com>
Co-authored-by: Colton Myers <colton@myers.fm>

CHANGELOG.asciidoc

@@ -33,9 +33,11 @@ endif::[]
 
 // Unreleased changes go here
 // When the next release happens, nest these changes under the "Python Agent version 6.x" heading
-//[float]
-//===== Features
-//
+[float]
+===== Features
+
+* Add redis query to context data for redis instrumentation {pull}1406[#1406]
+
 [float]
 ===== Bug fixes
 

elasticapm/instrumentation/packages/asyncio/aioredis.py

@@ -87,8 +87,14 @@ class RedisConnectionInstrumentation(AbstractInstrumentedModule):
 
     def call(self, module, method, wrapped, instance, args, kwargs):
         span = execution_context.get_span()
-        if span and span.subtype == "aioredis":
+        db_context = None
+        if args and len(args) > 1:
+            func_name = args[0].decode("utf-8") if isinstance(args[0], bytes) else args[0]
+            db_context = {"type": "query", "statement": "{} {}".format(func_name, args[1])}
+        if span and span.subtype == self.name:
             span.context["destination"] = _get_destination_info(instance)
+            if db_context:
+                span.context["db"] = db_context
         return wrapped(*args, **kwargs)
 
 

elasticapm/instrumentation/packages/redis.py

@@ -86,8 +86,14 @@ class RedisConnectionInstrumentation(AbstractInstrumentedModule):
 
     def call(self, module, method, wrapped, instance, args, kwargs):
         span = execution_context.get_span()
-        if span and span.subtype == "redis":
+        db_context = None
+        if args and len(args) > 1:
+            func_name = args[0].decode("utf-8") if isinstance(args[0], bytes) else args[0]
+            db_context = {"type": "query", "statement": "{} {}".format(func_name, args[1])}
+        if span and span.subtype == self.name:
             span.context["destination"] = get_destination_info(instance)
+            if db_context:
+                span.context["db"] = db_context
         return wrapped(*args, **kwargs)