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), attempt_count=( record.attempt_count + 1 if target_status == "extracting" else record.attempt_count ), failure_stage=( record.status if target_status == "failed" else None if target_status == "queued" else record.failure_stage ), 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")