ai-agent-admin/backend-fastapi/alembic/versions/f6g7h8i9j0k1_ensure_user_oauth_userid_columns.py

"""Ensure organization and workflow runtime columns exist.

Revision ID: f6g7h8i9j0k1
Revises: e5f6g7h8i9j0
Create Date: 2026-06-13 16:00:00.000000
"""
from typing import Sequence, Union

import sqlalchemy as sa
from alembic import op


revision: str = "f6g7h8i9j0k1"
down_revision: Union[str, None] = "e5f6g7h8i9j0"
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def _column_exists(table_name: str, column_name: str) -> bool:
    inspector = sa.inspect(op.get_bind())
    return column_name in {column["name"] for column in inspector.get_columns(table_name)}


def _index_exists(table_name: str, index_name: str) -> bool:
    inspector = sa.inspect(op.get_bind())
    return index_name in {index["name"] for index in inspector.get_indexes(table_name)}


def upgrade() -> None:
    if not _column_exists("core_dept", "dingtalk_dept_id"):
        op.add_column(
            "core_dept",
            sa.Column(
                "dingtalk_dept_id",
                sa.String(length=64),
                nullable=True,
                comment="DingTalk department ID",
            ),
        )

    if not _column_exists("core_dept", "wecom_dept_id"):
        op.add_column(
            "core_dept",
            sa.Column(
                "wecom_dept_id",
                sa.String(length=64),
                nullable=True,
                comment="WeCom department ID",
            ),
        )

    if not _column_exists("core_dept", "feishu_dept_id"):
        op.add_column(
            "core_dept",
            sa.Column(
                "feishu_dept_id",
                sa.String(length=128),
                nullable=True,
                comment="Feishu department ID",
            ),
        )

    if not _column_exists("core_user", "dingtalk_userid"):
        op.add_column(
            "core_user",
            sa.Column(
                "dingtalk_userid",
                sa.String(length=200),
                nullable=True,
                comment="DingTalk user ID",
            ),
        )

    if not _column_exists("core_user", "feishu_userid"):
        op.add_column(
            "core_user",
            sa.Column(
                "feishu_userid",
                sa.String(length=200),
                nullable=True,
                comment="Feishu user ID",
            ),
        )

    if not _index_exists("core_user", "ix_core_user_dingtalk_userid"):
        op.create_index(
            "ix_core_user_dingtalk_userid",
            "core_user",
            ["dingtalk_userid"],
            unique=True,
        )

    if not _index_exists("core_user", "ix_core_user_feishu_userid"):
        op.create_index(
            "ix_core_user_feishu_userid",
            "core_user",
            ["feishu_userid"],
            unique=True,
        )

    if not _index_exists("core_dept", "ix_core_dept_dingtalk_dept_id"):
        op.create_index(
            "ix_core_dept_dingtalk_dept_id",
            "core_dept",
            ["dingtalk_dept_id"],
            unique=True,
        )

    if not _index_exists("core_dept", "ix_core_dept_wecom_dept_id"):
        op.create_index(
            "ix_core_dept_wecom_dept_id",
            "core_dept",
            ["wecom_dept_id"],
            unique=True,
        )

    if not _index_exists("core_dept", "ix_core_dept_feishu_dept_id"):
        op.create_index(
            "ix_core_dept_feishu_dept_id",
            "core_dept",
            ["feishu_dept_id"],
            unique=True,
        )

    if not _column_exists("ai_workflow_run", "trigger_type"):
        op.add_column(
            "ai_workflow_run",
            sa.Column(
                "trigger_type",
                sa.String(length=30),
                nullable=True,
                comment="Workflow trigger source",
            ),
        )
    if not _column_exists("ai_workflow_run", "use_draft"):
        op.add_column(
            "ai_workflow_run",
            sa.Column(
                "use_draft",
                sa.Boolean(),
                nullable=True,
                comment="Whether the run uses draft definition",
            ),
        )
    if not _column_exists("ai_workflow_run", "workflow_version"):
        op.add_column(
            "ai_workflow_run",
            sa.Column(
                "workflow_version",
                sa.Integer(),
                nullable=True,
                comment="Published workflow version used for this run",
            ),
        )
    if not _column_exists("ai_workflow_run", "definition_snapshot"):
        op.add_column(
            "ai_workflow_run",
            sa.Column(
                "definition_snapshot",
                sa.JSON(),
                nullable=True,
                comment="Workflow definition snapshot at run start",
            ),
        )

    op.execute("UPDATE ai_workflow_run SET trigger_type = 'api' WHERE trigger_type IS NULL")
    op.execute("UPDATE ai_workflow_run SET use_draft = false WHERE use_draft IS NULL")
    op.execute("UPDATE ai_workflow_run SET definition_snapshot = '{}' WHERE definition_snapshot IS NULL")


def downgrade() -> None:
    # This migration repairs drifted databases. On a clean schema these columns
    # come from earlier revisions, so rollback must not remove them.
    pass