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
+146 -60
View File
@@ -20,55 +20,15 @@ down_revision: str | None = "c9d8e7f6a5b4"
branch_labels: str | Sequence[str] | None = None branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None depends_on: str | Sequence[str] | None = None
BACKFILL_BATCH_SIZE = 500
def column_exists(table_name: str, column_name: str) -> bool: _UPGRADE_BACKFILL_SQL = sa.text(
bind = op.get_bind()
inspector = inspect(bind)
columns = [c["name"] for c in inspector.get_columns(table_name)]
return column_name in columns
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:
batch_op.alter_column(
"total_tokens",
new_column_name="input_output_total_tokens",
existing_type=sa.Integer(),
existing_nullable=True,
)
with op.batch_alter_table("usage") as batch_op:
if not column_exists("usage", "input_context_tokens"):
batch_op.add_column(sa.Column("input_context_tokens", sa.Integer(), nullable=False, server_default="0"))
if not column_exists("usage", "total_tokens"):
batch_op.add_column(sa.Column("total_tokens", sa.Integer(), nullable=False, server_default="0"))
if not column_exists("usage", "cache_creation_cost_usd_5m"):
batch_op.add_column(sa.Column("cache_creation_cost_usd_5m", sa.Numeric(20, 8), nullable=False, server_default="0"))
if not column_exists("usage", "cache_creation_cost_usd_1h"):
batch_op.add_column(sa.Column("cache_creation_cost_usd_1h", sa.Numeric(20, 8), nullable=False, server_default="0"))
if not column_exists("usage", "actual_cache_creation_cost_usd_5m"):
batch_op.add_column(sa.Column("actual_cache_creation_cost_usd_5m", sa.Numeric(20, 8), nullable=False, server_default="0"))
if not column_exists("usage", "actual_cache_creation_cost_usd_1h"):
batch_op.add_column(sa.Column("actual_cache_creation_cost_usd_1h", sa.Numeric(20, 8), nullable=False, server_default="0"))
if not column_exists("usage", "actual_cache_cost_usd"):
batch_op.add_column(sa.Column("actual_cache_cost_usd", sa.Numeric(20, 8), nullable=False, server_default="0"))
if not column_exists("usage", "cache_creation_price_per_1m_5m"):
batch_op.add_column(sa.Column("cache_creation_price_per_1m_5m", sa.Numeric(20, 8), nullable=True))
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 UPDATE usage
SET SET
input_output_total_tokens = COALESCE(input_output_total_tokens, COALESCE(input_tokens, 0) + COALESCE(output_tokens, 0)), 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), input_context_tokens = COALESCE(input_tokens, 0) + COALESCE(cache_read_input_tokens, 0),
total_tokens = COALESCE(input_output_total_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_creation_input_tokens, 0)
+ COALESCE(cache_read_input_tokens, 0), + COALESCE(cache_read_input_tokens, 0),
cache_creation_cost_usd_5m = CASE cache_creation_cost_usd_5m = CASE
@@ -130,39 +90,165 @@ def upgrade() -> None:
END, END,
cache_cost_usd = COALESCE(cache_creation_cost_usd, 0) + COALESCE(cache_read_cost_usd, 0) cache_cost_usd = COALESCE(cache_creation_cost_usd, 0) + COALESCE(cache_read_cost_usd, 0)
WHERE id IN ( WHERE id IN (
SELECT id FROM usage SELECT id
WHERE input_context_tokens = 0 AND total_tokens = 0 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 LIMIT :batch_size
) )
""" """
),
{"batch_size": batch_size},
) )
if result.rowcount == 0:
break
_DOWNGRADE_BACKFILL_SQL = sa.text(
def downgrade() -> None:
conn = op.get_bind()
batch_size = 5000
while True:
result = conn.execute(
sa.text(
""" """
UPDATE usage UPDATE usage
SET total_tokens = COALESCE(input_output_total_tokens, COALESCE(input_tokens, 0) + COALESCE(output_tokens, 0)) SET total_tokens = COALESCE(input_output_total_tokens, COALESCE(input_tokens, 0) + COALESCE(output_tokens, 0))
WHERE id IN ( WHERE id IN (
SELECT id FROM usage SELECT id
FROM usage
WHERE total_tokens != COALESCE(input_output_total_tokens, COALESCE(input_tokens, 0) + COALESCE(output_tokens, 0)) WHERE total_tokens != COALESCE(input_output_total_tokens, COALESCE(input_tokens, 0) + COALESCE(output_tokens, 0))
ORDER BY id
LIMIT :batch_size LIMIT :batch_size
) )
""" """
),
{"batch_size": batch_size},
) )
if result.rowcount == 0:
def column_exists(table_name: str, column_name: str) -> bool:
bind = op.get_bind()
inspector = inspect(bind)
columns = [c["name"] for c in inspector.get_columns(table_name)]
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 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:
batch_op.alter_column(
"total_tokens",
new_column_name="input_output_total_tokens",
existing_type=sa.Integer(),
existing_nullable=True,
)
with op.batch_alter_table("usage") as batch_op:
if not column_exists("usage", "input_context_tokens"):
batch_op.add_column(sa.Column("input_context_tokens", sa.Integer(), nullable=False, server_default="0"))
if not column_exists("usage", "total_tokens"):
batch_op.add_column(sa.Column("total_tokens", sa.Integer(), nullable=False, server_default="0"))
if not column_exists("usage", "cache_creation_cost_usd_5m"):
batch_op.add_column(sa.Column("cache_creation_cost_usd_5m", sa.Numeric(20, 8), nullable=False, server_default="0"))
if not column_exists("usage", "cache_creation_cost_usd_1h"):
batch_op.add_column(sa.Column("cache_creation_cost_usd_1h", sa.Numeric(20, 8), nullable=False, server_default="0"))
if not column_exists("usage", "actual_cache_creation_cost_usd_5m"):
batch_op.add_column(sa.Column("actual_cache_creation_cost_usd_5m", sa.Numeric(20, 8), nullable=False, server_default="0"))
if not column_exists("usage", "actual_cache_creation_cost_usd_1h"):
batch_op.add_column(sa.Column("actual_cache_creation_cost_usd_1h", sa.Numeric(20, 8), nullable=False, server_default="0"))
if not column_exists("usage", "actual_cache_cost_usd"):
batch_op.add_column(sa.Column("actual_cache_cost_usd", sa.Numeric(20, 8), nullable=False, server_default="0"))
if not column_exists("usage", "cache_creation_price_per_1m_5m"):
batch_op.add_column(sa.Column("cache_creation_price_per_1m_5m", sa.Numeric(20, 8), nullable=True))
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))
run_backfill_in_batches(_UPGRADE_BACKFILL_SQL)
def downgrade() -> None:
run_backfill_in_batches(_DOWNGRADE_BACKFILL_SQL)
with op.batch_alter_table("usage") as batch_op: with op.batch_alter_table("usage") as batch_op:
if column_exists("usage", "cache_creation_price_per_1m_1h"): if column_exists("usage", "cache_creation_price_per_1m_1h"):
batch_op.drop_column("cache_creation_price_per_1m_1h") batch_op.drop_column("cache_creation_price_per_1m_1h")