"""确保本地 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 # 日志里可能含 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" 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 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, timeout=DOCKER_CMD_TIMEOUT, ) return True except (OSError, subprocess.CalledProcessError, subprocess.TimeoutExpired): 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("启动 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: 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 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 if not ensure_test_db(): return False return True if __name__ == "__main__": sys.exit(0 if ensure() else 1)