state_changes.py 6.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198
  1. # orm/state_changes.py
  2. # Copyright (C) 2005-2024 the SQLAlchemy authors and contributors
  3. # <see AUTHORS file>
  4. #
  5. # This module is part of SQLAlchemy and is released under
  6. # the MIT License: https://www.opensource.org/licenses/mit-license.php
  7. """State tracking utilities used by :class:`_orm.Session`.
  8. """
  9. from __future__ import annotations
  10. import contextlib
  11. from enum import Enum
  12. from typing import Any
  13. from typing import Callable
  14. from typing import cast
  15. from typing import Iterator
  16. from typing import NoReturn
  17. from typing import Optional
  18. from typing import Tuple
  19. from typing import TypeVar
  20. from typing import Union
  21. from .. import exc as sa_exc
  22. from .. import util
  23. from ..util.typing import Literal
  24. _F = TypeVar("_F", bound=Callable[..., Any])
  25. class _StateChangeState(Enum):
  26. pass
  27. class _StateChangeStates(_StateChangeState):
  28. ANY = 1
  29. NO_CHANGE = 2
  30. CHANGE_IN_PROGRESS = 3
  31. class _StateChange:
  32. """Supplies state assertion decorators.
  33. The current use case is for the :class:`_orm.SessionTransaction` class. The
  34. :class:`_StateChange` class itself is agnostic of the
  35. :class:`_orm.SessionTransaction` class so could in theory be generalized
  36. for other systems as well.
  37. """
  38. _next_state: _StateChangeState = _StateChangeStates.ANY
  39. _state: _StateChangeState = _StateChangeStates.NO_CHANGE
  40. _current_fn: Optional[Callable[..., Any]] = None
  41. def _raise_for_prerequisite_state(
  42. self, operation_name: str, state: _StateChangeState
  43. ) -> NoReturn:
  44. raise sa_exc.IllegalStateChangeError(
  45. f"Can't run operation '{operation_name}()' when Session "
  46. f"is in state {state!r}",
  47. code="isce",
  48. )
  49. @classmethod
  50. def declare_states(
  51. cls,
  52. prerequisite_states: Union[
  53. Literal[_StateChangeStates.ANY], Tuple[_StateChangeState, ...]
  54. ],
  55. moves_to: _StateChangeState,
  56. ) -> Callable[[_F], _F]:
  57. """Method decorator declaring valid states.
  58. :param prerequisite_states: sequence of acceptable prerequisite
  59. states. Can be the single constant _State.ANY to indicate no
  60. prerequisite state
  61. :param moves_to: the expected state at the end of the method, assuming
  62. no exceptions raised. Can be the constant _State.NO_CHANGE to
  63. indicate state should not change at the end of the method.
  64. """
  65. assert prerequisite_states, "no prequisite states sent"
  66. has_prerequisite_states = (
  67. prerequisite_states is not _StateChangeStates.ANY
  68. )
  69. prerequisite_state_collection = cast(
  70. "Tuple[_StateChangeState, ...]", prerequisite_states
  71. )
  72. expect_state_change = moves_to is not _StateChangeStates.NO_CHANGE
  73. @util.decorator
  74. def _go(fn: _F, self: Any, *arg: Any, **kw: Any) -> Any:
  75. current_state = self._state
  76. if (
  77. has_prerequisite_states
  78. and current_state not in prerequisite_state_collection
  79. ):
  80. self._raise_for_prerequisite_state(fn.__name__, current_state)
  81. next_state = self._next_state
  82. existing_fn = self._current_fn
  83. expect_state = moves_to if expect_state_change else current_state
  84. if (
  85. # destination states are restricted
  86. next_state is not _StateChangeStates.ANY
  87. # method seeks to change state
  88. and expect_state_change
  89. # destination state incorrect
  90. and next_state is not expect_state
  91. ):
  92. if existing_fn and next_state in (
  93. _StateChangeStates.NO_CHANGE,
  94. _StateChangeStates.CHANGE_IN_PROGRESS,
  95. ):
  96. raise sa_exc.IllegalStateChangeError(
  97. f"Method '{fn.__name__}()' can't be called here; "
  98. f"method '{existing_fn.__name__}()' is already "
  99. f"in progress and this would cause an unexpected "
  100. f"state change to {moves_to!r}",
  101. code="isce",
  102. )
  103. else:
  104. raise sa_exc.IllegalStateChangeError(
  105. f"Cant run operation '{fn.__name__}()' here; "
  106. f"will move to state {moves_to!r} where we are "
  107. f"expecting {next_state!r}",
  108. code="isce",
  109. )
  110. self._current_fn = fn
  111. self._next_state = _StateChangeStates.CHANGE_IN_PROGRESS
  112. try:
  113. ret_value = fn(self, *arg, **kw)
  114. except:
  115. raise
  116. else:
  117. if self._state is expect_state:
  118. return ret_value
  119. if self._state is current_state:
  120. raise sa_exc.IllegalStateChangeError(
  121. f"Method '{fn.__name__}()' failed to "
  122. "change state "
  123. f"to {moves_to!r} as expected",
  124. code="isce",
  125. )
  126. elif existing_fn:
  127. raise sa_exc.IllegalStateChangeError(
  128. f"While method '{existing_fn.__name__}()' was "
  129. "running, "
  130. f"method '{fn.__name__}()' caused an "
  131. "unexpected "
  132. f"state change to {self._state!r}",
  133. code="isce",
  134. )
  135. else:
  136. raise sa_exc.IllegalStateChangeError(
  137. f"Method '{fn.__name__}()' caused an unexpected "
  138. f"state change to {self._state!r}",
  139. code="isce",
  140. )
  141. finally:
  142. self._next_state = next_state
  143. self._current_fn = existing_fn
  144. return _go
  145. @contextlib.contextmanager
  146. def _expect_state(self, expected: _StateChangeState) -> Iterator[Any]:
  147. """called within a method that changes states.
  148. method must also use the ``@declare_states()`` decorator.
  149. """
  150. assert self._next_state is _StateChangeStates.CHANGE_IN_PROGRESS, (
  151. "Unexpected call to _expect_state outside of "
  152. "state-changing method"
  153. )
  154. self._next_state = expected
  155. try:
  156. yield
  157. except:
  158. raise
  159. else:
  160. if self._state is not expected:
  161. raise sa_exc.IllegalStateChangeError(
  162. f"Unexpected state change to {self._state!r}", code="isce"
  163. )
  164. finally:
  165. self._next_state = _StateChangeStates.CHANGE_IN_PROGRESS