DealDocumentScreening/migrations/versions/0006_prescreen.py

"""Add prescreen_results table and expand status/queue enums.

Revision ID: 0006
Revises: 0005
Create Date: 2026-08-16

Adds the prescreen stage between extract and analyze:
- documents.status gains 'prescreening'
- jobs.queue gains 'prescreen'
- new prescreen_results table stores deterministic regex extraction output
  and routing decision
- reports gains FK + meta column for downstream enrichment
"""

from __future__ import annotations

from collections.abc import Sequence

import sqlalchemy as sa
from alembic import op
from sqlalchemy.dialects import postgresql

# revision identifiers, used by Alembic.
revision: str = "0006"
down_revision: str | None = "0005"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None


def upgrade() -> None:
    # ── prescreen_results ────────────────────────────────────────────────────
    op.create_table(
        "prescreen_results",
        sa.Column(
            "id",
            postgresql.UUID(as_uuid=True),
            primary_key=True,
            server_default=sa.text("gen_random_uuid()"),
        ),
        sa.Column(
            "document_id",
            postgresql.UUID(as_uuid=True),
            sa.ForeignKey("documents.id", ondelete="CASCADE"),
            nullable=False,
            index=True,
        ),
        sa.Column(
            "correlation_id",
            postgresql.UUID(as_uuid=True),
            nullable=False,
        ),
        sa.Column("contract_type", sa.String(64), nullable=True),
        sa.Column("party_a", sa.Text, nullable=True),
        sa.Column("party_b", sa.Text, nullable=True),
        sa.Column("total_amount", sa.Numeric(18, 2), nullable=True),
        sa.Column("currency", sa.String(8), nullable=True),
        sa.Column("start_date", sa.Date, nullable=True),
        sa.Column("end_date", sa.Date, nullable=True),
        sa.Column("has_penalty_clause", sa.Boolean, nullable=True),
        sa.Column("has_termination_clause", sa.Boolean, nullable=True),
        sa.Column("has_arbitration", sa.Boolean, nullable=True),
        sa.Column(
            "confidence_score",
            sa.Numeric(4, 3),
            nullable=True,
            comment="Field-coverage ratio 0.000-1.000",
        ),
        sa.Column(
            "routing_decision",
            sa.String(32),
            nullable=False,
            server_default=sa.text("'manual_review'"),
        ),
        sa.Column(
            "prescreened_at",
            sa.DateTime(timezone=True),
            nullable=False,
            server_default=sa.func.now(),
        ),
        sa.Column("processing_ms", sa.Integer, nullable=True),
        sa.Column(
            "extractor_version",
            sa.String(32),
            nullable=False,
            server_default=sa.text("'regex-v1'"),
        ),
        sa.Column("auto_summary", sa.Text, nullable=True),
        sa.Column(
            "auto_findings", postgresql.JSONB, nullable=False, server_default=sa.text("'[]'")
        ),
        sa.Column("error_message", sa.Text, nullable=True),
        sa.Column(
            "retry_count",
            sa.Integer,
            nullable=False,
            server_default=sa.text("0"),
        ),
        sa.CheckConstraint(
            "routing_decision IN ('auto_approve','manual_review','deep_analysis')",
            name="prescreen_results_routing_decision_check",
        ),
        sa.CheckConstraint(
            "confidence_score IS NULL OR (confidence_score >= 0 AND confidence_score <= 1)",
            name="prescreen_results_confidence_check",
        ),
        sa.CheckConstraint("retry_count >= 0", name="prescreen_results_retry_count_nonneg"),
    )

    op.create_index("prescreen_results_routing_idx", "prescreen_results", ["routing_decision"])
    op.create_index("prescreen_results_confidence_idx", "prescreen_results", ["confidence_score"])

    # ── reports enrichment from prescreen ────────────────────────────────────
    op.add_column(
        "reports",
        sa.Column(
            "prescreen_result_id",
            postgresql.UUID(as_uuid=True),
            sa.ForeignKey("prescreen_results.id", ondelete="SET NULL"),
            nullable=True,
        ),
    )
    op.add_column(
        "reports",
        sa.Column(
            "prescreen_meta",
            postgresql.JSONB,
            nullable=True,
            comment="Serialized PrescreenCompleted metadata used by the analyzer",
        ),
    )

    # ── expand additive TEXT+CHECK enums ─────────────────────────────────────
    # Drop the old narrower constraints first so the new definitions can reuse
    # the same names.
    op.drop_constraint("documents_status_check", "documents", type_="check")
    op.create_check_constraint(
        "documents_status_check",
        "documents",
        "status IN ('queued','extracting','prescreening','ocr','analyzing','done','failed')",
    )
    op.drop_constraint("jobs_queue_check", "jobs", type_="check")
    op.create_check_constraint(
        "jobs_queue_check",
        "jobs",
        "queue IN ('extract','prescreen','analyze')",
    )


def downgrade() -> None:
    op.drop_constraint("jobs_queue_check", "jobs", type_="check")
    op.create_check_constraint(
        "jobs_queue_check",
        "jobs",
        "queue IN ('extract','analyze')",
    )

    op.drop_constraint("documents_status_check", "documents", type_="check")
    op.create_check_constraint(
        "documents_status_check",
        "documents",
        "status IN ('queued','extracting','ocr','analyzing','done','failed')",
    )

    op.drop_column("reports", "prescreen_meta")
    op.drop_column("reports", "prescreen_result_id")

    op.drop_index("prescreen_results_confidence_idx", table_name="prescreen_results")
    op.drop_index("prescreen_results_routing_idx", table_name="prescreen_results")
    op.drop_table("prescreen_results")