ai_v/venv/Lib/site-packages/redis/multidb/circuit.py

from abc import ABC, abstractmethod
from enum import Enum
from typing import Callable

import pybreaker

DEFAULT_GRACE_PERIOD = 60


class State(Enum):
    CLOSED = "closed"
    OPEN = "open"
    HALF_OPEN = "half-open"


class CircuitBreaker(ABC):
    @property
    @abstractmethod
    def grace_period(self) -> float:
        """The grace period in seconds when the circle should be kept open."""
        pass

    @grace_period.setter
    @abstractmethod
    def grace_period(self, grace_period: float):
        """Set the grace period in seconds."""

    @property
    @abstractmethod
    def state(self) -> State:
        """The current state of the circuit."""
        pass

    @state.setter
    @abstractmethod
    def state(self, state: State):
        """Set current state of the circuit."""
        pass

    @property
    @abstractmethod
    def database(self):
        """Database associated with this circuit."""
        pass

    @database.setter
    @abstractmethod
    def database(self, database):
        """Set database associated with this circuit."""
        pass

    @abstractmethod
    def on_state_changed(self, cb: Callable[["CircuitBreaker", State, State], None]):
        """Callback called when the state of the circuit changes."""
        pass


class BaseCircuitBreaker(CircuitBreaker):
    """
    Base implementation of Circuit Breaker interface.
    """

    def __init__(self, cb: pybreaker.CircuitBreaker):
        self._cb = cb
        self._state_pb_mapper = {
            State.CLOSED: self._cb.close,
            State.OPEN: self._cb.open,
            State.HALF_OPEN: self._cb.half_open,
        }
        self._database = None

    @property
    def grace_period(self) -> float:
        return self._cb.reset_timeout

    @grace_period.setter
    def grace_period(self, grace_period: float):
        self._cb.reset_timeout = grace_period

    @property
    def state(self) -> State:
        return State(value=self._cb.state.name)

    @state.setter
    def state(self, state: State):
        self._state_pb_mapper[state]()

    @property
    def database(self):
        return self._database

    @database.setter
    def database(self, database):
        self._database = database

    @abstractmethod
    def on_state_changed(self, cb: Callable[["CircuitBreaker", State, State], None]):
        """Callback called when the state of the circuit changes."""
        pass


class PBListener(pybreaker.CircuitBreakerListener):
    """Wrapper for callback to be compatible with pybreaker implementation."""

    def __init__(
        self,
        cb: Callable[[CircuitBreaker, State, State], None],
        database,
    ):
        """
        Initialize a PBListener instance.

        Args:
            cb: Callback function that will be called when the circuit breaker state changes.
            database: Database instance associated with this circuit breaker.
        """

        self._cb = cb
        self._database = database

    def state_change(self, cb, old_state, new_state):
        cb = PBCircuitBreakerAdapter(cb)
        cb.database = self._database
        old_state = State(value=old_state.name)
        new_state = State(value=new_state.name)
        self._cb(cb, old_state, new_state)


class PBCircuitBreakerAdapter(BaseCircuitBreaker):
    def __init__(self, cb: pybreaker.CircuitBreaker):
        """
        Initialize a PBCircuitBreakerAdapter instance.

        This adapter wraps pybreaker's CircuitBreaker implementation to make it compatible
        with our CircuitBreaker interface.

        Args:
            cb: A pybreaker CircuitBreaker instance to be adapted.
        """
        super().__init__(cb)

    def on_state_changed(self, cb: Callable[["CircuitBreaker", State, State], None]):
        listener = PBListener(cb, self.database)
        self._cb.add_listener(listener)