async common

Author: sentrivana

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,36 @@ 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",
+                attributes={
+                    "sentry.origin": SPAN_ORIGIN,
+                    "sentry.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 +84,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 +93,55 @@ 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"],
+                    attributes={
+                        "sentry.op": cache_properties["op"],
+                        "sentry.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"],
+                attributes={
+                    "sentry.op": db_properties["op"],
+                    "sentry.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)