diff --git a/scripts/ensure_pg.py b/scripts/ensure_pg.py new file mode 100644 index 0000000..d6d6e1c --- /dev/null +++ b/scripts/ensure_pg.py @@ -0,0 +1,208 @@ +"""确保本地 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 → 幂等确保测试库存在。全程无 SQLite 兜底。 + +生产用原生 PG(scripts/init_postgres.py),不走本模块。 +""" +from __future__ import annotations + +import os +import socket +import subprocess +import sys +import time +from pathlib import Path +from urllib.parse import urlsplit + +ROOT = Path(__file__).resolve().parent.parent + +APP_DB = "shaguabijia" +TEST_DB = "shaguabijia_test" +DB_USER = "shaguabijia_app" +COMPOSE_SERVICE = "postgres" + +DOCKER_START_TIMEOUT = int(os.environ.get("ENSURE_PG_DOCKER_TIMEOUT", "120")) +PG_READY_TIMEOUT = int(os.environ.get("ENSURE_PG_READY_TIMEOUT", "60")) +POLL_INTERVAL = 3.0 + +SQLITE_FIX_HINT = ( + "postgresql+psycopg://shaguabijia_app:shaguabijia_dev_pw@localhost:5432/shaguabijia" +) + + +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 _docker_desktop_cmd(platform: str, program_files: str) -> list[str] | None: + """按平台给出启动 Docker Desktop 的命令;Linux 返回 None(daemon 需 sudo,让用户手动)。""" + if platform.startswith("win"): + return [str(Path(program_files) / "Docker" / "Docker" / "Docker Desktop.exe")] + 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, + ) + return True + except (OSError, subprocess.CalledProcessError): + return False + + +def _docker_cli_ok() -> bool: + 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, os.environ.get("ProgramFiles", r"C:\Program Files") + ) + if cmd is None: + _log("Docker 守护进程未运行。Linux 请手动:sudo systemctl start docker,然后重试。") + return False + if sys.platform.startswith("win") and not Path(cmd[0]).exists(): + _log(f"找不到 Docker Desktop:{cmd[0]}。请手动启动 Docker Desktop 后重试。") + return False + _log(f"启动 Docker Desktop(首次冷启可能 30-60s)…") + 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 _compose_up() -> bool: + _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: + 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, + ) + 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 _ensure_test_db() -> None: + """幂等建测试库(兼容老 pgdata 卷首启没跑 initdb 的情况)。""" + check = 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, + ) + if check.returncode == 0 and check.stdout.strip() == "1": + return + _log(f"建测试库 {TEST_DB}…") + 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, check=False, + ) + + +def ensure(database_url: str | None = None) -> bool: + """确保 PG 就绪,返回 True/False。database_url 缺省从 settings 读(尊重 .env)。""" + if database_url is None: + from app.core.config import settings # 延迟导入,避免过早固化 settings + + database_url = settings.DATABASE_URL + + if _is_sqlite(database_url): + _log("检测到 DATABASE_URL 仍是 SQLite。本地开发/测试已切 PostgreSQL,请改成:") + _log(f" DATABASE_URL={SQLITE_FIX_HINT}") + return False + + 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/") + return False + if not _start_docker_daemon(): + return False + if not _compose_up(): + return False + if not _wait_pg_ready(host, port): + return False + + _ensure_test_db() + return True + + +if __name__ == "__main__": + sys.exit(0 if ensure() else 1) diff --git a/tests/test_ensure_pg.py b/tests/test_ensure_pg.py new file mode 100644 index 0000000..ee6e686 --- /dev/null +++ b/tests/test_ensure_pg.py @@ -0,0 +1,87 @@ +"""scripts/ensure_pg.py 纯函数单测(不需要 Docker/PG)。""" +from __future__ import annotations + +import socket + +from scripts.ensure_pg import ( + _docker_desktop_cmd, + _is_sqlite, + _parse_host_port, + _port_open, + ensure, +) + + +def test_is_sqlite(): + assert _is_sqlite("sqlite:///./data/app.db") + assert _is_sqlite(" SQLite:///x ") + assert not _is_sqlite("postgresql+psycopg://u:p@localhost:5432/db") + + +def test_parse_host_port_full(): + assert _parse_host_port( + "postgresql+psycopg://u:p@localhost:5432/shaguabijia" + ) == ("localhost", 5432) + + +def test_parse_host_port_defaults(): + # 缺端口 → 5432 + assert _parse_host_port("postgresql+psycopg://u:p@db.example/x")[1] == 5432 + # 缺 host → localhost + assert _parse_host_port("postgresql+psycopg:///x") == ("localhost", 5432) + + +def test_parse_host_port_testdb(): + assert _parse_host_port( + "postgresql+psycopg://u:p@localhost:5432/shaguabijia_test" + ) == ("localhost", 5432) + + +def test_port_open_true(): + srv = socket.socket() + srv.bind(("127.0.0.1", 0)) + srv.listen(1) + port = srv.getsockname()[1] + try: + assert _port_open("127.0.0.1", port, timeout=1.0) + finally: + srv.close() + + +def test_port_open_false(): + s = socket.socket() + s.bind(("127.0.0.1", 0)) + port = s.getsockname()[1] + s.close() # 释放端口,无人监听 → 连接应失败 + assert not _port_open("127.0.0.1", port, timeout=0.3) + + +def test_docker_desktop_cmd_windows(): + cmd = _docker_desktop_cmd("win32", r"C:\Program Files") + assert cmd is not None + assert cmd[0].endswith("Docker Desktop.exe") + assert "Docker" in cmd[0] + + +def test_docker_desktop_cmd_darwin(): + assert _docker_desktop_cmd("darwin", "") == ["open", "-a", "Docker"] + + +def test_docker_desktop_cmd_linux(): + assert _docker_desktop_cmd("linux", "") is None + + +def test_ensure_rejects_sqlite(): + # dev 守卫:sqlite 直接 False(不碰 Docker) + assert ensure("sqlite:///./data/app.db") is False + + +def test_ensure_shortcircuits_when_pg_up(monkeypatch): + # 端口通 → 直接 True,绝不触碰 docker + monkeypatch.setattr("scripts.ensure_pg._port_open", lambda *a, **k: True) + + def _boom(): + raise AssertionError("端口通时不应调用 docker") + + monkeypatch.setattr("scripts.ensure_pg._docker_cli_ok", _boom) + assert ensure("postgresql+psycopg://u:p@localhost:5432/shaguabijia") is True