Coverage for /Users/davegaeddert/Developer/dropseed/plain/plain-models/plain/models/transaction.py: 63%

1 from contextlib import ContextDecorator, contextmanager

2

3 from plain.models.db import (

4 DEFAULT_DB_ALIAS,

5 DatabaseError,

6 Error,

7 ProgrammingError,

8 connections,

9)

10

11

12 class TransactionManagementError(ProgrammingError):

13 """Transaction management is used improperly."""

14

15 pass

16

17

18 def get_connection(using=None):

19 """

20 Get a database connection by name, or the default database connection

21 if no name is provided. This is a private API.

22 """

23 if using is None:

24 using = DEFAULT_DB_ALIAS

25 return connections[using]

26

27

28 def get_autocommit(using=None):

29 """Get the autocommit status of the connection."""

30 return get_connection(using).get_autocommit()

31

32

33 def set_autocommit(autocommit, using=None):

34 """Set the autocommit status of the connection."""

35 return get_connection(using).set_autocommit(autocommit)

36

37

38 def commit(using=None):

39 """Commit a transaction."""

40 get_connection(using).commit()

41

42

43 def rollback(using=None):

44 """Roll back a transaction."""

45 get_connection(using).rollback()

46

47

48 def savepoint(using=None):

49 """

50 Create a savepoint (if supported and required by the backend) inside the

51 current transaction. Return an identifier for the savepoint that will be

52 used for the subsequent rollback or commit.

53 """

54 return get_connection(using).savepoint()

55

56

57 def savepoint_rollback(sid, using=None):

58 """

59 Roll back the most recent savepoint (if one exists). Do nothing if

60 savepoints are not supported.

61 """

62 get_connection(using).savepoint_rollback(sid)

63

64

65 def savepoint_commit(sid, using=None):

66 """

67 Commit the most recent savepoint (if one exists). Do nothing if

68 savepoints are not supported.

69 """

70 get_connection(using).savepoint_commit(sid)

71

72

73 def clean_savepoints(using=None):

74 """

75 Reset the counter used to generate unique savepoint ids in this thread.

76 """

77 get_connection(using).clean_savepoints()

78

79

80 def get_rollback(using=None):

81 """Get the "needs rollback" flag -- for *advanced use* only."""

82 return get_connection(using).get_rollback()

83

84

85 def set_rollback(rollback, using=None):

86 """

87 Set or unset the "needs rollback" flag -- for *advanced use* only.

88

89 When `rollback` is `True`, trigger a rollback when exiting the innermost

90 enclosing atomic block that has `savepoint=True` (that's the default). Use

91 this to force a rollback without raising an exception.

92

93 When `rollback` is `False`, prevent such a rollback. Use this only after

94 rolling back to a known-good state! Otherwise, you break the atomic block

95 and data corruption may occur.

96 """

97 return get_connection(using).set_rollback(rollback)

98

99

100@contextmanager

101 def mark_for_rollback_on_error(using=None):

102 """

103 Internal low-level utility to mark a transaction as "needs rollback" when

104 an exception is raised while not enforcing the enclosed block to be in a

105 transaction. This is needed by Model.save() and friends to avoid starting a

106 transaction when in autocommit mode and a single query is executed.

107

108 It's equivalent to:

109

110 connection = get_connection(using)

111 if connection.get_autocommit():

112 yield

113 else:

114 with transaction.atomic(using=using, savepoint=False):

115 yield

116

117 but it uses low-level utilities to avoid performance overhead.

118 """

119 try:

120 yield

121 except Exception as exc:

122 connection = get_connection(using)

123 if connection.in_atomic_block:

124 connection.needs_rollback = True

125 connection.rollback_exc = exc

126 raise

127

128

129 def on_commit(func, using=None, robust=False):

130 """

131 Register `func` to be called when the current transaction is committed.

132 If the current transaction is rolled back, `func` will not be called.

133 """

134 get_connection(using).on_commit(func, robust)

135

136

137#################################

138# Decorators / context managers #

139#################################

140

141

142 class Atomic(ContextDecorator):

143 """

144 Guarantee the atomic execution of a given block.

145

146 An instance can be used either as a decorator or as a context manager.

147

148 When it's used as a decorator, __call__ wraps the execution of the

149 decorated function in the instance itself, used as a context manager.

150

151 When it's used as a context manager, __enter__ creates a transaction or a

152 savepoint, depending on whether a transaction is already in progress, and

153 __exit__ commits the transaction or releases the savepoint on normal exit,

154 and rolls back the transaction or to the savepoint on exceptions.

155

156 It's possible to disable the creation of savepoints if the goal is to

157 ensure that some code runs within a transaction without creating overhead.

158

159 A stack of savepoints identifiers is maintained as an attribute of the

160 connection. None denotes the absence of a savepoint.

161

162 This allows reentrancy even if the same AtomicWrapper is reused. For

163 example, it's possible to define `oa = atomic('other')` and use `@oa` or

164 `with oa:` multiple times.

165

166 Since database connections are thread-local, this is thread-safe.

167

168 An atomic block can be tagged as durable. In this case, raise a

169 RuntimeError if it's nested within another atomic block. This guarantees

170 that database changes in a durable block are committed to the database when

171 the block exists without error.

172

173 This is a private API.

174 """

175

176 def __init__(self, using, savepoint, durable):

177 self.using = using

178 self.savepoint = savepoint

179 self.durable = durable

180 self._from_testcase = False

181

182 def __enter__(self):

183 connection = get_connection(self.using)

184

185 if (

186 self.durable

187 and connection.atomic_blocks

188 and not connection.atomic_blocks[-1]._from_testcase

189 ):

190 raise RuntimeError(

191 "A durable atomic block cannot be nested within another "

192 "atomic block."

193 )

194 if not connection.in_atomic_block:

195 # Reset state when entering an outermost atomic block.

196 connection.commit_on_exit = True

197 connection.needs_rollback = False

198 if not connection.get_autocommit():

199 # Pretend we're already in an atomic block to bypass the code

200 # that disables autocommit to enter a transaction, and make a

201 # note to deal with this case in __exit__.

202 connection.in_atomic_block = True

203 connection.commit_on_exit = False

204

205 if connection.in_atomic_block:

206 # We're already in a transaction; create a savepoint, unless we

207 # were told not to or we're already waiting for a rollback. The

208 # second condition avoids creating useless savepoints and prevents

209 # overwriting needs_rollback until the rollback is performed.

210 if self.savepoint and not connection.needs_rollback:

211 sid = connection.savepoint()

212 connection.savepoint_ids.append(sid)

213 else:

214 connection.savepoint_ids.append(None)

215 else:

216 connection.set_autocommit(

217 False, force_begin_transaction_with_broken_autocommit=True

218 )

219 connection.in_atomic_block = True

220

221 if connection.in_atomic_block:

222 connection.atomic_blocks.append(self)

223

224 def __exit__(self, exc_type, exc_value, traceback):

225 connection = get_connection(self.using)

226

227 if connection.in_atomic_block:

228 connection.atomic_blocks.pop()

229

230 if connection.savepoint_ids:

231 sid = connection.savepoint_ids.pop()

232 else:

233 # Prematurely unset this flag to allow using commit or rollback.

234 connection.in_atomic_block = False

235

236 try:

237 if connection.closed_in_transaction:

238 # The database will perform a rollback by itself.

239 # Wait until we exit the outermost block.

240 pass

241

242 elif exc_type is None and not connection.needs_rollback:

243 if connection.in_atomic_block:

244 # Release savepoint if there is one

245 if sid is not None:

246 try:

247 connection.savepoint_commit(sid)

248 except DatabaseError:

249 try:

250 connection.savepoint_rollback(sid)

251 # The savepoint won't be reused. Release it to

252 # minimize overhead for the database server.

253 connection.savepoint_commit(sid)

254 except Error:

255 # If rolling back to a savepoint fails, mark for

256 # rollback at a higher level and avoid shadowing

257 # the original exception.

258 connection.needs_rollback = True

259 raise

260 else:

261 # Commit transaction

262 try:

263 connection.commit()

264 except DatabaseError:

265 try:

266 connection.rollback()

267 except Error:

268 # An error during rollback means that something

269 # went wrong with the connection. Drop it.

270 connection.close()

271 raise

272 else:

273 # This flag will be set to True again if there isn't a savepoint

274 # allowing to perform the rollback at this level.

275 connection.needs_rollback = False

276 if connection.in_atomic_block:

277 # Roll back to savepoint if there is one, mark for rollback

278 # otherwise.

279 if sid is None:

280 connection.needs_rollback = True

281 else:

282 try:

283 connection.savepoint_rollback(sid)

284 # The savepoint won't be reused. Release it to

285 # minimize overhead for the database server.

286 connection.savepoint_commit(sid)

287 except Error:

288 # If rolling back to a savepoint fails, mark for

289 # rollback at a higher level and avoid shadowing

290 # the original exception.

291 connection.needs_rollback = True

292 else:

293 # Roll back transaction

294 try:

295 connection.rollback()

296 except Error:

297 # An error during rollback means that something

298 # went wrong with the connection. Drop it.

299 connection.close()

300

301 finally:

302 # Outermost block exit when autocommit was enabled.

303 if not connection.in_atomic_block:

304 if connection.closed_in_transaction:

305 connection.connection = None

306 else:

307 connection.set_autocommit(True)

308 # Outermost block exit when autocommit was disabled.

309 elif not connection.savepoint_ids and not connection.commit_on_exit:

310 if connection.closed_in_transaction:

311 connection.connection = None

312 else:

313 connection.in_atomic_block = False

314

315

316 def atomic(using=None, savepoint=True, durable=False):

317 # Bare decorator: @atomic -- although the first argument is called

318 # `using`, it's actually the function being decorated.

319 if callable(using):

320 return Atomic(DEFAULT_DB_ALIAS, savepoint, durable)(using)

321 # Decorator: @atomic(...) or context manager: with atomic(...): ...

322 else:

323 return Atomic(using, savepoint, durable)