135 lines
4.7 KiB
Python
135 lines
4.7 KiB
Python
"""One poll of one source, and the two housekeeping jobs around it.
|
|
|
|
The bookkeeping is the whole point of sharing this: every branch out of a poll
|
|
has to leave the source with a next_poll_at, a last_status and a counter, or the
|
|
dispatcher either re-runs it every minute forever or never runs it again. Both
|
|
happened while this logic lived in two places.
|
|
|
|
A service supplies a :class:`Poller`; the Celery tasks stay in the service so the
|
|
beat schedule keeps naming them.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import traceback
|
|
from dataclasses import dataclass, field
|
|
from datetime import timedelta
|
|
from typing import Any, Callable
|
|
|
|
from django.utils import timezone
|
|
|
|
from yksa_kit.choices import RunStatus
|
|
|
|
|
|
@dataclass
|
|
class Poller:
|
|
"""How one service polls.
|
|
|
|
``persist`` is called per fetched record and returns True when the record is
|
|
new. ``track`` is called per record with a dict of extra source fields to
|
|
write on success — tdas uses it to carry ``last_rx_at`` forward. Per record
|
|
and not per run, because a poll can yield more frames than fit in memory.
|
|
"""
|
|
|
|
source_model: type
|
|
run_model: type
|
|
adapters: dict
|
|
persist: Callable[[Any, Any], bool]
|
|
track: Callable[[Any, Any, dict], None] | None = None
|
|
#: Kinds that are fed from outside and must never be polled.
|
|
unpollable_kinds: tuple[str, ...] = field(default_factory=tuple)
|
|
|
|
|
|
def run_source(poller: Poller, source_id: int) -> dict:
|
|
"""Poll one source now. Synchronous, so the admin action can call it too."""
|
|
src = poller.source_model.objects.get(pk=source_id)
|
|
if not src.is_enabled:
|
|
return {"source_id": source_id, "skipped": True, "reason": "disabled"}
|
|
|
|
adapter_cls = poller.adapters.get(src.kind)
|
|
if adapter_cls is None:
|
|
return _skip(poller, src, f"no adapter registered for kind={src.kind}")
|
|
|
|
run = poller.run_model.objects.create(source=src, status=RunStatus.RUNNING)
|
|
fetched = 0
|
|
new_count = 0
|
|
extra: dict[str, Any] = {}
|
|
try:
|
|
for record in adapter_cls(src).fetch():
|
|
fetched += 1
|
|
if poller.persist(src, record):
|
|
new_count += 1
|
|
if poller.track:
|
|
poller.track(src, record, extra)
|
|
except NotImplementedError as exc:
|
|
# An adapter that cannot run for this configuration is not a failure;
|
|
# counting it as one would page someone for a source nobody enabled.
|
|
run.mark_skipped(str(exc))
|
|
_mark_source(poller, src, RunStatus.SKIPPED, error=str(exc))
|
|
return {"source_id": source_id, "status": RunStatus.SKIPPED, "fetched": 0, "new": 0}
|
|
except Exception as exc:
|
|
run.mark_failed(f"{exc}\n{traceback.format_exc()}")
|
|
_mark_source(poller, src, RunStatus.FAILED, error=str(exc), count_failure=True)
|
|
raise
|
|
|
|
run.mark_success(fetched=fetched, new=new_count)
|
|
_mark_source(poller, src, RunStatus.SUCCESS, count_success=True, extra=extra)
|
|
return {
|
|
"source_id": source_id,
|
|
"status": RunStatus.SUCCESS,
|
|
"fetched": fetched,
|
|
"new": new_count,
|
|
}
|
|
|
|
|
|
def due_source_ids(poller: Poller) -> list[int]:
|
|
queryset = poller.source_model.objects.filter(
|
|
is_enabled=True, next_poll_at__lte=timezone.now()
|
|
)
|
|
if poller.unpollable_kinds:
|
|
queryset = queryset.exclude(kind__in=poller.unpollable_kinds)
|
|
return list(queryset.values_list("id", flat=True))
|
|
|
|
|
|
def cleanup_runs(run_model: type, retention_days: int) -> dict:
|
|
cutoff = timezone.now() - timedelta(days=retention_days)
|
|
deleted, _ = run_model.objects.filter(started_at__lt=cutoff).delete()
|
|
return {"deleted": deleted, "retention_days": retention_days}
|
|
|
|
|
|
def _skip(poller: Poller, src, reason: str) -> dict:
|
|
poller.run_model.objects.create(
|
|
source=src, status=RunStatus.SKIPPED, error=reason, finished_at=timezone.now(),
|
|
)
|
|
_mark_source(poller, src, RunStatus.SKIPPED, error=reason)
|
|
return {"source_id": src.pk, "status": RunStatus.SKIPPED, "reason": reason}
|
|
|
|
|
|
def _mark_source(
|
|
poller: Poller,
|
|
src,
|
|
status: str,
|
|
*,
|
|
error: str = "",
|
|
count_success: bool = False,
|
|
count_failure: bool = False,
|
|
extra: dict | None = None,
|
|
) -> None:
|
|
"""Write the outcome and re-arm the schedule.
|
|
|
|
Written with an UPDATE rather than save(): a poll can take minutes, and the
|
|
in-memory source is stale by the time it ends.
|
|
"""
|
|
fields: dict[str, Any] = {
|
|
"last_run_at": timezone.now(),
|
|
"last_status": status,
|
|
"last_error": error[:4000],
|
|
}
|
|
if count_success:
|
|
fields["success_count"] = src.success_count + 1
|
|
if count_failure:
|
|
fields["failure_count"] = src.failure_count + 1
|
|
fields.update(extra or {})
|
|
poller.source_model.objects.filter(pk=src.pk).update(**fields)
|
|
src.refresh_from_db()
|
|
src.schedule_next_poll()
|