Connection attributes in redis database spans (#2398)

This adds db connection parameters like database host, database port, database name, database system ("redis" in this case) to all database spans that are created by our Redis integration. Works for async and sync connections to redis and redis cluster.

Author: antonpirker

sentry_sdk/integrations/redis/__init__.py

@@ -2,32 +2,31 @@
 
 from sentry_sdk import Hub
 from sentry_sdk.consts import OP, SPANDATA
+from sentry_sdk._compat import text_type
 from sentry_sdk.hub import _should_send_default_pii
+from sentry_sdk.integrations import Integration, DidNotEnable
+from sentry_sdk._types import TYPE_CHECKING
 from sentry_sdk.utils import (
     SENSITIVE_DATA_SUBSTITUTE,
     capture_internal_exceptions,
     logger,
 )
-from sentry_sdk.integrations import Integration, DidNotEnable
-
-from sentry_sdk._types import TYPE_CHECKING
 
 if TYPE_CHECKING:
-    from typing import Any, Sequence
+    from typing import Any, Dict, Sequence
     from sentry_sdk.tracing import Span
 
 _SINGLE_KEY_COMMANDS = frozenset(
-    ["decr", "decrby", "get", "incr", "incrby", "pttl", "set", "setex", "setnx", "ttl"]
+    ["decr", "decrby", "get", "incr", "incrby", "pttl", "set", "setex", "setnx", "ttl"],
+)
+_MULTI_KEY_COMMANDS = frozenset(
+    ["del", "touch", "unlink"],
 )
-_MULTI_KEY_COMMANDS = frozenset(["del", "touch", "unlink"])
-
 _COMMANDS_INCLUDING_SENSITIVE_DATA = [
     "auth",
 ]
-
 _MAX_NUM_ARGS = 10  # Trim argument lists to this many values
 _MAX_NUM_COMMANDS = 10  # Trim command lists to this many values
-
 _DEFAULT_MAX_DATA_SIZE = 1024
 
 
@@ -59,6 +58,26 @@ def _get_safe_command(name, args):
     return command
 
 
+def _get_span_description(name, *args):
+    # type: (str, *Any) -> str
+    description = name
+
+    with capture_internal_exceptions():
+        description = _get_safe_command(name, args)
+
+    return description
+
+
+def _get_redis_command_args(command):
+    # type: (Any) -> Sequence[Any]
+    return command[0]
+
+
+def _parse_rediscluster_command(command):
+    # type: (Any) -> Sequence[Any]
+    return command.args
+
+
 def _set_pipeline_data(
     span, is_cluster, get_command_args_fn, is_transaction, command_stack
 ):
@@ -84,6 +103,38 @@ def _set_pipeline_data(
     )
 
 
+def _set_client_data(span, is_cluster, name, *args):
+    # type: (Span, bool, str, *Any) -> None
+    span.set_tag("redis.is_cluster", is_cluster)
+    if name:
+        span.set_tag("redis.command", name)
+        span.set_tag(SPANDATA.DB_OPERATION, name)
+
+    if name and args:
+        name_low = name.lower()
+        if (name_low in _SINGLE_KEY_COMMANDS) or (
+            name_low in _MULTI_KEY_COMMANDS and len(args) == 1
+        ):
+            span.set_tag("redis.key", args[0])
+
+
+def _set_db_data(span, connection_params):
+    # type: (Span, Dict[str, Any]) -> None
+    span.set_data(SPANDATA.DB_SYSTEM, "redis")
+
+    db = connection_params.get("db")
+    if db is not None:
+        span.set_data(SPANDATA.DB_NAME, text_type(db))
+
+    host = connection_params.get("host")
+    if host is not None:
+        span.set_data(SPANDATA.SERVER_ADDRESS, host)
+
+    port = connection_params.get("port")
+    if port is not None:
+        span.set_data(SPANDATA.SERVER_PORT, port)
+
+
 def patch_redis_pipeline(pipeline_cls, is_cluster, get_command_args_fn):
     # type: (Any, bool, Any) -> None
     old_execute = pipeline_cls.execute
@@ -99,28 +150,51 @@ def sentry_patched_execute(self, *args, **kwargs):
             op=OP.DB_REDIS, description="redis.pipeline.execute"
         ) as span:
             with capture_internal_exceptions():
+                _set_db_data(span, self.connection_pool.connection_kwargs)
                 _set_pipeline_data(
                     span,
                     is_cluster,
                     get_command_args_fn,
                     self.transaction,
                     self.command_stack,
                 )
-                span.set_data(SPANDATA.DB_SYSTEM, "redis")
 
             return old_execute(self, *args, **kwargs)
 
     pipeline_cls.execute = sentry_patched_execute
 
 
-def _get_redis_command_args(command):
-    # type: (Any) -> Sequence[Any]
-    return command[0]
+def patch_redis_client(cls, is_cluster):
+    # type: (Any, bool) -> None
+    """
+    This function can be used to instrument custom redis client classes or
+    subclasses.
+    """
+    old_execute_command = cls.execute_command
 
+    def sentry_patched_execute_command(self, name, *args, **kwargs):
+        # type: (Any, str, *Any, **Any) -> Any
+        hub = Hub.current
+        integration = hub.get_integration(RedisIntegration)
 
-def _parse_rediscluster_command(command):
-    # type: (Any) -> Sequence[Any]
-    return command.args
+        if integration is None:
+            return old_execute_command(self, name, *args, **kwargs)
+
+        description = _get_span_description(name, *args)
+
+        data_should_be_truncated = (
+            integration.max_data_size and len(description) > integration.max_data_size
+        )
+        if data_should_be_truncated:
+            description = description[: integration.max_data_size - len("...")] + "..."
+
+        with hub.start_span(op=OP.DB_REDIS, description=description) as span:
+            _set_db_data(span, self.connection_pool.connection_kwargs)
+            _set_client_data(span, is_cluster, name, *args)
+
+            return old_execute_command(self, name, *args, **kwargs)
+
+    cls.execute_command = sentry_patched_execute_command
 
 
 def _patch_redis(StrictRedis, client):  # noqa: N803
@@ -206,61 +280,3 @@ def setup_once():
             _patch_rediscluster()
         except Exception:
             logger.exception("Error occurred while patching `rediscluster` library")
-
-
-def _get_span_description(name, *args):
-    # type: (str, *Any) -> str
-    description = name
-
-    with capture_internal_exceptions():
-        description = _get_safe_command(name, args)
-
-    return description
-
-
-def _set_client_data(span, is_cluster, name, *args):
-    # type: (Span, bool, str, *Any) -> None
-    span.set_data(SPANDATA.DB_SYSTEM, "redis")
-    span.set_tag("redis.is_cluster", is_cluster)
-    if name:
-        span.set_tag("redis.command", name)
-        span.set_tag(SPANDATA.DB_OPERATION, name)
-
-    if name and args:
-        name_low = name.lower()
-        if (name_low in _SINGLE_KEY_COMMANDS) or (
-            name_low in _MULTI_KEY_COMMANDS and len(args) == 1
-        ):
-            span.set_tag("redis.key", args[0])
-
-
-def patch_redis_client(cls, is_cluster):
-    # type: (Any, bool) -> None
-    """
-    This function can be used to instrument custom redis client classes or
-    subclasses.
-    """
-    old_execute_command = cls.execute_command
-
-    def sentry_patched_execute_command(self, name, *args, **kwargs):
-        # type: (Any, str, *Any, **Any) -> Any
-        hub = Hub.current
-        integration = hub.get_integration(RedisIntegration)
-
-        if integration is None:
-            return old_execute_command(self, name, *args, **kwargs)
-
-        description = _get_span_description(name, *args)
-
-        data_should_be_truncated = (
-            integration.max_data_size and len(description) > integration.max_data_size
-        )
-        if data_should_be_truncated:
-            description = description[: integration.max_data_size - len("...")] + "..."
-
-        with hub.start_span(op=OP.DB_REDIS, description=description) as span:
-            _set_client_data(span, is_cluster, name, *args)
-
-            return old_execute_command(self, name, *args, **kwargs)
-
-    cls.execute_command = sentry_patched_execute_command

sentry_sdk/integrations/redis/asyncio.py

@@ -2,19 +2,18 @@
 
 from sentry_sdk import Hub
 from sentry_sdk.consts import OP
-from sentry_sdk.utils import capture_internal_exceptions
 from sentry_sdk.integrations.redis import (
     RedisIntegration,
     _get_redis_command_args,
     _get_span_description,
     _set_client_data,
+    _set_db_data,
     _set_pipeline_data,
 )
+from sentry_sdk._types import TYPE_CHECKING
+from sentry_sdk.utils import capture_internal_exceptions
 
-
-from sentry_sdk._types import MYPY
-
-if MYPY:
+if TYPE_CHECKING:
     from typing import Any
 
 
@@ -33,6 +32,7 @@ async def _sentry_execute(self, *args, **kwargs):
             op=OP.DB_REDIS, description="redis.pipeline.execute"
         ) as span:
             with capture_internal_exceptions():
+                _set_db_data(span, self.connection_pool.connection_kwargs)
                 _set_pipeline_data(
                     span,
                     False,
@@ -60,6 +60,7 @@ async def _sentry_execute_command(self, name, *args, **kwargs):
         description = _get_span_description(name, *args)
 
         with hub.start_span(op=OP.DB_REDIS, description=description) as span:
+            _set_db_data(span, self.connection_pool.connection_kwargs)
             _set_client_data(span, False, name, *args)
 
             return await old_execute_command(self, name, *args, **kwargs)

tests/integrations/redis/asyncio/test_redis_asyncio.py

@@ -1,6 +1,7 @@
 import pytest
 
 from sentry_sdk import capture_message, start_transaction
+from sentry_sdk.consts import SPANDATA
 from sentry_sdk.integrations.redis import RedisIntegration
 
 from fakeredis.aioredis import FakeRedis
@@ -67,7 +68,13 @@ async def test_async_redis_pipeline(
         "redis.commands": {
             "count": 3,
             "first_ten": expected_first_ten,
-        }
+        },
+        SPANDATA.DB_SYSTEM: "redis",
+        SPANDATA.DB_NAME: "0",
+        SPANDATA.SERVER_ADDRESS: connection.connection_pool.connection_kwargs.get(
+            "host"
+        ),
+        SPANDATA.SERVER_PORT: 6379,
     }
     assert span["tags"] == {
         "redis.transaction": is_transaction,

tests/integrations/redis/test_redis.py

@@ -12,6 +12,14 @@
     import mock  # python < 3.3
 
 
+MOCK_CONNECTION_POOL = mock.MagicMock()
+MOCK_CONNECTION_POOL.connection_kwargs = {
+    "host": "localhost",
+    "port": 63791,
+    "db": 1,
+}
+
+
 def test_basic(sentry_init, capture_events):
     sentry_init(integrations=[RedisIntegration()])
     events = capture_events()
@@ -67,12 +75,10 @@ def test_redis_pipeline(
     (span,) = event["spans"]
     assert span["op"] == "db.redis"
     assert span["description"] == "redis.pipeline.execute"
-    assert span["data"] == {
-        "redis.commands": {
-            "count": 3,
-            "first_ten": expected_first_ten,
-        },
-        SPANDATA.DB_SYSTEM: "redis",
+    assert span["data"][SPANDATA.DB_SYSTEM] == "redis"
+    assert span["data"]["redis.commands"] == {
+        "count": 3,
+        "first_ten": expected_first_ten,
     }
     assert span["tags"] == {
         "redis.transaction": is_transaction,
@@ -242,3 +248,51 @@ def test_breadcrumbs(sentry_init, capture_events):
         },
         "timestamp": crumbs[1]["timestamp"],
     }
+
+
+def test_db_connection_attributes_client(sentry_init, capture_events):
+    sentry_init(
+        traces_sample_rate=1.0,
+        integrations=[RedisIntegration()],
+    )
+    events = capture_events()
+
+    with start_transaction():
+        connection = FakeStrictRedis(connection_pool=MOCK_CONNECTION_POOL)
+        connection.get("foobar")
+
+    (event,) = events
+    (span,) = event["spans"]
+
+    assert span["op"] == "db.redis"
+    assert span["description"] == "GET 'foobar'"
+    assert span["data"][SPANDATA.DB_SYSTEM] == "redis"
+    assert span["data"][SPANDATA.DB_NAME] == "1"
+    assert span["data"][SPANDATA.SERVER_ADDRESS] == "localhost"
+    assert span["data"][SPANDATA.SERVER_PORT] == 63791
+
+
+def test_db_connection_attributes_pipeline(sentry_init, capture_events):
+    sentry_init(
+        traces_sample_rate=1.0,
+        integrations=[RedisIntegration()],
+    )
+    events = capture_events()
+
+    with start_transaction():
+        connection = FakeStrictRedis(connection_pool=MOCK_CONNECTION_POOL)
+        pipeline = connection.pipeline(transaction=False)
+        pipeline.get("foo")
+        pipeline.set("bar", 1)
+        pipeline.set("baz", 2)
+        pipeline.execute()
+
+    (event,) = events
+    (span,) = event["spans"]
+
+    assert span["op"] == "db.redis"
+    assert span["description"] == "redis.pipeline.execute"
+    assert span["data"][SPANDATA.DB_SYSTEM] == "redis"
+    assert span["data"][SPANDATA.DB_NAME] == "1"
+    assert span["data"][SPANDATA.SERVER_ADDRESS] == "localhost"
+    assert span["data"][SPANDATA.SERVER_PORT] == 63791