"""为用户表添加实体关联字段

Revision ID: 20241129_001
Revises:
Create Date: 2024-11-29 15:00:00

"""
from alembic import op
import sqlalchemy as sa

# revision identifiers
revision = '20241129_001'
down_revision = None
branch_labels = None
depends_on = None


def upgrade():
    """
    为 sys_user 表添加 entity_id 和 entity_type 字段
    """
    # 添加 entity_id 字段
    op.add_column(
        'sys_user',
        sa.Column('entity_id', sa.String(50), nullable=True, comment='关联的企业实体ID（MCN机构或主播）')
    )

    # 添加 entity_type 字段
    op.add_column(
        'sys_user',
        sa.Column('entity_type', sa.String(20), nullable=True, comment='关联实体类型：mcn-机构，streamer-主播')
    )

    # 创建索引（可选）
    op.create_index(
        op.f('ix_sys_user_entity_id'),
        'sys_user',
        ['entity_id'],
        unique=False
    )

    op.create_index(
        op.f('ix_sys_user_entity_type'),
        'sys_user',
        ['entity_type'],
        unique=False
    )


def downgrade():
    """
    移除 entity_id 和 entity_type 字段
    """
    # 删除索引
    op.drop_index(op.f('ix_sys_user_entity_type'), table_name='sys_user')
    op.drop_index(op.f('ix_sys_user_entity_id'), table_name='sys_user')

    # 删除字段
    op.drop_column('sys_user', 'entity_type')
    op.drop_column('sys_user', 'entity_id')