feat(dev): scripts/ensure_pg.py 探测/拉起本地 Docker PostgreSQL
This commit is contained in:
guke
@@ -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)
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user