|
1 | 1 | from contextlib import ContextDecorator
|
2 | 2 |
|
3 |
| -from django.db import DEFAULT_DB_ALIAS, DatabaseError, Error |
| 3 | +from django.db import DEFAULT_DB_ALIAS, DatabaseError |
4 | 4 | from django.db.transaction import get_connection
|
5 | 5 |
|
6 | 6 |
|
@@ -30,82 +30,61 @@ class Atomic(ContextDecorator):
|
30 | 30 | `with oa:` multiple times.
|
31 | 31 |
|
32 | 32 | Since database connections are thread-local, this is thread-safe.
|
| 33 | +
|
| 34 | + Simplified from django.db.transaction. |
33 | 35 | """
|
34 | 36 |
|
35 | 37 | def __init__(self, using):
|
36 | 38 | self.using = using
|
37 | 39 |
|
38 | 40 | def __enter__(self):
|
39 | 41 | connection = get_connection(self.using)
|
40 |
| - if not connection.in_atomic_block_mongo: |
41 |
| - # Reset state when entering an outermost atomic block. |
42 |
| - connection.needs_rollback_mongo = False |
43 |
| - |
44 | 42 | if connection.in_atomic_block_mongo:
|
45 |
| - # We're already in a transaction. Increment the number of nested atomics. |
| 43 | + # If we're already in an atomic(), track the number of nested calls. |
46 | 44 | connection.nested_atomics += 1
|
47 | 45 | else:
|
| 46 | + # Start a transaction for the outermost atomic(). |
48 | 47 | connection._start_transaction()
|
49 | 48 | connection.in_atomic_block_mongo = True
|
50 | 49 |
|
51 |
| - if connection.in_atomic_block_mongo: |
52 |
| - connection.atomic_blocks_mongo.append(self) |
53 |
| - |
54 | 50 | def __exit__(self, exc_type, exc_value, traceback):
|
55 | 51 | connection = get_connection(self.using)
|
56 |
| - |
57 |
| - if connection.in_atomic_block_mongo: |
58 |
| - connection.atomic_blocks_mongo.pop() |
59 |
| - |
60 | 52 | if connection.nested_atomics:
|
61 | 53 | connection.nested_atomics -= 1
|
62 | 54 | else:
|
63 |
| - # Prematurely unset this flag to allow using commit or rollback. |
64 | 55 | connection.in_atomic_block_mongo = False
|
65 | 56 | try:
|
66 |
| - if exc_type is None and not connection.needs_rollback_mongo: |
| 57 | + if exc_type is None: |
| 58 | + # atomic() exited without an error. |
67 | 59 | if connection.in_atomic_block_mongo:
|
68 |
| - # Release savepoint if there is one |
| 60 | + # Do nothing for an inner atomic(). |
69 | 61 | pass
|
70 | 62 | else:
|
71 |
| - # Commit transaction |
| 63 | + # Commit transaction. |
72 | 64 | try:
|
73 | 65 | connection.commit_mongo()
|
74 | 66 | except DatabaseError:
|
75 |
| - try: |
76 |
| - connection.rollback_mongo() |
77 |
| - except Error: |
78 |
| - # An error during rollback means that something |
79 |
| - # went wrong with the connection. Drop it. |
80 |
| - connection.close() |
81 |
| - raise |
| 67 | + connection.rollback_mongo() |
82 | 68 | else:
|
83 |
| - # This flag will be set to True again if there isn't a savepoint |
84 |
| - # allowing to perform the rollback at this level. |
85 |
| - connection.needs_rollback_mongo = False |
| 69 | + # atomic() exited with an error. |
86 | 70 | if connection.in_atomic_block_mongo:
|
87 |
| - # Mark for rollback |
88 |
| - connection.needs_rollback_mongo = True |
| 71 | + # Do nothing for an inner atomic(). |
| 72 | + pass |
89 | 73 | else:
|
90 |
| - # Roll back transaction |
91 |
| - try: |
92 |
| - connection.rollback_mongo() |
93 |
| - except Error: |
94 |
| - # An error during rollback means that something |
95 |
| - # went wrong with the connection. Drop it. |
96 |
| - connection.close() |
| 74 | + # Rollback transaction. |
| 75 | + connection.rollback_mongo() |
97 | 76 | finally:
|
98 |
| - # Outermost block exit |
99 | 77 | if (
|
100 | 78 | not connection.in_atomic_block_mongo
|
101 | 79 | and connection.run_commit_hooks_on_set_autocommit_on
|
102 | 80 | ):
|
| 81 | + # Run on_commit() callbacks after outermost atomic() |
103 | 82 | connection.run_and_clear_commit_hooks()
|
104 | 83 |
|
105 | 84 |
|
106 | 85 | def atomic(using=None):
|
107 |
| - # Bare decorator: @atomic -- although the first argument is called |
108 |
| - # `using`, it's actually the function being decorated. |
| 86 | + # Bare decorator: @atomic -- although the first argument is called `using`, it's |
| 87 | + # actually the function being decorated. |
109 | 88 | if callable(using):
|
110 | 89 | return Atomic(DEFAULT_DB_ALIAS)(using)
|
111 | 90 | # Decorator: @atomic(...) or context manager: with atomic(...): ...
|
|
0 commit comments