118 lines
3.7 KiB
Python
118 lines
3.7 KiB
Python
"""
|
||
启动时执行 Alembic 迁移（与 CLI `uv run alembic upgrade head` 等价）。
|
||
|
||
- 对瞬时连接错误重试（数据库尚未就绪、网络抖动）。
|
||
- 可通过 settings 关闭迁移、或失败时中止进程（生产推荐）。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import time
|
||
from pathlib import Path
|
||
from typing import Final
|
||
|
||
from sqlalchemy.exc import OperationalError
|
||
|
||
from app.core.config import settings
|
||
from app.core.logging import get_logger
|
||
|
||
logger = get_logger(__name__)
|
||
|
||
_API_DIR: Final[Path] = Path(__file__).resolve().parent.parent.parent
|
||
|
||
|
||
def _repair_withdrawn_revision_stamp_if_needed() -> None:
|
||
from sqlalchemy import create_engine
|
||
|
||
from app.core.alembic_revision_repair import try_repair_withdrawn_0020_revision
|
||
from app.core.db import _database_url
|
||
|
||
engine = create_engine(_database_url())
|
||
with engine.connect() as conn:
|
||
if try_repair_withdrawn_0020_revision(conn):
|
||
conn.commit()
|
||
logger.warning(
|
||
"alembic_version 曾为已撤回的 0020_*，已回退到 0018；"
|
||
"将重新执行 0019_align_legacy_schema"
|
||
)
|
||
|
||
|
||
def _run_alembic_upgrade_once() -> None:
|
||
from alembic.command import upgrade
|
||
from alembic.config import Config
|
||
|
||
_repair_withdrawn_revision_stamp_if_needed()
|
||
cfg = Config(str(_API_DIR / "alembic.ini"))
|
||
upgrade(cfg, "head")
|
||
|
||
|
||
def _is_retryable_migration_error(err: BaseException) -> bool:
|
||
if isinstance(err, OperationalError):
|
||
return True
|
||
msg = str(err).lower()
|
||
needles = (
|
||
"connection refused",
|
||
"could not connect",
|
||
"connection reset",
|
||
"timeout",
|
||
"server closed the connection",
|
||
"the database system is starting",
|
||
"database system is shutting down",
|
||
"too many connections",
|
||
)
|
||
return any(n in msg for n in needles)
|
||
|
||
|
||
def run_alembic_upgrade_at_startup() -> None:
|
||
"""
|
||
同步执行迁移；失败时按 settings 记录日志或抛出。
|
||
|
||
在 asyncio 中请使用 ``asyncio.to_thread(run_alembic_upgrade_at_startup)``，
|
||
避免阻塞事件循环。
|
||
"""
|
||
if not settings.alembic_run_on_startup:
|
||
logger.info("跳过 Alembic 迁移（alembic_run_on_startup=False）")
|
||
return
|
||
|
||
max_tries = max(1, settings.alembic_startup_max_retries)
|
||
base_delay = float(settings.alembic_startup_retry_base_seconds)
|
||
last: BaseException | None = None
|
||
|
||
for attempt in range(max_tries):
|
||
try:
|
||
_run_alembic_upgrade_once()
|
||
if attempt > 0:
|
||
logger.info(
|
||
"Alembic 迁移成功（第 {} 次尝试）",
|
||
attempt + 1,
|
||
)
|
||
else:
|
||
logger.info("Alembic 迁移已就绪")
|
||
return
|
||
except Exception as e:
|
||
last = e
|
||
will_retry = attempt < max_tries - 1 and _is_retryable_migration_error(e)
|
||
if will_retry:
|
||
delay = base_delay * (2**attempt)
|
||
logger.warning(
|
||
"Alembic 迁移失败（将重试）: {} — {:.1f}s 后重试 ({}/{})",
|
||
e,
|
||
delay,
|
||
attempt + 1,
|
||
max_tries,
|
||
)
|
||
time.sleep(delay)
|
||
else:
|
||
break
|
||
|
||
assert last is not None
|
||
logger.exception("Alembic 迁移失败: {}", last)
|
||
if settings.alembic_startup_fail_fast:
|
||
raise RuntimeError(
|
||
"Alembic migration failed; set ALEMBIC_STARTUP_FAIL_FAST=false "
|
||
"for dev-only tolerate mode"
|
||
) from last
|
||
logger.error(
|
||
"应用将继续启动，但数据库结构可能未更新；请检查 DATABASE_URL 与 PostgreSQL 是否可达"
|
||
)
|