DataOps-platform/app/core/data_research/ingestion.py

from __future__ import annotations

import hashlib
import json
import re
from dataclasses import replace
from typing import Any, Callable

from app.core.common.identifiers import new_governance_uid
from app.core.common.timezone_utils import now_china_naive
from app.core.data_research.errors import (
    IngestionJobNotFound,
    IngestionPayloadInvalid,
    InvalidJobTransition,
)
from app.core.data_research.models import IngestionJobRecord, IngestionJobSpec


ALLOWED_TRANSITIONS = {
    "created": frozenset({"queued", "cancelled"}),
    "queued": frozenset({"extracting", "failed", "cancelled"}),
    "extracting": frozenset({"normalizing", "failed", "cancelled"}),
    "normalizing": frozenset({"matching", "failed", "cancelled"}),
    "matching": frozenset(
        {"awaiting_review", "partial", "failed", "cancelled"}
    ),
    "awaiting_review": frozenset({"published", "partial", "cancelled"}),
    "partial": frozenset({"queued", "cancelled"}),
    "failed": frozenset({"queued", "cancelled"}),
    "published": frozenset(),
    "cancelled": frozenset(),
}
TERMINAL_STATUSES = frozenset({"published", "cancelled"})
_SECRET_ASSIGNMENT = re.compile(
    r"(?i)\b(password|token|secret|authorization)\s*[:=]\s*[^\s,;]+"
)


def _required_text(payload: dict[str, Any], name: str) -> str:
    value = str(payload.get(name) or "").strip()
    if not value:
        raise IngestionPayloadInvalid(f"{name} is required")
    return value


def _sanitize_error(error: Any) -> str | None:
    if error is None:
        return None
    value = _SECRET_ASSIGNMENT.sub(lambda match: f"{match.group(1)}=[redacted]", str(error))
    return value[:1000]


class IngestionService:
    def __init__(
        self,
        repository,
        *,
        uid_factory: Callable[[], str] = new_governance_uid,
        nonce_factory: Callable[[], str] = new_governance_uid,
        clock: Callable[[], Any] = now_china_naive,
        commit: Callable[[], Any] = lambda: None,
        rollback: Callable[[], Any] = lambda: None,
    ):
        self.repository = repository
        self.uid_factory = uid_factory
        self.nonce_factory = nonce_factory
        self.clock = clock
        self.commit = commit
        self.rollback = rollback

    @staticmethod
    def _spec(payload: dict[str, Any]) -> IngestionJobSpec:
        if not isinstance(payload, dict):
            raise IngestionPayloadInvalid("payload must be an object")
        parameters = payload.get("parameters", {})
        if parameters is None:
            parameters = {}
        if not isinstance(parameters, dict):
            raise IngestionPayloadInvalid("parameters must be an object")
        artifact_uid = str(payload.get("artifact_uid") or "").strip() or None
        return IngestionJobSpec(
            source_uid=_required_text(payload, "source_uid"),
            artifact_uid=artifact_uid,
            job_type=_required_text(payload, "job_type"),
            parser_version=_required_text(payload, "parser_version"),
            parameters=dict(parameters),
            force_rerun=bool(payload.get("force_rerun", False)),
        )

    def _idempotency_key(self, spec: IngestionJobSpec) -> str:
        value = {
            "source_uid": spec.source_uid,
            "artifact_uid": spec.artifact_uid,
            "job_type": spec.job_type,
            "parser_version": spec.parser_version,
            "parameters": spec.parameters,
        }
        if spec.force_rerun:
            value["rerun_nonce"] = self.nonce_factory()
        canonical = json.dumps(
            value,
            sort_keys=True,
            ensure_ascii=False,
            separators=(",", ":"),
        )
        return hashlib.sha256(canonical.encode("utf-8")).hexdigest()

    def create_job(self, payload: dict[str, Any], actor_uid: str | None):
        spec = self._spec(payload)
        idempotency_key = self._idempotency_key(spec)
        if not spec.force_rerun:
            existing = self.repository.get_by_idempotency_key(idempotency_key)
            if existing is not None:
                return existing, False
        now = self.clock()
        record = IngestionJobRecord(
            uid=self.uid_factory(),
            source_uid=spec.source_uid,
            artifact_uid=spec.artifact_uid,
            job_type=spec.job_type,
            parser_version=spec.parser_version,
            idempotency_key=idempotency_key,
            actor_uid=actor_uid,
            parameters=spec.parameters,
            force_rerun=spec.force_rerun,
            created_at=now,
            updated_at=now,
        )
        try:
            saved = self.repository.add(record)
            self.commit()
            return saved, True
        except Exception:
            self.rollback()
            raise

    def _get(self, uid: str) -> IngestionJobRecord:
        record = self.repository.get(str(uid))
        if record is None:
            raise IngestionJobNotFound(f"ingestion job {uid} was not found")
        return record

    def transition(
        self,
        uid: str,
        target_status: str,
        statistics: dict[str, Any] | None = None,
        error: Any = None,
    ) -> IngestionJobRecord:
        record = self._get(uid)
        target_status = str(target_status or "").strip()
        if target_status not in ALLOWED_TRANSITIONS.get(record.status, frozenset()):
            raise InvalidJobTransition(f"{record.status} -> {target_status} is not allowed")
        now = self.clock()
        updated = replace(
            record,
            status=target_status,
            statistics=dict(statistics or record.statistics),
            last_error=_sanitize_error(error),
            updated_at=now,
            started_at=(now if target_status == "extracting" and record.started_at is None else record.started_at),
            finished_at=(now if target_status in TERMINAL_STATUSES | {"failed", "partial"} else None),
        )
        try:
            saved = self.repository.save(updated)
            self.commit()
            return saved
        except Exception:
            self.rollback()
            raise

    def get_job(self, uid: str) -> IngestionJobRecord:
        return self._get(uid)

    def list_jobs(self, filters: dict[str, Any] | None = None):
        return self.repository.list(filters or {})

    def retry(self, uid: str) -> IngestionJobRecord:
        return self.transition(uid, "queued")

    def cancel(self, uid: str) -> IngestionJobRecord:
        return self.transition(uid, "cancelled")