|
2 | 2 | import traceback |
3 | 3 | from contextlib import ExitStack |
4 | 4 |
|
5 | | -from django.db import transaction |
| 5 | +from django.db import router, transaction |
| 6 | +from django.db import DEFAULT_DB_ALIAS |
6 | 7 | from django.utils.translation import gettext as _ |
7 | 8 |
|
8 | 9 | from core.signals import clear_events |
| 10 | +from dcim.models import Device |
9 | 11 | from extras.models import Script as ScriptModel |
| 12 | +from netbox.context_managers import event_tracking |
10 | 13 | from netbox.jobs import JobRunner |
11 | 14 | from netbox.registry import registry |
12 | 15 | from utilities.exceptions import AbortScript, AbortTransaction |
@@ -42,10 +45,21 @@ def run_script(self, script, request, data, commit): |
42 | 45 | # A script can modify multiple models so need to do an atomic lock on |
43 | 46 | # both the default database (for non ChangeLogged models) and potentially |
44 | 47 | # any other database (for ChangeLogged models) |
45 | | - with transaction.atomic(): |
46 | | - script.output = script.run(data, commit) |
47 | | - if not commit: |
48 | | - raise AbortTransaction() |
| 48 | + changeloged_db = router.db_for_write(Device) |
| 49 | + with transaction.atomic(using=DEFAULT_DB_ALIAS): |
| 50 | + # If branch database is different from default, wrap in a second atomic transaction |
| 51 | + # Note: Don't add any extra code between the two atomic transactions, |
| 52 | + # otherwise the changes might get committed to the default database |
| 53 | + # if there are any raised exceptions. |
| 54 | + if changeloged_db != DEFAULT_DB_ALIAS: |
| 55 | + with transaction.atomic(using=changeloged_db): |
| 56 | + script.output = script.run(data, commit) |
| 57 | + if not commit: |
| 58 | + raise AbortTransaction() |
| 59 | + else: |
| 60 | + script.output = script.run(data, commit) |
| 61 | + if not commit: |
| 62 | + raise AbortTransaction() |
49 | 63 | except AbortTransaction: |
50 | 64 | script.log_info(message=_("Database changes have been reverted automatically.")) |
51 | 65 | if script.failed: |
@@ -108,14 +122,14 @@ def run(self, data, request=None, commit=True, **kwargs): |
108 | 122 | script.request = request |
109 | 123 | self.logger.debug(f"Request ID: {request.id if request else None}") |
110 | 124 |
|
111 | | - # Execute the script. If commit is True, wrap it with the event_tracking context manager to ensure we process |
112 | | - # change logging, event rules, etc. |
113 | 125 | if commit: |
114 | 126 | self.logger.info("Executing script (commit enabled)") |
115 | | - with ExitStack() as stack: |
116 | | - for request_processor in registry['request_processors']: |
117 | | - stack.enter_context(request_processor(request)) |
118 | | - self.run_script(script, request, data, commit) |
119 | 127 | else: |
120 | 128 | self.logger.warning("Executing script (commit disabled)") |
| 129 | + |
| 130 | + with ExitStack() as stack: |
| 131 | + for request_processor in registry['request_processors']: |
| 132 | + if not commit and request_processor is event_tracking: |
| 133 | + continue |
| 134 | + stack.enter_context(request_processor(request)) |
121 | 135 | self.run_script(script, request, data, commit) |
0 commit comments