image-match/python/py/Lib/site-packages/sentry_sdk/integrations/redis/redis_cluster.py

"""
Instrumentation for RedisCluster
This is part of the main redis-py client.

https://github.com/redis/redis-py/blob/master/redis/cluster.py
"""

from sentry_sdk.integrations.redis._sync_common import (
    patch_redis_client,
    patch_redis_pipeline,
)
from sentry_sdk.integrations.redis.modules.queries import _set_db_data_on_span
from sentry_sdk.integrations.redis.utils import _parse_rediscluster_command

from sentry_sdk.utils import capture_internal_exceptions

from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from typing import Any
    from redis import RedisCluster
    from redis.asyncio.cluster import (
        RedisCluster as AsyncRedisCluster,
        ClusterPipeline as AsyncClusterPipeline,
    )
    from sentry_sdk.tracing import Span


def _set_async_cluster_db_data(
    span: "Span", async_redis_cluster_instance: "AsyncRedisCluster[Any]"
) -> None:
    default_node = async_redis_cluster_instance.get_default_node()
    if default_node is not None and default_node.connection_kwargs is not None:
        _set_db_data_on_span(span, default_node.connection_kwargs)


def _set_async_cluster_pipeline_db_data(
    span: "Span", async_redis_cluster_pipeline_instance: "AsyncClusterPipeline[Any]"
) -> None:
    with capture_internal_exceptions():
        client = getattr(async_redis_cluster_pipeline_instance, "cluster_client", None)
        if client is None:
            # In older redis-py versions, the AsyncClusterPipeline had a `_client`
            # attr but it is private so potentially problematic and mypy does not
            # recognize it - see
            # https://github.com/redis/redis-py/blame/v5.0.0/redis/asyncio/cluster.py#L1386
            client = (
                async_redis_cluster_pipeline_instance._client  # type: ignore[attr-defined]
            )

        _set_async_cluster_db_data(
            span,
            client,
        )


def _set_cluster_db_data(
    span: "Span", redis_cluster_instance: "RedisCluster[Any]"
) -> None:
    default_node = redis_cluster_instance.get_default_node()

    if default_node is not None:
        connection_params = {
            "host": default_node.host,
            "port": default_node.port,
        }
        _set_db_data_on_span(span, connection_params)


def _patch_redis_cluster() -> None:
    """Patches the cluster module on redis SDK (as opposed to rediscluster library)"""
    try:
        from redis import RedisCluster, cluster
    except ImportError:
        pass
    else:
        patch_redis_client(
            RedisCluster,
            is_cluster=True,
            set_db_data_fn=_set_cluster_db_data,
        )
        patch_redis_pipeline(
            cluster.ClusterPipeline,
            is_cluster=True,
            get_command_args_fn=_parse_rediscluster_command,
            set_db_data_fn=_set_cluster_db_data,
        )

    try:
        from redis.asyncio import cluster as async_cluster
    except ImportError:
        pass
    else:
        from sentry_sdk.integrations.redis._async_common import (
            patch_redis_async_client,
            patch_redis_async_pipeline,
        )

        patch_redis_async_client(
            async_cluster.RedisCluster,
            is_cluster=True,
            set_db_data_fn=_set_async_cluster_db_data,
        )
        patch_redis_async_pipeline(
            async_cluster.ClusterPipeline,
            is_cluster=True,
            get_command_args_fn=_parse_rediscluster_command,
            set_db_data_fn=_set_async_cluster_pipeline_db_data,
        )