|
| 1 | +"""启动前的数据库就绪检查（含本地免安装 PostgreSQL 的自动拉起）。 |
| 2 | +
|
| 3 | +`npm run dev`（tools/dev-local.mjs）起 API 前会先探 5433、没起就跑 `pg-dev.ps1 start`； |
| 4 | +但按文档「分开启动」直接跑 uvicorn 的人拿不到这层照顾——PG 一停（开机后没起、或哪个 |
| 5 | +会话收尾时顺手 stop 了），API 就只剩百行 psycopg 超时 traceback，前端全部接口跟着 |
| 6 | +ECONNREFUSED。这里把同一判断搬进 API 启动路径： |
| 7 | +
|
| 8 | +1. 先 `SELECT 1` 探库； |
| 9 | +2. 连不上且目标就是 tools/pg-dev.ps1 管的本地实例（Windows、127.0.0.1/localhost:5433、 |
| 10 | + 脚本在场）→ 自动拉起一次再探；Docker 5432、远端库、非 Windows 一律不插手； |
| 11 | +3. 仍连不上 → 抛 `DatabaseUnavailableError`，一行说清「连不上哪、怎么起」。 |
| 12 | +""" |
| 13 | + |
| 14 | +from __future__ import annotations |
| 15 | + |
| 16 | +import locale |
| 17 | +import logging |
| 18 | +import subprocess |
| 19 | +import sys |
| 20 | +from pathlib import Path |
| 21 | + |
| 22 | +from sqlalchemy.engine import make_url |
| 23 | + |
| 24 | +from .config import SERVICE_ROOT, Settings |
| 25 | +from .db import Database |
| 26 | + |
| 27 | +logger = logging.getLogger(__name__) |
| 28 | + |
| 29 | +LOCAL_PG_PORT = 5433 |
| 30 | +_LOCAL_HOSTS = frozenset({"127.0.0.1", "localhost", "::1"}) |
| 31 | +# tools/pg-dev.ps1 在仓库根；SERVICE_ROOT = backend/api |
| 32 | +PG_DEV_SCRIPT = SERVICE_ROOT.parents[1] / "tools" / "pg-dev.ps1" |
| 33 | +# pg_ctl start -w -t 60 的等待上限，再留 PowerShell 启动开销 |
| 34 | +START_TIMEOUT_SECONDS = 90.0 |
| 35 | + |
| 36 | + |
| 37 | +class DatabaseUnavailableError(RuntimeError): |
| 38 | + """启动时数据库不可达。消息即用户可照着做的提示，不夹带驱动 traceback。""" |
| 39 | + |
| 40 | + |
| 41 | +def describe_target(database_url: str) -> str: |
| 42 | + """连接串的可读目标（不含凭据）：backend://host:port/database。""" |
| 43 | + try: |
| 44 | + url = make_url(database_url) |
| 45 | + except Exception: # 非法连接串也要能原样报出去 |
| 46 | + return database_url |
| 47 | + backend = url.get_backend_name() |
| 48 | + if backend == "sqlite": |
| 49 | + return f"sqlite:{url.database or ':memory:'}" |
| 50 | + host = url.host or "localhost" |
| 51 | + port = f":{url.port}" if url.port else "" |
| 52 | + return f"{backend}://{host}{port}/{url.database or ''}" |
| 53 | + |
| 54 | + |
| 55 | +def manages_local_pg( |
| 56 | + database_url: str, |
| 57 | + *, |
| 58 | + script: Path = PG_DEV_SCRIPT, |
| 59 | + platform: str = sys.platform, |
| 60 | +) -> bool: |
| 61 | + """连接串是否指向 tools/pg-dev.ps1 管理的本地实例（且脚本在、本机是 Windows）。""" |
| 62 | + if platform != "win32" or not script.is_file(): |
| 63 | + return False |
| 64 | + try: |
| 65 | + url = make_url(database_url) |
| 66 | + except Exception: |
| 67 | + return False |
| 68 | + if url.get_backend_name() != "postgresql": |
| 69 | + return False |
| 70 | + return (url.host or "") in _LOCAL_HOSTS and url.port == LOCAL_PG_PORT |
| 71 | + |
| 72 | + |
| 73 | +def start_local_pg( |
| 74 | + script: Path = PG_DEV_SCRIPT, timeout_seconds: float = START_TIMEOUT_SECONDS |
| 75 | +) -> bool: |
| 76 | + """跑一次 `pg-dev.ps1 start`；返回是否成功退出。脚本输出并入本进程日志。""" |
| 77 | + command = [ |
| 78 | + "powershell", |
| 79 | + "-NoProfile", |
| 80 | + "-ExecutionPolicy", |
| 81 | + "Bypass", |
| 82 | + "-File", |
| 83 | + str(script), |
| 84 | + "start", |
| 85 | + ] |
| 86 | + try: |
| 87 | + completed = subprocess.run( |
| 88 | + command, capture_output=True, timeout=timeout_seconds, check=False |
| 89 | + ) |
| 90 | + except (OSError, subprocess.TimeoutExpired) as exc: |
| 91 | + logger.error("自动拉起本地 PostgreSQL 失败：%s", exc) |
| 92 | + return False |
| 93 | + output = (_decode_console(completed.stdout) + _decode_console(completed.stderr)).strip() |
| 94 | + if output: |
| 95 | + logger.info("pg-dev.ps1 start（exit=%s）：%s", completed.returncode, output) |
| 96 | + return completed.returncode == 0 |
| 97 | + |
| 98 | + |
| 99 | +def _decode_console(raw: bytes) -> str: |
| 100 | + """PowerShell 子进程的输出编码随控制台代码页变（UTF-8 或 GBK 都可能）： |
| 101 | + 先按 UTF-8 严格解，解不出再按本机首选编码兜底，别把脚本的中文提示解成乱码。""" |
| 102 | + try: |
| 103 | + return raw.decode("utf-8") |
| 104 | + except UnicodeDecodeError: |
| 105 | + return raw.decode(locale.getpreferredencoding(False), errors="replace") |
| 106 | + |
| 107 | + |
| 108 | +def ensure_database_ready(db: Database, settings: Settings) -> None: |
| 109 | + """探库；本地 pg-dev 实例没起就拉起一次；仍不可达则抛 DatabaseUnavailableError。""" |
| 110 | + error = db.ping() |
| 111 | + if error is None: |
| 112 | + return |
| 113 | + target = describe_target(settings.database_url) |
| 114 | + if settings.local_pg_autostart and manages_local_pg(settings.database_url): |
| 115 | + logger.warning("数据库连不上（%s：%s），尝试自动拉起本地 PostgreSQL…", target, error) |
| 116 | + if start_local_pg(): |
| 117 | + error = db.ping() |
| 118 | + if error is None: |
| 119 | + logger.info("本地 PostgreSQL 已拉起，数据库就绪：%s", target) |
| 120 | + return |
| 121 | + message = ( |
| 122 | + f"数据库连不上：{target}（{error}）。" |
| 123 | + "本地开发请先起 PostgreSQL：.\\tools\\pg-dev.ps1 start（首次先 init），" |
| 124 | + "或直接 npm run dev（会自动拉起）；连接串由 OMM_DATABASE_URL / backend/api/.env 决定。" |
| 125 | + ) |
| 126 | + logger.error(message) |
| 127 | + raise DatabaseUnavailableError(message) |
0 commit comments