From 868f3349e5763833e7b3113d4451256f1cce93a1 Mon Sep 17 00:00:00 2001
From: fawney19 <elky0401@gmail.com>
Date: Tue, 23 Dec 2025 00:18:11 +0800
Subject: [PATCH] fix: use AUTOCOMMIT mode for CREATE INDEX CONCURRENTLY in
 migration
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

PostgreSQL 不允许在事务块内执行 CREATE INDEX CONCURRENTLY，
通过创建独立连接并设置 AUTOCOMMIT 隔离级别来解决此问题。
---
 ...251220_1500_add_usage_composite_indexes.py | 55 ++++++++++---------
 1 file changed, 29 insertions(+), 26 deletions(-)

diff --git a/alembic/versions/20251220_1500_add_usage_composite_indexes.py b/alembic/versions/20251220_1500_add_usage_composite_indexes.py
index 122da57..a0221f5 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
+from sqlalchemy import inspect, text
 
 # revision identifiers, used by Alembic.
 revision = 'b2c3d4e5f6g7'
@@ -24,33 +24,36 @@ def index_exists(table_name: str, index_name: str) -> bool:
 
 
 def upgrade() -> None:
-    """为 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
-        )
+    """为 usage 表添加复合索引以优化常见查询
 
-    # 2. api_key_id + created_at 复合索引 (API Key 用量查询)
-    if not index_exists('usage', 'idx_usage_apikey_created'):
-        op.create_index(
-            'idx_usage_apikey_created',
-            'usage',
-            ['api_key_id', 'created_at'],
-            postgresql_concurrently=True
-        )
+    使用 CONCURRENTLY 创建索引以避免锁表，
+    但需要在 AUTOCOMMIT 模式下执行（不能在事务内）
+    """
+    conn = op.get_bind()
+    engine = conn.engine
 
-    # 3. provider + model + created_at 复合索引 (模型统计查询)
-    if not index_exists('usage', 'idx_usage_provider_model_created'):
-        op.create_index(
-            'idx_usage_provider_model_created',
-            'usage',
-            ['provider', 'model', 'created_at'],
-            postgresql_concurrently=True
-        )
+    # 使用新连接并设置 AUTOCOMMIT 模式以支持 CREATE INDEX CONCURRENTLY
+    with engine.connect().execution_options(isolation_level="AUTOCOMMIT") as autocommit_conn:
+        # 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)"
+            ))
+
+        # 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: