i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-10-13 04:29:50 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix(migration): 修复 usage token 语义迁移脚本

Browse Source 
- 将内联 SQL 提取为模块级常量 _UPGRADE_BACKFILL_SQL / _DOWNGRADE_BACKFILL_SQL
- 抽取 run_backfill_in_batches() 统一批量更新逻辑，通过 autocommit_block 释放 ALTER TABLE 锁
- 修正 upgrade 中 total_tokens 计算逻辑：用 COALESCE(input_tokens,0)+COALESCE(output_tokens,0) 补全 input_output_total_tokens
- WHERE 条件改为按实际字段差异筛选待回填行，替换原先仅过滤 input_context_tokens=0 AND total_tokens=0 的不完整条件
- 批量大小由 5000 调整为 500，减少单批锁持有时间
This commit is contained in:
fawney19
2026-03-24 01:56:36 +08:00
parent bd4e5f3a5d
commit 7c678b715f
1 changed files with 187 additions and 101 deletions
Download Patch File Download Diff File Expand all files Collapse all files
alembic/versions/20260324_1400_c3d4e5f6a7b8_usage_token_semantics_v2.py
+187 -101
View File
@@ -20,6 +20,180 @@ down_revision: str | None = "c9d8e7f6a5b4"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None
BACKFILL_BATCH_SIZE = 500
_UPGRADE_BACKFILL_SQL = sa.text(
"""
UPDATE usage
SET
input_output_total_tokens = COALESCE(input_output_total_tokens, COALESCE(input_tokens, 0) + COALESCE(output_tokens, 0)),
input_context_tokens = COALESCE(input_tokens, 0) + COALESCE(cache_read_input_tokens, 0),
total_tokens = COALESCE(input_output_total_tokens, COALESCE(input_tokens, 0) + COALESCE(output_tokens, 0))
+ COALESCE(cache_creation_input_tokens, 0)
+ COALESCE(cache_read_input_tokens, 0),
cache_creation_cost_usd_5m = CASE
WHEN COALESCE(cache_creation_input_tokens_5m, 0) > 0
AND COALESCE(cache_creation_input_tokens_1h, 0) = 0
THEN COALESCE(cache_creation_cost_usd, 0)
WHEN COALESCE(cache_creation_input_tokens_5m, 0) > 0
AND COALESCE(cache_creation_input_tokens, 0) > 0
THEN COALESCE(cache_creation_cost_usd, 0)
* (COALESCE(cache_creation_input_tokens_5m, 0) * 1.0
/ GREATEST(COALESCE(cache_creation_input_tokens, 0), 1))
ELSE 0
END,
cache_creation_cost_usd_1h = CASE
WHEN COALESCE(cache_creation_input_tokens_1h, 0) > 0
AND COALESCE(cache_creation_input_tokens_5m, 0) = 0
THEN COALESCE(cache_creation_cost_usd, 0)
WHEN COALESCE(cache_creation_input_tokens_1h, 0) > 0
AND COALESCE(cache_creation_input_tokens, 0) > 0
THEN COALESCE(cache_creation_cost_usd, 0)
* (COALESCE(cache_creation_input_tokens_1h, 0) * 1.0
/ GREATEST(COALESCE(cache_creation_input_tokens, 0), 1))
ELSE 0
END,
actual_cache_creation_cost_usd_5m = CASE
WHEN COALESCE(cache_creation_input_tokens_5m, 0) > 0
AND COALESCE(cache_creation_input_tokens_1h, 0) = 0
THEN COALESCE(actual_cache_creation_cost_usd, 0)
WHEN COALESCE(cache_creation_input_tokens_5m, 0) > 0
AND COALESCE(cache_creation_input_tokens, 0) > 0
THEN COALESCE(actual_cache_creation_cost_usd, 0)
* (COALESCE(cache_creation_input_tokens_5m, 0) * 1.0
/ GREATEST(COALESCE(cache_creation_input_tokens, 0), 1))
ELSE 0
END,
actual_cache_creation_cost_usd_1h = CASE
WHEN COALESCE(cache_creation_input_tokens_1h, 0) > 0
AND COALESCE(cache_creation_input_tokens_5m, 0) = 0
THEN COALESCE(actual_cache_creation_cost_usd, 0)
WHEN COALESCE(cache_creation_input_tokens_1h, 0) > 0
AND COALESCE(cache_creation_input_tokens, 0) > 0
THEN COALESCE(actual_cache_creation_cost_usd, 0)
* (COALESCE(cache_creation_input_tokens_1h, 0) * 1.0
/ GREATEST(COALESCE(cache_creation_input_tokens, 0), 1))
ELSE 0
END,
actual_cache_cost_usd = COALESCE(actual_cache_creation_cost_usd, 0) + COALESCE(actual_cache_read_cost_usd, 0),
cache_creation_price_per_1m_5m = CASE
WHEN COALESCE(cache_creation_input_tokens_5m, 0) > 0
AND COALESCE(cache_creation_input_tokens_1h, 0) = 0
THEN cache_creation_price_per_1m
ELSE NULL
END,
cache_creation_price_per_1m_1h = CASE
WHEN COALESCE(cache_creation_input_tokens_1h, 0) > 0
AND COALESCE(cache_creation_input_tokens_5m, 0) = 0
THEN cache_creation_price_per_1m
ELSE NULL
END,
cache_cost_usd = COALESCE(cache_creation_cost_usd, 0) + COALESCE(cache_read_cost_usd, 0)
WHERE id IN (
SELECT id
FROM usage
WHERE
input_output_total_tokens IS DISTINCT FROM COALESCE(input_output_total_tokens, COALESCE(input_tokens, 0) + COALESCE(output_tokens, 0))
OR input_context_tokens IS DISTINCT FROM COALESCE(input_tokens, 0) + COALESCE(cache_read_input_tokens, 0)
OR total_tokens IS DISTINCT FROM (
COALESCE(input_output_total_tokens, COALESCE(input_tokens, 0) + COALESCE(output_tokens, 0))
+ COALESCE(cache_creation_input_tokens, 0)
+ COALESCE(cache_read_input_tokens, 0)
)
OR cache_creation_cost_usd_5m IS DISTINCT FROM (
CASE
WHEN COALESCE(cache_creation_input_tokens_5m, 0) > 0
AND COALESCE(cache_creation_input_tokens_1h, 0) = 0
THEN COALESCE(cache_creation_cost_usd, 0)
WHEN COALESCE(cache_creation_input_tokens_5m, 0) > 0
AND COALESCE(cache_creation_input_tokens, 0) > 0
THEN COALESCE(cache_creation_cost_usd, 0)
* (COALESCE(cache_creation_input_tokens_5m, 0) * 1.0
/ GREATEST(COALESCE(cache_creation_input_tokens, 0), 1))
ELSE 0
END
)
OR cache_creation_cost_usd_1h IS DISTINCT FROM (
CASE
WHEN COALESCE(cache_creation_input_tokens_1h, 0) > 0
AND COALESCE(cache_creation_input_tokens_5m, 0) = 0
THEN COALESCE(cache_creation_cost_usd, 0)
WHEN COALESCE(cache_creation_input_tokens_1h, 0) > 0
AND COALESCE(cache_creation_input_tokens, 0) > 0
THEN COALESCE(cache_creation_cost_usd, 0)
* (COALESCE(cache_creation_input_tokens_1h, 0) * 1.0
/ GREATEST(COALESCE(cache_creation_input_tokens, 0), 1))
ELSE 0
END
)
OR actual_cache_creation_cost_usd_5m IS DISTINCT FROM (
CASE
WHEN COALESCE(cache_creation_input_tokens_5m, 0) > 0
AND COALESCE(cache_creation_input_tokens_1h, 0) = 0
THEN COALESCE(actual_cache_creation_cost_usd, 0)
WHEN COALESCE(cache_creation_input_tokens_5m, 0) > 0
AND COALESCE(cache_creation_input_tokens, 0) > 0
THEN COALESCE(actual_cache_creation_cost_usd, 0)
* (COALESCE(cache_creation_input_tokens_5m, 0) * 1.0
/ GREATEST(COALESCE(cache_creation_input_tokens, 0), 1))
ELSE 0
END
)
OR actual_cache_creation_cost_usd_1h IS DISTINCT FROM (
CASE
WHEN COALESCE(cache_creation_input_tokens_1h, 0) > 0
AND COALESCE(cache_creation_input_tokens_5m, 0) = 0
THEN COALESCE(actual_cache_creation_cost_usd, 0)
WHEN COALESCE(cache_creation_input_tokens_1h, 0) > 0
AND COALESCE(cache_creation_input_tokens, 0) > 0
THEN COALESCE(actual_cache_creation_cost_usd, 0)
* (COALESCE(cache_creation_input_tokens_1h, 0) * 1.0
/ GREATEST(COALESCE(cache_creation_input_tokens, 0), 1))
ELSE 0
END
)
OR actual_cache_cost_usd IS DISTINCT FROM (
COALESCE(actual_cache_creation_cost_usd, 0) + COALESCE(actual_cache_read_cost_usd, 0)
)
OR cache_creation_price_per_1m_5m IS DISTINCT FROM (
CASE
WHEN COALESCE(cache_creation_input_tokens_5m, 0) > 0
AND COALESCE(cache_creation_input_tokens_1h, 0) = 0
THEN cache_creation_price_per_1m
ELSE NULL
END
)
OR cache_creation_price_per_1m_1h IS DISTINCT FROM (
CASE
WHEN COALESCE(cache_creation_input_tokens_1h, 0) > 0
AND COALESCE(cache_creation_input_tokens_5m, 0) = 0
THEN cache_creation_price_per_1m
ELSE NULL
END
)
OR cache_cost_usd IS DISTINCT FROM (
COALESCE(cache_creation_cost_usd, 0) + COALESCE(cache_read_cost_usd, 0)
)
ORDER BY id
LIMIT :batch_size
)
"""
)
_DOWNGRADE_BACKFILL_SQL = sa.text(
"""
UPDATE usage
SET total_tokens = COALESCE(input_output_total_tokens, COALESCE(input_tokens, 0) + COALESCE(output_tokens, 0))
WHERE id IN (
SELECT id
FROM usage
WHERE total_tokens != COALESCE(input_output_total_tokens, COALESCE(input_tokens, 0) + COALESCE(output_tokens, 0))
ORDER BY id
LIMIT :batch_size
)
"""
)
def column_exists(table_name: str, column_name: str) -> bool:
bind = op.get_bind()
@@ -28,6 +202,17 @@ def column_exists(table_name: str, column_name: str) -> bool:
return column_name in columns
def run_backfill_in_batches(sql: sa.TextClause, batch_size: int = BACKFILL_BATCH_SIZE) -> None:
context = op.get_context()
while True:
# Commit the preceding schema transaction before each batch so PostgreSQL
# does not keep ALTER TABLE locks for the entire data backfill.
with context.autocommit_block():
rowcount = op.get_bind().execute(sql, {"batch_size": batch_size}).rowcount
if rowcount == 0:
break
def upgrade() -> None:
if column_exists("usage", "total_tokens") and not column_exists("usage", "input_output_total_tokens"):
with op.batch_alter_table("usage") as batch_op:
@@ -58,110 +243,11 @@ def upgrade() -> None:
if not column_exists("usage", "cache_creation_price_per_1m_1h"):
batch_op.add_column(sa.Column("cache_creation_price_per_1m_1h", sa.Numeric(20, 8), nullable=True))
conn = op.get_bind()
batch_size = 5000
while True:
result = conn.execute(
sa.text(
"""
UPDATE usage
SET
input_output_total_tokens = COALESCE(input_output_total_tokens, COALESCE(input_tokens, 0) + COALESCE(output_tokens, 0)),
input_context_tokens = COALESCE(input_tokens, 0) + COALESCE(cache_read_input_tokens, 0),
total_tokens = COALESCE(input_output_total_tokens, 0)
+ COALESCE(cache_creation_input_tokens, 0)
+ COALESCE(cache_read_input_tokens, 0),
cache_creation_cost_usd_5m = CASE
WHEN COALESCE(cache_creation_input_tokens_5m, 0) > 0
AND COALESCE(cache_creation_input_tokens_1h, 0) = 0
THEN COALESCE(cache_creation_cost_usd, 0)
WHEN COALESCE(cache_creation_input_tokens_5m, 0) > 0
AND COALESCE(cache_creation_input_tokens, 0) > 0
THEN COALESCE(cache_creation_cost_usd, 0)
* (COALESCE(cache_creation_input_tokens_5m, 0) * 1.0
/ GREATEST(COALESCE(cache_creation_input_tokens, 0), 1))
ELSE 0
END,
cache_creation_cost_usd_1h = CASE
WHEN COALESCE(cache_creation_input_tokens_1h, 0) > 0
AND COALESCE(cache_creation_input_tokens_5m, 0) = 0
THEN COALESCE(cache_creation_cost_usd, 0)
WHEN COALESCE(cache_creation_input_tokens_1h, 0) > 0
AND COALESCE(cache_creation_input_tokens, 0) > 0
THEN COALESCE(cache_creation_cost_usd, 0)
* (COALESCE(cache_creation_input_tokens_1h, 0) * 1.0
/ GREATEST(COALESCE(cache_creation_input_tokens, 0), 1))
ELSE 0
END,
actual_cache_creation_cost_usd_5m = CASE
WHEN COALESCE(cache_creation_input_tokens_5m, 0) > 0
AND COALESCE(cache_creation_input_tokens_1h, 0) = 0
THEN COALESCE(actual_cache_creation_cost_usd, 0)
WHEN COALESCE(cache_creation_input_tokens_5m, 0) > 0
AND COALESCE(cache_creation_input_tokens, 0) > 0
THEN COALESCE(actual_cache_creation_cost_usd, 0)
* (COALESCE(cache_creation_input_tokens_5m, 0) * 1.0
/ GREATEST(COALESCE(cache_creation_input_tokens, 0), 1))
ELSE 0
END,
actual_cache_creation_cost_usd_1h = CASE
WHEN COALESCE(cache_creation_input_tokens_1h, 0) > 0
AND COALESCE(cache_creation_input_tokens_5m, 0) = 0
THEN COALESCE(actual_cache_creation_cost_usd, 0)
WHEN COALESCE(cache_creation_input_tokens_1h, 0) > 0
AND COALESCE(cache_creation_input_tokens, 0) > 0
THEN COALESCE(actual_cache_creation_cost_usd, 0)
* (COALESCE(cache_creation_input_tokens_1h, 0) * 1.0
/ GREATEST(COALESCE(cache_creation_input_tokens, 0), 1))
ELSE 0
END,
actual_cache_cost_usd = COALESCE(actual_cache_creation_cost_usd, 0) + COALESCE(actual_cache_read_cost_usd, 0),
cache_creation_price_per_1m_5m = CASE
WHEN COALESCE(cache_creation_input_tokens_5m, 0) > 0
AND COALESCE(cache_creation_input_tokens_1h, 0) = 0
THEN cache_creation_price_per_1m
ELSE NULL
END,
cache_creation_price_per_1m_1h = CASE
WHEN COALESCE(cache_creation_input_tokens_1h, 0) > 0
AND COALESCE(cache_creation_input_tokens_5m, 0) = 0
THEN cache_creation_price_per_1m
ELSE NULL
END,
cache_cost_usd = COALESCE(cache_creation_cost_usd, 0) + COALESCE(cache_read_cost_usd, 0)
WHERE id IN (
SELECT id FROM usage
WHERE input_context_tokens = 0 AND total_tokens = 0
LIMIT :batch_size
)
"""
),
{"batch_size": batch_size},
)
if result.rowcount == 0:
break
run_backfill_in_batches(_UPGRADE_BACKFILL_SQL)
def downgrade() -> None:
conn = op.get_bind()
batch_size = 5000
while True:
result = conn.execute(
sa.text(
"""
UPDATE usage
SET total_tokens = COALESCE(input_output_total_tokens, COALESCE(input_tokens, 0) + COALESCE(output_tokens, 0))
WHERE id IN (
SELECT id FROM usage
WHERE total_tokens != COALESCE(input_output_total_tokens, COALESCE(input_tokens, 0) + COALESCE(output_tokens, 0))
LIMIT :batch_size
)
"""
),
{"batch_size": batch_size},
)
if result.rowcount == 0:
break
run_backfill_in_batches(_DOWNGRADE_BACKFILL_SQL)
with op.batch_alter_table("usage") as batch_op:
if column_exists("usage", "cache_creation_price_per_1m_1h"):