Coverage for netbox_data_import/jobs.py: 94.25%
261 statements
« prev ^ index » next coverage.py v7.17.0, created at 2026-10-30 09:47 +0000
1# SPDX-License-Identifier: Apache-2.0
2# Copyright (C) 2025 Marcin Zieba <marcinpsk@gmail.com>
3"""Native NetBox background jobs for data imports."""
5import logging
6from dataclasses import dataclass, replace
7from datetime import timedelta
8from functools import partial
9from typing import Any, NoReturn
11from django.core.exceptions import ValidationError
12from django.db import DatabaseError, connection, transaction
13from django.utils import timezone
14from django_pg_utils import advisory_lock
15from redis.exceptions import ConnectionError as RedisConnectionError
16from redis.exceptions import TimeoutError as RedisTimeoutError
17from rq import get_current_job
18from rq.exceptions import InvalidJobOperation, NoSuchJobError
19from rq.job import Job as RQJob, JobStatus
20from rq.utils import as_text
21from rq.worker import Worker, WorkerStatus
23from core.exceptions import JobFailed
24from netbox.context_managers import event_tracking
25from netbox.jobs import JobRunner, system_job
26from utilities.request import NetBoxFakeRequest
28from . import branching
29from .adapters import SourceUnreadable, UnknownSourceAdapter
30from .import_engine import (
31 EngineConfigurationError,
32 ImportEngine,
33 PreconditionFailed,
34 SelectionError,
35 StalePlan,
36 StaleSourceDocument,
37 operator_failure_message,
38)
39from .models import (
40 ExecutionOutcome,
41 ImportExecution,
42 ImportProfile,
43 ProposalFailureReason,
44 SourceDocument,
45 validate_registered_adapter,
46)
47from .netbox_reader import PlanningTargetUnavailable
48from .object_permissions import ObjectPermissionDenied
49from .plan import PlanError
52_PROGRESS_REPORT_INTERVAL = 25
53logger = logging.getLogger(__name__)
55IMPORT_TASK_LOST = "The import task is no longer queued or running. Re-read its preview to recover."
56PROPOSAL_RUN_ABANDONED = "An earlier delivery of this Job claimed its proposal and ended without a result."
57QUEUE_PUSH_GRACE = timedelta(minutes=1)
58QUEUE_UNREADABLE = "The job queue cannot be read now."
61def _job_lock(job):
62 """Name the lock a worker holds while it delivers a Job, which recovery must get before it releases one."""
63 with connection.cursor() as cursor:
64 cursor.execute("SELECT hashtextextended(%s, 0)", [f"netbox-data-import-job:{job.job_id}"])
65 return cursor.fetchone()[0]
68def import_queue_task(job):
69 """Fetch queue evidence without removing a missing task's ID from Redis."""
70 import django_rq
72 try:
73 queue = django_rq.get_queue(job.queue_name)
74 except KeyError:
75 return None
76 try:
77 return RQJob.fetch(str(job.job_id), connection=queue.connection, serializer=queue.serializer)
78 except NoSuchJobError:
79 return None
82def _task_lost(job, rq_job) -> bool:
83 """Return whether an active Job's queue task can no longer run it."""
84 from core.choices import JobStatusChoices
86 if rq_job is None:
87 # The Job commits before its queue push, so a young pending Job may have no task yet.
88 return not (job.status == JobStatusChoices.STATUS_PENDING and job.created > timezone.now() - QUEUE_PUSH_GRACE)
89 try:
90 status = rq_job.get_status(refresh=True)
91 except InvalidJobOperation:
92 return True
93 return status in (
94 None,
95 JobStatus.FINISHED,
96 JobStatus.FAILED,
97 JobStatus.CANCELED,
98 JobStatus.STOPPED,
99 )
102def import_job_abandoned(job) -> bool:
103 """Read whether a native active import has lost its queue task, without changing either."""
104 from core.choices import JobStatusChoices
106 return job.status in JobStatusChoices.ENQUEUED_STATE_CHOICES and _task_lost(job, import_queue_task(job))
109QUEUED, RUNNING, LOST, ENDED = "queued", "running", "lost", "ended"
112@dataclass(frozen=True)
113class ImportJobStatus:
114 """What one import Job and its queue task say now, for a page that shows the Job."""
116 job: Any
117 state: str
118 phase: str = ""
119 processed: int = 0
120 total: int = 0
121 # Only for a queued Job whose task is in its queue: its place there (1 is next), and its queue's workers.
122 position: int | None = None
123 workers: int | None = None
124 busy: int = 0
125 queue_readable: bool = True
127 @property
128 def active(self) -> bool:
129 """Return whether the Job still waits in its queue or runs."""
130 return self.state in (QUEUED, RUNNING)
132 @property
133 def queue_note(self) -> str:
134 """Return what a page says when Redis did not answer, or "" when it did."""
135 return "" if self.queue_readable else QUEUE_UNREADABLE
138def import_job_status(job) -> ImportJobStatus:
139 """Read one import Job, its queue task and its queue, without changing any of them."""
140 from core.choices import JobStatusChoices
142 data = job.data or {}
143 status = ImportJobStatus(
144 job,
145 ENDED,
146 phase=str(data.get("phase") or ""),
147 processed=int(data.get("processed") or 0),
148 total=int(data.get("total") or 0),
149 )
150 if job.status not in JobStatusChoices.ENQUEUED_STATE_CHOICES:
151 return status
152 try:
153 return _with_queue_evidence(status)
154 except (RedisConnectionError, RedisTimeoutError):
155 # An outage says nothing about the task, so the Job row alone decides the state.
156 state = RUNNING if job.status == JobStatusChoices.STATUS_RUNNING else QUEUED
157 return replace(status, state=state, queue_readable=False)
160def _with_queue_evidence(status: ImportJobStatus) -> ImportJobStatus:
161 """Add what the Job's queue task and its queue say; Redis errors reach the caller."""
162 from core.choices import JobStatusChoices
164 job = status.job
165 rq_job = import_queue_task(job)
166 if _task_lost(job, rq_job):
167 return replace(status, state=LOST)
168 if rq_job is None:
169 # The queue push follows the commit, so the queue knows nothing of this Job yet.
170 return replace(status, state=QUEUED)
171 # The worker publishes row progress to the task, not to the Job row.
172 status = replace(
173 status,
174 phase=str(rq_job.meta.get("phase") or status.phase),
175 processed=int(rq_job.meta.get("processed") or status.processed),
176 total=int(rq_job.meta.get("total") or status.total),
177 )
178 if job.status == JobStatusChoices.STATUS_RUNNING:
179 return replace(status, state=RUNNING)
180 return _with_queue_place(replace(status, state=QUEUED), rq_job)
183def _with_queue_place(status: ImportJobStatus, rq_job) -> ImportJobStatus:
184 """Add the task's place in its own queue, and the live workers of that queue; a read that writes nothing."""
185 import django_rq
187 queue = django_rq.get_queue(status.job.queue_name)
188 # Worker.all() removes the registration of an expired worker, so read each worker hash directly.
189 states = [queue.connection.hget(key, "state") for key in Worker.all_keys(queue=queue)]
190 live = [as_text(state) for state in states if state is not None]
191 position = rq_job.get_position()
192 return replace(
193 status,
194 position=None if position is None else position + 1,
195 workers=len(live),
196 busy=live.count(WorkerStatus.BUSY),
197 )
200def _try_job_lock(job) -> bool:
201 """Take the Job lock until the transaction ends, or return False at once when a worker holds it."""
202 if not connection.in_atomic_block:
203 raise RuntimeError("The Job lock requires the transaction of the recovery that takes it.")
204 # Never wait for a worker while holding other locks. PostgreSQL holds this lock through commit or rollback.
205 with connection.cursor() as cursor:
206 cursor.execute("SELECT pg_try_advisory_xact_lock(%s)", [_job_lock(job)])
207 return cursor.fetchone()[0]
210def recover_abandoned_import_job(job) -> None:
211 """Fail a lost task under the profile lock, unless a worker still owns its delivery."""
212 from core.choices import JobStatusChoices
213 from core.models import Job
215 if _try_job_lock(job):
216 job.refresh_from_db()
217 if import_job_abandoned(job):
218 Job.objects.filter(pk=job.pk, status__in=JobStatusChoices.ENQUEUED_STATE_CHOICES).update(
219 status=JobStatusChoices.STATUS_ERRORED,
220 completed=timezone.now(),
221 error=IMPORT_TASK_LOST,
222 data={**(job.data or {}), "phase": "failed", "message": IMPORT_TASK_LOST},
223 )
226def cancel_queued_import_job(job, message: str) -> bool:
227 """Fail a Job that no worker has started, and cancel its queue task after the commit.
229 Return False, and change nothing, when a worker holds the Job lock or the Job has left the queue.
230 """
231 from core.choices import JobStatusChoices
232 from core.models import Job
234 if not _try_job_lock(job):
235 return False
236 job.refresh_from_db()
237 waiting = (JobStatusChoices.STATUS_PENDING, JobStatusChoices.STATUS_SCHEDULED)
238 if job.status not in waiting or job.started is not None:
239 return False
240 # NetBox has no cancelled status; a failed Job with this phase is one the operator stopped.
241 Job.objects.filter(pk=job.pk).update(
242 status=JobStatusChoices.STATUS_FAILED,
243 completed=timezone.now(),
244 error=message,
245 data={**(job.data or {}), "phase": "cancelled", "message": message},
246 )
247 transaction.on_commit(partial(_cancel_queue_task, job))
248 return True
251def _cancel_queue_task(job) -> None:
252 """Take a cancelled Job's task out of its queue; a worker that already took it skips the Job."""
253 try:
254 rq_job = import_queue_task(job)
255 if rq_job is not None and rq_job.get_status(refresh=True) in (
256 JobStatus.QUEUED,
257 JobStatus.SCHEDULED,
258 JobStatus.DEFERRED,
259 ):
260 rq_job.cancel()
261 except (InvalidJobOperation, NoSuchJobError):
262 # Another actor, such as NetBox Job.delete(), removed the task first, which is the goal.
263 logger.debug("The queue task of cancelled import Job %s was already gone.", job.pk)
264 except (RedisConnectionError, RedisTimeoutError):
265 # The Job row is the record: ImportJobRunner.handle runs only a Job that is still enqueued.
266 logger.warning("The queue task of cancelled import Job %s stays in its queue.", job.pk)
269class OwnedJobRunner(JobRunner):
270 """A runner whose worker holds the Job lock for the whole delivery, so recovery cannot release a live Job."""
272 @classmethod
273 def handle(cls, job, *args, **kwargs):
274 """Skip late or duplicate delivery without adding a transaction around the durable audit."""
275 from core.choices import JobStatusChoices
276 from core.models import Job
278 with advisory_lock(_job_lock(job)):
279 try:
280 job.refresh_from_db()
281 except Job.DoesNotExist:
282 return
283 if job.status in JobStatusChoices.ENQUEUED_STATE_CHOICES:
284 super().handle(job, *args, **kwargs)
287class ImportJobRunner(OwnedJobRunner):
288 """Validate and execute one import while publishing row progress to RQ."""
290 job_type = "netbox_data_import.import"
292 class Meta:
293 name = "Data Import"
295 def _save_data(self, **values):
296 """Merge values into the native Job data."""
297 data = {**(self.job.data or {}), **values}
298 data.pop("accepted_plan", None)
299 self.job.data = data
300 self.job.save(update_fields=["data"])
302 def _fail(self, message) -> NoReturn:
303 """Record a recoverable failure and stop the native Job."""
304 values = {"phase": "failed", "message": message}
305 execution = ImportExecution.objects.filter(job=self.job).first()
306 if execution is not None:
307 values["import_execution_id"] = execution.pk
308 self._save_data(**values)
309 raise JobFailed
311 @staticmethod
312 def _publish_progress(processed, total):
313 """Publish progress outside the database transaction through RQ metadata."""
314 if processed not in (0, total) and processed % _PROGRESS_REPORT_INTERVAL:
315 return
316 rq_job = get_current_job()
317 if rq_job is None:
318 return
319 rq_job.meta.update({"processed": processed, "total": total, "phase": "importing"})
320 rq_job.save_meta()
322 def _change_logging_request(self, user) -> NetBoxFakeRequest:
323 """Return the request NetBox records this Job's changes and events under, with the Job UUID as its id."""
324 return NetBoxFakeRequest(
325 {
326 "META": {},
327 "COOKIES": {},
328 "POST": {},
329 "GET": {},
330 "FILES": {},
331 "user": user,
332 "method": "POST",
333 "path": "",
334 "path_info": "",
335 "id": self.job.job_id,
336 }
337 )
339 def run(self, profile_id, source_document_id, accepted_plan, selection, idempotency_key):
340 """Execute one accepted Import Plan as the Job's actor."""
341 try:
342 branching.fail_job_in_branch(self)
343 except JobFailed as exc:
344 self._save_data(phase="failed", message=str(exc))
345 raise
346 user = self.job.user
347 if user is None:
348 self._fail("The user who started this import is no longer available.")
350 profile = ImportProfile.objects.restrict(user, "change").filter(pk=profile_id).first()
351 if profile is None:
352 self._fail("The import profile is no longer available.")
353 try:
354 validate_registered_adapter(profile)
355 except ValidationError as exc:
356 self._fail(operator_failure_message(exc))
357 source_document = SourceDocument.objects.filter(pk=source_document_id, profile=profile).first()
358 if source_document is None:
359 self._fail("The stored source is no longer available. Upload it again.")
361 self._save_data(phase="validating")
362 progress = {"processed": 0, "total": 0}
364 def publish_progress(processed, total):
365 """Remember final progress and publish the bounded RQ updates."""
366 progress.update(processed=processed, total=total)
367 self._publish_progress(processed, total)
369 # Only event_tracking: the other request processors (netbox-branching) must not run in a worker.
370 with event_tracking(self._change_logging_request(user)):
371 try:
372 execution = ImportEngine.execute(
373 profile,
374 source_document,
375 accepted_plan,
376 selection,
377 idempotency_key,
378 user,
379 job=self.job,
380 progress_callback=publish_progress,
381 )
382 except ImportProfile.DoesNotExist:
383 self._fail("The import profile is no longer available.")
384 except DatabaseError as exc:
385 logger.exception("Import execution failed with a database error")
386 self._fail(operator_failure_message(exc))
387 except (EngineConfigurationError, SourceUnreadable, UnknownSourceAdapter) as exc:
388 logger.exception("Import execution failed before its source could be planned")
389 self._fail(operator_failure_message(exc))
390 except (
391 ObjectPermissionDenied,
392 PlanError,
393 PlanningTargetUnavailable,
394 PreconditionFailed,
395 SelectionError,
396 StalePlan,
397 StaleSourceDocument,
398 ValidationError,
399 ) as exc:
400 self._fail(operator_failure_message(exc))
401 # Raising here leaves the block before event_tracking flushes the queued events.
402 if execution.outcome != ExecutionOutcome.SUCCEEDED:
403 reason = (execution.failure_detail or {}).get("reason") or execution.outcome or "unknown"
404 self._fail(f"The accepted import execution did not succeed ({reason}).")
405 self._save_data(
406 phase="completed",
407 processed=progress["processed"],
408 total=progress["total"],
409 import_execution_id=execution.pk,
410 )
413def retained_sync_jobs(user, profile_id, document_id):
414 """Return the active per-trace sync Jobs of one source, whichever preview queued them.
416 The Job rows are the record, so a request that lost a race to the coordinator still sees the sync.
417 """
418 from core.choices import JobStatusChoices
420 return ImportJobRunner.get_jobs().filter(
421 user=user,
422 data__job_type=ImportJobRunner.job_type,
423 data__keeps_preview=True,
424 data__profile_id=profile_id,
425 data__source_document_id=document_id,
426 status__in=JobStatusChoices.ENQUEUED_STATE_CHOICES,
427 )
430@system_job(interval=60 * 24)
431class SourceDocumentRetentionJob(JobRunner):
432 """Reclaim stored uploads no Import Execution references (section 9.1), and expire old previews."""
434 class Meta:
435 name = "Data Import source document retention"
437 @staticmethod
438 def purge() -> int:
439 """Apply the retention rules and return the number of deleted documents."""
440 return SourceDocument.purge_unreferenced()
442 def run(self, *args, **kwargs):
443 """Run one retention pass; a preview expires with the 30 days its Source Document is kept."""
444 from .preview_coordinator import expire_previews
446 branching.fail_job_in_branch(self)
447 expire_previews()
448 return self.purge()
451class ResolutionProposalJob(OwnedJobRunner):
452 """Run inference for one proposal id through the worker service."""
454 job_type = "netbox_data_import.resolution_proposal"
456 class Meta:
457 name = "Resolution Proposal"
459 @classmethod
460 def handle(cls, job, *args, **kwargs):
461 """Run the proposal, then give its request slot to the next waiting proposal, however the run ended."""
462 from .proposal_dispatch import dispatch_proposals
464 try:
465 super().handle(job, *args, **kwargs)
466 finally:
467 # The Job is terminal and its lock is released here, so the dispatcher no longer counts its slot.
468 try:
469 dispatch_proposals()
470 except (RedisConnectionError, RedisTimeoutError):
471 logger.warning("The proposal queue did not answer; the next card read dispatches the next proposal.")
473 def run(self, proposal_id):
474 """Resolve backend configuration on the worker after claiming the proposal."""
475 from .proposal_jobs import run_proposal
476 from .resolution_proposals import fail_abandoned_proposal, fail_proposal
478 try:
479 branching.fail_job_in_branch(self)
480 except JobFailed:
481 # The row is main-only, so this write lands in main and frees the field for a new request.
482 fail_proposal(proposal_id, reason=ProposalFailureReason.BRANCH_ACTIVE)
483 raise
484 if run_proposal(proposal_id):
485 return
486 # The Job lock makes this the only delivery, so a row this Job claimed earlier has no worker left.
487 if fail_abandoned_proposal(proposal_id, job=self.job):
488 self.logger.error(PROPOSAL_RUN_ABANDONED)
489 raise JobFailed(PROPOSAL_RUN_ABANDONED)
492__all__ = (
493 "ImportJobRunner",
494 "ResolutionProposalJob",
495 "SourceDocumentRetentionJob",
496)