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)