Jens
@@ -0,0 +1,253 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import imaplib
|
||||
from contextlib import suppress
|
||||
from datetime import timedelta
|
||||
from uuid import uuid4
|
||||
|
||||
from celery import shared_task
|
||||
from django.conf import settings
|
||||
from django.db.models import Q
|
||||
from django.utils import timezone
|
||||
|
||||
from apps.jobs.services.pipeline import process_raw_document
|
||||
from apps.sources.models import EmailMessageRecord, RawDocument, Source, SourceRun
|
||||
from apps.sources.services.email_import import ingest_email, message_identity
|
||||
from apps.sources.services.fetcher import (
|
||||
FetchError,
|
||||
FetchTimeoutError,
|
||||
PolicyBlockedError,
|
||||
RateLimitedError,
|
||||
fetch_url,
|
||||
)
|
||||
from apps.sources.services.health import canary_recovery_sources, evaluate_source_health, start_health_canary
|
||||
from apps.sources.services.policy import assess_url
|
||||
from apps.sources.services.scheduling import (
|
||||
acquire_source_lease,
|
||||
calculate_failure_backoff_seconds,
|
||||
calculate_success_jitter_seconds,
|
||||
release_source_lease,
|
||||
)
|
||||
|
||||
|
||||
def _kind_for(content_type: str) -> str:
|
||||
content_type = (content_type or "").lower()
|
||||
if "html" in content_type:
|
||||
return RawDocument.Kind.HTML
|
||||
if "xml" in content_type or "rss" in content_type or "atom" in content_type:
|
||||
return RawDocument.Kind.XML
|
||||
if "json" in content_type:
|
||||
return RawDocument.Kind.JSON
|
||||
return RawDocument.Kind.TEXT
|
||||
|
||||
|
||||
@shared_task(
|
||||
bind=True,
|
||||
name="apps.sources.tasks.fetch_source",
|
||||
)
|
||||
def fetch_source(self, source_id: int, force: bool = False) -> dict[str, object]:
|
||||
source = Source.objects.get(pk=source_id)
|
||||
run = SourceRun.objects.create(source=source)
|
||||
now = timezone.now()
|
||||
lease_token = str(self.request.id or uuid4())
|
||||
lease = None
|
||||
decision = assess_url(source.base_url, source=source)
|
||||
if source.status not in {Source.Status.ACTIVE, Source.Status.TRIAL, Source.Status.CANDIDATE}:
|
||||
run.finish(SourceRun.Status.SKIPPED, error_category="status")
|
||||
return {"status": "skipped", "reason": "status"}
|
||||
if not decision.allowed:
|
||||
run.finish(
|
||||
SourceRun.Status.SKIPPED,
|
||||
error_category="policy",
|
||||
error_message=decision.reason,
|
||||
)
|
||||
return {"status": "skipped", "reason": "policy"}
|
||||
if not force and source.next_run_at and source.next_run_at > now:
|
||||
run.finish(SourceRun.Status.SKIPPED, error_category="not_due")
|
||||
return {"status": "skipped", "reason": "not_due"}
|
||||
|
||||
lease = acquire_source_lease(source_id=source.pk, worker_token=lease_token, now=now)
|
||||
if lease is None:
|
||||
run.finish(
|
||||
SourceRun.Status.SKIPPED,
|
||||
error_category="lease",
|
||||
error_message="Bron staat al in verwerking.",
|
||||
)
|
||||
return {"status": "skipped", "reason": "lease"}
|
||||
headers: dict[str, str] = {}
|
||||
if source.etag:
|
||||
headers["If-None-Match"] = source.etag
|
||||
if source.last_modified:
|
||||
headers["If-Modified-Since"] = source.last_modified
|
||||
try:
|
||||
fetched = fetch_url(source.base_url, source=source, conditional_headers=headers)
|
||||
if fetched.status_code == 304:
|
||||
source.schedule_after_success(
|
||||
jitter_seconds=calculate_success_jitter_seconds(source)
|
||||
)
|
||||
run.finish(SourceRun.Status.SUCCESS, http_status=304)
|
||||
return {"status": "not_modified"}
|
||||
content_type = fetched.headers.get("content-type", "")
|
||||
document = RawDocument.objects.create(
|
||||
source=source,
|
||||
source_run=run,
|
||||
url=fetched.requested_url,
|
||||
final_url=fetched.final_url,
|
||||
kind=_kind_for(content_type),
|
||||
content_type=content_type[:200],
|
||||
http_status=fetched.status_code,
|
||||
response_headers={
|
||||
key: value
|
||||
for key, value in fetched.headers.items()
|
||||
if key.lower() in {"etag", "last-modified", "content-type", "cache-control"}
|
||||
},
|
||||
content_hash=fetched.sha256,
|
||||
body_text=fetched.text,
|
||||
byte_length=len(fetched.content),
|
||||
retain_until=timezone.now() + timedelta(days=settings.RAW_DOCUMENT_RETENTION_DAYS),
|
||||
)
|
||||
metrics = process_raw_document(document)
|
||||
source.etag = fetched.headers.get("etag", source.etag)
|
||||
source.last_modified = fetched.headers.get("last-modified", source.last_modified)
|
||||
source.save(update_fields=["etag", "last_modified", "updated_at"])
|
||||
source.schedule_after_success(
|
||||
jitter_seconds=calculate_success_jitter_seconds(source)
|
||||
)
|
||||
run.finish(
|
||||
SourceRun.Status.SUCCESS,
|
||||
http_status=fetched.status_code,
|
||||
extracted_count=int(metrics["extracted"]),
|
||||
created_count=int(metrics["created"]),
|
||||
updated_count=int(metrics["updated"]),
|
||||
duplicate_count=int(metrics["duplicates"]),
|
||||
metrics={"parser": metrics["parser"], "warnings": metrics["warnings"]},
|
||||
)
|
||||
return metrics
|
||||
except PolicyBlockedError as exc:
|
||||
source.status = Source.Status.QUARANTINED
|
||||
source.policy = Source.Policy.DENY
|
||||
source.policy_reason = str(exc)
|
||||
source.save(update_fields=["status", "policy", "policy_reason", "updated_at"])
|
||||
run.finish(
|
||||
SourceRun.Status.SKIPPED,
|
||||
error_category="policy",
|
||||
error_message=str(exc)[:1000],
|
||||
)
|
||||
return {"status": "blocked", "reason": str(exc)}
|
||||
except RateLimitedError as exc:
|
||||
backoff = calculate_failure_backoff_seconds(
|
||||
source,
|
||||
failure_count=source.failure_count + 1,
|
||||
retry_after_seconds=exc.retry_after_seconds,
|
||||
)
|
||||
source.schedule_after_failure(now=timezone.now(), backoff_seconds=backoff)
|
||||
run.finish(
|
||||
SourceRun.Status.FAILED,
|
||||
error_category="rate_limited",
|
||||
error_message=str(exc)[:1000],
|
||||
)
|
||||
return {"status": "rate_limited", "backoff": backoff}
|
||||
except FetchTimeoutError as exc:
|
||||
backoff = calculate_failure_backoff_seconds(
|
||||
source, failure_count=source.failure_count + 1, timeout=True
|
||||
)
|
||||
source.schedule_after_failure(now=timezone.now(), backoff_seconds=backoff)
|
||||
run.finish(
|
||||
SourceRun.Status.FAILED,
|
||||
error_category="timeout",
|
||||
error_message=str(exc)[:1000],
|
||||
)
|
||||
return {"status": "timeout", "backoff": backoff}
|
||||
except FetchError as exc:
|
||||
backoff = calculate_failure_backoff_seconds(
|
||||
source, failure_count=source.failure_count + 1
|
||||
)
|
||||
source.schedule_after_failure(now=timezone.now(), backoff_seconds=backoff)
|
||||
run.finish(
|
||||
SourceRun.Status.FAILED,
|
||||
error_category=exc.__class__.__name__,
|
||||
error_message=str(exc)[:1000],
|
||||
)
|
||||
return {"status": "failed", "error_category": exc.__class__.__name__, "backoff": backoff}
|
||||
except Exception as exc:
|
||||
source.schedule_after_failure()
|
||||
run.finish(
|
||||
SourceRun.Status.FAILED,
|
||||
error_category="unexpected",
|
||||
error_message=str(exc)[:1000],
|
||||
)
|
||||
return {"status": "failed", "error_category": "unexpected"}
|
||||
finally:
|
||||
if lease is not None:
|
||||
release_source_lease(source_id=source.pk, worker_token=lease_token)
|
||||
|
||||
|
||||
@shared_task(name="apps.sources.tasks.schedule_due_sources")
|
||||
def schedule_due_sources() -> dict[str, int]:
|
||||
now = timezone.now()
|
||||
due = Source.objects.filter(
|
||||
status__in=[Source.Status.ACTIVE, Source.Status.TRIAL],
|
||||
).filter(Q(next_run_at__isnull=True) | Q(next_run_at__lte=now))
|
||||
ids = list(due.values_list("id", flat=True).distinct()[:200])
|
||||
for source_id in ids:
|
||||
fetch_source.delay(source_id)
|
||||
return {"scheduled": len(ids)}
|
||||
|
||||
|
||||
@shared_task(name="apps.sources.tasks.poll_imap_inbox")
|
||||
def poll_imap_inbox() -> dict[str, int | str]:
|
||||
if not settings.IMAP_ENABLED:
|
||||
return {"status": "disabled", "imported": 0}
|
||||
if not settings.IMAP_HOST or not settings.IMAP_USER or not settings.IMAP_PASSWORD:
|
||||
return {"status": "not_configured", "imported": 0}
|
||||
client_cls = imaplib.IMAP4_SSL if settings.IMAP_USE_SSL else imaplib.IMAP4
|
||||
client = client_cls(settings.IMAP_HOST, settings.IMAP_PORT)
|
||||
imported = 0
|
||||
try:
|
||||
client.login(settings.IMAP_USER, settings.IMAP_PASSWORD)
|
||||
status, _ = client.select(settings.IMAP_MAILBOX, readonly=not settings.IMAP_MARK_SEEN)
|
||||
if status != "OK":
|
||||
raise RuntimeError("IMAP-mailbox kon niet worden geopend")
|
||||
status, data = client.uid("search", None, "ALL")
|
||||
if status != "OK":
|
||||
raise RuntimeError("IMAP-zoekopdracht mislukt")
|
||||
uids = (data[0] or b"").split()[-200:]
|
||||
for uid in uids:
|
||||
status, payload = client.uid("fetch", uid, "(RFC822)")
|
||||
if status != "OK" or not payload:
|
||||
continue
|
||||
raw = next((part[1] for part in payload if isinstance(part, tuple)), None)
|
||||
if not raw:
|
||||
continue
|
||||
message_id = message_identity(raw)
|
||||
was_known = EmailMessageRecord.objects.filter(message_id=message_id).exists()
|
||||
ingest_email(raw, mailbox=settings.IMAP_MAILBOX)
|
||||
if not was_known:
|
||||
imported += 1
|
||||
return {"status": "ok", "imported": imported}
|
||||
finally:
|
||||
with suppress(Exception):
|
||||
client.logout()
|
||||
|
||||
|
||||
@shared_task(name="apps.sources.tasks.cleanup_raw_documents")
|
||||
def cleanup_raw_documents() -> dict[str, int]:
|
||||
deleted, _ = RawDocument.objects.filter(
|
||||
retain_until__lt=timezone.now(), quarantined=False
|
||||
).delete()
|
||||
return {"deleted": deleted}
|
||||
|
||||
|
||||
@shared_task(name="apps.sources.tasks.enforce_source_health")
|
||||
def enforce_source_health() -> dict[str, int]:
|
||||
return evaluate_source_health()
|
||||
|
||||
|
||||
@shared_task(name="apps.sources.tasks.recover_source_health")
|
||||
def recover_source_health() -> dict[str, int]:
|
||||
started = 0
|
||||
for source in canary_recovery_sources():
|
||||
start_health_canary(source)
|
||||
fetch_source.delay(source.pk, force=True)
|
||||
started += 1
|
||||
return {"canary_started": started}
|
||||
Reference in New Issue
Block a user