commit fe5ef9132ab95566af13a9dd53d85253cb711f61 Author: ThePetrovich Date: Tue Aug 18 22:05:57 2026 +0800 Initial commit diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..5e70718 --- /dev/null +++ b/.gitignore @@ -0,0 +1,155 @@ + +# Created by https://www.toptal.com/developers/gitignore/api/python +# Edit at https://www.toptal.com/developers/gitignore?templates=python + +### Python ### +# Byte-compiled / optimized / DLL files +__pycache__/ +*.py[cod] +*$py.class + +# C extensions +*.so + +# Distribution / packaging +.Python +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib64/ +parts/ +sdist/ +var/ +wheels/ +pip-wheel-metadata/ +share/python-wheels/ +*.egg-info/ +.installed.cfg +*.egg +MANIFEST + +# PyInstaller +# Usually these files are written by a python script from a template +# before PyInstaller builds the exe, so as to inject date/other infos into it. +*.manifest +*.spec + +# Installer logs +pip-log.txt +pip-delete-this-directory.txt + +# Unit test / coverage reports +htmlcov/ +.tox/ +.nox/ +.coverage +.coverage.* +.cache +nosetests.xml +coverage.xml +*.cover +*.py,cover +.hypothesis/ +.pytest_cache/ +pytestdebug.log + +# Translations +*.mo +*.pot + +# Django stuff: +*.log +local_settings.py +db.sqlite3 +db.sqlite3-journal + +# Flask stuff: +instance/ +.webassets-cache + +# Scrapy stuff: +.scrapy + +# Sphinx documentation +docs/_build/ +doc/_build/ + +# PyBuilder +target/ + +# Jupyter Notebook +.ipynb_checkpoints + +# IPython +profile_default/ +ipython_config.py + +# pyenv +.python-version + +# pipenv +# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. +# However, in case of collaboration, if having platform-specific dependencies or dependencies +# having no cross-platform support, pipenv may install dependencies that don't work, or not +# install all needed dependencies. +#Pipfile.lock + +# PEP 582; used by e.g. github.com/David-OConnor/pyflow +__pypackages__/ + +# Celery stuff +celerybeat-schedule +celerybeat.pid + +# SageMath parsed files +*.sage.py + +# Environments +.env +.venv +env/ +venv/ +ENV/ +env.bak/ +venv.bak/ +pythonenv* + +# Spyder project settings +.spyderproject +.spyproject + +# Rope project settings +.ropeproject + +# mkdocs documentation +/site + +# mypy +.mypy_cache/ +.dmypy.json +dmypy.json + +# Pyre type checker +.pyre/ + +# pytype static type analyzer +.pytype/ + +# profiling data +.prof + +# Db and static files +*.sqlite3 +/media +/static +/postgres +/EXAMPLE_* + +# Docker +docker-compose.override.yml +docker-compose.override + +# End of https://www.toptal.com/developers/gitignore/api/python diff --git a/README.md b/README.md new file mode 100644 index 0000000..6cb69b3 --- /dev/null +++ b/README.md @@ -0,0 +1,89 @@ +# yksa-source-poller + +The machine that polls upstreams on a schedule: a source model, a run log, the +dispatcher and the bookkeeping around one attempt. The cargo — what a fetched +record *is* and what happens to it — stays in the service. + +Extracted from `tdas/yksa_tdas/sources/` and `tle/yksa_tle/sources/`, which were +the same ~500 lines with different payload types. + +## Install + +``` +yksa-source-poller @ git+https://git.intra.yksa.space/web/yksa-source-poller.git@v0.1.0 +``` + +Requires [yksa-django-kit](../yksa-django-kit) for `RunStatus` and `read_secret`. +It is **not** added to `INSTALLED_APPS`: every model in it is abstract, so it +ships no tables and needs no app config. + +## Use + +```python +# sources/models.py +class TelemetrySource(AbstractPollingSource): + SUBSCRIPTION_ID_FIELDS = {"norad": "norad_cat_id", "satnogs": "satnogs_id"} + + kind = models.CharField(max_length=32, choices=SourceKind.choices) + subscription_id_kind = models.CharField(..., default=SubscriptionIdKind.SATNOGS) + tracked_satellites = models.ManyToManyField("satellites.Satellite", blank=True) + + +class SourceRun(AbstractSourceRun): + source = models.ForeignKey(TelemetrySource, on_delete=models.CASCADE, + related_name="runs") +``` + +```python +# sources/tasks.py +POLLER = Poller( + source_model=TelemetrySource, + run_model=SourceRun, + adapters=ADAPTERS, + persist=_persist_frame, # -> True when the record is new + track=_track_latest_rx, # optional: extra source fields on success + unpollable_kinds=(SourceKind.INTERNAL_PUSH,), +) + + +@shared_task(bind=True, max_retries=3, default_retry_delay=120) +def run_source_task(self, source_id): + try: + return run_source(POLLER, source_id) + except Exception as exc: + raise self.retry(exc=exc) +``` + +The Celery tasks stay in the service on purpose: `CELERY_BEAT_SCHEDULE` names +them by dotted path, and a task that moved into this package would silently stop +being scheduled. + +## Why the models are abstract + +The two concrete source models differ where they have to — which satellite model +the M2M points at, which identifiers the upstream accepts, and (in tdas) a +`last_rx_at` watermark. Everything else was identical. Abstract bases share the +identical part without pretending the rest is the same. + +Field definitions here reproduce what the services already had, so adopting the +package needs **no migration**. Where a service's own definition differed — a +different `poll_interval_sec` default, different wording — the service overrides +the field, which Django allows for abstract inheritance. Check with +`manage.py makemigrations --check --dry-run` after wiring it up. + +## What the runner guarantees + +Every exit path — success, failure, unconfigured adapter, missing adapter — +writes `last_status`, updates a counter and re-arms `next_poll_at`. That is the +reason to share it: a path that forgets leaves the dispatcher either hammering +the source once a minute or ignoring it forever, and both shipped once. + +`NotImplementedError` from an adapter is a *skip*, not a failure: a half-configured +source must not look like an outage. + +## Tests + +```sh +pip install -e ".[test]" ../yksa-django-kit +python -m pytest +``` diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..43e48ae --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,28 @@ +[build-system] +requires = ["setuptools>=68"] +build-backend = "setuptools.build_meta" + +[project] +name = "yksa-source-poller" +version = "0.1.0" +description = "Polling framework for YKSA ingest services: sources, runs, dispatch" +requires-python = ">=3.13" +license = { text = "Proprietary" } +dependencies = [ + "Django>=5.2", + "celery>=5.4", + "yksa-django-kit>=0.1", +] + +[project.optional-dependencies] +test = [ + "pytest>=8.0", + "pytest-django>=4.8", +] + +[tool.setuptools.packages.find] +include = ["yksa_poller*"] + +[tool.pytest.ini_options] +DJANGO_SETTINGS_MODULE = "tests.settings" +python_files = ["test_*.py"] diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/settings.py b/tests/settings.py new file mode 100644 index 0000000..c85720e --- /dev/null +++ b/tests/settings.py @@ -0,0 +1,19 @@ +SECRET_KEY = "test-only" +DEBUG = True +USE_TZ = True +PRODUCTION = False + +YKSA_SERVICE = "poller" + +INSTALLED_APPS = [ + "django.contrib.contenttypes", + "django.contrib.auth", + "yksa_kit", + "tests.testapp", +] + +DATABASES = {"default": {"ENGINE": "django.db.backends.sqlite3", "NAME": ":memory:"}} + +CACHES = {"default": {"BACKEND": "django.core.cache.backends.locmem.LocMemCache"}} + +CELERY_TASK_ALWAYS_EAGER = True diff --git a/tests/test_runner.py b/tests/test_runner.py new file mode 100644 index 0000000..73337d8 --- /dev/null +++ b/tests/test_runner.py @@ -0,0 +1,169 @@ +"""What the runner has to guarantee: every exit leaves the source schedulable. + +The bugs this suite exists for are all the same shape — a poll that ends without +re-arming next_poll_at, so the dispatcher either hammers the source every minute +or drops it forever. +""" + +from __future__ import annotations + +from datetime import timedelta + +import pytest +from django.utils import timezone + +from yksa_kit.choices import RunStatus +from yksa_poller.adapters import AdapterRegistry, BaseAdapter +from yksa_poller.runner import Poller, cleanup_runs, due_source_ids, run_source + +from tests.testapp.models import Kind, Run, Satellite, Source + +ADAPTERS = AdapterRegistry() + + +@ADAPTERS.register(Kind.HTTP) +class TwoRecords(BaseAdapter): + def fetch(self): + yield {"id": 1} + yield {"id": 2} + + +class Exploding(BaseAdapter): + def fetch(self): + raise RuntimeError("upstream on fire") + yield # pragma: no cover + + +class Unconfigured(BaseAdapter): + def fetch(self): + raise NotImplementedError("no credentials for this source") + yield # pragma: no cover + + +def poller(adapters=None, **kwargs): + return Poller( + source_model=Source, + run_model=Run, + adapters=adapters if adapters is not None else ADAPTERS, + persist=kwargs.pop("persist", lambda src, rec: True), + **kwargs, + ) + + +@pytest.fixture +def source(db): + return Source.objects.create(name="Upstream", slug="upstream", kind=Kind.HTTP) + + +@pytest.mark.django_db +def test_a_successful_poll_counts_and_reschedules(source): + before = source.next_poll_at + result = run_source(poller(), source.pk) + + assert result["fetched"] == 2 + assert result["new"] == 2 + source.refresh_from_db() + assert source.last_status == RunStatus.SUCCESS + assert source.success_count == 1 + assert source.next_poll_at > before + assert source.runs.get().status == RunStatus.SUCCESS + + +@pytest.mark.django_db +def test_only_new_records_are_counted_as_new(source): + result = run_source(poller(persist=lambda src, rec: rec["id"] == 1), source.pk) + assert (result["fetched"], result["new"]) == (2, 1) + + +@pytest.mark.django_db +def test_a_failing_poll_still_reschedules_and_reraises(source): + before = source.next_poll_at + with pytest.raises(RuntimeError): + run_source(poller(adapters={Kind.HTTP: Exploding}), source.pk) + + source.refresh_from_db() + assert source.last_status == RunStatus.FAILED + assert source.failure_count == 1 + assert source.next_poll_at > before + assert "upstream on fire" in source.runs.get().error + + +@pytest.mark.django_db +def test_an_unconfigured_adapter_is_skipped_not_failed(source): + """A source nobody finished configuring must not look like an outage.""" + result = run_source(poller(adapters={Kind.HTTP: Unconfigured}), source.pk) + + assert result["status"] == RunStatus.SKIPPED + source.refresh_from_db() + assert source.failure_count == 0 + assert source.runs.get().status == RunStatus.SKIPPED + + +@pytest.mark.django_db +def test_a_kind_with_no_adapter_is_skipped(source): + source.kind = Kind.PUSH + source.save() + result = run_source(poller(), source.pk) + assert result["status"] == RunStatus.SKIPPED + assert "no adapter" in source.runs.get().error + + +@pytest.mark.django_db +def test_a_disabled_source_does_nothing_at_all(source): + source.is_enabled = False + source.save() + assert run_source(poller(), source.pk)["skipped"] is True + assert source.runs.count() == 0 + + +@pytest.mark.django_db +def test_track_carries_extra_fields_onto_the_source(source): + """tdas uses this for last_rx_at, the lower bound of its next query.""" + stamp = timezone.now() + + def track(src, record, extra): + extra["last_run_at"] = stamp + + run_source(poller(track=track), source.pk) + source.refresh_from_db() + assert source.last_run_at == stamp + + +@pytest.mark.django_db +def test_dispatch_picks_only_due_enabled_pollable_sources(): + past = timezone.now() - timedelta(minutes=1) + future = timezone.now() + timedelta(hours=1) + due = Source.objects.create(name="due", slug="due", kind=Kind.HTTP, next_poll_at=past) + Source.objects.create(name="later", slug="later", kind=Kind.HTTP, next_poll_at=future) + Source.objects.create(name="off", slug="off", kind=Kind.HTTP, next_poll_at=past, is_enabled=False) + Source.objects.create(name="pushed", slug="pushed", kind=Kind.PUSH, next_poll_at=past) + + ids = due_source_ids(poller(unpollable_kinds=(Kind.PUSH,))) + assert ids == [due.pk] + + +@pytest.mark.django_db +def test_cleanup_deletes_only_old_runs(source): + fresh = Run.objects.create(source=source, status=RunStatus.SUCCESS) + old = Run.objects.create(source=source, status=RunStatus.SUCCESS) + # started_at is auto_now_add, so age has to be forced. + Run.objects.filter(pk=old.pk).update( + started_at=timezone.now() - timedelta(days=40) + ) + + assert cleanup_runs(Run, retention_days=30)["deleted"] == 1 + assert list(Run.objects.values_list("pk", flat=True)) == [fresh.pk] + + +@pytest.mark.django_db +def test_subscription_ids_skip_satellites_missing_the_identifier(source): + """An upstream query built from a blank identifier returns either nothing or + everything, and both are wrong.""" + with_norad = Satellite.objects.create(norad_cat_id="25544", internal_id="iss") + Satellite.objects.create(norad_cat_id="", internal_id="unlaunched") + source.tracked_satellites.set(Satellite.objects.all()) + + assert source.subscription_id_values() == [with_norad.norad_cat_id] + + source.subscription_id_kind = "internal" + assert sorted(source.subscription_id_values()) == ["iss", "unlaunched"] diff --git a/tests/testapp/__init__.py b/tests/testapp/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/testapp/migrations/__init__.py b/tests/testapp/migrations/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/testapp/models.py b/tests/testapp/models.py new file mode 100644 index 0000000..0d89051 --- /dev/null +++ b/tests/testapp/models.py @@ -0,0 +1,27 @@ +"""A minimal service, standing in for tdas and odms.""" + +from django.db import models + +from yksa_poller.models import AbstractPollingSource, AbstractSourceRun + + +class Kind(models.TextChoices): + HTTP = "http", "HTTP" + PUSH = "push", "Push (fed from outside)" + + +class Satellite(models.Model): + norad_cat_id = models.CharField(max_length=16, blank=True) + internal_id = models.SlugField(max_length=64) + + +class Source(AbstractPollingSource): + SUBSCRIPTION_ID_FIELDS = {"norad": "norad_cat_id", "internal": "internal_id"} + + kind = models.CharField(max_length=32, choices=Kind.choices) + subscription_id_kind = models.CharField(max_length=16, default="norad") + tracked_satellites = models.ManyToManyField(Satellite, blank=True) + + +class Run(AbstractSourceRun): + source = models.ForeignKey(Source, on_delete=models.CASCADE, related_name="runs") diff --git a/yksa_poller/__init__.py b/yksa_poller/__init__.py new file mode 100644 index 0000000..3dc1f76 --- /dev/null +++ b/yksa_poller/__init__.py @@ -0,0 +1 @@ +__version__ = "0.1.0" diff --git a/yksa_poller/adapters.py b/yksa_poller/adapters.py new file mode 100644 index 0000000..0b5bf38 --- /dev/null +++ b/yksa_poller/adapters.py @@ -0,0 +1,62 @@ +"""What an adapter is, and how the runner finds one. + +An adapter does I/O and parsing and nothing else: it yields records, and the +runner matches, dedups and writes them. Keeping the split means an adapter can be +tested against a captured payload with no database. +""" + +from __future__ import annotations + +from typing import Iterator + + +class BaseAdapter: + """Subclasses implement :meth:`fetch`. + + The payload type is the service's own — a frame in tdas, an element set in + odms. That is the cargo, not the machine, so it is not defined here. + """ + + def __init__(self, source) -> None: + self.source = source + + @property + def config(self) -> dict: + return self.source.config or {} + + def api_key(self) -> str | None: + """This source's API key, or None. + + An inline ``api_key`` in the config wins; otherwise ``api_key_secret`` + names a Docker secret or env var. Production keeps the key out of the + database that way. + """ + direct = self.config.get("api_key") + if direct: + return str(direct).strip() + secret_name = self.config.get("api_key_secret") + if secret_name: + from yksa_kit.settings import read_secret + + value = read_secret(str(secret_name)) + return value.strip() if value else None + return None + + def fetch(self) -> Iterator: + raise NotImplementedError + + +class AdapterRegistry(dict): + """`kind` -> adapter class, with a decorator for declaring one. + + A kind with no adapter is legitimate — tdas's `internal_push` is fed by an + API, odms's `internal_corrected` by a pipeline — and the runner records those + as skipped rather than failed. + """ + + def register(self, kind: str): + def decorator(cls): + self[kind] = cls + return cls + + return decorator diff --git a/yksa_poller/admin.py b/yksa_poller/admin.py new file mode 100644 index 0000000..acddc70 --- /dev/null +++ b/yksa_poller/admin.py @@ -0,0 +1,58 @@ +"""Admin building blocks. A service still registers its own models.""" + +from __future__ import annotations + +from django.contrib import admin, messages +from django.utils.translation import gettext_lazy as _ + + +class SourceAdminMixin: + """The columns, filters and the run-now action every source list needs. + + ``run_task`` is the service's Celery task. It is queued rather than run + inline so a slow upstream cannot hold the admin request open. + """ + + run_task = None + + list_filter = ("kind", "is_enabled", "is_public", "last_status") + search_fields = ("name", "slug", "description") + prepopulated_fields = {"slug": ("name",)} + filter_horizontal = ("tracked_satellites",) + readonly_fields = ( + "last_run_at", "last_status", "last_error", + "success_count", "failure_count", "created_at", "updated_at", + ) + + @admin.action(description=_("Run selected sources now")) + def run_now(self, request, queryset): + if self.run_task is None: # pragma: no cover - configuration error + raise NotImplementedError("set run_task on the ModelAdmin") + queued = 0 + for src in queryset: + self.run_task.delay(src.pk) + queued += 1 + messages.success(request, _("Queued %(n)d source(s) for fetch.") % {"n": queued}) + + actions = ["run_now"] + + +class SourceRunAdminMixin: + """Run history is a log: readable, filterable, never editable.""" + + list_display = ( + "source", "started_at", "finished_at", "status", "fetched_count", "new_count", + ) + list_filter = ("status", "source") + search_fields = ("source__name", "source__slug") + readonly_fields = ( + "source", "started_at", "finished_at", + "status", "fetched_count", "new_count", "error", + ) + date_hierarchy = "started_at" + + def has_add_permission(self, request): + return False + + def has_change_permission(self, request, obj=None): + return False diff --git a/yksa_poller/apps.py b/yksa_poller/apps.py new file mode 100644 index 0000000..0730e84 --- /dev/null +++ b/yksa_poller/apps.py @@ -0,0 +1,7 @@ +from django.apps import AppConfig + + +class YksaPollerConfig(AppConfig): + name = "yksa_poller" + label = "yksa_poller" + verbose_name = "YKSA source poller" diff --git a/yksa_poller/models.py b/yksa_poller/models.py new file mode 100644 index 0000000..3faaba8 --- /dev/null +++ b/yksa_poller/models.py @@ -0,0 +1,138 @@ +"""Abstract models. This package ships no tables of its own. + +A source is a place to fetch from on a schedule; a run is one attempt. Both are +abstract because the concrete models differ where they must — which satellite +model the M2M points at, which identifier the upstream subscribes by — and are +identical everywhere else. + +Field definitions here match what the services already had, so adopting this +package needs no migration. A service that wants a different default or wording +overrides the field; Django permits that for fields inherited from an abstract +base, and it is the intended escape hatch. +""" + +from __future__ import annotations + +from datetime import timedelta + +from django.db import models +from django.utils import timezone +from django.utils.translation import gettext_lazy as _ + +from yksa_kit.choices import RunStatus +from yksa_kit.states import health_state + + +class AbstractPollingSource(models.Model): + #: subscription_id_kind value -> Satellite field to read. Set by the service; + #: subscription_id_values() is a no-op without it. + SUBSCRIPTION_ID_FIELDS: dict[str, str] = {} + + name = models.CharField(max_length=128, unique=True) + slug = models.SlugField(max_length=64, unique=True) + config = models.JSONField( + default=dict, blank=True, + help_text=_("Adapter-specific config."), + ) + description = models.TextField(blank=True) + + poll_interval_sec = models.PositiveIntegerField(default=3600) + rate_limit_host = models.CharField( + max_length=128, blank=True, + help_text=_("Optional shared-bucket key for outbound rate limiting."), + ) + rate_limit_per_min = models.PositiveIntegerField(default=20) + + is_enabled = models.BooleanField(default=True) + is_public = models.BooleanField( + default=True, + help_text=_("Show on the public source-status page."), + ) + + next_poll_at = models.DateTimeField(default=timezone.now, db_index=True) + last_run_at = models.DateTimeField(null=True, blank=True) + last_status = models.CharField(max_length=16, choices=RunStatus.choices, blank=True) + last_error = models.TextField(blank=True) + success_count = models.PositiveIntegerField(default=0) + failure_count = models.PositiveIntegerField(default=0) + + created_at = models.DateTimeField(auto_now_add=True) + updated_at = models.DateTimeField(auto_now=True) + + class Meta: + abstract = True + + def __str__(self) -> str: + return self.name + + @property + def ui_state(self) -> str: + """For yksa/ui/_state.html. A source that has never run reads as unknown, + not as an error: nothing has gone wrong yet.""" + return health_state(self.last_status) + + def schedule_next_poll(self, *, now=None) -> None: + ref = now or timezone.now() + self.next_poll_at = ref + timedelta(seconds=self.poll_interval_sec) + self.save(update_fields=["next_poll_at", "updated_at"]) + + def subscription_id_values(self) -> list: + """Identifier values for the satellites this source subscribes to. + + Reads the field named by ``subscription_id_kind`` through + ``SUBSCRIPTION_ID_FIELDS``. Satellites whose chosen field is empty are + dropped: an upstream query built from a blank identifier returns either + nothing or everything, and both are wrong. + """ + if not self.SUBSCRIPTION_ID_FIELDS: + return [] + default_field = next(iter(self.SUBSCRIPTION_ID_FIELDS.values())) + field_name = self.SUBSCRIPTION_ID_FIELDS.get( + self.subscription_id_kind, default_field + ) + values = ( + self.tracked_satellites + .exclude(**{f"{field_name}__isnull": True}) + .values_list(field_name, flat=True) + ) + return [v for v in values if v not in ("", None)] + + +class AbstractSourceRun(models.Model): + started_at = models.DateTimeField(auto_now_add=True, db_index=True) + finished_at = models.DateTimeField(null=True, blank=True) + status = models.CharField(max_length=16, choices=RunStatus.choices) + fetched_count = models.PositiveIntegerField(default=0) + new_count = models.PositiveIntegerField(default=0) + error = models.TextField(blank=True) + + class Meta: + abstract = True + + def __str__(self) -> str: + return f"{self.source_id} @ {self.started_at:%Y-%m-%d %H:%M:%S} ({self.status})" + + @property + def ui_state(self) -> str: + return health_state(self.status) + + def mark_success(self, *, fetched: int, new: int) -> None: + self.status = RunStatus.SUCCESS + self.fetched_count = fetched + self.new_count = new + self.finished_at = timezone.now() + self.save(update_fields=["status", "fetched_count", "new_count", "finished_at"]) + + def mark_failed(self, error: str) -> None: + # Truncated: a traceback from a broken upstream can be megabytes, and the + # column is read on a status page. + self.status = RunStatus.FAILED + self.error = error[:4000] + self.finished_at = timezone.now() + self.save(update_fields=["status", "error", "finished_at"]) + + def mark_skipped(self, reason: str) -> None: + self.status = RunStatus.SKIPPED + self.error = reason + self.finished_at = timezone.now() + self.save(update_fields=["status", "error", "finished_at"]) diff --git a/yksa_poller/runner.py b/yksa_poller/runner.py new file mode 100644 index 0000000..f1e88ee --- /dev/null +++ b/yksa_poller/runner.py @@ -0,0 +1,135 @@ +"""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()