1from __future__ import annotations
2
3from collections.abc import Callable, Generator
4from contextlib import ContextDecorator, contextmanager
5from types import TracebackType
6from typing import Any
7
8import psycopg
9from plain.postgres.db import get_connection
10
11
12class TransactionManagementError(psycopg.ProgrammingError):
13 """Transaction management is used improperly."""
14
15
16@contextmanager
17def mark_for_rollback_on_error() -> Generator[None]:
18 """
19 Internal low-level utility to mark a transaction as "needs rollback" when
20 an exception is raised while not enforcing the enclosed block to be in a
21 transaction. This is needed by Model.create()/update() and friends to avoid
22 starting a transaction when in autocommit mode and a single query is executed.
23
24 It's equivalent to:
25
26 if get_connection().get_autocommit():
27 yield
28 else:
29 with transaction.atomic(savepoint=False):
30 yield
31
32 but it uses low-level utilities to avoid performance overhead.
33 """
34 try:
35 yield
36 except Exception as exc:
37 conn = get_connection()
38 if conn.in_atomic_block:
39 conn.needs_rollback = True
40 conn.rollback_exc = exc
41 raise
42
43
44def on_commit(func: Callable[[], Any], robust: bool = False) -> None:
45 """
46 Register `func` to be called when the current transaction is committed.
47 If the current transaction is rolled back, `func` will not be called.
48 """
49 get_connection().on_commit(func, robust)
50
51
52#################################
53# Decorators / context managers #
54#################################
55
56
57class Atomic(ContextDecorator):
58 """
59 Guarantee the atomic execution of a given block.
60
61 An instance can be used either as a decorator or as a context manager.
62
63 When it's used as a decorator, __call__ wraps the execution of the
64 decorated function in the instance itself, used as a context manager.
65
66 When it's used as a context manager, __enter__ creates a transaction or a
67 savepoint, depending on whether a transaction is already in progress, and
68 __exit__ commits the transaction or releases the savepoint on normal exit,
69 and rolls back the transaction or to the savepoint on exceptions.
70
71 It's possible to disable the creation of savepoints if the goal is to
72 ensure that some code runs within a transaction without creating overhead.
73
74 A stack of savepoints identifiers is maintained as an attribute of the
75 connection. None denotes the absence of a savepoint.
76
77 This allows reentrancy even if the same AtomicWrapper is reused. For
78 example, it's possible to define `oa = atomic('other')` and use `@oa` or
79 `with oa:` multiple times.
80
81 Since database connections are stored per-context (ContextVar), this is thread-safe.
82
83 An atomic block can be tagged as durable. In this case, raise a
84 RuntimeError if it's nested within another atomic block. This guarantees
85 that database changes in a durable block are committed to the database when
86 the block exists without error.
87
88 This is a private API.
89 """
90
91 def __init__(self, savepoint: bool, durable: bool) -> None:
92 self.savepoint = savepoint
93 self.durable = durable
94 self._from_testcase = False
95
96 def __enter__(self) -> None:
97 conn = get_connection()
98 if (
99 self.durable
100 and conn.atomic_blocks
101 and not conn.atomic_blocks[-1]._from_testcase
102 ):
103 raise RuntimeError(
104 "A durable atomic block cannot be nested within another atomic block."
105 )
106 if not conn.in_atomic_block:
107 # Reset state when entering an outermost atomic block. Clearing
108 # rollback_exc keeps a stale value from the reused connection wrapper
109 # from being misattributed as the cause of a later block's
110 # broken-transaction error.
111 conn.needs_rollback = False
112 conn.rollback_exc = None
113 if conn.in_atomic_block:
114 # We're already in a transaction; create a savepoint, unless we
115 # were told not to or we're already waiting for a rollback. The
116 # second condition avoids creating useless savepoints and prevents
117 # overwriting needs_rollback until the rollback is performed.
118 if self.savepoint and not conn.needs_rollback:
119 sid = conn.savepoint()
120 conn.savepoint_ids.append(sid)
121 else:
122 conn.savepoint_ids.append(None)
123 else:
124 conn.set_autocommit(False)
125 conn.in_atomic_block = True
126
127 if conn.in_atomic_block:
128 conn.atomic_blocks.append(self)
129
130 def __exit__(
131 self,
132 exc_type: type[BaseException] | None,
133 exc_value: BaseException | None,
134 traceback: TracebackType | None,
135 ) -> None:
136 conn = get_connection()
137 if conn.in_atomic_block:
138 conn.atomic_blocks.pop()
139
140 if conn.savepoint_ids:
141 sid = conn.savepoint_ids.pop()
142 else:
143 # Prematurely unset this flag to allow using commit or rollback.
144 conn.in_atomic_block = False
145
146 try:
147 if exc_type is None and not conn.needs_rollback:
148 if conn.in_atomic_block:
149 # Release savepoint if there is one
150 if sid is not None:
151 try:
152 conn.savepoint_commit(sid)
153 except psycopg.DatabaseError:
154 try:
155 conn.savepoint_rollback(sid)
156 # The savepoint won't be reused. Release it to
157 # minimize overhead for the database server.
158 conn.savepoint_commit(sid)
159 except psycopg.Error:
160 # If rolling back to a savepoint fails, mark for
161 # rollback at a higher level and avoid shadowing
162 # the original exception.
163 conn.needs_rollback = True
164 raise
165 else:
166 # Commit transaction
167 try:
168 conn.commit()
169 except psycopg.DatabaseError:
170 try:
171 conn.rollback()
172 except psycopg.Error:
173 # An error during rollback means that something
174 # went wrong with the connection. Drop it.
175 conn.close()
176 raise
177 else:
178 # This flag will be set to True again if there isn't a savepoint
179 # allowing to perform the rollback at this level.
180 conn.needs_rollback = False
181 if conn.in_atomic_block:
182 # Roll back to savepoint if there is one, mark for rollback
183 # otherwise.
184 if sid is None:
185 conn.needs_rollback = True
186 # Record what broke the transaction (when we can see it)
187 # so validate_no_broken_transaction() chains from the
188 # real cause rather than a stale or absent one. Mirror
189 # mark_for_rollback_on_error() and only capture Exception
190 # (not BaseException like KeyboardInterrupt/SystemExit).
191 if isinstance(exc_value, Exception):
192 conn.rollback_exc = exc_value
193 else:
194 try:
195 conn.savepoint_rollback(sid)
196 # The savepoint won't be reused. Release it to
197 # minimize overhead for the database server.
198 conn.savepoint_commit(sid)
199 except psycopg.Error:
200 # If rolling back to a savepoint fails, mark for
201 # rollback at a higher level and avoid shadowing
202 # the original exception.
203 conn.needs_rollback = True
204 else:
205 # Roll back transaction
206 try:
207 conn.rollback()
208 except psycopg.Error:
209 # An error during rollback means that something
210 # went wrong with the connection. Drop it.
211 conn.close()
212
213 finally:
214 # Outermost block exit when autocommit was enabled. Skip when
215 # the connection was dropped during rollback/commit failure —
216 # ensure_connection() would otherwise acquire a fresh pool
217 # connection just to flip autocommit while the original error
218 # is still propagating.
219 if not conn.in_atomic_block and conn.connection is not None:
220 conn.set_autocommit(True)
221
222
223def atomic[F: Callable[..., Any]](
224 func: F | None = None, *, savepoint: bool = True, durable: bool = False
225) -> F | Atomic:
226 """Create an atomic transaction context or decorator."""
227 if callable(func):
228 return Atomic(savepoint, durable)(func)
229 return Atomic(savepoint, durable)