101 lines
3.1 KiB
Python
101 lines
3.1 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 _run_alembic_upgrade_once() -> None:
|
|||
|
|
from alembic.command import upgrade
|
|||
|
|
from alembic.config import Config
|
|||
|
|
|
|||
|
|
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 是否可达"
|
|||
|
|
)
|