Files
Aether/alembic/versions/20251220_1500_add_usage_composite_indexes.py
fawney19 03ad16ea8a fix: 修复迁移脚本在全新安装时报错及改进统计回填逻辑
迁移脚本修复:
- 移除 AUTOCOMMIT 模式,改为在同一事务中创建索引
- 分别检查每个索引是否存在,只创建缺失的索引
- 修复全新安装时 AUTOCOMMIT 连接看不到未提交表的问题 (#46)

统计回填改进:
- 分别检查 StatsDaily 和 StatsDailyModel 的缺失日期
- 只回填实际缺失的数据而非连续区间
- 添加失败统计计数和 rollback 错误日志
2025-12-24 21:50:05 +08:00

66 lines
1.9 KiB
Python

"""add usage table composite indexes for query optimization
Revision ID: b2c3d4e5f6g7
Revises: a1b2c3d4e5f6
Create Date: 2025-12-20 15:00:00.000000+00:00
"""
from alembic import op
from sqlalchemy import text
# revision identifiers, used by Alembic.
revision = 'b2c3d4e5f6g7'
down_revision = 'a1b2c3d4e5f6'
branch_labels = None
depends_on = None
def upgrade() -> None:
"""为 usage 表添加复合索引以优化常见查询
注意:这些索引已经在 baseline 迁移中创建。
此迁移仅用于从旧版本升级的场景,新安装会跳过。
"""
conn = op.get_bind()
# 检查 usage 表是否存在
result = conn.execute(text(
"SELECT EXISTS (SELECT FROM information_schema.tables WHERE table_name = 'usage')"
))
if not result.scalar():
# 表不存在,跳过
return
# 定义需要创建的索引
indexes = [
("idx_usage_user_created", "ON usage (user_id, created_at)"),
("idx_usage_apikey_created", "ON usage (api_key_id, created_at)"),
("idx_usage_provider_model_created", "ON usage (provider, model, created_at)"),
]
# 分别检查并创建每个索引
for index_name, index_def in indexes:
result = conn.execute(text(
f"SELECT EXISTS (SELECT 1 FROM pg_indexes WHERE indexname = '{index_name}')"
))
if result.scalar():
continue # 索引已存在,跳过
conn.execute(text(f"CREATE INDEX {index_name} {index_def}"))
def downgrade() -> None:
"""删除复合索引"""
conn = op.get_bind()
# 使用 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"
))