|
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