From 7185818724359b432779f0e4a55e9a6d4fcb9151 Mon Sep 17 00:00:00 2001 From: fawney19 Date: Tue, 23 Dec 2025 00:21:03 +0800 Subject: [PATCH] fix: remove index_exists check to avoid transaction conflict in migration - Remove index_exists function that used op.get_bind() within transaction - Use IF NOT EXISTS / IF EXISTS SQL syntax instead - Fixes CREATE INDEX CONCURRENTLY error in Docker migration --- ...251220_1500_add_usage_composite_indexes.py | 57 +++++++++---------- 1 file changed, 26 insertions(+), 31 deletions(-) diff --git a/alembic/versions/20251220_1500_add_usage_composite_indexes.py b/alembic/versions/20251220_1500_add_usage_composite_indexes.py index a0221f5..024c817 100644 --- a/alembic/versions/20251220_1500_add_usage_composite_indexes.py +++ b/alembic/versions/20251220_1500_add_usage_composite_indexes.py @@ -6,7 +6,7 @@ Create Date: 2025-12-20 15:00:00.000000+00:00 """ from alembic import op -from sqlalchemy import inspect, text +from sqlalchemy import text # revision identifiers, used by Alembic. revision = 'b2c3d4e5f6g7' @@ -15,14 +15,6 @@ branch_labels = None depends_on = None -def index_exists(table_name: str, index_name: str) -> bool: - """检查索引是否存在""" - bind = op.get_bind() - inspector = inspect(bind) - indexes = [idx['name'] for idx in inspector.get_indexes(table_name)] - return index_name in indexes - - def upgrade() -> None: """为 usage 表添加复合索引以优化常见查询 @@ -34,35 +26,38 @@ def upgrade() -> None: # 使用新连接并设置 AUTOCOMMIT 模式以支持 CREATE INDEX CONCURRENTLY with engine.connect().execution_options(isolation_level="AUTOCOMMIT") as autocommit_conn: + # 使用 IF NOT EXISTS 避免重复创建，无需单独检查索引是否存在 + # 1. user_id + created_at 复合索引 (用户用量查询) - if not index_exists('usage', 'idx_usage_user_created'): - autocommit_conn.execute(text( - "CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_usage_user_created " - "ON usage (user_id, created_at)" - )) + autocommit_conn.execute(text( + "CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_usage_user_created " + "ON usage (user_id, created_at)" + )) # 2. api_key_id + created_at 复合索引 (API Key 用量查询) - if not index_exists('usage', 'idx_usage_apikey_created'): - autocommit_conn.execute(text( - "CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_usage_apikey_created " - "ON usage (api_key_id, created_at)" - )) + autocommit_conn.execute(text( + "CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_usage_apikey_created " + "ON usage (api_key_id, created_at)" + )) # 3. provider + model + created_at 复合索引 (模型统计查询) - if not index_exists('usage', 'idx_usage_provider_model_created'): - autocommit_conn.execute(text( - "CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_usage_provider_model_created " - "ON usage (provider, model, created_at)" - )) + autocommit_conn.execute(text( + "CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_usage_provider_model_created " + "ON usage (provider, model, created_at)" + )) def downgrade() -> None: """删除复合索引""" - if index_exists('usage', 'idx_usage_provider_model_created'): - op.drop_index('idx_usage_provider_model_created', table_name='usage') + conn = op.get_bind() - if index_exists('usage', 'idx_usage_apikey_created'): - op.drop_index('idx_usage_apikey_created', table_name='usage') - - if index_exists('usage', 'idx_usage_user_created'): - op.drop_index('idx_usage_user_created', table_name='usage') + # 使用 IF EXISTS 避免索引不存在时报错 + conn.execute(text( + "DROP INDEX IF EXISTS idx_usage_provider_model_created" + )) + conn.execute(text( + "DROP INDEX IF EXISTS idx_usage_apikey_created" + )) + conn.execute(text( + "DROP INDEX IF EXISTS idx_usage_user_created" + ))