from abc import abstractmethod from asyncio import iscoroutinefunction from datetime import datetime from typing import Any, Awaitable, Callable, List, Optional, Union from redis.asyncio import RedisCluster from redis.asyncio.client import Pipeline, PubSub from redis.asyncio.multidb.database import AsyncDatabase, Database, Databases from redis.asyncio.multidb.event import ( AsyncActiveDatabaseChanged, CloseConnectionOnActiveDatabaseChanged, RegisterCommandFailure, ResubscribeOnActiveDatabaseChanged, ) from redis.asyncio.multidb.failover import ( DEFAULT_FAILOVER_ATTEMPTS, DEFAULT_FAILOVER_DELAY, AsyncFailoverStrategy, DefaultFailoverStrategyExecutor, FailoverStrategyExecutor, ) from redis.asyncio.multidb.failure_detector import AsyncFailureDetector from redis.asyncio.retry import Retry from redis.event import AsyncOnCommandsFailEvent, EventDispatcherInterface from redis.multidb.circuit import State as CBState from redis.multidb.command_executor import BaseCommandExecutor, CommandExecutor from redis.multidb.config import DEFAULT_AUTO_FALLBACK_INTERVAL from redis.typing import KeyT class AsyncCommandExecutor(CommandExecutor): @property @abstractmethod def databases(self) -> Databases: """Returns a list of databases.""" pass @property @abstractmethod def failure_detectors(self) -> List[AsyncFailureDetector]: """Returns a list of failure detectors.""" pass @abstractmethod def add_failure_detector(self, failure_detector: AsyncFailureDetector) -> None: """Adds a new failure detector to the list of failure detectors.""" pass @property @abstractmethod def active_database(self) -> Optional[AsyncDatabase]: """Returns currently active database.""" pass @abstractmethod async def set_active_database(self, database: AsyncDatabase) -> None: """Sets the currently active database. Args: database: The new active database. """ pass @property @abstractmethod def active_pubsub(self) -> Optional[PubSub]: """Returns currently active pubsub.""" pass @active_pubsub.setter @abstractmethod def active_pubsub(self, pubsub: PubSub) -> None: """Sets currently active pubsub.""" pass @property @abstractmethod def failover_strategy_executor(self) -> FailoverStrategyExecutor: """Returns failover strategy executor.""" pass @property @abstractmethod def command_retry(self) -> Retry: """Returns command retry object.""" pass @abstractmethod async def pubsub(self, **kwargs): """Initializes a PubSub object on a currently active database""" pass @abstractmethod async def execute_command(self, *args, **options): """Executes a command and returns the result.""" pass @abstractmethod async def execute_pipeline(self, command_stack: tuple): """Executes a stack of commands in pipeline.""" pass @abstractmethod async def execute_transaction( self, transaction: Callable[[Pipeline], None], *watches, **options ): """Executes a transaction block wrapped in callback.""" pass @abstractmethod async def execute_pubsub_method(self, method_name: str, *args, **kwargs): """Executes a given method on active pub/sub.""" pass @abstractmethod async def execute_pubsub_run(self, sleep_time: float, **kwargs) -> Any: """Executes pub/sub run in a thread.""" pass class DefaultCommandExecutor(BaseCommandExecutor, AsyncCommandExecutor): def __init__( self, failure_detectors: List[AsyncFailureDetector], databases: Databases, command_retry: Retry, failover_strategy: AsyncFailoverStrategy, event_dispatcher: EventDispatcherInterface, failover_attempts: int = DEFAULT_FAILOVER_ATTEMPTS, failover_delay: float = DEFAULT_FAILOVER_DELAY, auto_fallback_interval: float = DEFAULT_AUTO_FALLBACK_INTERVAL, ): """ Initialize the DefaultCommandExecutor instance. Args: failure_detectors: List of failure detector instances to monitor database health databases: Collection of available databases to execute commands on command_retry: Retry policy for failed command execution failover_strategy: Strategy for handling database failover event_dispatcher: Interface for dispatching events failover_attempts: Number of failover attempts failover_delay: Delay between failover attempts auto_fallback_interval: Time interval in seconds between attempts to fall back to a primary database """ super().__init__(auto_fallback_interval) for fd in failure_detectors: fd.set_command_executor(command_executor=self) self._databases = databases self._failure_detectors = failure_detectors self._command_retry = command_retry self._failover_strategy_executor = DefaultFailoverStrategyExecutor( failover_strategy, failover_attempts, failover_delay ) self._event_dispatcher = event_dispatcher self._active_database: Optional[Database] = None self._active_pubsub: Optional[PubSub] = None self._active_pubsub_kwargs = {} self._setup_event_dispatcher() self._schedule_next_fallback() @property def databases(self) -> Databases: return self._databases @property def failure_detectors(self) -> List[AsyncFailureDetector]: return self._failure_detectors def add_failure_detector(self, failure_detector: AsyncFailureDetector) -> None: self._failure_detectors.append(failure_detector) @property def active_database(self) -> Optional[AsyncDatabase]: return self._active_database async def set_active_database(self, database: AsyncDatabase) -> None: old_active = self._active_database self._active_database = database if old_active is not None and old_active is not database: await self._event_dispatcher.dispatch_async( AsyncActiveDatabaseChanged( old_active, self._active_database, self, **self._active_pubsub_kwargs, ) ) @property def active_pubsub(self) -> Optional[PubSub]: return self._active_pubsub @active_pubsub.setter def active_pubsub(self, pubsub: PubSub) -> None: self._active_pubsub = pubsub @property def failover_strategy_executor(self) -> FailoverStrategyExecutor: return self._failover_strategy_executor @property def command_retry(self) -> Retry: return self._command_retry def pubsub(self, **kwargs): if self._active_pubsub is None: if isinstance(self._active_database.client, RedisCluster): raise ValueError("PubSub is not supported for RedisCluster") self._active_pubsub = self._active_database.client.pubsub(**kwargs) self._active_pubsub_kwargs = kwargs async def execute_command(self, *args, **options): async def callback(): response = await self._active_database.client.execute_command( *args, **options ) await self._register_command_execution(args) return response return await self._execute_with_failure_detection(callback, args) async def execute_pipeline(self, command_stack: tuple): async def callback(): async with self._active_database.client.pipeline() as pipe: for command, options in command_stack: pipe.execute_command(*command, **options) response = await pipe.execute() await self._register_command_execution(command_stack) return response return await self._execute_with_failure_detection(callback, command_stack) async def execute_transaction( self, func: Callable[["Pipeline"], Union[Any, Awaitable[Any]]], *watches: KeyT, shard_hint: Optional[str] = None, value_from_callable: bool = False, watch_delay: Optional[float] = None, ): async def callback(): response = await self._active_database.client.transaction( func, *watches, shard_hint=shard_hint, value_from_callable=value_from_callable, watch_delay=watch_delay, ) await self._register_command_execution(()) return response return await self._execute_with_failure_detection(callback) async def execute_pubsub_method(self, method_name: str, *args, **kwargs): async def callback(): method = getattr(self.active_pubsub, method_name) if iscoroutinefunction(method): response = await method(*args, **kwargs) else: response = method(*args, **kwargs) await self._register_command_execution(args) return response return await self._execute_with_failure_detection(callback, *args) async def execute_pubsub_run( self, sleep_time: float, exception_handler=None, pubsub=None ) -> Any: async def callback(): return await self._active_pubsub.run( poll_timeout=sleep_time, exception_handler=exception_handler, pubsub=pubsub, ) return await self._execute_with_failure_detection(callback) async def _execute_with_failure_detection( self, callback: Callable, cmds: tuple = () ): """ Execute a commands execution callback with failure detection. """ async def wrapper(): # On each retry we need to check active database as it might change. await self._check_active_database() return await callback() return await self._command_retry.call_with_retry( lambda: wrapper(), lambda error: self._on_command_fail(error, *cmds), ) async def _check_active_database(self): """ Checks if active a database needs to be updated. """ if ( self._active_database is None or self._active_database.circuit.state != CBState.CLOSED or ( self._auto_fallback_interval > 0 and self._next_fallback_attempt <= datetime.now() ) ): await self.set_active_database( await self._failover_strategy_executor.execute() ) self._schedule_next_fallback() async def _on_command_fail(self, error, *args): await self._event_dispatcher.dispatch_async( AsyncOnCommandsFailEvent(args, error) ) async def _register_command_execution(self, cmd: tuple): for detector in self._failure_detectors: await detector.register_command_execution(cmd) def _setup_event_dispatcher(self): """ Registers necessary listeners. """ failure_listener = RegisterCommandFailure(self._failure_detectors) resubscribe_listener = ResubscribeOnActiveDatabaseChanged() close_connection_listener = CloseConnectionOnActiveDatabaseChanged() self._event_dispatcher.register_listeners( { AsyncOnCommandsFailEvent: [failure_listener], AsyncActiveDatabaseChanged: [ close_connection_listener, resubscribe_listener, ], } )