A class for managing persistent recoverable state using a Pydantic model. This class facilitates state persistence to a `KeyValueStore`, allowing data to be saved and retrieved across migrations or restarts. It manages the loading, saving, and resetting of state data, with optional pers
| 17 | |
| 18 | |
| 19 | class RecoverableState(Generic[TStateModel]): |
| 20 | """A class for managing persistent recoverable state using a Pydantic model. |
| 21 | |
| 22 | This class facilitates state persistence to a `KeyValueStore`, allowing data to be saved and retrieved |
| 23 | across migrations or restarts. It manages the loading, saving, and resetting of state data, |
| 24 | with optional persistence capabilities. |
| 25 | |
| 26 | The state is represented by a Pydantic model that can be serialized to and deserialized from JSON. |
| 27 | The class automatically hooks into the event system to persist state when needed. |
| 28 | |
| 29 | Type Parameters: |
| 30 | TStateModel: A Pydantic BaseModel type that defines the structure of the state data. |
| 31 | Typically, it should be inferred from the `default_state` constructor parameter. |
| 32 | """ |
| 33 | |
| 34 | def __init__( |
| 35 | self, |
| 36 | *, |
| 37 | default_state: TStateModel, |
| 38 | persist_state_key: str, |
| 39 | persistence_enabled: Literal[True, False, 'explicit_only'] = False, |
| 40 | persist_state_kvs_name: str | None = None, |
| 41 | persist_state_kvs_id: str | None = None, |
| 42 | persist_state_kvs_factory: Callable[[], Coroutine[None, None, KeyValueStore]] | None = None, |
| 43 | logger: logging.Logger, |
| 44 | ) -> None: |
| 45 | """Initialize a new recoverable state object. |
| 46 | |
| 47 | Args: |
| 48 | default_state: The default state model instance to use when no persisted state is found. |
| 49 | A deep copy is made each time the state is used. |
| 50 | persist_state_key: The key under which the state is stored in the KeyValueStore |
| 51 | persistence_enabled: Flag to enable or disable state persistence. Use 'explicit_only' if you want to be able |
| 52 | to save the state manually, but without any automatic persistence. |
| 53 | persist_state_kvs_name: The name of the KeyValueStore to use for persistence. |
| 54 | If neither a name nor and id are supplied, the default store will be used. |
| 55 | persist_state_kvs_id: The identifier of the KeyValueStore to use for persistence. |
| 56 | If neither a name nor and id are supplied, the default store will be used. |
| 57 | persist_state_kvs_factory: Factory that can be awaited to create KeyValueStore to use for persistence. If |
| 58 | not provided, a system-wide KeyValueStore will be used, based on service locator configuration. |
| 59 | logger: A logger instance for logging operations related to state persistence |
| 60 | """ |
| 61 | raise_if_too_many_kwargs( |
| 62 | persist_state_kvs_name=persist_state_kvs_name, |
| 63 | persist_state_kvs_id=persist_state_kvs_id, |
| 64 | persist_state_kvs_factory=persist_state_kvs_factory, |
| 65 | ) |
| 66 | if not persist_state_kvs_factory: |
| 67 | logger.debug( |
| 68 | 'No explicit key_value_store set for recoverable state. Recovery will use a system-wide KeyValueStore ' |
| 69 | 'based on service_locator configuration, potentially calling service_locator.set_storage_client in the ' |
| 70 | 'process. It is recommended to initialize RecoverableState with explicit key_value_store to avoid ' |
| 71 | 'global side effects.' |
| 72 | ) |
| 73 | |
| 74 | self._default_state = default_state |
| 75 | self._state_type: type[TStateModel] = self._default_state.__class__ |
| 76 | self._state: TStateModel | None = None |