Initial commit (Clean history)

path/to/venv/lib/python3.12/site-packages/redis/multidb/circuit.py

@@ -0,0 +1,144 @@
+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)