Coverage for netbox_interface_name_rules / transactions.py: 98.45%
143 statements
« prev ^ index » next coverage.py v7.13.4, created at 2026-10-16 10:13 +0000
1# SPDX-License-Identifier: Apache-2.0
2# Copyright (C) 2025 Marcin Zieba <marcinpsk@gmail.com>
3"""The one owner of the plugin's database connections and transaction state.
5Every plugin write runs in a write scope. On main the scope holds ``default``. In a netbox-branching
6branch it holds ``default`` and the branch alias: netbox-branching writes the ChangeDiff rows of a
7branch change on ``default``, so a block opens both, ``default`` outside the branch.
8"""
10from contextlib import ExitStack, contextmanager
11from contextvars import ContextVar
12from dataclasses import dataclass
13from functools import partial
15from django.db import DEFAULT_DB_ALIAS, DatabaseError, connections, router, transaction
16from netbox.context import events_queue
18# PostgreSQL does not detect a lock cycle through the two sessions of one request in a branch.
19LOCK_TIMEOUT = "10s"
20READ_LOCK_TIMEOUT = "SHOW lock_timeout"
21# Sets the session value, which the setup and each restore use.
22SET_LOCK_TIMEOUT = "SELECT set_config('lock_timeout', %s, false)"
24# The aliases of the open write scope, or None outside one.
25_scope_aliases = ContextVar("write_scope_aliases", default=None)
26# The event queues that enclose the current block, outermost first.
27_enclosing_queues = ContextVar("atomic_with_events_enclosing_queues", default=())
30def write_alias() -> str:
31 """Return the alias that NetBox's router gives for a write of an interface."""
32 from dcim.models import Interface
34 return router.db_for_write(Interface)
37@contextmanager
38def write_scope(expected_alias=None):
39 """Pin the aliases of one plugin operation and yield them: ``default``, then the branch alias in a branch.
41 It raises before any query when the write alias is not *expected_alias*. A nested scope joins the
42 open one, and raises when its write alias differs. In a branch the outermost scope sets
43 ``lock_timeout`` on both connections and restores each value when it exits; on main it runs no query.
44 """
45 alias = write_alias()
46 if expected_alias is not None and alias != expected_alias:
47 raise RuntimeError(f"The write alias is {alias!r}, but the operation expects {expected_alias!r}.")
48 enclosing = _scope_aliases.get()
49 if enclosing is not None:
50 if alias != enclosing[-1]:
51 raise RuntimeError(f"The write alias is {alias!r}, but the open write scope writes to {enclosing[-1]!r}.")
52 yield enclosing
53 return
54 aliases = (DEFAULT_DB_ALIAS,) if alias == DEFAULT_DB_ALIAS else (DEFAULT_DB_ALIAS, alias)
55 token = _scope_aliases.set(aliases)
56 try:
57 if len(aliases) == 1:
58 yield aliases
59 else:
60 with _lock_timeout(aliases):
61 yield aliases
62 finally:
63 _scope_aliases.reset(token)
66@contextmanager
67def _lock_timeout(aliases):
68 """Set ``lock_timeout`` on the connection of each of *aliases*, and restore each value when the block exits.
70 A caller's open transaction holds both the set and the restore, so either outcome of it keeps the
71 value from before it. A setup that fails restores the connections it changed.
72 """
73 previous = {}
74 try:
75 for alias in aliases:
76 connection = connections[alias]
77 with connection.cursor() as cursor:
78 cursor.execute(READ_LOCK_TIMEOUT)
79 (value,) = cursor.fetchone()
80 cursor.execute(SET_LOCK_TIMEOUT, [LOCK_TIMEOUT])
81 previous[alias] = value
82 yield
83 finally:
84 _restore_lock_timeouts(previous)
87def _restore_lock_timeouts(previous):
88 """Restore each connection's ``lock_timeout`` from *previous*; discard each connection whose restore fails."""
89 failures = []
90 for alias, value in previous.items():
91 connection = connections[alias]
92 try:
93 with connection.cursor() as cursor:
94 cursor.execute(SET_LOCK_TIMEOUT, [value])
95 except DatabaseError as error:
96 _discard(connection)
97 failures.append((alias, error))
98 if failures:
99 names = ", ".join(repr(alias) for alias, _ in failures)
100 raise RuntimeError(f"Could not restore lock_timeout on {names}; the connection is closed.") from failures[0][1]
103def _discard(connection):
104 """Close the session of *connection*, so that no pool can hand out its changed setting."""
105 session = connection.connection
106 try:
107 if session is not None: 107 ↛ 110line 107 didn't jump to line 110 because the condition on line 107 was always true
108 session.close()
109 finally:
110 connection.close()
113class _CommitJoin:
114 """Run a callback once, when every alias of the scope acknowledged the commit of its transaction."""
116 def __init__(self, callback, aliases):
117 self._callback = callback
118 # The whole set is pending before the first registration: a connection in autocommit acknowledges at once.
119 self._pending = set(aliases)
120 self._fired = False
122 def acknowledge(self, alias):
123 self._pending.discard(alias)
124 if self._pending or self._fired:
125 return
126 self._fired = True
127 self._callback()
130def on_commit(callback):
131 """Run *callback* once, after the transactions open at the call commit on every alias of the write scope.
133 Each alias acknowledges through Django's ``on_commit``, which drops the acknowledgement when its
134 transaction, or a savepoint around the call, rolls back; *callback* then never runs.
135 """
136 aliases = _scope_aliases.get()
137 if aliases is None:
138 raise RuntimeError("on_commit() runs only inside a write scope.")
139 if len(aliases) == 1:
140 transaction.on_commit(callback, using=aliases[0])
141 return
142 join = _CommitJoin(callback, aliases)
143 for alias in aliases:
144 transaction.on_commit(partial(join.acknowledge, alias), using=alias)
147@dataclass(frozen=True)
148class Block:
149 """An open ``atomic_with_events()`` block, with one atomic block on each alias of its write scope."""
151 aliases: tuple[str, ...]
153 def set_rollback(self):
154 """Mark the block for rollback on every alias."""
155 for alias in self.aliases:
156 transaction.set_rollback(True, using=alias)
159@contextmanager
160def atomic_with_events():
161 """Run the block atomically on every alias of the write scope, with its own NetBox event queue.
163 It joins the open write scope or opens one. It opens ``default`` first and the branch alias inside
164 it, so the branch commits first. The block keeps its queued events when the branch block commits
165 (outermost) or is released (nested) with no alias marked for rollback, and drops them otherwise.
166 """
167 with write_scope() as aliases:
168 enclosing = (*_enclosing_queues.get(), events_queue.get())
169 enclosing_token = _enclosing_queues.set(enclosing)
170 queue_token = events_queue.set({})
171 write = aliases[-1]
172 outermost = transaction.get_autocommit(using=write)
173 commit_marks = []
174 released = False
175 try:
176 with ExitStack() as blocks:
177 for alias in aliases:
178 blocks.enter_context(transaction.atomic(using=alias))
179 if outermost:
180 # The first commit callback of the write connection: it runs after COMMIT, before a callback can raise.
181 transaction.on_commit(lambda: commit_marks.append(True), using=write)
182 yield Block(aliases)
183 rollback_marked = any(transaction.get_rollback(using=alias) for alias in aliases)
184 released = not rollback_marked
185 finally:
186 events = events_queue.get()
187 events_queue.reset(queue_token)
188 _enclosing_queues.reset(enclosing_token)
189 # COMMIT runs only at the exit of the outermost block, and it can fail there or in a callback after it.
190 if commit_marks if outermost else released:
191 _keep(events)
192 else:
193 _point_at_saved_rows(events, enclosing)
196def _point_at_saved_rows(dropped, enclosing):
197 """Point each enclosing event of an object that the rolled-back block queued at a new copy of its saved row.
199 NetBox 4.5 and later serialize the queued instance at the flush, and the block can have changed it.
200 """
201 from core.events import OBJECT_DELETED
203 for event in (queue[key] for key in dropped for queue in enclosing if key in queue):
204 if event["event_type"] == OBJECT_DELETED or "object" not in event:
205 continue
206 model = event["object_type"].model_class()
207 row = model._base_manager.filter(pk=event["object_id"]).first()
208 if row is None:
209 raise RuntimeError(
210 f"{model._meta.label} {event['object_id']} has a queued event but no row after a rollback"
211 )
212 if hasattr(event, "refresh_serialization_source"):
213 event.refresh_serialization_source(row)
214 else:
215 event["object"] = row
216 if "data" in event: 216 ↛ 217line 216 didn't jump to line 217 because the condition on line 216 was never true
217 del event["data"]
220def _keep(events):
221 """Add the events of a committed block to the enclosing queue, coalesced as NetBox's enqueue_event() does."""
222 if not events:
223 return
224 from core.events import OBJECT_DELETED
226 queue = events_queue.get()
227 for key, event in events.items():
228 earlier = queue.get(key)
229 if earlier is not None:
230 # The object keeps its first before-state and event type, unless the block deleted it.
231 event["snapshots"]["prechange"] = earlier["snapshots"]["prechange"]
232 if event["event_type"] != OBJECT_DELETED:
233 event["event_type"] = earlier["event_type"]
234 queue[key] = event
235 events_queue.set(queue)