feat: Migrate redis integration to span first

sentrivana · sentrivana · commit 021f3f7e7ccd · 2026-02-25T11:02:49.000+01:00
diff --git a/sentry_sdk/integrations/redis/_async_common.py b/sentry_sdk/integrations/redis/_async_common.py
@@ -11,13 +11,15 @@
     _set_pipeline_data,
 )
 from sentry_sdk.tracing import Span
+from sentry_sdk.tracing_utils import has_span_streaming_enabled
 from sentry_sdk.utils import capture_internal_exceptions
 
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
     from collections.abc import Callable
-    from typing import Any, Union
+    from typing import Any, Optional, Union
+    from sentry_sdk.traces import StreamedSpan
     from redis.asyncio.client import Pipeline, StrictRedis
     from redis.asyncio.cluster import ClusterPipeline, RedisCluster
 
@@ -26,21 +28,32 @@ def patch_redis_async_pipeline(
     pipeline_cls: "Union[type[Pipeline[Any]], type[ClusterPipeline[Any]]]",
     is_cluster: bool,
     get_command_args_fn: "Any",
-    set_db_data_fn: "Callable[[Span, Any], None]",
+    set_db_data_fn: "Callable[[Union[Span, StreamedSpan], Any], None]",
 ) -> None:
     old_execute = pipeline_cls.execute
 
     from sentry_sdk.integrations.redis import RedisIntegration
 
     async def _sentry_execute(self: "Any", *args: "Any", **kwargs: "Any") -> "Any":
-        if sentry_sdk.get_client().get_integration(RedisIntegration) is None:
+        client = sentry_sdk.get_client()
+        if client.get_integration(RedisIntegration) is None:
             return await old_execute(self, *args, **kwargs)
 
-        with sentry_sdk.start_span(
-            op=OP.DB_REDIS,
-            name="redis.pipeline.execute",
-            origin=SPAN_ORIGIN,
-        ) as span:
+        span_streaming = has_span_streaming_enabled(client.options)
+
+        span: "Union[Span, StreamedSpan]"
+        if span_streaming:
+            span = sentry_sdk.traces.start_span(name="redis.pipeline.execute")
+            span.set_origin(SPAN_ORIGIN)
+            span.set_op(OP.DB_REDIS)
+        else:
+            span = sentry_sdk.start_span(
+                op=OP.DB_REDIS,
+                name="redis.pipeline.execute",
+                origin=SPAN_ORIGIN,
+            )
+
+        with span:
             with capture_internal_exceptions():
                 try:
                     command_seq = self._execution_strategy._command_queue
@@ -67,7 +80,7 @@ async def _sentry_execute(self: "Any", *args: "Any", **kwargs: "Any") -> "Any":
 def patch_redis_async_client(
     cls: "Union[type[StrictRedis[Any]], type[RedisCluster[Any]]]",
     is_cluster: bool,
-    set_db_data_fn: "Callable[[Span, Any], None]",
+    set_db_data_fn: "Callable[[Union[Span, StreamedSpan], Any], None]",
 ) -> None:
     old_execute_command = cls.execute_command
 
@@ -76,33 +89,49 @@ def patch_redis_async_client(
     async def _sentry_execute_command(
         self: "Any", name: str, *args: "Any", **kwargs: "Any"
     ) -> "Any":
-        integration = sentry_sdk.get_client().get_integration(RedisIntegration)
+        client = sentry_sdk.get_client()
+        integration = client.get_integration(RedisIntegration)
         if integration is None:
             return await old_execute_command(self, name, *args, **kwargs)
 
+        span_streaming = has_span_streaming_enabled(client.options)
+
         cache_properties = _compile_cache_span_properties(
             name,
             args,
             kwargs,
             integration,
         )
 
-        cache_span = None
+        cache_span: "Optional[Union[Span, StreamedSpan]]" = None
         if cache_properties["is_cache_key"] and cache_properties["op"] is not None:
-            cache_span = sentry_sdk.start_span(
-                op=cache_properties["op"],
-                name=cache_properties["description"],
-                origin=SPAN_ORIGIN,
-            )
+            if span_streaming:
+                cache_span = sentry_sdk.traces.start_span(
+                    name=cache_properties["description"]
+                )
+                cache_span.set_op(cache_properties["op"])
+                cache_span.set_origin(SPAN_ORIGIN)
+            else:
+                cache_span = sentry_sdk.start_span(
+                    op=cache_properties["op"],
+                    name=cache_properties["description"],
+                    origin=SPAN_ORIGIN,
+                )
             cache_span.__enter__()
 
         db_properties = _compile_db_span_properties(integration, name, args)
 
-        db_span = sentry_sdk.start_span(
-            op=db_properties["op"],
-            name=db_properties["description"],
-            origin=SPAN_ORIGIN,
-        )
+        db_span: "Union[Span, StreamedSpan]"
+        if span_streaming:
+            db_span = sentry_sdk.traces.start_span(name=db_properties["description"])
+            db_span.set_op(db_properties["op"])
+            db_span.set_origin(SPAN_ORIGIN)
+        else:
+            db_span = sentry_sdk.start_span(
+                op=db_properties["op"],
+                name=db_properties["description"],
+                origin=SPAN_ORIGIN,
+            )
         db_span.__enter__()
 
         set_db_data_fn(db_span, self)
diff --git a/sentry_sdk/integrations/redis/_sync_common.py b/sentry_sdk/integrations/redis/_sync_common.py
@@ -11,34 +11,47 @@
     _set_pipeline_data,
 )
 from sentry_sdk.tracing import Span
+from sentry_sdk.tracing_utils import has_span_streaming_enabled
 from sentry_sdk.utils import capture_internal_exceptions
 
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
     from collections.abc import Callable
-    from typing import Any
+    from typing import Any, Optional, Union
+    from sentry_sdk.traces import StreamedSpan
 
 
 def patch_redis_pipeline(
     pipeline_cls: "Any",
     is_cluster: bool,
     get_command_args_fn: "Any",
-    set_db_data_fn: "Callable[[Span, Any], None]",
+    set_db_data_fn: "Callable[[Union[Span, StreamedSpan], Any], None]",
 ) -> None:
     old_execute = pipeline_cls.execute
 
     from sentry_sdk.integrations.redis import RedisIntegration
 
     def sentry_patched_execute(self: "Any", *args: "Any", **kwargs: "Any") -> "Any":
-        if sentry_sdk.get_client().get_integration(RedisIntegration) is None:
+        client = sentry_sdk.get_client()
+        if client.get_integration(RedisIntegration) is None:
             return old_execute(self, *args, **kwargs)
 
-        with sentry_sdk.start_span(
-            op=OP.DB_REDIS,
-            name="redis.pipeline.execute",
-            origin=SPAN_ORIGIN,
-        ) as span:
+        span_streaming = has_span_streaming_enabled(client.options)
+
+        span: "Union[Span, StreamedSpan]"
+        if span_streaming:
+            span = sentry_sdk.traces.start_span(name="redis.pipeline.execute")
+            span.set_origin(SPAN_ORIGIN)
+            span.set_op(OP.DB_REDIS)
+        else:
+            span = sentry_sdk.start_span(
+                op=OP.DB_REDIS,
+                name="redis.pipeline.execute",
+                origin=SPAN_ORIGIN,
+            )
+
+        with span:
             with capture_internal_exceptions():
                 command_seq = None
                 try:
@@ -61,7 +74,9 @@ def sentry_patched_execute(self: "Any", *args: "Any", **kwargs: "Any") -> "Any":
 
 
 def patch_redis_client(
-    cls: "Any", is_cluster: bool, set_db_data_fn: "Callable[[Span, Any], None]"
+    cls: "Any",
+    is_cluster: bool,
+    set_db_data_fn: "Callable[[Union[Span, StreamedSpan], Any], None]",
 ) -> None:
     """
     This function can be used to instrument custom redis client classes or
@@ -74,33 +89,49 @@ def patch_redis_client(
     def sentry_patched_execute_command(
         self: "Any", name: str, *args: "Any", **kwargs: "Any"
     ) -> "Any":
-        integration = sentry_sdk.get_client().get_integration(RedisIntegration)
+        client = sentry_sdk.get_client()
+        integration = client.get_integration(RedisIntegration)
         if integration is None:
             return old_execute_command(self, name, *args, **kwargs)
 
+        span_streaming = has_span_streaming_enabled(client.options)
+
         cache_properties = _compile_cache_span_properties(
             name,
             args,
             kwargs,
             integration,
         )
 
-        cache_span = None
+        cache_span: "Optional[Union[Span, StreamedSpan]]" = None
         if cache_properties["is_cache_key"] and cache_properties["op"] is not None:
-            cache_span = sentry_sdk.start_span(
-                op=cache_properties["op"],
-                name=cache_properties["description"],
-                origin=SPAN_ORIGIN,
-            )
+            if span_streaming:
+                cache_span = sentry_sdk.traces.start_span(
+                    name=cache_properties["description"]
+                )
+                cache_span.set_op(cache_properties["op"])
+                cache_span.set_origin(SPAN_ORIGIN)
+            else:
+                cache_span = sentry_sdk.start_span(
+                    op=cache_properties["op"],
+                    name=cache_properties["description"],
+                    origin=SPAN_ORIGIN,
+                )
             cache_span.__enter__()
 
         db_properties = _compile_db_span_properties(integration, name, args)
 
-        db_span = sentry_sdk.start_span(
-            op=db_properties["op"],
-            name=db_properties["description"],
-            origin=SPAN_ORIGIN,
-        )
+        db_span: "Union[Span, StreamedSpan]"
+        if span_streaming:
+            db_span = sentry_sdk.traces.start_span(name=db_properties["description"])
+            db_span.set_op(db_properties["op"])
+            db_span.set_origin(SPAN_ORIGIN)
+        else:
+            db_span = sentry_sdk.start_span(
+                op=db_properties["op"],
+                name=db_properties["description"],
+                origin=SPAN_ORIGIN,
+            )
         db_span.__enter__()
 
         set_db_data_fn(db_span, self)
diff --git a/sentry_sdk/integrations/redis/modules/caches.py b/sentry_sdk/integrations/redis/modules/caches.py
@@ -4,6 +4,7 @@
 
 from sentry_sdk.consts import OP, SPANDATA
 from sentry_sdk.integrations.redis.utils import _get_safe_key, _key_as_string
+from sentry_sdk.traces import StreamedSpan
 from sentry_sdk.utils import capture_internal_exceptions
 
 GET_COMMANDS = ("get", "mget")
@@ -14,7 +15,7 @@
 if TYPE_CHECKING:
     from sentry_sdk.integrations.redis import RedisIntegration
     from sentry_sdk.tracing import Span
-    from typing import Any, Optional
+    from typing import Any, Optional, Union
 
 
 def _get_op(name: str) -> "Optional[str]":
@@ -80,25 +81,30 @@ def _get_cache_span_description(
 
 
 def _set_cache_data(
-    span: "Span",
+    span: "Union[Span, StreamedSpan]",
     redis_client: "Any",
     properties: "dict[str, Any]",
     return_value: "Optional[Any]",
 ) -> None:
+    if isinstance(span, StreamedSpan):
+        set_on_span = span.set_attribute
+    else:
+        set_on_span = span.set_data
+
     with capture_internal_exceptions():
-        span.set_data(SPANDATA.CACHE_KEY, properties["key"])
+        set_on_span(SPANDATA.CACHE_KEY, properties["key"])
 
         if properties["redis_command"] in GET_COMMANDS:
             if return_value is not None:
-                span.set_data(SPANDATA.CACHE_HIT, True)
+                set_on_span(SPANDATA.CACHE_HIT, True)
                 size = (
                     len(str(return_value).encode("utf-8"))
                     if not isinstance(return_value, bytes)
                     else len(return_value)
                 )
-                span.set_data(SPANDATA.CACHE_ITEM_SIZE, size)
+                set_on_span(SPANDATA.CACHE_ITEM_SIZE, size)
             else:
-                span.set_data(SPANDATA.CACHE_HIT, False)
+                set_on_span(SPANDATA.CACHE_HIT, False)
 
         elif properties["redis_command"] in SET_COMMANDS:
             if properties["value"] is not None:
@@ -107,7 +113,7 @@ def _set_cache_data(
                     if not isinstance(properties["value"], bytes)
                     else len(properties["value"])
                 )
-                span.set_data(SPANDATA.CACHE_ITEM_SIZE, size)
+                set_on_span(SPANDATA.CACHE_ITEM_SIZE, size)
 
         try:
             connection_params = redis_client.connection_pool.connection_kwargs
@@ -122,8 +128,8 @@ def _set_cache_data(
 
         host = connection_params.get("host")
         if host is not None:
-            span.set_data(SPANDATA.NETWORK_PEER_ADDRESS, host)
+            set_on_span(SPANDATA.NETWORK_PEER_ADDRESS, host)
 
         port = connection_params.get("port")
         if port is not None:
-            span.set_data(SPANDATA.NETWORK_PEER_PORT, port)
+            set_on_span(SPANDATA.NETWORK_PEER_PORT, port)
diff --git a/sentry_sdk/integrations/redis/modules/queries.py b/sentry_sdk/integrations/redis/modules/queries.py
@@ -4,6 +4,7 @@
 
 from sentry_sdk.consts import OP, SPANDATA
 from sentry_sdk.integrations.redis.utils import _get_safe_command
+from sentry_sdk.traces import StreamedSpan
 from sentry_sdk.utils import capture_internal_exceptions
 
 from typing import TYPE_CHECKING
@@ -12,7 +13,7 @@
     from redis import Redis
     from sentry_sdk.integrations.redis import RedisIntegration
     from sentry_sdk.tracing import Span
-    from typing import Any
+    from typing import Any, Union
 
 
 def _compile_db_span_properties(
@@ -42,23 +43,32 @@ def _get_db_span_description(
     return description
 
 
-def _set_db_data_on_span(span: "Span", connection_params: "dict[str, Any]") -> None:
-    span.set_data(SPANDATA.DB_SYSTEM, "redis")
+def _set_db_data_on_span(
+    span: "Union[Span, StreamedSpan]", connection_params: "dict[str, Any]"
+) -> None:
+    if isinstance(span, StreamedSpan):
+        set_on_span = span.set_attribute
+    else:
+        set_on_span = span.set_data
+
+    set_on_span(SPANDATA.DB_SYSTEM, "redis")
 
     db = connection_params.get("db")
     if db is not None:
-        span.set_data(SPANDATA.DB_NAME, str(db))
+        set_on_span(SPANDATA.DB_NAME, str(db))
 
     host = connection_params.get("host")
     if host is not None:
-        span.set_data(SPANDATA.SERVER_ADDRESS, host)
+        set_on_span(SPANDATA.SERVER_ADDRESS, host)
 
     port = connection_params.get("port")
     if port is not None:
-        span.set_data(SPANDATA.SERVER_PORT, port)
+        set_on_span(SPANDATA.SERVER_PORT, port)
 
 
-def _set_db_data(span: "Span", redis_instance: "Redis[Any]") -> None:
+def _set_db_data(
+    span: "Union[Span, StreamedSpan]", redis_instance: "Redis[Any]"
+) -> None:
     try:
         _set_db_data_on_span(span, redis_instance.connection_pool.connection_kwargs)
     except AttributeError:
diff --git a/sentry_sdk/integrations/redis/redis_cluster.py b/sentry_sdk/integrations/redis/redis_cluster.py
diff --git a/sentry_sdk/integrations/redis/utils.py b/sentry_sdk/integrations/redis/utils.py
