vent_ai_platform/alembic/versions/0005_extraction_judgement_status.py

"""add extraction judgement status

Revision ID: 0005_extraction_judgement_status
Revises: 0004_extraction_edit_accuracy
Create Date: 2026-10-12
"""

from __future__ import annotations

import sqlalchemy as sa
from alembic import op


revision = "0005_extraction_judgement_status"
down_revision = "0004_extraction_edit_accuracy"
branch_labels = None
depends_on = None


def upgrade() -> None:
    op.add_column(
        "document_extractions",
        sa.Column(
            "edit_judgement_status",
            sa.String(length=32),
            nullable=False,
            server_default="not_required",
        ),
    )
    op.execute(
        """
        UPDATE document_extractions
        SET edit_judgement_status = CASE
            WHEN edit_judged_at IS NOT NULL THEN 'completed'
            WHEN edited = 1 THEN 'pending'
            ELSE 'not_required'
        END
        """
    )
    op.create_index(
        "ix_document_extractions_edit_judgement_status",
        "document_extractions",
        ["edit_judgement_status"],
    )


def downgrade() -> None:
    op.drop_index(
        "ix_document_extractions_edit_judgement_status",
        table_name="document_extractions",
    )
    op.drop_column("document_extractions", "edit_judgement_status")