geointel/backend/app/api/routes/source_registry.py

from __future__ import annotations

from uuid import UUID

from fastapi import APIRouter, Depends
from sqlalchemy import func, or_
from sqlalchemy.orm import Session

from app.core.errors import AppError
from app.db.session import get_db
from app.models import Dataset, DatasetLineageEdge, DatasetQuarantine, Project, SourceRegistry, SourceSnapshot
from app.schemas import (
    DatasetLineageEdgeRead,
    DatasetProvenanceRead,
    DatasetQuarantineRead,
    Envelope,
    ItemList,
    SourceRegistryDetailRead,
    SourceRegistryRead,
    SourceSnapshotRead,
)
from app.utils.response import envelope


router = APIRouter(tags=["source-registry"])


def _source_read(source: SourceRegistry, *, snapshot_count: int = 0) -> SourceRegistryRead:
    return SourceRegistryRead.model_validate(source).model_copy(update={"snapshot_count": int(snapshot_count)})


@router.get("/source-registry", response_model=Envelope[ItemList[SourceRegistryRead]])
def list_source_registry(
    classification: str | None = None,
    db: Session = Depends(get_db),
) -> dict:
    query = (
        db.query(SourceRegistry, func.count(SourceSnapshot.id).label("snapshot_count"))
        .outerjoin(SourceSnapshot, SourceSnapshot.source_registry_id == SourceRegistry.id)
    )
    if classification:
        query = query.filter(SourceRegistry.classification == classification.strip().lower())
    rows = (
        query.group_by(SourceRegistry.id)
        .order_by(SourceRegistry.classification.asc(), SourceRegistry.display_name.asc())
        .all()
    )
    items = [_source_read(source, snapshot_count=count) for source, count in rows]
    return envelope({"items": items, "total": len(items)})


@router.get("/source-registry/{source_key}", response_model=Envelope[SourceRegistryDetailRead])
def get_source_registry_entry(source_key: str, db: Session = Depends(get_db)) -> dict:
    normalized_key = source_key.strip().lower()
    source = db.query(SourceRegistry).filter(SourceRegistry.source_key == normalized_key).one_or_none()
    if source is None:
        raise AppError(code="SOURCE_REGISTRY_ENTRY_NOT_FOUND", message="Source registry entry was not found", status_code=404)
    snapshots = (
        db.query(SourceSnapshot)
        .filter(SourceSnapshot.source_registry_id == source.id)
        .order_by(SourceSnapshot.fetched_at.desc(), SourceSnapshot.created_at.desc())
        .all()
    )
    detail = SourceRegistryDetailRead(
        source=_source_read(source, snapshot_count=len(snapshots)),
        snapshots=[SourceSnapshotRead.model_validate(snapshot) for snapshot in snapshots],
    )
    return envelope(detail)


@router.get(
    "/projects/{project_id}/datasets/{dataset_id}/provenance",
    response_model=Envelope[DatasetProvenanceRead],
)
def get_dataset_provenance(
    project_id: UUID,
    dataset_id: UUID,
    db: Session = Depends(get_db),
) -> dict:
    if db.get(Project, project_id) is None:
        raise AppError(code="PROJECT_NOT_FOUND", message="Project not found", status_code=404)
    dataset = db.get(Dataset, dataset_id)
    if dataset is None or dataset.project_id != project_id:
        raise AppError(code="DATASET_NOT_FOUND", message="Dataset not found", status_code=404)

    source = db.get(SourceRegistry, dataset.source_registry_id) if dataset.source_registry_id else None
    snapshot = db.get(SourceSnapshot, dataset.source_snapshot_id) if dataset.source_snapshot_id else None
    lineage = (
        db.query(DatasetLineageEdge)
        .filter(
            or_(
                DatasetLineageEdge.parent_dataset_id == dataset.id,
                DatasetLineageEdge.child_dataset_id == dataset.id,
            )
        )
        .order_by(DatasetLineageEdge.created_at.asc(), DatasetLineageEdge.id.asc())
        .all()
    )
    quarantines = (
        db.query(DatasetQuarantine)
        .filter(DatasetQuarantine.dataset_id == dataset.id)
        .order_by(DatasetQuarantine.created_at.desc(), DatasetQuarantine.id.desc())
        .all()
    )
    result = DatasetProvenanceRead(
        dataset_id=dataset.id,
        source=_source_read(source) if source else None,
        snapshot=SourceSnapshotRead.model_validate(snapshot) if snapshot else None,
        data_contract_key=dataset.data_contract_key,
        data_contract_version=dataset.data_contract_version,
        validation_status=dataset.validation_status,
        validation_report_json=dataset.validation_report_json,
        provenance_status=dataset.provenance_status,
        lineage_status=dataset.lineage_status,
        quarantine_status=dataset.quarantine_status,
        lineage=[DatasetLineageEdgeRead.model_validate(item) for item in lineage],
        quarantines=[DatasetQuarantineRead.model_validate(item) for item in quarantines],
    )
    return envelope(result)