yksa-source-poller/yksa_poller/models.py
2026-08-18 22:05:57 +08:00

138 lines
5.2 KiB
Python

"""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"])