225 lines
8.5 KiB
Python
225 lines
8.5 KiB
Python
"""202512051405
|
||
|
||
Revision ID: 20a742ef1d93
|
||
Revises: fc9664b8e4a1
|
||
Create Date: 2025-12-05 14:24:33.545865
|
||
|
||
"""
|
||
from typing import Sequence, Union
|
||
|
||
from alembic import op
|
||
import sqlalchemy as sa
|
||
from sqlalchemy import inspect
|
||
from sqlalchemy.dialects import postgresql
|
||
|
||
# revision identifiers, used by Alembic.
|
||
revision: str = '20a742ef1d93'
|
||
down_revision: Union[str, None] = 'fc9664b8e4a1'
|
||
branch_labels: Union[str, Sequence[str], None] = None
|
||
depends_on: Union[str, Sequence[str], None] = None
|
||
|
||
|
||
def table_exists(table_name: str) -> bool:
|
||
"""检查表是否存在"""
|
||
bind = op.get_bind()
|
||
inspector = inspect(bind)
|
||
return table_name in inspector.get_table_names()
|
||
|
||
|
||
def column_exists(table_name: str, column_name: str) -> bool:
|
||
"""检查列是否存在"""
|
||
bind = op.get_bind()
|
||
inspector = inspect(bind)
|
||
if not table_exists(table_name):
|
||
return False
|
||
columns = [col['name'] for col in inspector.get_columns(table_name)]
|
||
return column_name in columns
|
||
|
||
|
||
def index_exists(table_name: str, index_name: str) -> bool:
|
||
"""检查索引是否存在"""
|
||
bind = op.get_bind()
|
||
inspector = inspect(bind)
|
||
if not table_exists(table_name):
|
||
return False
|
||
indexes = [idx['name'] for idx in inspector.get_indexes(table_name)]
|
||
return index_name in indexes
|
||
|
||
|
||
def constraint_exists(table_name: str, constraint_name: str) -> bool:
|
||
"""检查约束是否存在（外键、唯一约束等）"""
|
||
bind = op.get_bind()
|
||
inspector = inspect(bind)
|
||
if not table_exists(table_name):
|
||
return False
|
||
|
||
# 检查外键约束
|
||
foreign_keys = [fk['name'] for fk in inspector.get_foreign_keys(table_name) if fk['name']]
|
||
if constraint_name in foreign_keys:
|
||
return True
|
||
|
||
# 检查唯一约束
|
||
unique_constraints = [uc['name'] for uc in inspector.get_unique_constraints(table_name) if uc['name']]
|
||
if constraint_name in unique_constraints:
|
||
return True
|
||
|
||
# 检查检查约束
|
||
check_constraints = [cc['name'] for cc in inspector.get_check_constraints(table_name) if cc['name']]
|
||
if constraint_name in check_constraints:
|
||
return True
|
||
|
||
return False
|
||
|
||
|
||
def trigger_exists(trigger_name: str) -> bool:
|
||
"""检查触发器是否存在（PostgreSQL）"""
|
||
bind = op.get_bind()
|
||
result = bind.execute(sa.text(
|
||
"SELECT EXISTS (SELECT 1 FROM pg_trigger WHERE tgname = :trigger_name)"
|
||
), {"trigger_name": trigger_name})
|
||
return result.scalar()
|
||
|
||
|
||
def sequence_exists(sequence_name: str) -> bool:
|
||
"""检查序列是否存在（PostgreSQL）"""
|
||
bind = op.get_bind()
|
||
result = bind.execute(sa.text(
|
||
"SELECT EXISTS (SELECT 1 FROM pg_class WHERE relkind = 'S' AND relname = :sequence_name)"
|
||
), {"sequence_name": sequence_name})
|
||
return result.scalar()
|
||
|
||
|
||
def enum_exists(enum_name: str) -> bool:
|
||
"""检查枚举类型是否存在（PostgreSQL）"""
|
||
bind = op.get_bind()
|
||
result = bind.execute(sa.text(
|
||
"SELECT EXISTS (SELECT 1 FROM pg_type WHERE typname = :enum_name)"
|
||
), {"enum_name": enum_name})
|
||
return result.scalar()
|
||
|
||
|
||
def upgrade() -> None:
|
||
# ### commands auto generated by Alembic - please adjust! ###
|
||
with op.batch_alter_table('data_config', schema=None) as batch_op:
|
||
batch_op.alter_column('llm',
|
||
existing_type=sa.VARCHAR(length=255),
|
||
comment='LLM模型配置ID',
|
||
existing_comment='LS',
|
||
existing_nullable=True)
|
||
|
||
with op.batch_alter_table('end_users', schema=None) as batch_op:
|
||
batch_op.alter_column('id',
|
||
existing_type=sa.UUID(),
|
||
server_default=None,
|
||
existing_nullable=False)
|
||
|
||
with op.batch_alter_table('knowledges', schema=None) as batch_op:
|
||
batch_op.alter_column('permission_id',
|
||
existing_type=sa.VARCHAR(),
|
||
comment='permission ID:Private|Share|Memory',
|
||
existing_comment='permission ID:Private|Share',
|
||
existing_nullable=True)
|
||
|
||
with op.batch_alter_table('memory_increments', schema=None) as batch_op:
|
||
batch_op.alter_column('id',
|
||
existing_type=sa.UUID(),
|
||
server_default=None,
|
||
existing_nullable=False)
|
||
batch_op.alter_column('created_at',
|
||
existing_type=postgresql.TIMESTAMP(precision=0),
|
||
server_default=None,
|
||
existing_nullable=True)
|
||
batch_op.alter_column('updated_at',
|
||
existing_type=postgresql.TIMESTAMP(precision=0),
|
||
server_default=None,
|
||
existing_nullable=True)
|
||
|
||
# 为 model_configs 添加 tenant_id 列和相关约束
|
||
if table_exists('model_configs'):
|
||
# 添加列（分步骤处理 NOT NULL 约束）
|
||
if not column_exists('model_configs', 'tenant_id'):
|
||
# 步骤1: 先添加可空列
|
||
with op.batch_alter_table('model_configs', schema=None) as batch_op:
|
||
batch_op.add_column(sa.Column('tenant_id', sa.UUID(), nullable=True, comment='租户ID'))
|
||
|
||
# 步骤2: 为现有数据填充默认值
|
||
# 注意：这里需要根据实际业务逻辑填充，这里假设获取第一个租户ID
|
||
bind = op.get_bind()
|
||
result = bind.execute(sa.text("SELECT id FROM tenants LIMIT 1"))
|
||
first_tenant = result.scalar()
|
||
|
||
if first_tenant:
|
||
# 更新现有记录
|
||
bind.execute(
|
||
sa.text("UPDATE model_configs SET tenant_id = :tenant_id WHERE tenant_id IS NULL"),
|
||
{"tenant_id": first_tenant}
|
||
)
|
||
|
||
# 步骤3: 将列改为 NOT NULL
|
||
with op.batch_alter_table('model_configs', schema=None) as batch_op:
|
||
batch_op.alter_column('tenant_id', nullable=False)
|
||
|
||
# 创建索引
|
||
if not index_exists('model_configs', 'ix_model_configs_tenant_id'):
|
||
op.create_index('ix_model_configs_tenant_id', 'model_configs', ['tenant_id'], unique=False)
|
||
|
||
# 创建外键
|
||
if not constraint_exists('model_configs', 'model_configs_tenant_id_fkey'):
|
||
op.create_foreign_key('model_configs_tenant_id_fkey', 'model_configs', 'tenants', ['tenant_id'], ['id'])
|
||
|
||
# ### end Alembic commands ###
|
||
|
||
|
||
def downgrade() -> None:
|
||
# ### commands auto generated by Alembic - please adjust! ###
|
||
# 回滚 model_configs 的修改
|
||
if table_exists('model_configs'):
|
||
# 删除外键
|
||
if constraint_exists('model_configs', 'model_configs_tenant_id_fkey'):
|
||
op.drop_constraint('model_configs_tenant_id_fkey', 'model_configs', type_='foreignkey')
|
||
|
||
# 删除索引
|
||
if index_exists('model_configs', 'ix_model_configs_tenant_id'):
|
||
op.drop_index('ix_model_configs_tenant_id', table_name='model_configs')
|
||
|
||
# 删除列
|
||
if column_exists('model_configs', 'tenant_id'):
|
||
with op.batch_alter_table('model_configs', schema=None) as batch_op:
|
||
batch_op.drop_column('tenant_id')
|
||
|
||
with op.batch_alter_table('memory_increments', schema=None) as batch_op:
|
||
batch_op.alter_column('updated_at',
|
||
existing_type=postgresql.TIMESTAMP(precision=0),
|
||
server_default=sa.text('CURRENT_TIMESTAMP'),
|
||
existing_nullable=True)
|
||
batch_op.alter_column('created_at',
|
||
existing_type=postgresql.TIMESTAMP(precision=0),
|
||
server_default=sa.text('CURRENT_TIMESTAMP'),
|
||
existing_nullable=True)
|
||
batch_op.alter_column('id',
|
||
existing_type=sa.UUID(),
|
||
server_default=sa.text('gen_random_uuid()'),
|
||
existing_nullable=False)
|
||
|
||
with op.batch_alter_table('knowledges', schema=None) as batch_op:
|
||
batch_op.alter_column('permission_id',
|
||
existing_type=sa.VARCHAR(),
|
||
comment='permission ID:Private|Share',
|
||
existing_comment='permission ID:Private|Share|Memory',
|
||
existing_nullable=True)
|
||
|
||
with op.batch_alter_table('end_users', schema=None) as batch_op:
|
||
batch_op.alter_column('id',
|
||
existing_type=sa.UUID(),
|
||
server_default=sa.text('gen_random_uuid()'),
|
||
existing_nullable=False)
|
||
|
||
with op.batch_alter_table('data_config', schema=None) as batch_op:
|
||
batch_op.alter_column('llm',
|
||
existing_type=sa.VARCHAR(length=255),
|
||
comment='LS',
|
||
existing_comment='LLM模型配置ID',
|
||
existing_nullable=True)
|
||
|
||
# ### end Alembic commands ###
|