Repository: elastic/apm-agent-python

Issue:
Redis: Support publish and subscribe methods
While running the examples from the blog post [How to instrument a polyglot microservices application with Elastic APM](https://www.elastic.co/blog/how-to-instrument-a-polyglot-microservices-application-with-elastic-apm), I noticed Redis doesn't show up on the service map as being connected to the Python service in this example.

It looks like that's because [according to our documentation we don't have the `publish` and `subscribe` methods instrumented](https://www.elastic.co/guide/en/apm/agent/python/5.x/supported-technologies.html#automatic-instrumentation-db-redis).

If these methods were instrumented we would be able to see Redis on service maps for applications that are using it for pub/sub.


Discussion:
Thanks for the request!

Resolving patch:
diff --git a/elasticapm/instrumentation/packages/asyncio/aioredis.py b/elasticapm/instrumentation/packages/asyncio/aioredis.py
--- a/elasticapm/instrumentation/packages/asyncio/aioredis.py
+++ b/elasticapm/instrumentation/packages/asyncio/aioredis.py
@@ -38,7 +38,8 @@
 class RedisConnectionPoolInstrumentation(AbstractInstrumentedModule):
     name = "aioredis"
 
-    instrument_list = [("aioredis.pool", "ConnectionsPool.execute")]
+    instrument_list = [("aioredis.pool", "ConnectionsPool.execute"),
+                       ("aioredis.pool", "ConnectionsPool.execute_pubsub")]
 
     def call(self, module, method, wrapped, instance, args, kwargs):
         if len(args) > 0:
@@ -73,7 +74,8 @@ def call(self, module, method, wrapped, instance, args, kwargs):
 class RedisConnectionInstrumentation(AbstractInstrumentedModule):
     name = "aioredis"
 
-    instrument_list = (("aioredis.connection", "RedisConnection.execute"),)
+    instrument_list = (("aioredis.connection", "RedisConnection.execute"),
+                       ("aioredis.pool", "ConnectionsPool.execute_pubsub"))
 
     def call(self, module, method, wrapped, instance, args, kwargs):
         span = execution_context.get_span()
diff --git a/elasticapm/instrumentation/packages/redis.py b/elasticapm/instrumentation/packages/redis.py
--- a/elasticapm/instrumentation/packages/redis.py
+++ b/elasticapm/instrumentation/packages/redis.py
@@ -53,7 +53,7 @@ class RedisInstrumentation(Redis3CheckMixin, AbstractInstrumentedModule):
     name = "redis"
 
     # no need to instrument StrictRedis in redis-py >= 3.0
-    instrument_list_3 = [("redis.client", "Redis.execute_command")]
+    instrument_list_3 = [("redis.client", "Redis.execute_command"), ("redis.client", "PubSub.execute_command")]
     instrument_list = [("redis.client", "Redis.execute_command"), ("redis.client", "StrictRedis.execute_command")]
 
     def call(self, module, method, wrapped, instance, args, kwargs):


Test patch:
diff --git a/tests/instrumentation/asyncio_tests/aioredis_tests.py b/tests/instrumentation/asyncio_tests/aioredis_tests.py
--- a/tests/instrumentation/asyncio_tests/aioredis_tests.py
+++ b/tests/instrumentation/asyncio_tests/aioredis_tests.py
@@ -123,3 +123,48 @@ async def test_redis_client(instrument, elasticapm_client, redis_conn):
     assert spans[2]["type"] == "test"
 
     assert len(spans) == 3
+
+
+@pytest.mark.integrationtest
+async def test_publish_subscribe(instrument, elasticapm_client, redis_conn):
+    elasticapm_client.begin_transaction("transaction.test")
+    with capture_span("test_publish_subscribe", "test"):
+        # publish
+        await redis_conn.publish("mykey", "a")
+
+        #subscribe
+        await redis_conn.subscribe("mykey")
+
+    elasticapm_client.end_transaction("MyView")
+
+    transactions = elasticapm_client.events[TRANSACTION]
+    spans = elasticapm_client.spans_for_transaction(transactions[0])
+
+    expected_signatures = {"test_publish_subscribe", "PUBLISH", "SUBSCRIBE"}
+
+    assert {t["name"] for t in spans} == expected_signatures
+
+    assert spans[0]["name"] == "PUBLISH"
+    assert spans[0]["type"] == "db"
+    assert spans[0]["subtype"] == "redis"
+    assert spans[0]["action"] == "query"
+    assert spans[0]["context"]["destination"] == {
+        "address": os.environ.get("REDIS_HOST", "localhost"),
+        "port": int(os.environ.get("REDIS_PORT", 6379)),
+        "service": {"name": "aioredis", "resource": "redis", "type": "db"},
+    }
+
+    assert spans[1]["name"] == "SUBSCRIBE"
+    assert spans[1]["type"] == "db"
+    assert spans[1]["subtype"] == "redis"
+    assert spans[1]["action"] == "query"
+    assert spans[1]["context"]["destination"] == {
+        "address": os.environ.get("REDIS_HOST", "localhost"),
+        "port": int(os.environ.get("REDIS_PORT", 6379)),
+        "service": {"name": "aioredis", "resource": "redis", "type": "db"},
+    }
+
+    assert spans[2]["name"] == "test_publish_subscribe"
+    assert spans[2]["type"] == "test"
+
+    assert len(spans) == 3
diff --git a/tests/instrumentation/redis_tests.py b/tests/instrumentation/redis_tests.py
--- a/tests/instrumentation/redis_tests.py
+++ b/tests/instrumentation/redis_tests.py
@@ -163,3 +163,49 @@ def test_unix_domain_socket_connection_destination_info():
     destination_info = get_destination_info(conn)
     assert destination_info["port"] is None
     assert destination_info["address"] == "unix:///some/path"
+
+
+@pytest.mark.integrationtest
+def test_publish_subscribe(instrument, elasticapm_client, redis_conn):
+    elasticapm_client.begin_transaction("transaction.test")
+    with capture_span("test_publish_subscribe", "test"):
+        # publish
+        redis_conn.publish("mykey", "a")
+
+        # subscribe
+        with redis_conn.pubsub() as channel:
+            channel.subscribe("mykey")
+
+    elasticapm_client.end_transaction("MyView")
+
+    transactions = elasticapm_client.events[TRANSACTION]
+    spans = elasticapm_client.spans_for_transaction(transactions[0])
+
+    expected_signatures = {"test_publish_subscribe", "PUBLISH", "SUBSCRIBE"}
+
+    assert {t["name"] for t in spans} == expected_signatures
+
+    assert spans[0]["name"] == "PUBLISH"
+    assert spans[0]["type"] == "db"
+    assert spans[0]["subtype"] == "redis"
+    assert spans[0]["action"] == "query"
+    assert spans[0]["context"]["destination"] == {
+        "address": os.environ.get("REDIS_HOST", "localhost"),
+        "port": int(os.environ.get("REDIS_PORT", 6379)),
+        "service": {"name": "redis", "resource": "redis", "type": "db"},
+    }
+
+    assert spans[1]["name"] == "SUBSCRIBE"
+    assert spans[1]["type"] == "db"
+    assert spans[1]["subtype"] == "redis"
+    assert spans[1]["action"] == "query"
+    assert spans[1]["context"]["destination"] == {
+        "address": os.environ.get("REDIS_HOST", "localhost"),
+        "port": int(os.environ.get("REDIS_PORT", 6379)),
+        "service": {"name": "redis", "resource": "redis", "type": "db"},
+    }
+
+    assert spans[2]["name"] == "test_publish_subscribe"
+    assert spans[2]["type"] == "test"
+
+    assert len(spans) == 3