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