guke
3f449d3600
原先把 Docker Desktop.exe 写死在 %ProgramFiles%\Docker\Docker\,装在 D 盘就自动拉起 失败(明明 docker CLI 已在 PATH 上可用、只是 daemon 没启)。改为按可靠性排序的多路探测: DOCKER_DESKTOP_EXE 显式指定 → 从 PATH 上 docker CLI 反推安装目录(跟随实际盘符)→ 注册表 → 各 Program Files 变体兜底;找不到时提示 DOCKER_DESKTOP_EXE / sqlite 逃生舱。 .env.example 记一笔用法。 顺带修一个先前就红、与现行"sqlite→降级返 True"契约矛盾的测试 (test_ensure_rejects_sqlite → test_ensure_sqlite_escape_hatch)。 Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
395 lines
16 KiB
Python
395 lines
16 KiB
Python
"""确保本地 PostgreSQL 就绪(开发/测试统一用 Docker PG)。
|
|
|
|
被三处复用:
|
|
- run.sh / run.bat:`python -m scripts.ensure_pg`(CLI,失败退非 0)
|
|
- tests/conftest.py:`from scripts.ensure_pg import ensure; ensure(test_url)`
|
|
|
|
流程:读 DATABASE_URL → TCP 探测 → 没起就(必要时启 Docker Desktop)→
|
|
`docker compose up -d` → 等 PG ready → 幂等确保测试库存在。
|
|
|
|
运行时(run.sh/run.bat)支持【显式】SQLite 逃生舱:DATABASE_URL 设为 sqlite → 放行并打印
|
|
醒目降级横幅(绝不静默替你切库);无 docker CLI 时报错里也指路该逃生舱。测试侧
|
|
(conftest 传 PG URL)不降级——sqlite 分支不触发,没 PG 直接 raise。详见设计文档 §10。
|
|
|
|
生产用原生 PG(scripts/init_postgres.py),不走本模块。
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
import shutil
|
|
import socket
|
|
import subprocess
|
|
import sys
|
|
import time
|
|
from collections.abc import Mapping
|
|
from pathlib import Path, PureWindowsPath
|
|
from urllib.parse import urlsplit
|
|
|
|
ROOT = Path(__file__).resolve().parent.parent
|
|
|
|
# 日志里可能含 emoji(如 ✅);Windows GBK 控制台(cmd.exe)无法编码会抛 UnicodeEncodeError → 脚本崩、
|
|
# run.bat 误判 ensure_pg 失败。用 backslashreplace 保底:中文仍正常,仅不可编码字符被转义,不崩。
|
|
for _stream in (sys.stdout, sys.stderr):
|
|
try:
|
|
_stream.reconfigure(errors="backslashreplace")
|
|
except (AttributeError, ValueError):
|
|
pass
|
|
|
|
APP_DB = "shaguabijia"
|
|
TEST_DB = "shaguabijia_test"
|
|
DB_USER = "shaguabijia_app"
|
|
COMPOSE_SERVICE = "postgres"
|
|
CONTAINER_NAME = "shaguabijia-pg" # 必须与 docker-compose.yml 的 container_name 一致
|
|
|
|
# 钉死 compose 项目名:否则它默认取运行目录 basename,在不同目录/worktree 之间切会各自
|
|
# 成一个项目 → 同一个固定 container_name 撞名报错、pgdata 卷还会按项目名分裂成多份。
|
|
# 钉成 app 名后,无论从哪个目录/worktree 跑都是同一个项目、同一个卷。setdefault:尊重外部覆盖。
|
|
os.environ.setdefault("COMPOSE_PROJECT_NAME", "shaguabijia")
|
|
|
|
DOCKER_START_TIMEOUT = int(os.environ.get("ENSURE_PG_DOCKER_TIMEOUT", "120"))
|
|
PG_READY_TIMEOUT = int(os.environ.get("ENSURE_PG_READY_TIMEOUT", "60"))
|
|
# 单条 docker 探测/exec 命令的超时:防 Docker 守护进程半死(尤其 Windows 冷启)时
|
|
# docker info / exec 无限挂起、绕过上面的总超时。
|
|
DOCKER_CMD_TIMEOUT = int(os.environ.get("ENSURE_PG_CMD_TIMEOUT", "15"))
|
|
POLL_INTERVAL = 3.0
|
|
|
|
PG_URL = (
|
|
"postgresql+psycopg://shaguabijia_app:shaguabijia_dev_pw@localhost:5432/shaguabijia"
|
|
)
|
|
SQLITE_URL = "sqlite:///./data/app.db" # 无 Docker 时的显式降级逃生舱(仅 run 运行时)
|
|
|
|
|
|
def _log(msg: str) -> None:
|
|
print(f"[ensure_pg] {msg}", flush=True)
|
|
|
|
|
|
def _is_sqlite(url: str) -> bool:
|
|
return url.strip().lower().startswith("sqlite")
|
|
|
|
|
|
def _parse_host_port(url: str) -> tuple[str, int]:
|
|
"""从 SQLAlchemy URL 取 host/port,缺省 localhost:5432。"""
|
|
parts = urlsplit(url)
|
|
return (parts.hostname or "localhost"), (parts.port or 5432)
|
|
|
|
|
|
def _port_open(host: str, port: int, timeout: float = 1.0) -> bool:
|
|
try:
|
|
with socket.create_connection((host, port), timeout=timeout):
|
|
return True
|
|
except OSError:
|
|
return False
|
|
|
|
|
|
def _win_docker_desktop_candidates(
|
|
env: Mapping[str, str], docker_cli: str | None
|
|
) -> list[PureWindowsPath]:
|
|
"""Windows 上 Docker Desktop.exe 的候选路径(按优先级)。纯函数:不碰文件系统、不读注册表。
|
|
|
|
用 PureWindowsPath 解析,故在任意 OS 上跑单测都按 Windows 语义(反斜杠分隔),行为确定。
|
|
优先级:
|
|
① 环境变量 DOCKER_DESKTOP_EXE 显式指定(终极逃生舱,盘符随你);
|
|
② 从 PATH 上的 docker CLI 反推——Docker Desktop 的 CLI 在
|
|
<安装目录>\\resources\\bin\\docker.exe,往上几级即安装目录,天然跟随实际盘符
|
|
(装在 D 盘就反推出 D 盘,不再写死 C 盘);多取几级容忍未来目录布局微调;
|
|
③ 各 Program Files 变体下的标准安装路径兜底(覆盖常规 C 盘装)。
|
|
调用方按序取第一个真实存在的。
|
|
"""
|
|
out: list[PureWindowsPath] = []
|
|
override = (env.get("DOCKER_DESKTOP_EXE") or "").strip().strip('"')
|
|
if override:
|
|
out.append(PureWindowsPath(override))
|
|
if docker_cli:
|
|
for parent in list(PureWindowsPath(docker_cli).parents)[:4]:
|
|
out.append(parent / "Docker Desktop.exe")
|
|
for var in ("ProgramFiles", "ProgramW6432", "ProgramFiles(x86)"):
|
|
root = env.get(var)
|
|
if root:
|
|
out.append(PureWindowsPath(root) / "Docker" / "Docker" / "Docker Desktop.exe")
|
|
return out
|
|
|
|
|
|
def _docker_desktop_from_registry() -> Path | None:
|
|
"""从注册表尽力取 Docker Desktop.exe 位置(best-effort;非 Windows / 任何异常都当没找到)。
|
|
|
|
比路径猜测更权威且完全跟随实际盘符。探两处:
|
|
- App Paths\\Docker Desktop.exe 的默认值(通常就是 exe 全路径);
|
|
- Uninstall\\Docker Desktop 的 InstallLocation(安装目录,需再拼 exe 名)。
|
|
"""
|
|
try:
|
|
import winreg
|
|
except ImportError: # 非 Windows
|
|
return None
|
|
probes = (
|
|
(winreg.HKEY_LOCAL_MACHINE,
|
|
r"SOFTWARE\Microsoft\Windows\CurrentVersion\App Paths\Docker Desktop.exe", "", False),
|
|
(winreg.HKEY_LOCAL_MACHINE,
|
|
r"SOFTWARE\Microsoft\Windows\CurrentVersion\Uninstall\Docker Desktop",
|
|
"InstallLocation", True),
|
|
)
|
|
for hive, subkey, value_name, join_exe in probes:
|
|
try:
|
|
with winreg.OpenKey(hive, subkey) as key:
|
|
val, _ = winreg.QueryValueEx(key, value_name)
|
|
except OSError:
|
|
continue # 键不存在/无权限 → 下一个
|
|
if not val:
|
|
continue
|
|
exe = Path(val) / "Docker Desktop.exe" if join_exe else Path(val)
|
|
if exe.exists():
|
|
return exe
|
|
return None
|
|
|
|
|
|
def _find_docker_desktop_exe() -> Path | None:
|
|
"""Windows 上尽力定位【真实存在】的 Docker Desktop.exe;遍历候选 + 注册表兜底,找不到返回 None。"""
|
|
for cand in _win_docker_desktop_candidates(os.environ, shutil.which("docker")):
|
|
if Path(cand).exists():
|
|
return Path(cand)
|
|
return _docker_desktop_from_registry()
|
|
|
|
|
|
def _docker_desktop_cmd(platform: str) -> list[str] | None:
|
|
"""按平台给出启动 Docker Desktop 的命令。
|
|
|
|
Windows:智能定位 exe(见 _find_docker_desktop_exe),找不到 → None。
|
|
macOS:交给 `open -a Docker`。Linux:None(daemon 需 sudo,让用户手动)。
|
|
"""
|
|
if platform.startswith("win"):
|
|
exe = _find_docker_desktop_exe()
|
|
return [str(exe)] if exe else None
|
|
if platform == "darwin":
|
|
return ["open", "-a", "Docker"]
|
|
return None
|
|
|
|
|
|
def _docker_ok(subcmd: str) -> bool:
|
|
"""`docker --version`(CLI 在不在,纯客户端)/`docker info`(daemon 起没起)成功与否。"""
|
|
try:
|
|
subprocess.run(
|
|
["docker", subcmd],
|
|
cwd=ROOT,
|
|
stdout=subprocess.DEVNULL,
|
|
stderr=subprocess.DEVNULL,
|
|
check=True,
|
|
timeout=DOCKER_CMD_TIMEOUT,
|
|
)
|
|
return True
|
|
except (OSError, subprocess.CalledProcessError, subprocess.TimeoutExpired):
|
|
return False
|
|
|
|
|
|
def _docker_cli_ok() -> bool:
|
|
# 必须用 `docker --version`(纯客户端,不连 daemon)而非 `docker version`
|
|
# (后者要连 daemon,守护进程没起时退非零)——否则「装了 Docker 但没启动」
|
|
# 会被误判成「没装 CLI」,直接绕过下面 _start_docker_daemon() 的自动拉起(需求②)。
|
|
return _docker_ok("--version")
|
|
|
|
|
|
def _docker_daemon_ok() -> bool:
|
|
return _docker_ok("info")
|
|
|
|
|
|
def _start_docker_daemon() -> bool:
|
|
"""守护进程没起时按平台拉起,轮询到就绪。返回是否成功。"""
|
|
if _docker_daemon_ok():
|
|
return True
|
|
cmd = _docker_desktop_cmd(sys.platform)
|
|
if cmd is None:
|
|
if sys.platform.startswith("win"):
|
|
# docker CLI 在 PATH 上(否则走不到这)、却定位不到 Docker Desktop.exe:多为非标准安装位置
|
|
_log("找不到 Docker Desktop.exe(已试:PATH 上 docker CLI 反推、注册表、常见安装目录)。")
|
|
_log(" 确已安装 → 设环境变量 DOCKER_DESKTOP_EXE=<Docker Desktop.exe 全路径> 再重试,"
|
|
"或先手动启动 Docker Desktop。")
|
|
_log(f" 不想折腾 → 把 .env 的 DATABASE_URL 改成 {SQLITE_URL} 可降级用 SQLite 跑(仅救急)。")
|
|
else:
|
|
_log("Docker 守护进程未运行。Linux 请手动:sudo systemctl start docker,然后重试。")
|
|
return False
|
|
_log(f"启动 Docker Desktop(首次冷启可能 30-60s):{cmd[0]}")
|
|
try:
|
|
subprocess.Popen(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
|
|
except OSError as e:
|
|
_log(f"启动 Docker Desktop 失败:{e}")
|
|
return False
|
|
deadline = time.monotonic() + DOCKER_START_TIMEOUT
|
|
while time.monotonic() < deadline:
|
|
if _docker_daemon_ok():
|
|
_log("Docker 守护进程已就绪。")
|
|
return True
|
|
_log("等待 Docker 守护进程…")
|
|
time.sleep(POLL_INTERVAL)
|
|
_log(f"等待 Docker 守护进程超时({DOCKER_START_TIMEOUT}s)。")
|
|
return False
|
|
|
|
|
|
def _ps_names(*filters: str) -> str:
|
|
"""docker ps -a 按 filter 查容器名(每行一个);失败返回空串。"""
|
|
args = ["docker", "ps", "-a", "--format", "{{.Names}}"]
|
|
for f in filters:
|
|
args += ["--filter", f]
|
|
try:
|
|
r = subprocess.run(
|
|
args, cwd=ROOT, capture_output=True, text=True, timeout=DOCKER_CMD_TIMEOUT,
|
|
)
|
|
except (OSError, subprocess.TimeoutExpired):
|
|
return ""
|
|
return r.stdout if r.returncode == 0 else ""
|
|
|
|
|
|
def _remove_stale_container() -> None:
|
|
"""删掉「同名但不属于本 compose 项目」的残留容器(旧目录/worktree 建的)。
|
|
|
|
固定的 container_name 是全局唯一名:若旧项目留下一个同名容器,`docker compose up`
|
|
会因撞名报 "container name already in use" 而失败。这里在 up 之前主动清掉它。
|
|
数据在命名卷(<project>_pgdata)里,删容器不删卷、不丢数据。
|
|
"""
|
|
project = os.environ.get("COMPOSE_PROJECT_NAME", "")
|
|
name_filter = f"name=^{CONTAINER_NAME}$"
|
|
if CONTAINER_NAME not in _ps_names(name_filter).split():
|
|
return # 没有同名容器
|
|
ours = _ps_names(name_filter, f"label=com.docker.compose.project={project}")
|
|
if CONTAINER_NAME in ours.split():
|
|
return # 就是本项目的容器,compose 会自己 start/复用,别删
|
|
_log(f"发现残留同名容器 {CONTAINER_NAME}(非本项目 '{project}'),删除以避免撞名"
|
|
f"(数据在卷里,不丢)…")
|
|
try:
|
|
subprocess.run(
|
|
["docker", "rm", "-f", CONTAINER_NAME], cwd=ROOT,
|
|
stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, timeout=DOCKER_CMD_TIMEOUT,
|
|
)
|
|
except (OSError, subprocess.TimeoutExpired):
|
|
_log(f"⚠️ 删除残留容器失败,可手动: docker rm -f {CONTAINER_NAME}")
|
|
|
|
|
|
def _compose_up() -> bool:
|
|
_remove_stale_container()
|
|
_log("docker compose up -d(镜像缺失会自动拉取,首用约几十秒)…")
|
|
try:
|
|
subprocess.run(["docker", "compose", "up", "-d"], cwd=ROOT, check=True)
|
|
return True
|
|
except (OSError, subprocess.CalledProcessError) as e:
|
|
_log(f"docker compose up 失败:{e}")
|
|
return False
|
|
|
|
|
|
def _pg_isready() -> bool:
|
|
try:
|
|
r = subprocess.run(
|
|
["docker", "compose", "exec", "-T", COMPOSE_SERVICE,
|
|
"pg_isready", "-U", DB_USER, "-d", APP_DB],
|
|
cwd=ROOT, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
|
|
timeout=DOCKER_CMD_TIMEOUT,
|
|
)
|
|
except (OSError, subprocess.TimeoutExpired):
|
|
return False
|
|
return r.returncode == 0
|
|
|
|
|
|
def _wait_pg_ready(host: str, port: int) -> bool:
|
|
deadline = time.monotonic() + PG_READY_TIMEOUT
|
|
while time.monotonic() < deadline:
|
|
if _port_open(host, port) and _pg_isready():
|
|
_log("PostgreSQL 已就绪。")
|
|
return True
|
|
_log("等待 PostgreSQL 就绪…")
|
|
time.sleep(POLL_INTERVAL)
|
|
_log(f"等待 PostgreSQL 就绪超时({PG_READY_TIMEOUT}s)。")
|
|
return False
|
|
|
|
|
|
def _test_db_exists() -> bool:
|
|
"""测试库是否已存在(连业务库 shaguabijia 查 pg_database)。"""
|
|
try:
|
|
r = subprocess.run(
|
|
["docker", "compose", "exec", "-T", COMPOSE_SERVICE,
|
|
"psql", "-U", DB_USER, "-d", APP_DB, "-tAc",
|
|
f"SELECT 1 FROM pg_database WHERE datname='{TEST_DB}'"],
|
|
cwd=ROOT, capture_output=True, text=True, timeout=DOCKER_CMD_TIMEOUT,
|
|
)
|
|
except (OSError, subprocess.TimeoutExpired):
|
|
return False
|
|
return r.returncode == 0 and r.stdout.strip() == "1"
|
|
|
|
|
|
def ensure_test_db() -> bool:
|
|
"""幂等建测试库(兼容老 pgdata 卷首启没跑 initdb 的情况)。返回测试库是否就绪。
|
|
|
|
公开给 conftest 单独调用:ensure() 在「端口已通」时会短路返回、不建测试库,
|
|
所以测试侧需在 ensure() 之后再显式补一刀(best-effort)。
|
|
"""
|
|
if _test_db_exists():
|
|
return True
|
|
_log(f"建测试库 {TEST_DB}…")
|
|
try:
|
|
create = subprocess.run(
|
|
["docker", "compose", "exec", "-T", COMPOSE_SERVICE,
|
|
"psql", "-U", DB_USER, "-d", APP_DB, "-c",
|
|
f"CREATE DATABASE {TEST_DB} OWNER {DB_USER}"],
|
|
cwd=ROOT, capture_output=True, text=True, timeout=DOCKER_CMD_TIMEOUT,
|
|
)
|
|
except (OSError, subprocess.TimeoutExpired) as e:
|
|
_log(f"⚠️ 建测试库 {TEST_DB} 失败:{e}")
|
|
return False
|
|
# returncode==0=建成功;非 0 但库已存在=与并发创建者竞争失败(42P04),仍算就绪
|
|
if create.returncode == 0 or _test_db_exists():
|
|
return True
|
|
_log(f"⚠️ 建测试库 {TEST_DB} 失败:{(create.stderr or '').strip()}")
|
|
return False
|
|
|
|
|
|
def _warn_sqlite_degraded() -> None:
|
|
"""DATABASE_URL 是 SQLite 时打印醒目降级横幅(显式逃生舱,非静默切库)。"""
|
|
for line in (
|
|
"⚠️ ================= 降级模式(SQLite) =================",
|
|
"⚠️ DATABASE_URL 是 SQLite,不是 PostgreSQL。",
|
|
"⚠️ PG 专有 SQL(窗口函数/FILTER/JSONB)与严格类型在此模式【不被验证】。",
|
|
"⚠️ 提交前请在装了 Docker 的机器上用 PG 复跑;装好后把 DATABASE_URL 改回:",
|
|
f"⚠️ {PG_URL}",
|
|
"⚠️ ===================================================",
|
|
):
|
|
_log(line)
|
|
|
|
|
|
def ensure(database_url: str | None = None) -> bool:
|
|
"""确保 PG 就绪,返回 True/False。database_url 缺省从 settings 读(尊重 .env)。
|
|
|
|
运行时若 DATABASE_URL 是 SQLite → 打印降级横幅并返回 True(显式逃生舱);
|
|
conftest 传的是 PG URL,故测试侧永不走此分支。
|
|
"""
|
|
if database_url is None:
|
|
from app.core.config import settings # 延迟导入,避免过早固化 settings
|
|
|
|
database_url = settings.DATABASE_URL
|
|
|
|
if _is_sqlite(database_url):
|
|
_warn_sqlite_degraded()
|
|
return True
|
|
|
|
host, port = _parse_host_port(database_url)
|
|
|
|
if _port_open(host, port):
|
|
_log(f"✅ PostgreSQL 已在 {host}:{port} 运行,跳过 Docker。")
|
|
return True
|
|
|
|
_log(f"{host}:{port} 无 PostgreSQL,准备用 Docker 拉起…")
|
|
|
|
if not _docker_cli_ok():
|
|
_log("未检测到 docker 命令。请先安装 Docker Desktop:")
|
|
_log(" https://www.docker.com/products/docker-desktop/")
|
|
_log(f"装不了 Docker?把 .env 的 DATABASE_URL 改成 {SQLITE_URL} 可降级用 SQLite 跑")
|
|
_log(" (PG 专有 SQL/严格性不被验证,仅救急);改完重跑 run.sh/run.bat。")
|
|
return False
|
|
if not _start_docker_daemon():
|
|
return False
|
|
if not _compose_up():
|
|
return False
|
|
if not _wait_pg_ready(host, port):
|
|
return False
|
|
if not ensure_test_db():
|
|
return False
|
|
return True
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(0 if ensure() else 1)
|