1# engine/util.py
2# Copyright (C) 2005-2026 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
8from __future__ import annotations
9
10from typing import Any
11from typing import Callable
12from typing import Optional
13from typing import Protocol
14from typing import TypeVar
15
16from ._util_cy import _distill_params_20 as _distill_params_20 # noqa: F401
17from ._util_cy import _distill_raw_params as _distill_raw_params # noqa: F401
18from .. import exc
19from .. import util
20from ..util.typing import Self
21
22_C = TypeVar("_C", bound=Callable[[], Any])
23
24
25def connection_memoize(key: str) -> Callable[[_C], _C]:
26 """Decorator, memoize a function in a connection.info stash.
27
28 Only applicable to functions which take no arguments other than a
29 connection. The memo will be stored in ``connection.info[key]``.
30 """
31
32 @util.decorator
33 def decorated(fn, self, connection): # type: ignore
34 connection = connection.connect()
35 try:
36 return connection.info[key]
37 except KeyError:
38 connection.info[key] = val = fn(self, connection)
39 return val
40
41 return decorated
42
43
44class _TConsSubject(Protocol):
45 _trans_context_manager: Optional[TransactionalContext]
46
47
48class TransactionalContext:
49 """Apply Python context manager behavior to transaction objects.
50
51 Performs validation to ensure the subject of the transaction is not
52 used if the transaction were ended prematurely.
53
54 """
55
56 __slots__ = ("_outer_trans_ctx", "_trans_subject", "__weakref__")
57
58 _trans_subject: Optional[_TConsSubject]
59 _rollback_exception: Optional[BaseException] = None
60
61 def _transaction_is_active(self) -> bool:
62 raise NotImplementedError()
63
64 def _transaction_is_closed(self) -> bool:
65 raise NotImplementedError()
66
67 def _rollback_can_be_called(self) -> bool:
68 """indicates the object is in a state that is known to be acceptable
69 for rollback() to be called.
70
71 This does not necessarily mean rollback() will succeed or not raise
72 an error, just that there is currently no state detected that indicates
73 rollback() would fail or emit warnings.
74
75 It also does not mean that there's a transaction in progress, as
76 it is usually safe to call rollback() even if no transaction is
77 present.
78
79 .. versionadded:: 1.4.28
80
81 """
82 raise NotImplementedError()
83
84 def _get_subject(self) -> _TConsSubject:
85 raise NotImplementedError()
86
87 def commit(self) -> None:
88 raise NotImplementedError()
89
90 def rollback(self) -> None:
91 raise NotImplementedError()
92
93 def close(self) -> None:
94 raise NotImplementedError()
95
96 @classmethod
97 def _trans_ctx_check(cls, subject: _TConsSubject) -> None:
98 trans_context = subject._trans_context_manager
99 if trans_context and not trans_context._transaction_is_active():
100 rollback_exc = trans_context._rollback_exception
101 raise exc.InvalidRequestError(
102 "Can't operate on closed transaction inside context "
103 "manager. "
104 + (
105 "The transaction was rolled back due to an "
106 f"exception: {rollback_exc}. "
107 if rollback_exc is not None
108 else ""
109 )
110 + "Please complete the context manager before "
111 "emitting further commands."
112 )
113
114 def __enter__(self) -> Self:
115 subject = self._get_subject()
116
117 # none for outer transaction, may be non-None for nested
118 # savepoint, legacy nesting cases
119 trans_context = subject._trans_context_manager
120 self._outer_trans_ctx = trans_context
121
122 self._trans_subject = subject
123 subject._trans_context_manager = self
124 return self
125
126 def __exit__(self, type_: Any, value: Any, traceback: Any) -> None:
127 subject = getattr(self, "_trans_subject", None)
128
129 # simplistically we could assume that
130 # "subject._trans_context_manager is self". However, any calling
131 # code that is manipulating __exit__ directly would break this
132 # assumption. alembic context manager
133 # is an example of partial use that just calls __exit__ and
134 # not __enter__ at the moment. it's safe to assume this is being done
135 # in the wild also
136 out_of_band_exit = (
137 subject is None or subject._trans_context_manager is not self
138 )
139
140 if type_ is None and self._transaction_is_active():
141 try:
142 self.commit()
143 except:
144 with util.safe_reraise():
145 if self._rollback_can_be_called():
146 self.rollback()
147 finally:
148 if not out_of_band_exit:
149 assert subject is not None
150 subject._trans_context_manager = self._outer_trans_ctx
151 self._trans_subject = self._outer_trans_ctx = None
152 else:
153 try:
154 if not self._transaction_is_active():
155 if not self._transaction_is_closed():
156 self.close()
157 else:
158 if self._rollback_can_be_called():
159 self.rollback()
160 finally:
161 if not out_of_band_exit:
162 assert subject is not None
163 subject._trans_context_manager = self._outer_trans_ctx
164 self._trans_subject = self._outer_trans_ctx = None