vent_ai_platform/alembic/versions/0002_document_processing.py

"""document processing workflow

Revision ID: 0002_document_processing
Revises: 0001_initial
Create Date: 2026-10-09
"""

from __future__ import annotations

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


revision = "0002_document_processing"
down_revision = "0001_initial"
branch_labels = None
depends_on = None


def upgrade() -> None:
    op.add_column(
        "business_documents",
        sa.Column(
            "document_type",
            sa.String(length=32),
            nullable=False,
            server_default="",
        ),
    )
    op.create_index(
        "ix_business_documents_document_type",
        "business_documents",
        ["document_type"],
    )

    op.create_table(
        "document_extraction_templates",
        sa.Column("code", sa.String(length=64), nullable=False),
        sa.Column("name", sa.String(length=128), nullable=False),
        sa.Column("document_type", sa.String(length=32), nullable=False),
        sa.Column("description", sa.Text(), nullable=False),
        sa.Column("version", sa.String(length=32), nullable=False),
        sa.Column("prompt", sa.Text(), nullable=False),
        sa.Column("extraction_schema", sa.JSON(), nullable=False),
        sa.Column("enabled", sa.Boolean(), nullable=False),
        sa.Column("id", sa.String(length=32), nullable=False),
        sa.Column("created_at", sa.DateTime(timezone=True), nullable=False),
        sa.Column("updated_at", sa.DateTime(timezone=True), nullable=False),
        sa.PrimaryKeyConstraint("id", name="pk_document_extraction_templates"),
        sa.UniqueConstraint("code", name="uq_document_extraction_templates_code"),
    )
    op.create_index(
        "ix_document_extraction_templates_code",
        "document_extraction_templates",
        ["code"],
    )
    op.create_index(
        "ix_document_extraction_templates_document_type",
        "document_extraction_templates",
        ["document_type"],
    )
    op.create_index(
        "ix_document_extraction_templates_enabled",
        "document_extraction_templates",
        ["enabled"],
    )

    op.create_table(
        "document_processing_runs",
        sa.Column("document_id", sa.String(length=32), nullable=False),
        sa.Column("template_id", sa.String(length=32), nullable=True),
        sa.Column("status", sa.String(length=32), nullable=False),
        sa.Column("stage", sa.String(length=32), nullable=False),
        sa.Column("progress", sa.Integer(), nullable=False),
        sa.Column("document_type", sa.String(length=32), nullable=False),
        sa.Column("parser_engine", sa.String(length=32), nullable=False),
        sa.Column("parsed_markdown_path", sa.String(length=512), nullable=False),
        sa.Column(
            "parsed_markdown",
            sa.Text().with_variant(mysql.LONGTEXT(), "mysql"),
            nullable=False,
        ),
        sa.Column("model_output", sa.JSON(), nullable=False),
        sa.Column("error", sa.Text(), nullable=False),
        sa.Column("started_at", sa.DateTime(timezone=True), nullable=True),
        sa.Column("finished_at", sa.DateTime(timezone=True), nullable=True),
        sa.Column("metadata", sa.JSON(), nullable=False),
        sa.Column("id", sa.String(length=32), nullable=False),
        sa.Column("created_at", sa.DateTime(timezone=True), nullable=False),
        sa.Column("updated_at", sa.DateTime(timezone=True), nullable=False),
        sa.ForeignKeyConstraint(
            ["document_id"],
            ["business_documents.id"],
            name="fk_document_processing_runs_document",
            ondelete="CASCADE",
        ),
        sa.ForeignKeyConstraint(
            ["template_id"],
            ["document_extraction_templates.id"],
            name="fk_document_processing_runs_template",
            ondelete="SET NULL",
        ),
        sa.PrimaryKeyConstraint("id", name="pk_document_processing_runs"),
    )
    op.create_index(
        "ix_document_processing_runs_document_id",
        "document_processing_runs",
        ["document_id"],
    )
    op.create_index(
        "ix_document_processing_runs_template_id",
        "document_processing_runs",
        ["template_id"],
    )
    op.create_index(
        "ix_document_processing_runs_status",
        "document_processing_runs",
        ["status"],
    )
    op.create_index(
        "ix_document_processing_runs_document_type",
        "document_processing_runs",
        ["document_type"],
    )


def downgrade() -> None:
    op.drop_index(
        "ix_document_processing_runs_document_type",
        table_name="document_processing_runs",
    )
    op.drop_index(
        "ix_document_processing_runs_status",
        table_name="document_processing_runs",
    )
    op.drop_index(
        "ix_document_processing_runs_template_id",
        table_name="document_processing_runs",
    )
    op.drop_index(
        "ix_document_processing_runs_document_id",
        table_name="document_processing_runs",
    )
    op.drop_table("document_processing_runs")

    op.drop_index(
        "ix_document_extraction_templates_enabled",
        table_name="document_extraction_templates",
    )
    op.drop_index(
        "ix_document_extraction_templates_document_type",
        table_name="document_extraction_templates",
    )
    op.drop_index(
        "ix_document_extraction_templates_code",
        table_name="document_extraction_templates",
    )
    op.drop_table("document_extraction_templates")

    op.drop_index(
        "ix_business_documents_document_type",
        table_name="business_documents",
    )
    op.drop_column("business_documents", "document_type")