redis 7.0.0b2__py3-none-any.whl → 7.0.1__py3-none-any.whl

redis/maint_notifications.py

@@ -32,9 +32,8 @@ class EndpointType(enum.Enum):
 
 if TYPE_CHECKING:
     from redis.connection import (
-        BlockingConnectionPool,
-        ConnectionInterface,
-        ConnectionPool,
+        MaintNotificationsAbstractConnection,
+        MaintNotificationsAbstractConnectionPool,
     )
 
 
@@ -501,7 +500,7 @@ class MaintNotificationsConfig:
         return self.relaxed_timeout != -1
 
     def get_endpoint_type(
-        self, host: str, connection: "ConnectionInterface"
+        self, host: str, connection: "MaintNotificationsAbstractConnection"
     ) -> EndpointType:
         """
         Determine the appropriate endpoint type for CLIENT MAINT_NOTIFICATIONS command.
@@ -558,7 +557,7 @@ class MaintNotificationsConfig:
 class MaintNotificationsPoolHandler:
     def __init__(
         self,
-        pool: Union["ConnectionPool", "BlockingConnectionPool"],
+        pool: "MaintNotificationsAbstractConnectionPool",
         config: MaintNotificationsConfig,
     ) -> None:
         self.pool = pool
@@ -567,9 +566,19 @@ class MaintNotificationsPoolHandler:
         self._lock = threading.RLock()
         self.connection = None
 
-    def set_connection(self, connection: "ConnectionInterface"):
+    def set_connection(self, connection: "MaintNotificationsAbstractConnection"):
         self.connection = connection
 
+    def get_handler_for_connection(self):
+        # Copy all data that should be shared between connections
+        # but each connection should have its own pool handler
+        # since each connection can be in a different state
+        copy = MaintNotificationsPoolHandler(self.pool, self.config)
+        copy._processed_notifications = self._processed_notifications
+        copy._lock = self._lock
+        copy.connection = None
+        return copy
+
     def remove_expired_notifications(self):
         with self._lock:
             for notification in tuple(self._processed_notifications):
@@ -751,7 +760,9 @@ class MaintNotificationsConnectionHandler:
     }
 
     def __init__(
-        self, connection: "ConnectionInterface", config: MaintNotificationsConfig
+        self,
+        connection: "MaintNotificationsAbstractConnection",
+        config: MaintNotificationsConfig,
     ) -> None:
         self.connection = connection
         self.config = config

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)