i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-09-01 17:00:21 +08:00
Code Issues Packages Projects Releases Wiki Activity

perf: 同步 DB 操作迁移到 asyncio.to_thread，避免阻塞事件循环

Browse Source 
failover/stream_telemetry/recording/stream 中的同步 DB 操作（commit/execute/query）
会阻塞 asyncio 事件循环，导致 Hub PING 心跳无法发送、worker idle timeout 断连。
将这些操作包装到 asyncio.to_thread() 中执行。

同时提取 Alembic 迁移脚本中重复的幂等性辅助函数到 alembic/helpers.py，
用批量查询缓存替代逐条 information_schema 查询，backfill SQL 合并为 LEFT JOIN。

新增 failover 中客户端断连的快速终止路径，避免继续无意义的重试。
This commit is contained in:
fawney19
2026-03-09 11:49:03 +08:00
parent 4dbfeb87b8
commit 0bddc7965b
10 changed files with 923 additions and 741 deletions
Download Patch File Download Diff File Expand all files Collapse all files

171
alembic/helpers.py Normal file
View File

@@ -0,0 +1,171 @@
"""Shared idempotent helpers for Alembic migrations.
All metadata lookups are batched: one query loads an entire table's column info
or all FK delete rules, then results are cached for the migration's lifetime.
"""
from __future__ import annotations
from collections import defaultdict
from collections.abc import Callable
import sqlalchemy as sa
from alembic import op
# ---------------------------------------------------------------------------
# Batch metadata cache
# ---------------------------------------------------------------------------
class _SchemaCache:
"""Lazy, per-migration cache for information_schema lookups.
Call ``load_columns(tables)`` / ``load_fk_rules(tables)`` once at the top
of ``upgrade()`` or ``downgrade()`` to prime the cache. Subsequent
``column_exists`` / ``column_type`` / ``fk_ondelete`` calls are pure
dict lookups -- zero extra DB round-trips.
"""
def __init__(self) -> None:
# {table_name: {column_name: data_type}}
self._columns: dict[str, dict[str, str]] = {}
# {(table_name, constraint_name): delete_rule}
self._fk_rules: dict[tuple[str, str], str] = {}
self._fk_loaded_tables: set[str] = set()
# -- loaders (one query per call) --------------------------------------
def load_columns(self, tables: list[str]) -> None:
"""Fetch column names + data types for *tables* in one query."""
need = [t for t in tables if t not in self._columns]
if not need:
return
bind = op.get_bind()
rows = bind.execute(
sa.text(
"SELECT table_name, column_name, data_type "
"FROM information_schema.columns "
"WHERE table_name = ANY(:tables)"
),
{"tables": need},
).fetchall()
# Initialise even empty tables so we don't re-query
for t in need:
self._columns.setdefault(t, {})
for table, col, dtype in rows:
self._columns[table][col] = dtype
def load_fk_rules(self, tables: list[str]) -> None:
"""Fetch FK delete rules for *tables* in one query."""
need = [t for t in tables if t not in self._fk_loaded_tables]
if not need:
return
bind = op.get_bind()
rows = bind.execute(
sa.text(
"SELECT tc.table_name, tc.constraint_name, rc.delete_rule "
"FROM information_schema.referential_constraints rc "
"JOIN information_schema.table_constraints tc "
" ON rc.constraint_name = tc.constraint_name "
"WHERE tc.table_name = ANY(:tables)"
),
{"tables": need},
).fetchall()
for table, name, rule in rows:
self._fk_rules[(table, name)] = rule
self._fk_loaded_tables.update(need)
# -- lookups (pure dict, zero DB) --------------------------------------
def column_exists(self, table: str, column: str) -> bool:
return column in self._columns.get(table, {})
def column_type(self, table: str, column: str) -> str | None:
return self._columns.get(table, {}).get(column)
def is_numeric(self, table: str, column: str) -> bool:
return self.column_type(table, column) == "numeric"
def fk_ondelete(self, table: str, constraint: str) -> str | None:
return self._fk_rules.get((table, constraint))
def invalidate_columns(self, table: str) -> None:
"""Force re-load on next load_columns() for *table* (after ADD/DROP COLUMN)."""
self._columns.pop(table, None)
def new_cache() -> _SchemaCache:
"""Create a fresh schema cache for a single migration run."""
return _SchemaCache()
# ---------------------------------------------------------------------------
# Idempotent DDL helpers
# ---------------------------------------------------------------------------
def replace_fk_if_needed(
cache: _SchemaCache,
constraint_name: str,
table_name: str,
ref_table: str,
local_cols: list[str],
remote_cols: list[str],
desired_ondelete: str,
) -> None:
"""Drop and recreate a FK only if the current ON DELETE rule differs."""
current = cache.fk_ondelete(table_name, constraint_name)
if current and current.upper() == desired_ondelete.upper():
return
if current:
op.drop_constraint(constraint_name, table_name, type_="foreignkey")
op.create_foreign_key(
constraint_name,
table_name,
ref_table,
local_cols,
remote_cols,
ondelete=desired_ondelete,
)
def index_exists(index_name: str) -> bool:
bind = op.get_bind()
result = bind.execute(
sa.text("SELECT 1 FROM pg_indexes WHERE indexname = :name"),
{"name": index_name},
)
return result.scalar() is not None
# ---------------------------------------------------------------------------
# Batch ALTER TYPE helper
# ---------------------------------------------------------------------------
def batch_alter_type(
cache: _SchemaCache,
columns: list[tuple[str, str, bool, str | None]],
cast_suffix: str,
type_fn: Callable[[str], str],
) -> None:
"""Group columns by table and issue ONE ``ALTER TABLE`` per table.
Skips columns that don't exist. Each tuple is
``(table_name, column_name, nullable, server_default)``.
"""
by_table: dict[str, list[tuple[str, str]]] = defaultdict(list)
for table, col, _nullable, _default in columns:
if not cache.column_exists(table, col):
continue
by_table[table].append((col, type_fn(col)))
bind = op.get_bind()
for table, col_types in by_table.items():
parts = [
f"ALTER COLUMN {col} TYPE {target} USING {col}::{cast_suffix}"
for col, target in col_types
]
if parts:
bind.execute(sa.text(f"ALTER TABLE {table} " + ", ".join(parts)))

219
alembic/versions/20260308_0348_45b118150a78_usage_stats_retention_set_null_on_.py
View File

@@ -5,137 +5,93 @@ Revises: 2d932114930d
Create Date: 2026-03-08 03:48:49.622091+00:00
"""
from alembic import op
import sqlalchemy as sa
from alembic import op
from alembic.helpers import new_cache, replace_fk_if_needed
# revision identifiers, used by Alembic.
revision = '45b118150a78'
down_revision = '2d932114930d'
revision = "45b118150a78"
down_revision = "2d932114930d"
branch_labels = None
depends_on = None
# ---------------------------------------------------------------------------
# Idempotent helpers
# ---------------------------------------------------------------------------
def _column_exists(table_name: str, column_name: str) -> bool:
bind = op.get_bind()
result = bind.execute(
sa.text(
"SELECT 1 FROM information_schema.columns "
"WHERE table_name = :table AND column_name = :col"
),
{"table": table_name, "col": column_name},
)
return result.scalar() is not None
def _constraint_exists(table_name: str, constraint_name: str) -> bool:
bind = op.get_bind()
result = bind.execute(
sa.text(
"SELECT 1 FROM information_schema.table_constraints "
"WHERE table_name = :table AND constraint_name = :name"
),
{"table": table_name, "name": constraint_name},
)
return result.scalar() is not None
def _fk_ondelete(table_name: str, constraint_name: str) -> str | None:
"""Return the ON DELETE action for a foreign key, or None if it doesn't exist."""
bind = op.get_bind()
result = bind.execute(
sa.text(
"SELECT rc.delete_rule "
"FROM information_schema.referential_constraints rc "
"JOIN information_schema.table_constraints tc "
" ON rc.constraint_name = tc.constraint_name "
"WHERE tc.table_name = :table AND tc.constraint_name = :name"
),
{"table": table_name, "name": constraint_name},
)
row = result.first()
return row[0] if row else None
def _replace_fk_if_needed(
constraint_name: str,
table_name: str,
ref_table: str,
local_cols: list[str],
remote_cols: list[str],
desired_ondelete: str,
) -> None:
"""Drop and recreate a FK only if the current ON DELETE rule differs."""
current = _fk_ondelete(table_name, constraint_name)
if current and current.upper() == desired_ondelete.upper():
return # already correct
if current:
op.drop_constraint(constraint_name, table_name, type_='foreignkey')
op.create_foreign_key(
constraint_name, table_name, ref_table,
local_cols, remote_cols, ondelete=desired_ondelete,
)
_TABLES = ["usage", "stats_user_daily", "stats_daily_api_key"]
def upgrade() -> None:
c = new_cache()
c.load_columns(_TABLES)
c.load_fk_rules(["stats_user_daily", "stats_daily_api_key"])
# --- Usage: add name snapshot columns ---
if not _column_exists('usage', 'username'):
op.add_column('usage', sa.Column('username', sa.String(100), nullable=True,
comment='用户名快照'))
if not _column_exists('usage', 'api_key_name'):
op.add_column('usage', sa.Column('api_key_name', sa.String(200), nullable=True,
comment='API Key 名称快照'))
if not c.column_exists("usage", "username"):
op.add_column(
"usage", sa.Column("username", sa.String(100), nullable=True, comment="用户名快照")
)
if not c.column_exists("usage", "api_key_name"):
op.add_column(
"usage",
sa.Column("api_key_name", sa.String(200), nullable=True, comment="API Key 名称快照"),
)
# --- StatsUserDaily: CASCADE -> SET NULL, add username snapshot ---
_replace_fk_if_needed(
'stats_user_daily_user_id_fkey', 'stats_user_daily',
'users', ['user_id'], ['id'], 'SET NULL',
replace_fk_if_needed(
c,
"stats_user_daily_user_id_fkey",
"stats_user_daily",
"users",
["user_id"],
["id"],
"SET NULL",
)
op.alter_column('stats_user_daily', 'user_id', existing_type=sa.String(36), nullable=True)
if not _column_exists('stats_user_daily', 'username'):
op.add_column('stats_user_daily', sa.Column('username', sa.String(100), nullable=True,
comment='用户名快照（删除用户后仍可追溯）'))
op.alter_column("stats_user_daily", "user_id", existing_type=sa.String(36), nullable=True)
if not c.column_exists("stats_user_daily", "username"):
op.add_column(
"stats_user_daily",
sa.Column(
"username",
sa.String(100),
nullable=True,
comment="用户名快照（删除用户后仍可追溯）",
),
)
# --- StatsDailyApiKey: CASCADE -> SET NULL, add api_key_name snapshot ---
_replace_fk_if_needed(
'stats_daily_api_key_api_key_id_fkey', 'stats_daily_api_key',
'api_keys', ['api_key_id'], ['id'], 'SET NULL',
replace_fk_if_needed(
c,
"stats_daily_api_key_api_key_id_fkey",
"stats_daily_api_key",
"api_keys",
["api_key_id"],
["id"],
"SET NULL",
)
op.alter_column('stats_daily_api_key', 'api_key_id', existing_type=sa.String(36),
nullable=True)
if not _column_exists('stats_daily_api_key', 'api_key_name'):
op.add_column('stats_daily_api_key', sa.Column('api_key_name', sa.String(200),
nullable=True,
comment='API Key 名称快照（删除 Key 后仍可追溯）'))
op.alter_column("stats_daily_api_key", "api_key_id", existing_type=sa.String(36), nullable=True)
if not c.column_exists("stats_daily_api_key", "api_key_name"):
op.add_column(
"stats_daily_api_key",
sa.Column(
"api_key_name",
sa.String(200),
nullable=True,
comment="API Key 名称快照（删除 Key 后仍可追溯）",
),
)
# --- Backfill: populate snapshots from existing FK joins ---
# Single UPDATE per table using LEFT JOINs to fill both columns in one pass.
# (WHERE ... IS NULL makes these inherently idempotent)
# One LEFT JOIN UPDATE per table covers all rows regardless of which FK is present.
op.execute("""
UPDATE usage u
SET username = COALESCE(u.username, usr.username),
api_key_name = COALESCE(u.api_key_name, ak.name)
FROM users usr, api_keys ak
WHERE usr.id = u.user_id
AND ak.id = u.api_key_id
FROM usage u2
LEFT JOIN users usr ON usr.id = u2.user_id
LEFT JOIN api_keys ak ON ak.id = u2.api_key_id
WHERE u.id = u2.id
AND (u.username IS NULL OR u.api_key_name IS NULL)
""")
# Catch rows that have user_id but no api_key_id (or vice versa)
op.execute("""
UPDATE usage u
SET username = usr.username
FROM users usr
WHERE u.user_id = usr.id AND u.username IS NULL
""")
op.execute("""
UPDATE usage u
SET api_key_name = ak.name
FROM api_keys ak
WHERE u.api_key_id = ak.id AND u.api_key_name IS NULL
""")
op.execute("""
UPDATE stats_user_daily s
SET username = usr.username
@@ -151,27 +107,42 @@ def upgrade() -> None:
def downgrade() -> None:
c = new_cache()
c.load_columns(["stats_daily_api_key", "stats_user_daily", "usage"])
c.load_fk_rules(["stats_daily_api_key", "stats_user_daily"])
# --- Remove snapshot columns ---
if _column_exists('stats_daily_api_key', 'api_key_name'):
op.drop_column('stats_daily_api_key', 'api_key_name')
if _column_exists('stats_user_daily', 'username'):
op.drop_column('stats_user_daily', 'username')
if _column_exists('usage', 'api_key_name'):
op.drop_column('usage', 'api_key_name')
if _column_exists('usage', 'username'):
op.drop_column('usage', 'username')
if c.column_exists("stats_daily_api_key", "api_key_name"):
op.drop_column("stats_daily_api_key", "api_key_name")
if c.column_exists("stats_user_daily", "username"):
op.drop_column("stats_user_daily", "username")
if c.column_exists("usage", "api_key_name"):
op.drop_column("usage", "api_key_name")
if c.column_exists("usage", "username"):
op.drop_column("usage", "username")
# --- StatsDailyApiKey: SET NULL -> CASCADE ---
_replace_fk_if_needed(
'stats_daily_api_key_api_key_id_fkey', 'stats_daily_api_key',
'api_keys', ['api_key_id'], ['id'], 'CASCADE',
replace_fk_if_needed(
c,
"stats_daily_api_key_api_key_id_fkey",
"stats_daily_api_key",
"api_keys",
["api_key_id"],
["id"],
"CASCADE",
)
op.alter_column(
"stats_daily_api_key", "api_key_id", existing_type=sa.String(36), nullable=False
)
op.alter_column('stats_daily_api_key', 'api_key_id', existing_type=sa.String(36),
nullable=False)
# --- StatsUserDaily: SET NULL -> CASCADE ---
_replace_fk_if_needed(
'stats_user_daily_user_id_fkey', 'stats_user_daily',
'users', ['user_id'], ['id'], 'CASCADE',
replace_fk_if_needed(
c,
"stats_user_daily_user_id_fkey",
"stats_user_daily",
"users",
["user_id"],
["id"],
"CASCADE",
)
op.alter_column('stats_user_daily', 'user_id', existing_type=sa.String(36), nullable=False)
op.alter_column("stats_user_daily", "user_id", existing_type=sa.String(36), nullable=False)

247
alembic/versions/20260308_1215_13a4c8f6d9e0_request_candidate_video_task_retention.py
View File

@@ -6,89 +6,32 @@ Create Date: 2026-03-08 12:15:00.000000+00:00
"""
from alembic import op
import sqlalchemy as sa
from alembic import op
from alembic.helpers import new_cache, replace_fk_if_needed
# revision identifiers, used by Alembic.
revision = "13a4c8f6d9e0"
down_revision = "45b118150a78"
branch_labels = None
depends_on = None
# ---------------------------------------------------------------------------
# Idempotent helpers
# ---------------------------------------------------------------------------
def _column_exists(table_name: str, column_name: str) -> bool:
bind = op.get_bind()
result = bind.execute(
sa.text(
"SELECT 1 FROM information_schema.columns "
"WHERE table_name = :table AND column_name = :col"
),
{"table": table_name, "col": column_name},
)
return result.scalar() is not None
def _constraint_exists(table_name: str, constraint_name: str) -> bool:
bind = op.get_bind()
result = bind.execute(
sa.text(
"SELECT 1 FROM information_schema.table_constraints "
"WHERE table_name = :table AND constraint_name = :name"
),
{"table": table_name, "name": constraint_name},
)
return result.scalar() is not None
def _fk_ondelete(table_name: str, constraint_name: str) -> str | None:
"""Return the ON DELETE action for a foreign key, or None if it doesn't exist."""
bind = op.get_bind()
result = bind.execute(
sa.text(
"SELECT rc.delete_rule "
"FROM information_schema.referential_constraints rc "
"JOIN information_schema.table_constraints tc "
" ON rc.constraint_name = tc.constraint_name "
"WHERE tc.table_name = :table AND tc.constraint_name = :name"
),
{"table": table_name, "name": constraint_name},
)
row = result.first()
return row[0] if row else None
def _replace_fk_if_needed(
constraint_name: str,
table_name: str,
ref_table: str,
local_cols: list[str],
remote_cols: list[str],
desired_ondelete: str,
) -> None:
"""Drop and recreate a FK only if the current ON DELETE rule differs."""
current = _fk_ondelete(table_name, constraint_name)
if current and current.upper() == desired_ondelete.upper():
return # already correct
if current:
op.drop_constraint(constraint_name, table_name, type_="foreignkey")
op.create_foreign_key(
constraint_name, table_name, ref_table,
local_cols, remote_cols, ondelete=desired_ondelete,
)
_TABLES = ["request_candidates", "video_tasks"]
def upgrade() -> None:
c = new_cache()
c.load_columns(_TABLES)
c.load_fk_rules(_TABLES)
# --- request_candidates: add snapshot columns ---
if not _column_exists("request_candidates", "username"):
if not c.column_exists("request_candidates", "username"):
op.add_column(
"request_candidates",
sa.Column("username", sa.String(length=100), nullable=True, comment="用户名快照"),
)
if not _column_exists("request_candidates", "api_key_name"):
if not c.column_exists("request_candidates", "api_key_name"):
op.add_column(
"request_candidates",
sa.Column(
@@ -100,22 +43,32 @@ def upgrade() -> None:
)
# --- request_candidates: CASCADE -> SET NULL ---
_replace_fk_if_needed(
"request_candidates_user_id_fkey", "request_candidates",
"users", ["user_id"], ["id"], "SET NULL",
replace_fk_if_needed(
c,
"request_candidates_user_id_fkey",
"request_candidates",
"users",
["user_id"],
["id"],
"SET NULL",
)
_replace_fk_if_needed(
"request_candidates_api_key_id_fkey", "request_candidates",
"api_keys", ["api_key_id"], ["id"], "SET NULL",
replace_fk_if_needed(
c,
"request_candidates_api_key_id_fkey",
"request_candidates",
"api_keys",
["api_key_id"],
["id"],
"SET NULL",
)
# --- video_tasks: add snapshot columns ---
if not _column_exists("video_tasks", "username"):
if not c.column_exists("video_tasks", "username"):
op.add_column(
"video_tasks",
sa.Column("username", sa.String(length=100), nullable=True, comment="用户名快照"),
)
if not _column_exists("video_tasks", "api_key_name"):
if not c.column_exists("video_tasks", "api_key_name"):
op.add_column(
"video_tasks",
sa.Column(
@@ -128,106 +81,100 @@ def upgrade() -> None:
# --- video_tasks: CASCADE -> SET NULL, user_id nullable ---
op.alter_column("video_tasks", "user_id", existing_type=sa.String(length=36), nullable=True)
_replace_fk_if_needed(
"video_tasks_user_id_fkey", "video_tasks",
"users", ["user_id"], ["id"], "SET NULL",
replace_fk_if_needed(
c,
"video_tasks_user_id_fkey",
"video_tasks",
"users",
["user_id"],
["id"],
"SET NULL",
)
_replace_fk_if_needed(
"video_tasks_api_key_id_fkey", "video_tasks",
"api_keys", ["api_key_id"], ["id"], "SET NULL",
replace_fk_if_needed(
c,
"video_tasks_api_key_id_fkey",
"video_tasks",
"api_keys",
["api_key_id"],
["id"],
"SET NULL",
)
# --- Backfill: populate snapshots from existing FK joins ---
# Single UPDATE per table using JOINs to fill both columns in one pass.
# (WHERE ... IS NULL makes these inherently idempotent)
# request_candidates: fill both columns where both FKs exist
op.execute(
"""
# One LEFT JOIN UPDATE per table covers all rows regardless of which FK is present.
op.execute("""
UPDATE request_candidates rc
SET username = COALESCE(rc.username, usr.username),
api_key_name = COALESCE(rc.api_key_name, ak.name)
FROM users usr, api_keys ak
WHERE usr.id = rc.user_id
AND ak.id = rc.api_key_id
FROM request_candidates rc2
LEFT JOIN users usr ON usr.id = rc2.user_id
LEFT JOIN api_keys ak ON ak.id = rc2.api_key_id
WHERE rc.id = rc2.id
AND (rc.username IS NULL OR rc.api_key_name IS NULL)
"""
)
# Catch rows with only user_id or only api_key_id
op.execute(
"""
UPDATE request_candidates rc
SET username = usr.username
FROM users usr
WHERE usr.id = rc.user_id AND rc.username IS NULL
"""
)
op.execute(
"""
UPDATE request_candidates rc
SET api_key_name = ak.name
FROM api_keys ak
WHERE ak.id = rc.api_key_id AND rc.api_key_name IS NULL
"""
)
""")
# video_tasks: fill both columns where both FKs exist
op.execute(
"""
op.execute("""
UPDATE video_tasks vt
SET username = COALESCE(vt.username, usr.username),
api_key_name = COALESCE(vt.api_key_name, ak.name)
FROM users usr, api_keys ak
WHERE usr.id = vt.user_id
AND ak.id = vt.api_key_id
FROM video_tasks vt2
LEFT JOIN users usr ON usr.id = vt2.user_id
LEFT JOIN api_keys ak ON ak.id = vt2.api_key_id
WHERE vt.id = vt2.id
AND (vt.username IS NULL OR vt.api_key_name IS NULL)
"""
)
# Catch rows with only user_id or only api_key_id
op.execute(
"""
UPDATE video_tasks vt
SET username = usr.username
FROM users usr
WHERE usr.id = vt.user_id AND vt.username IS NULL
"""
)
op.execute(
"""
UPDATE video_tasks vt
SET api_key_name = ak.name
FROM api_keys ak
WHERE ak.id = vt.api_key_id AND vt.api_key_name IS NULL
"""
)
""")
def downgrade() -> None:
c = new_cache()
c.load_columns(_TABLES)
c.load_fk_rules(_TABLES)
# --- video_tasks: SET NULL -> default (no action), restore NOT NULL ---
_replace_fk_if_needed(
"video_tasks_api_key_id_fkey", "video_tasks",
"api_keys", ["api_key_id"], ["id"], "NO ACTION",
replace_fk_if_needed(
c,
"video_tasks_api_key_id_fkey",
"video_tasks",
"api_keys",
["api_key_id"],
["id"],
"NO ACTION",
)
_replace_fk_if_needed(
"video_tasks_user_id_fkey", "video_tasks",
"users", ["user_id"], ["id"], "NO ACTION",
replace_fk_if_needed(
c,
"video_tasks_user_id_fkey",
"video_tasks",
"users",
["user_id"],
["id"],
"NO ACTION",
)
op.alter_column("video_tasks", "user_id", existing_type=sa.String(length=36), nullable=False)
if _column_exists("video_tasks", "api_key_name"):
if c.column_exists("video_tasks", "api_key_name"):
op.drop_column("video_tasks", "api_key_name")
if _column_exists("video_tasks", "username"):
if c.column_exists("video_tasks", "username"):
op.drop_column("video_tasks", "username")
# --- request_candidates: SET NULL -> CASCADE ---
_replace_fk_if_needed(
"request_candidates_api_key_id_fkey", "request_candidates",
"api_keys", ["api_key_id"], ["id"], "CASCADE",
replace_fk_if_needed(
c,
"request_candidates_api_key_id_fkey",
"request_candidates",
"api_keys",
["api_key_id"],
["id"],
"CASCADE",
)
_replace_fk_if_needed(
"request_candidates_user_id_fkey", "request_candidates",
"users", ["user_id"], ["id"], "CASCADE",
replace_fk_if_needed(
c,
"request_candidates_user_id_fkey",
"request_candidates",
"users",
["user_id"],
["id"],
"CASCADE",
)
if _column_exists("request_candidates", "api_key_name"):
if c.column_exists("request_candidates", "api_key_name"):
op.drop_column("request_candidates", "api_key_name")
if _column_exists("request_candidates", "username"):
if c.column_exists("request_candidates", "username"):
op.drop_column("request_candidates", "username")

102
alembic/versions/20260308_1530_2053ab8ed764_cost_fields_float_to_numeric.py
View File

@@ -6,9 +6,8 @@ Create Date: 2026-03-08 15:30:00.000000+00:00
"""
import sqlalchemy as sa
from alembic import op
from alembic.helpers import batch_alter_type, index_exists, new_cache
# revision identifiers, used by Alembic.
revision = "2053ab8ed764"
@@ -77,44 +76,7 @@ _COST_COLUMNS: list[tuple[str, str, bool, str | None]] = [
("stats_user_daily", "total_cost", False, "0.0"),
]
# rate_multiplier uses a smaller precision
_RATE_MULTIPLIER_TYPE = sa.Numeric(10, 6)
_COST_TYPE = sa.Numeric(20, 8)
def _column_exists(table_name: str, column_name: str) -> bool:
bind = op.get_bind()
result = bind.execute(
sa.text(
"SELECT 1 FROM information_schema.columns "
"WHERE table_name = :table AND column_name = :col"
),
{"table": table_name, "col": column_name},
)
return result.scalar() is not None
def _index_exists(index_name: str) -> bool:
bind = op.get_bind()
result = bind.execute(
sa.text("SELECT 1 FROM pg_indexes WHERE indexname = :name"),
{"name": index_name},
)
return result.scalar() is not None
def _is_numeric_type(table_name: str, column_name: str) -> bool:
"""Check if a column is already numeric type (not float/double precision)."""
bind = op.get_bind()
result = bind.execute(
sa.text(
"SELECT data_type FROM information_schema.columns "
"WHERE table_name = :table AND column_name = :col"
),
{"table": table_name, "col": column_name},
)
data_type = result.scalar()
return data_type == "numeric"
_ALL_TABLES = list({t for t, *_ in _COST_COLUMNS})
def _type_spec(col: str) -> str:
@@ -122,50 +84,20 @@ def _type_spec(col: str) -> str:
return "NUMERIC(10,6)" if col == "rate_multiplier" else "NUMERIC(20,8)"
def _batch_alter_type(
columns: list[tuple[str, str, bool, str | None]],
cast_suffix: str,
type_fn=None,
) -> None:
"""Group columns by table and issue ONE ALTER TABLE per table.
This avoids rewriting the same table N times (once per column).
"""
from collections import defaultdict
by_table: dict[str, list[tuple[str, str, bool, str | None]]] = defaultdict(list)
for table, col, nullable, default in columns:
if not _column_exists(table, col):
continue
by_table[table].append((table, col, nullable, default))
bind = op.get_bind()
for table, cols in by_table.items():
# Build a single ALTER TABLE with multiple ALTER COLUMN clauses
parts: list[str] = []
for _t, col, _nullable, _default in cols:
target_type = type_fn(col) if type_fn else _type_spec(col)
parts.append(
f"ALTER COLUMN {col} TYPE {target_type} USING {col}::{cast_suffix}"
)
if not parts:
continue
sql = f"ALTER TABLE {table} " + ", ".join(parts)
bind.execute(sa.text(sql))
def upgrade() -> None:
c = new_cache()
c.load_columns(_ALL_TABLES)
# -- 1. cost fields: Float -> Numeric (batched per table)
# Filter out columns that are already numeric
cols_to_convert = [
(t, c, n, d)
for t, c, n, d in _COST_COLUMNS
if _column_exists(t, c) and not _is_numeric_type(t, c)
(t, col, n, d)
for t, col, n, d in _COST_COLUMNS
if c.column_exists(t, col) and not c.is_numeric(t, col)
]
_batch_alter_type(cols_to_convert, cast_suffix="numeric", type_fn=_type_spec)
batch_alter_type(c, cols_to_convert, cast_suffix="numeric", type_fn=_type_spec)
# -- 2. provider_api_keys composite index
if not _index_exists("idx_provider_api_keys_provider_active"):
if not index_exists("idx_provider_api_keys_provider_active"):
op.create_index(
"idx_provider_api_keys_provider_active",
"provider_api_keys",
@@ -175,19 +107,23 @@ def upgrade() -> None:
def downgrade() -> None:
# -- 2. drop composite index
if _index_exists("idx_provider_api_keys_provider_active"):
if index_exists("idx_provider_api_keys_provider_active"):
op.drop_index(
"idx_provider_api_keys_provider_active",
table_name="provider_api_keys",
)
# -- 1. Numeric -> Float (batched per table)
c = new_cache()
c.load_columns(_ALL_TABLES)
cols_to_revert = [
(t, c, n, d)
for t, c, n, d in _COST_COLUMNS
if _column_exists(t, c) and _is_numeric_type(t, c)
(t, col, n, d)
for t, col, n, d in _COST_COLUMNS
if c.column_exists(t, col) and c.is_numeric(t, col)
]
_batch_alter_type(
batch_alter_type(
c,
cols_to_revert,
cast_suffix="double precision",
type_fn=lambda _col: "DOUBLE PRECISION",

53
alembic/versions/20260309_0100_d7649c1f8e21_video_task_key_id_set_null_on_delete.py
View File

@@ -6,9 +6,7 @@ Create Date: 2026-03-09 01:00:00.000000+00:00
"""
import sqlalchemy as sa
from alembic import op
from alembic.helpers import new_cache, replace_fk_if_needed
# revision identifiers, used by Alembic.
revision = "d7649c1f8e21"
@@ -20,47 +18,11 @@ _TABLE = "video_tasks"
_FK_NAME = "video_tasks_key_id_fkey"
def _fk_ondelete(table_name: str, constraint_name: str) -> str | None:
bind = op.get_bind()
result = bind.execute(
sa.text(
"SELECT rc.delete_rule "
"FROM information_schema.referential_constraints rc "
"JOIN information_schema.table_constraints tc "
" ON rc.constraint_name = tc.constraint_name "
"WHERE tc.table_name = :table AND tc.constraint_name = :name"
),
{"table": table_name, "name": constraint_name},
)
row = result.first()
return row[0] if row else None
def _replace_fk_if_needed(
constraint_name: str,
table_name: str,
ref_table: str,
local_cols: list[str],
remote_cols: list[str],
desired_ondelete: str,
) -> None:
current = _fk_ondelete(table_name, constraint_name)
if current and current.upper() == desired_ondelete.upper():
return
if current:
op.drop_constraint(constraint_name, table_name, type_="foreignkey")
op.create_foreign_key(
constraint_name,
table_name,
ref_table,
local_cols,
remote_cols,
ondelete=desired_ondelete,
)
def upgrade() -> None:
_replace_fk_if_needed(
c = new_cache()
c.load_fk_rules([_TABLE])
replace_fk_if_needed(
c,
_FK_NAME,
_TABLE,
"provider_api_keys",
@@ -71,7 +33,10 @@ def upgrade() -> None:
def downgrade() -> None:
_replace_fk_if_needed(
c = new_cache()
c.load_fk_rules([_TABLE])
replace_fk_if_needed(
c,
_FK_NAME,
_TABLE,
"provider_api_keys",