two version of R2R are here HEAD master

1 files changed, 108 insertions, 0 deletions

diff --git a/.venv/lib/python3.12/site-packages/sentry_sdk/integrations/redis/_async_common.py b/.venv/lib/python3.12/site-packages/sentry_sdk/integrations/redis/_async_common.py
new file mode 100644
index 00000000..196e85e7
--- /dev/null
+++ b/.venv/lib/python3.12/site-packages/sentry_sdk/integrations/redis/_async_common.py
@@ -0,0 +1,108 @@
+import sentry_sdk
+from sentry_sdk.consts import OP
+from sentry_sdk.integrations.redis.consts import SPAN_ORIGIN
+from sentry_sdk.integrations.redis.modules.caches import (
+    _compile_cache_span_properties,
+    _set_cache_data,
+)
+from sentry_sdk.integrations.redis.modules.queries import _compile_db_span_properties
+from sentry_sdk.integrations.redis.utils import (
+    _set_client_data,
+    _set_pipeline_data,
+)
+from sentry_sdk.tracing import Span
+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 redis.asyncio.client import Pipeline, StrictRedis
+    from redis.asyncio.cluster import ClusterPipeline, RedisCluster
+
+
+def patch_redis_async_pipeline(
+    pipeline_cls, is_cluster, get_command_args_fn, set_db_data_fn
+):
+    # type: (Union[type[Pipeline[Any]], type[ClusterPipeline[Any]]], bool, Any, Callable[[Span, Any], None]) -> None
+    old_execute = pipeline_cls.execute
+
+    from sentry_sdk.integrations.redis import RedisIntegration
+
+    async def _sentry_execute(self, *args, **kwargs):
+        # type: (Any, *Any, **Any) -> Any
+        if sentry_sdk.get_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:
+            with capture_internal_exceptions():
+                set_db_data_fn(span, self)
+                _set_pipeline_data(
+                    span,
+                    is_cluster,
+                    get_command_args_fn,
+                    False if is_cluster else self.is_transaction,
+                    self._command_stack if is_cluster else self.command_stack,
+                )
+
+            return await old_execute(self, *args, **kwargs)
+
+    pipeline_cls.execute = _sentry_execute  # type: ignore
+
+
+def patch_redis_async_client(cls, is_cluster, set_db_data_fn):
+    # type: (Union[type[StrictRedis[Any]], type[RedisCluster[Any]]], bool, Callable[[Span, Any], None]) -> None
+    old_execute_command = cls.execute_command
+
+    from sentry_sdk.integrations.redis import RedisIntegration
+
+    async def _sentry_execute_command(self, name, *args, **kwargs):
+        # type: (Any, str, *Any, **Any) -> Any
+        integration = sentry_sdk.get_client().get_integration(RedisIntegration)
+        if integration is None:
+            return await old_execute_command(self, name, *args, **kwargs)
+
+        cache_properties = _compile_cache_span_properties(
+            name,
+            args,
+            kwargs,
+            integration,
+        )
+
+        cache_span = 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,
+            )
+            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.__enter__()
+
+        set_db_data_fn(db_span, self)
+        _set_client_data(db_span, is_cluster, name, *args)
+
+        value = await old_execute_command(self, name, *args, **kwargs)
+
+        db_span.__exit__(None, None, None)
+
+        if cache_span:
+            _set_cache_data(cache_span, self, cache_properties, value)
+            cache_span.__exit__(None, None, None)
+
+        return value
+
+    cls.execute_command = _sentry_execute_command  # type: ignore