alien_py_cloud/alien_util/redis_client.py

from redis import Redis
from redis.asyncio import Redis as AsyncRedis
from redis.asyncio.sentinel import Sentinel as AsyncSentinel
from redis.sentinel import Sentinel

from alien_gateway.config import settings


def _sentinel_kwargs() -> dict:
    kwargs = {}
    if settings.REDIS_SENTINEL_USERNAME:
        kwargs["username"] = settings.REDIS_SENTINEL_USERNAME
    if settings.REDIS_SENTINEL_PASSWORD:
        kwargs["password"] = settings.REDIS_SENTINEL_PASSWORD
    return kwargs


def get_sentinel_client() -> Sentinel:
    return Sentinel(
        settings.REDIS_SENTINEL_NODES,
        socket_timeout=0.5,
        sentinel_kwargs=_sentinel_kwargs(),
    )


def get_async_sentinel_client() -> AsyncSentinel:
    return AsyncSentinel(
        settings.REDIS_SENTINEL_NODES,
        socket_timeout=0.5,
        sentinel_kwargs=_sentinel_kwargs(),
    )


def get_redis_master() -> Redis:
    return get_sentinel_client().master_for(
        service_name=settings.REDIS_MASTER_NAME,
        password=settings.REDIS_PASSWORD or None,
        db=settings.REDIS_DB,
    )


def get_redis_slave() -> Redis:
    return get_sentinel_client().slave_for(
        service_name=settings.REDIS_MASTER_NAME,
        password=settings.REDIS_PASSWORD or None,
        db=settings.REDIS_DB,
    )


def get_async_redis_master() -> AsyncRedis:
    return get_async_sentinel_client().master_for(
        service_name=settings.REDIS_MASTER_NAME,
        password=settings.REDIS_PASSWORD or None,
        db=settings.REDIS_DB,
    )


def get_async_redis_slave() -> AsyncRedis:
    return get_async_sentinel_client().slave_for(
        service_name=settings.REDIS_MASTER_NAME,
        password=settings.REDIS_PASSWORD or None,
        db=settings.REDIS_DB,
    )