fix: use AUTOCOMMIT mode for CREATE INDEX CONCURRENTLY in migration

PostgreSQL 不允许在事务块内执行 CREATE INDEX CONCURRENTLY,
通过创建独立连接并设置 AUTOCOMMIT 隔离级别来解决此问题。
This commit is contained in:
fawney19
2025-12-23 00:18:11 +08:00
parent d7384e69d9
commit 868f3349e5

View File

@@ -6,7 +6,7 @@ Create Date: 2025-12-20 15:00:00.000000+00:00
""" """
from alembic import op from alembic import op
from sqlalchemy import inspect from sqlalchemy import inspect, text
# revision identifiers, used by Alembic. # revision identifiers, used by Alembic.
revision = 'b2c3d4e5f6g7' revision = 'b2c3d4e5f6g7'
@@ -24,33 +24,36 @@ def index_exists(table_name: str, index_name: str) -> bool:
def upgrade() -> None: def upgrade() -> None:
"""为 usage 表添加复合索引以优化常见查询""" """为 usage 表添加复合索引以优化常见查询
# 1. user_id + created_at 复合索引 (用户用量查询)
if not index_exists('usage', 'idx_usage_user_created'):
op.create_index(
'idx_usage_user_created',
'usage',
['user_id', 'created_at'],
postgresql_concurrently=True
)
# 2. api_key_id + created_at 复合索引 (API Key 用量查询) 使用 CONCURRENTLY 创建索引以避免锁表,
if not index_exists('usage', 'idx_usage_apikey_created'): 但需要在 AUTOCOMMIT 模式下执行(不能在事务内)
op.create_index( """
'idx_usage_apikey_created', conn = op.get_bind()
'usage', engine = conn.engine
['api_key_id', 'created_at'],
postgresql_concurrently=True
)
# 3. provider + model + created_at 复合索引 (模型统计查询) # 使用新连接并设置 AUTOCOMMIT 模式以支持 CREATE INDEX CONCURRENTLY
if not index_exists('usage', 'idx_usage_provider_model_created'): with engine.connect().execution_options(isolation_level="AUTOCOMMIT") as autocommit_conn:
op.create_index( # 1. user_id + created_at 复合索引 (用户用量查询)
'idx_usage_provider_model_created', if not index_exists('usage', 'idx_usage_user_created'):
'usage', autocommit_conn.execute(text(
['provider', 'model', 'created_at'], "CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_usage_user_created "
postgresql_concurrently=True "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)"
))
# 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)"
))
def downgrade() -> None: def downgrade() -> None: