ax-backend/migrations/versions/003_update_documents_table.py

"""update documents table - remove title, template_id, make session_id required

Revision ID: 003
Revises: 002
Create Date: 2026-06-24
"""

from alembic import op
import sqlalchemy as sa

revision = "003"
down_revision = "002"
branch_labels = None
depends_on = None


def upgrade() -> None:
    # 移除 title 和 template_id 列
    op.drop_column("documents", "title")
    op.drop_column("documents", "template_id")
    
    # 将 session_id 改为 NOT NULL（先设置默认值，再改约束）
    op.execute("UPDATE documents SET session_id = 'unknown' WHERE session_id IS NULL")
    op.alter_column(
        "documents",
        "session_id",
        existing_type=sa.String(128),
        nullable=False,
    )


def downgrade() -> None:
    # 恢复 title 和 template_id
    op.add_column("documents", sa.Column("title", sa.String(255), nullable=False, server_default=""))
    op.add_column("documents", sa.Column("template_id", sa.String(64), nullable=True))
    
    # 恢复 session_id 可为空
    op.alter_column(
        "documents",
        "session_id",
        existing_type=sa.String(128),
        nullable=True,
    )