DataOps-platform/app/models/data_research.py

from __future__ import annotations

from typing import Any

from sqlalchemy.dialects.postgresql import JSONB, UUID

from app import db
from app.core.common.identifiers import new_governance_uid
from app.core.common.timezone_utils import now_china_naive


JOB_STATUSES = (
    "created",
    "queued",
    "extracting",
    "normalizing",
    "matching",
    "awaiting_review",
    "published",
    "partial",
    "failed",
    "cancelled",
)


def _iso(value: Any) -> str | None:
    return value.isoformat() if value is not None else None


class IngestionSource(db.Model):
    __tablename__ = "ingestion_sources"
    __table_args__ = (
        db.CheckConstraint(
            "source_type IN ('database','file','ddl')",
            name="ck_ingestion_source_type",
        ),
        db.CheckConstraint(
            "status IN ('active','disabled')",
            name="ck_ingestion_source_status",
        ),
        {"schema": "public"},
    )

    uid = db.Column(UUID(as_uuid=False), primary_key=True, default=new_governance_uid)
    source_type = db.Column(db.String(20), nullable=False)
    name = db.Column(db.String(300), nullable=False)
    config = db.Column(JSONB, nullable=False, default=dict)
    permission_scope = db.Column(JSONB, nullable=False, default=dict)
    status = db.Column(db.String(20), nullable=False, default="active")
    created_by = db.Column(db.String(100))
    created_at = db.Column(db.DateTime, nullable=False, default=now_china_naive)
    updated_at = db.Column(db.DateTime, nullable=False, default=now_china_naive)

    def to_dict(self) -> dict[str, Any]:
        return {
            "uid": str(self.uid),
            "source_type": self.source_type,
            "name": self.name,
            "permission_scope": dict(self.permission_scope or {}),
            "status": self.status,
            "created_by": self.created_by,
            "created_at": _iso(self.created_at),
            "updated_at": _iso(self.updated_at),
        }


class SourceArtifact(db.Model):
    __tablename__ = "source_artifacts"
    __table_args__ = (
        db.CheckConstraint("size_bytes >= 0", name="ck_source_artifact_size"),
        db.UniqueConstraint(
            "source_uid",
            "content_hash",
            "parser_version",
            name="uq_source_artifact_parser_content",
        ),
        {"schema": "public"},
    )

    uid = db.Column(UUID(as_uuid=False), primary_key=True, default=new_governance_uid)
    source_uid = db.Column(
        UUID(as_uuid=False),
        db.ForeignKey("public.ingestion_sources.uid"),
        nullable=False,
    )
    filename = db.Column(db.String(500), nullable=False)
    media_type = db.Column(db.String(200), nullable=False)
    size_bytes = db.Column(db.BigInteger, nullable=False)
    content_hash = db.Column(db.String(64), nullable=False)
    storage_ref = db.Column(db.Text, nullable=False)
    parser_version = db.Column(db.String(100), nullable=False)
    created_at = db.Column(db.DateTime, nullable=False, default=now_china_naive)

    def to_dict(self) -> dict[str, Any]:
        return {
            "uid": str(self.uid),
            "source_uid": str(self.source_uid),
            "filename": self.filename,
            "media_type": self.media_type,
            "size_bytes": int(self.size_bytes),
            "content_hash": self.content_hash,
            "parser_version": self.parser_version,
            "created_at": _iso(self.created_at),
        }


class IngestionJob(db.Model):
    __tablename__ = "ingestion_jobs"
    __table_args__ = (
        db.CheckConstraint(
            "status IN (" + ",".join(f"'{value}'" for value in JOB_STATUSES) + ")",
            name="ck_ingestion_job_status",
        ),
        db.UniqueConstraint(
            "idempotency_key",
            name="uq_ingestion_job_idempotency_key",
        ),
        {"schema": "public"},
    )

    uid = db.Column(UUID(as_uuid=False), primary_key=True, default=new_governance_uid)
    source_uid = db.Column(
        UUID(as_uuid=False),
        db.ForeignKey("public.ingestion_sources.uid"),
        nullable=False,
    )
    artifact_uid = db.Column(
        UUID(as_uuid=False),
        db.ForeignKey("public.source_artifacts.uid"),
    )
    job_type = db.Column(db.String(30), nullable=False)
    status = db.Column(db.String(30), nullable=False, default="created")
    idempotency_key = db.Column(db.String(64), nullable=False, unique=True)
    parser_version = db.Column(db.String(100), nullable=False)
    parameters = db.Column(JSONB, nullable=False, default=dict)
    statistics = db.Column(JSONB, nullable=False, default=dict)
    last_error = db.Column(db.String(1000))
    actor_uid = db.Column(db.String(100))
    force_rerun = db.Column(db.Boolean, nullable=False, default=False)
    created_at = db.Column(db.DateTime, nullable=False, default=now_china_naive)
    updated_at = db.Column(db.DateTime, nullable=False, default=now_china_naive)
    started_at = db.Column(db.DateTime)
    finished_at = db.Column(db.DateTime)

    def to_dict(self) -> dict[str, Any]:
        return {
            "uid": str(self.uid),
            "source_uid": str(self.source_uid),
            "artifact_uid": str(self.artifact_uid) if self.artifact_uid else None,
            "job_type": self.job_type,
            "status": self.status,
            "parser_version": self.parser_version,
            "parameters": dict(self.parameters or {}),
            "statistics": dict(self.statistics or {}),
            "last_error": self.last_error,
            "actor_uid": self.actor_uid,
            "created_at": _iso(self.created_at),
            "updated_at": _iso(self.updated_at),
            "started_at": _iso(self.started_at),
            "finished_at": _iso(self.finished_at),
        }


class EvidenceFragment(db.Model):
    __tablename__ = "evidence_fragments"
    __table_args__ = (
        db.CheckConstraint(
            "confidence IS NULL OR (confidence >= 0 AND confidence <= 1)",
            name="ck_evidence_confidence",
        ),
        {"schema": "public"},
    )

    uid = db.Column(UUID(as_uuid=False), primary_key=True, default=new_governance_uid)
    job_uid = db.Column(
        UUID(as_uuid=False),
        db.ForeignKey("public.ingestion_jobs.uid", ondelete="CASCADE"),
        nullable=False,
    )
    artifact_uid = db.Column(
        UUID(as_uuid=False),
        db.ForeignKey("public.source_artifacts.uid"),
    )
    locator = db.Column(JSONB, nullable=False)
    excerpt = db.Column(db.Text)
    confidence = db.Column(db.Float)
    created_at = db.Column(db.DateTime, nullable=False, default=now_china_naive)

    def to_dict(self) -> dict[str, Any]:
        return {
            "uid": str(self.uid),
            "job_uid": str(self.job_uid),
            "artifact_uid": str(self.artifact_uid) if self.artifact_uid else None,
            "locator": dict(self.locator or {}),
            "excerpt": self.excerpt,
            "confidence": self.confidence,
            "created_at": _iso(self.created_at),
        }


class ExtractionCandidate(db.Model):
    __tablename__ = "extraction_candidates"
    __table_args__ = (
        db.CheckConstraint(
            "confidence >= 0 AND confidence <= 1",
            name="ck_candidate_confidence",
        ),
        db.CheckConstraint(
            "status IN ('candidate','accepted','rejected','ignored')",
            name="ck_candidate_status",
        ),
        {"schema": "public"},
    )

    uid = db.Column(UUID(as_uuid=False), primary_key=True, default=new_governance_uid)
    job_uid = db.Column(
        UUID(as_uuid=False),
        db.ForeignKey("public.ingestion_jobs.uid", ondelete="CASCADE"),
        nullable=False,
    )
    candidate_type = db.Column(db.String(40), nullable=False)
    normalized_data = db.Column(JSONB, nullable=False)
    evidence_uids = db.Column(JSONB, nullable=False, default=list)
    confidence = db.Column(db.Float, nullable=False)
    status = db.Column(db.String(20), nullable=False, default="candidate")
    created_at = db.Column(db.DateTime, nullable=False, default=now_china_naive)
    updated_at = db.Column(db.DateTime, nullable=False, default=now_china_naive)

    def to_dict(self) -> dict[str, Any]:
        return {
            "uid": str(self.uid),
            "job_uid": str(self.job_uid),
            "candidate_type": self.candidate_type,
            "normalized_data": dict(self.normalized_data or {}),
            "evidence_uids": list(self.evidence_uids or []),
            "confidence": float(self.confidence),
            "status": self.status,
            "created_at": _iso(self.created_at),
            "updated_at": _iso(self.updated_at),
        }