278 lines
10 KiB
Python
278 lines
10 KiB
Python
from __future__ import annotations
|
|
|
|
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 MailboxConnection, RawDocument, Source, SourceRun
|
|
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.imap_import import poll_imap_mailbox
|
|
from apps.sources.services.mailbox_connections import (
|
|
MailboxCredentialError,
|
|
claim_mailbox_connection,
|
|
finish_mailbox_poll,
|
|
)
|
|
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.schedule_mailbox_polls")
|
|
def schedule_mailbox_polls() -> dict[str, int]:
|
|
now = timezone.now()
|
|
due_ids = list(
|
|
MailboxConnection.objects.filter(enabled=True, next_poll_at__lte=now)
|
|
.filter(Q(lease_expires_at__isnull=True) | Q(lease_expires_at__lte=now))
|
|
.values_list("id", flat=True)[:100]
|
|
)
|
|
for connection_id in due_ids:
|
|
poll_mailbox.delay(connection_id)
|
|
return {"scheduled": len(due_ids)}
|
|
|
|
|
|
@shared_task(bind=True, name="apps.sources.tasks.poll_mailbox")
|
|
def poll_mailbox(self, connection_id: int, force: bool = False) -> dict[str, int | str]:
|
|
worker_token = str(self.request.id or uuid4())
|
|
claimed = claim_mailbox_connection(
|
|
connection_id=connection_id, worker_token=worker_token, force=force
|
|
)
|
|
if claimed is None:
|
|
return {"status": "skipped", "imported": 0}
|
|
connection, token = claimed
|
|
try:
|
|
result = poll_imap_mailbox(connection)
|
|
finish_mailbox_poll(
|
|
connection_id=connection.pk,
|
|
worker_token=token,
|
|
success=result["status"] == "ok",
|
|
)
|
|
return result
|
|
except MailboxCredentialError as exc:
|
|
finish_mailbox_poll(
|
|
connection_id=connection.pk,
|
|
worker_token=token,
|
|
success=False,
|
|
error_category="credentials",
|
|
error_message=str(exc),
|
|
)
|
|
return {"status": "failed", "error_category": "credentials", "imported": 0}
|
|
except Exception as exc:
|
|
error_category = "imap"
|
|
if exc.__class__.__name__.lower().startswith("unsafe"):
|
|
error_category = "host_policy"
|
|
finish_mailbox_poll(
|
|
connection_id=connection.pk,
|
|
worker_token=token,
|
|
success=False,
|
|
error_category=error_category,
|
|
error_message="IMAP-verbinding of authenticatie mislukt.",
|
|
)
|
|
return {"status": "failed", "error_category": error_category, "imported": 0}
|
|
|
|
|
|
@shared_task(name="apps.sources.tasks.poll_imap_inbox")
|
|
def poll_imap_inbox() -> dict[str, int]:
|
|
"""Compatibility alias for old deployments; it now schedules due per-platform mailboxes."""
|
|
return schedule_mailbox_polls()
|
|
|
|
|
|
@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}
|