caches

Author: sentrivana

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)