from __future__ import annotations import json from datetime import UTC, datetime from pathlib import Path from typing import Any from sqlalchemy import DateTime, String, Text, create_engine from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column, sessionmaker from datapilot.config import Settings class Base(DeclarativeBase): pass class AnalysisRun(Base): __tablename__ = "analysis_runs" run_id: Mapped[str] = mapped_column(String(64), primary_key=True) dataset_name: Mapped[str] = mapped_column(String(255)) status: Mapped[str] = mapped_column(String(32), index=True) created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True)) updated_at: Mapped[datetime] = mapped_column(DateTime(timezone=True)) payload_json: Mapped[str] = mapped_column(Text) class RunStore: """SQLite locally; set DATABASE_URL to PostgreSQL in production.""" def __init__(self, settings: Settings): self.engine = create_engine(settings.database_url, future=True) Base.metadata.create_all(self.engine) self.sessions = sessionmaker(self.engine, expire_on_commit=False) def save(self, run_id: str, dataset_name: str, status: str, payload: dict[str, Any]) -> None: now = datetime.now(UTC) with self.sessions.begin() as session: record = session.get(AnalysisRun, run_id) if record is None: record = AnalysisRun( run_id=run_id, dataset_name=dataset_name, status=status, created_at=now, updated_at=now, payload_json=json.dumps(payload, default=str), ) session.add(record) else: record.status = status record.updated_at = now record.payload_json = json.dumps(payload, default=str) def get(self, run_id: str) -> dict[str, Any] | None: with self.sessions() as session: record = session.get(AnalysisRun, run_id) return json.loads(record.payload_json) if record else None def list_recent(self, limit: int = 20) -> list[dict[str, Any]]: from sqlalchemy import select with self.sessions() as session: records = session.scalars( select(AnalysisRun).order_by(AnalysisRun.created_at.desc()).limit(limit) ) return [ { "run_id": record.run_id, "dataset_name": record.dataset_name, "status": record.status, "created_at": record.created_at.isoformat(), } for record in records ] class ArtifactStore: """Local artifact storage with an interface that can be replaced by S3/MinIO.""" def __init__(self, root: Path): self.root = root self.root.mkdir(parents=True, exist_ok=True) def run_directory(self, run_id: str) -> Path: directory = (self.root / run_id).resolve() if self.root.resolve() not in directory.parents: raise ValueError("Invalid run identifier.") directory.mkdir(parents=True, exist_ok=True) return directory