After gathering all the data, set it on the span

Author: antonpirker

sentry_sdk/integrations/redis/_async_common.py

@@ -9,6 +9,7 @@
 from sentry_sdk.integrations.redis.utils import (
     _get_client_data,
     _get_pipeline_data,
+    _update_span,
 )
 from sentry_sdk.tracing import Span
 from sentry_sdk.utils import capture_internal_exceptions
@@ -48,6 +49,7 @@ async def _sentry_execute(self, *args, **kwargs):
                     is_transaction=False if is_cluster else self.is_transaction,
                     command_stack=self._command_stack if is_cluster else self.command_stack,
                 )
+                _update_span(span, span_data, pipeline_data)
 
             return await old_execute(self, *args, **kwargs)
 
@@ -92,14 +94,16 @@ async def _sentry_execute_command(self, name, *args, **kwargs):
         db_span.__enter__()
 
         db_span_data = get_db_data_fn(self)
-        db_client_span_data = _get_client_data(db_span, is_cluster, name, *args)
+        db_client_span_data = _get_client_data(is_cluster, name, *args)
+        _update_span(db_span, db_span_data, db_client_span_data)
 
         value = await old_execute_command(self, name, *args, **kwargs)
 
         db_span.__exit__(None, None, None)
 
         if cache_span:
             cache_span_data = _get_cache_data(self, cache_properties, value)
+            _update_span(cache_span, cache_span_data)
             cache_span.__exit__(None, None, None)
 
         return value

sentry_sdk/integrations/redis/_sync_common.py

@@ -9,6 +9,7 @@
 from sentry_sdk.integrations.redis.utils import (
     _get_client_data,
     _get_pipeline_data,
+    _update_span,
 )
 from sentry_sdk.tracing import Span
 from sentry_sdk.utils import capture_internal_exceptions
@@ -50,6 +51,7 @@ def sentry_patched_execute(self, *args, **kwargs):
                     is_transaction=False if is_cluster else self.transaction,
                     command_stack=self.command_stack,
                 )
+                _update_span(span, span_data, pipeline_data)
 
             return old_execute(self, *args, **kwargs)
 
@@ -99,13 +101,15 @@ def sentry_patched_execute_command(self, name, *args, **kwargs):
 
         db_span_data = get_db_data_fn(self)
         db_client_span_data = _get_client_data(db_span, is_cluster, name, *args)
+        _update_span(db_span, db_span_data, db_client_span_data)
 
         value = old_execute_command(self, name, *args, **kwargs)
 
         db_span.__exit__(None, None, None)
 
         if cache_span:
             cache_span_data = _get_cache_data(self, cache_properties, value)
+            _update_span(cache_span, cache_span_data)
             cache_span.__exit__(None, None, None)
 
         return value

sentry_sdk/integrations/redis/utils.py

@@ -16,6 +16,25 @@
     from sentry_sdk.tracing import Span
 
 
+TAG_KEYS = [
+    "redis.commands",
+    "redis.is_cluster",
+    "redis.key",
+    "redis.transaction",
+    SPANDATA.DB_OPERATION,
+]
+
+
+def _update_span(span, *data_bags):
+    # type: (Span, *dict[str, Any]) -> None
+    for data in data_bags:
+        for key, value in data.items():
+            if key in TAG_KEYS:
+                span.set_tag(key, value)
+            else:
+                span.set_data(key, value)
+
+
 def _get_safe_command(name, args):
     # type: (str, Sequence[Any]) -> str
     command_parts = [name]