Add Redis pub/sub instrumentation

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()

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):

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

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