120 lines
4.7 KiB
Python
120 lines
4.7 KiB
Python
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)
|