456 lines
16 KiB
Python
456 lines
16 KiB
Python
"""TeamHUB AssetManifest ingest orchestration.
|
|
|
|
This API boundary accepts module-owned object storage references, verifies the
|
|
original bytes, records an idempotent ingest job, and queues the existing
|
|
per-document OCR/extraction/indexing pipeline.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import shutil
|
|
import tempfile
|
|
import uuid
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Protocol
|
|
|
|
from sqlalchemy import select
|
|
from sqlalchemy.exc import IntegrityError
|
|
|
|
from app.api.schemas import (
|
|
AssetManifest,
|
|
AssetManifestEnvelope,
|
|
KnowledgeIngestOptions,
|
|
KnowledgeIngestResponse,
|
|
)
|
|
from app.common.json_logger import redact_secrets
|
|
from app.config import settings
|
|
from app.db.audit import record_audit
|
|
from app.db.models import (
|
|
ArtifactType,
|
|
AssetIngestJob,
|
|
Document,
|
|
DocumentStatus,
|
|
ProcessingEvent,
|
|
)
|
|
from app.db.session import session_scope
|
|
from app.logging_config import get_logger
|
|
from app.storage.artifacts import ensure_artifact
|
|
from app.storage.local_paths import work_dir_for
|
|
from app.storage.minio_client import MinioStorage, get_storage
|
|
from app.utils.hashing import sha256_file
|
|
|
|
logger = get_logger(__name__)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class VerifiedAssetObject:
|
|
local_path: Path
|
|
size_bytes: int
|
|
sha256: str
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ExistingKnowledgeIngestJob:
|
|
job_id: uuid.UUID
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class KnowledgeIngestRegistration:
|
|
job_id: uuid.UUID
|
|
document_id: uuid.UUID | None
|
|
duplicate: bool
|
|
|
|
|
|
class KnowledgeIngestError(Exception):
|
|
def __init__(self, status_code: int, reason_code: str, message: str) -> None:
|
|
super().__init__(message)
|
|
self.status_code = status_code
|
|
self.reason_code = reason_code
|
|
self.message = message
|
|
|
|
|
|
class KnowledgeIngestRepository(Protocol):
|
|
def find_existing_job(self, manifest: AssetManifest) -> ExistingKnowledgeIngestJob | None:
|
|
...
|
|
|
|
def register_verified_manifest(
|
|
self,
|
|
manifest: AssetManifest,
|
|
ingest_options: KnowledgeIngestOptions,
|
|
verified: VerifiedAssetObject,
|
|
idempotency_key: str | None,
|
|
) -> KnowledgeIngestRegistration:
|
|
...
|
|
|
|
|
|
class SqlAlchemyKnowledgeIngestRepository:
|
|
def find_existing_job(self, manifest: AssetManifest) -> ExistingKnowledgeIngestJob | None:
|
|
with session_scope() as db:
|
|
job = db.execute(_identity_query(manifest)).scalar_one_or_none()
|
|
if job is None:
|
|
return None
|
|
return ExistingKnowledgeIngestJob(job_id=job.id)
|
|
|
|
def register_verified_manifest(
|
|
self,
|
|
manifest: AssetManifest,
|
|
ingest_options: KnowledgeIngestOptions,
|
|
verified: VerifiedAssetObject,
|
|
idempotency_key: str | None,
|
|
) -> KnowledgeIngestRegistration:
|
|
try:
|
|
with session_scope() as db:
|
|
existing_job = db.execute(_identity_query(manifest)).scalar_one_or_none()
|
|
if existing_job is not None:
|
|
return KnowledgeIngestRegistration(
|
|
job_id=existing_job.id,
|
|
document_id=existing_job.document_id,
|
|
duplicate=True,
|
|
)
|
|
|
|
existing_doc = db.execute(
|
|
select(Document).where(Document.sha256 == manifest.asset.sha256.lower())
|
|
).scalar_one_or_none()
|
|
|
|
if existing_doc is None:
|
|
document_id = uuid.uuid4()
|
|
doc = Document(
|
|
id=document_id,
|
|
source_path=_storage_uri(manifest),
|
|
original_file_name=_original_filename(manifest),
|
|
sha256=manifest.asset.sha256.lower(),
|
|
file_size_bytes=verified.size_bytes,
|
|
mime_type=manifest.asset.content_type,
|
|
status=DocumentStatus.STORED_ORIGINAL,
|
|
)
|
|
db.add(doc)
|
|
db.add(
|
|
ProcessingEvent(
|
|
document_id=document_id,
|
|
stage=DocumentStatus.DISCOVERED,
|
|
level="INFO",
|
|
message="Asset manifest accepted",
|
|
data=_event_data(manifest, idempotency_key),
|
|
)
|
|
)
|
|
else:
|
|
document_id = existing_doc.id
|
|
if existing_doc.status == DocumentStatus.DISCOVERED:
|
|
existing_doc.status = DocumentStatus.STORED_ORIGINAL
|
|
|
|
# Flush the Document row before FK-dependent inserts: there is no
|
|
# ORM relationship between AssetIngestJob/DocumentArtifact and
|
|
# Document, so the unit of work does not order these INSERTs and
|
|
# PostgreSQL rejects the job row with a ForeignKeyViolation
|
|
# (SQLite-based tests do not enforce FKs, hence unseen there).
|
|
db.flush()
|
|
|
|
ensure_artifact(
|
|
db,
|
|
document_id=document_id,
|
|
artifact_type=ArtifactType.ORIGINAL_PDF,
|
|
bucket=manifest.storage.bucket,
|
|
key=manifest.storage.object_key,
|
|
checksum=manifest.asset.sha256.lower(),
|
|
)
|
|
|
|
job_id = uuid.uuid4()
|
|
db.add(
|
|
AssetIngestJob(
|
|
id=job_id,
|
|
document_id=document_id,
|
|
asset_id=manifest.asset.asset_id,
|
|
manifest_version=manifest.manifest_version,
|
|
sha256=manifest.asset.sha256.lower(),
|
|
owner_module=manifest.asset.owner_module,
|
|
owner_record_type=manifest.asset.owner_record_type,
|
|
owner_record_id=manifest.asset.owner_record_id,
|
|
status="QUEUED",
|
|
idempotency_key=idempotency_key,
|
|
storage_bucket=manifest.storage.bucket,
|
|
storage_key=manifest.storage.object_key,
|
|
storage_version_id=manifest.storage.version_id,
|
|
manifest_json=manifest.model_dump(mode="json"),
|
|
ingest_options=ingest_options.model_dump(mode="json"),
|
|
)
|
|
)
|
|
db.add(
|
|
ProcessingEvent(
|
|
run_id=job_id,
|
|
document_id=document_id,
|
|
stage=DocumentStatus.STORED_ORIGINAL,
|
|
level="INFO",
|
|
message="Original referenced from owner object storage",
|
|
data=_event_data(manifest, idempotency_key),
|
|
)
|
|
)
|
|
|
|
record_audit(
|
|
db,
|
|
action="asset.ingest.registered",
|
|
actor=manifest.asset.created_by,
|
|
app_code=manifest.asset.owner_module,
|
|
entity_type="document",
|
|
entity_id=str(document_id),
|
|
details={
|
|
"asset_id": manifest.asset.asset_id,
|
|
"job_id": str(job_id),
|
|
"sha256": manifest.asset.sha256.lower(),
|
|
"owner_record_type": manifest.asset.owner_record_type,
|
|
"owner_record_id": manifest.asset.owner_record_id,
|
|
},
|
|
)
|
|
|
|
return KnowledgeIngestRegistration(
|
|
job_id=job_id,
|
|
document_id=document_id,
|
|
duplicate=False,
|
|
)
|
|
except IntegrityError:
|
|
existing = self.find_existing_job(manifest)
|
|
if existing is not None:
|
|
return KnowledgeIngestRegistration(
|
|
job_id=existing.job_id,
|
|
document_id=None,
|
|
duplicate=True,
|
|
)
|
|
raise
|
|
|
|
|
|
def accept_knowledge_ingest(
|
|
req: AssetManifestEnvelope,
|
|
*,
|
|
idempotency_key: str | None = None,
|
|
storage: MinioStorage | None = None,
|
|
repository: KnowledgeIngestRepository | None = None,
|
|
enqueue_document=None,
|
|
cache_original=None,
|
|
temp_dir: Path | None = None,
|
|
) -> KnowledgeIngestResponse:
|
|
manifest = req.manifest
|
|
if manifest.security.gate_status.lower() != "approved":
|
|
return KnowledgeIngestResponse(
|
|
status="rejected",
|
|
asset_id=manifest.asset.asset_id,
|
|
idempotency_key=idempotency_key,
|
|
reason_code="security_gate_not_approved",
|
|
)
|
|
|
|
repo = repository or SqlAlchemyKnowledgeIngestRepository()
|
|
existing = repo.find_existing_job(manifest)
|
|
if existing is not None:
|
|
return KnowledgeIngestResponse(
|
|
status="duplicate",
|
|
ingest_job_id=existing.job_id,
|
|
asset_id=manifest.asset.asset_id,
|
|
idempotency_key=idempotency_key,
|
|
)
|
|
|
|
object_storage = storage or get_storage()
|
|
enqueue = enqueue_document or enqueue_process_document
|
|
cache = cache_original or cache_verified_original
|
|
|
|
with _DownloadTarget(temp_dir) as local_path:
|
|
try:
|
|
object_storage.get_to_path(
|
|
manifest.storage.bucket,
|
|
manifest.storage.object_key,
|
|
local_path,
|
|
version_id=manifest.storage.version_id,
|
|
)
|
|
except Exception as exc:
|
|
raise KnowledgeIngestError(
|
|
400,
|
|
"object_read_failed",
|
|
f"Could not read object from storage: {manifest.storage.bucket}/{manifest.storage.object_key}",
|
|
) from exc
|
|
|
|
verified = _verify_downloaded_object(local_path, manifest)
|
|
registration = repo.register_verified_manifest(
|
|
manifest,
|
|
req.ingest_options,
|
|
verified,
|
|
idempotency_key,
|
|
)
|
|
|
|
if registration.duplicate:
|
|
return KnowledgeIngestResponse(
|
|
status="duplicate",
|
|
ingest_job_id=registration.job_id,
|
|
asset_id=manifest.asset.asset_id,
|
|
idempotency_key=idempotency_key,
|
|
)
|
|
if registration.document_id is None:
|
|
raise KnowledgeIngestError(500, "document_registration_failed", "Document was not registered")
|
|
|
|
cache(registration.document_id, verified.sha256, verified.local_path)
|
|
enqueue(registration.document_id, registration.job_id)
|
|
|
|
logger.info(
|
|
"knowledge_ingest.accepted",
|
|
asset_id=manifest.asset.asset_id,
|
|
job_id=str(registration.job_id),
|
|
document_id=str(registration.document_id),
|
|
)
|
|
return KnowledgeIngestResponse(
|
|
status="accepted",
|
|
ingest_job_id=registration.job_id,
|
|
asset_id=manifest.asset.asset_id,
|
|
idempotency_key=idempotency_key,
|
|
)
|
|
|
|
|
|
def enqueue_process_document(document_id: uuid.UUID, run_id: uuid.UUID) -> None:
|
|
from app.workers.tasks import process_document # noqa: PLC0415
|
|
|
|
process_document.delay(str(document_id), str(run_id))
|
|
|
|
|
|
def cache_verified_original(document_id: uuid.UUID, sha256: str, source_path: Path) -> None:
|
|
dest = work_dir_for(document_id) / f"{sha256}.pdf"
|
|
if dest.exists() and sha256_file(dest) == sha256:
|
|
return
|
|
shutil.copyfile(source_path, dest)
|
|
|
|
|
|
def load_asset_metadata_for_document(document_id: uuid.UUID) -> dict[str, object]:
|
|
with session_scope() as db:
|
|
job = (
|
|
db.execute(
|
|
select(AssetIngestJob)
|
|
.where(AssetIngestJob.document_id == document_id)
|
|
.order_by(AssetIngestJob.created_at.desc())
|
|
)
|
|
.scalars()
|
|
.first()
|
|
)
|
|
if job is None:
|
|
return {}
|
|
return _metadata_from_manifest(job.manifest_json)
|
|
|
|
|
|
def load_asset_ingest_options(document_id: uuid.UUID) -> dict[str, object]:
|
|
"""Return the latest asset ingest job's options (e.g. ``emit_events``)."""
|
|
with session_scope() as db:
|
|
job = (
|
|
db.execute(
|
|
select(AssetIngestJob)
|
|
.where(AssetIngestJob.document_id == document_id)
|
|
.order_by(AssetIngestJob.created_at.desc())
|
|
)
|
|
.scalars()
|
|
.first()
|
|
)
|
|
if job is None:
|
|
return {}
|
|
return dict(job.ingest_options or {})
|
|
|
|
|
|
def _identity_query(manifest: AssetManifest):
|
|
return select(AssetIngestJob).where(
|
|
AssetIngestJob.asset_id == manifest.asset.asset_id,
|
|
AssetIngestJob.sha256 == manifest.asset.sha256.lower(),
|
|
AssetIngestJob.manifest_version == manifest.manifest_version,
|
|
)
|
|
|
|
|
|
def _verify_downloaded_object(path: Path, manifest: AssetManifest) -> VerifiedAssetObject:
|
|
size = path.stat().st_size
|
|
if size != manifest.asset.size_bytes:
|
|
raise KnowledgeIngestError(
|
|
400,
|
|
"size_mismatch",
|
|
f"Object size {size} does not match manifest size {manifest.asset.size_bytes}",
|
|
)
|
|
sha = sha256_file(path)
|
|
if sha != manifest.asset.sha256.lower():
|
|
raise KnowledgeIngestError(
|
|
400,
|
|
"sha256_mismatch",
|
|
"Object SHA256 does not match manifest SHA256",
|
|
)
|
|
return VerifiedAssetObject(local_path=path, size_bytes=size, sha256=sha)
|
|
|
|
|
|
def _storage_uri(manifest: AssetManifest) -> str:
|
|
uri = f"s3://{manifest.storage.bucket}/{manifest.storage.object_key}"
|
|
if manifest.storage.version_id:
|
|
return f"{uri}?versionId={manifest.storage.version_id}"
|
|
return uri
|
|
|
|
|
|
def _original_filename(manifest: AssetManifest) -> str:
|
|
if manifest.asset.original_filename:
|
|
return manifest.asset.original_filename
|
|
if manifest.asset.title:
|
|
return manifest.asset.title
|
|
return manifest.storage.object_key.rsplit("/", 1)[-1] or manifest.asset.asset_id
|
|
|
|
|
|
def _event_data(manifest: AssetManifest, idempotency_key: str | None) -> dict[str, object]:
|
|
# redact_secrets masks any secret-bearing field (e.g. storage.kms_key_id)
|
|
# before the payload is persisted to processing_events.data.
|
|
return redact_secrets(
|
|
{
|
|
"asset_id": manifest.asset.asset_id,
|
|
"manifest_version": manifest.manifest_version,
|
|
"owner_module": manifest.asset.owner_module,
|
|
"owner_record_type": manifest.asset.owner_record_type,
|
|
"owner_record_id": manifest.asset.owner_record_id,
|
|
"idempotency_key": idempotency_key,
|
|
"storage": manifest.storage.model_dump(mode="json"),
|
|
"security": manifest.security.model_dump(mode="json"),
|
|
"retention": manifest.retention.model_dump(mode="json"),
|
|
}
|
|
)
|
|
|
|
|
|
def _metadata_from_manifest(manifest: dict[str, object]) -> dict[str, object]:
|
|
asset = manifest.get("asset") if isinstance(manifest, dict) else {}
|
|
storage = manifest.get("storage") if isinstance(manifest, dict) else {}
|
|
retention = manifest.get("retention") if isinstance(manifest, dict) else {}
|
|
security = manifest.get("security") if isinstance(manifest, dict) else {}
|
|
if not isinstance(asset, dict):
|
|
asset = {}
|
|
if not isinstance(storage, dict):
|
|
storage = {}
|
|
if not isinstance(retention, dict):
|
|
retention = {}
|
|
if not isinstance(security, dict):
|
|
security = {}
|
|
return {
|
|
"asset_id": asset.get("asset_id"),
|
|
"owner_module": asset.get("owner_module"),
|
|
"owner_record_type": asset.get("owner_record_type"),
|
|
"owner_record_id": asset.get("owner_record_id"),
|
|
"sha256": asset.get("sha256"),
|
|
"manifest_version": manifest.get("manifest_version"),
|
|
"retention": retention,
|
|
"legal_hold": retention.get("legal_hold"),
|
|
"security_gate_status": security.get("gate_status"),
|
|
"source_storage": storage,
|
|
}
|
|
|
|
|
|
class _DownloadTarget:
|
|
def __init__(self, temp_dir: Path | None) -> None:
|
|
self.temp_dir = temp_dir
|
|
self._managed_dir: tempfile.TemporaryDirectory[str] | None = None
|
|
self.path: Path | None = None
|
|
|
|
def __enter__(self) -> Path:
|
|
if self.temp_dir is None:
|
|
parent = Path(settings.app_work_dir) / "_knowledge_ingest"
|
|
parent.mkdir(parents=True, exist_ok=True)
|
|
self._managed_dir = tempfile.TemporaryDirectory(dir=parent)
|
|
self.path = Path(self._managed_dir.name) / "original"
|
|
else:
|
|
self.temp_dir.mkdir(parents=True, exist_ok=True)
|
|
self.path = self.temp_dir / f"{uuid.uuid4()}.original"
|
|
return self.path
|
|
|
|
def __exit__(self, exc_type, exc, tb) -> None:
|
|
if self._managed_dir is not None:
|
|
self._managed_dir.cleanup()
|