190 lines
7.2 KiB
Python
190 lines
7.2 KiB
Python
#!/usr/bin/env python3
|
||
"""Check the backend's PostgreSQL target before init is allowed to change its schema."""
|
||
|
||
import asyncio
|
||
import json
|
||
import os
|
||
from pathlib import Path
|
||
import re
|
||
import socket
|
||
import subprocess
|
||
import sys
|
||
|
||
from sqlalchemy import text
|
||
from sqlalchemy.engine import URL, make_url
|
||
from sqlalchemy.ext.asyncio import create_async_engine
|
||
from sqlalchemy.pool import NullPool
|
||
|
||
ROOT = Path(__file__).resolve().parents[1]
|
||
CONNECT_TIMEOUT_SECONDS = 10
|
||
DEFAULT_POSTGRES_PORT = 5432
|
||
DEFAULT_REDIS_PORT = 6379
|
||
POSTGRES_CONTAINER = "planet_postgres"
|
||
POSTGRES_CONTAINER_PORT = "5432/tcp"
|
||
LOOPBACK_HOSTS = {"localhost", "127.0.0.1", "::1"}
|
||
|
||
|
||
class DatabaseReadinessError(RuntimeError):
|
||
"""A safe, actionable diagnostic that does not include connection credentials."""
|
||
|
||
|
||
class MissingPortBindingError(DatabaseReadinessError):
|
||
"""Exit 2 asks planet.sh to reconcile the managed container once, preserving its volume."""
|
||
|
||
|
||
def backend_database_url() -> str:
|
||
sys.path.insert(0, str(ROOT / "backend"))
|
||
from app.core.config import settings
|
||
|
||
return settings.DATABASE_URL
|
||
|
||
|
||
def configured_ports() -> tuple[int, int]:
|
||
database = make_url(backend_database_url())
|
||
from app.core.config import settings
|
||
|
||
redis = make_url(settings.REDIS_URL)
|
||
postgres_port = database.port or DEFAULT_POSTGRES_PORT
|
||
redis_port = redis.port or DEFAULT_REDIS_PORT
|
||
targets = []
|
||
if database.host in LOOPBACK_HOSTS:
|
||
targets.append(("PostgreSQL", POSTGRES_CONTAINER, postgres_port))
|
||
if redis.host in LOOPBACK_HOSTS:
|
||
targets.append(("Redis", "planet_redis", redis_port))
|
||
check_port_conflicts(targets)
|
||
return postgres_port, redis_port
|
||
|
||
|
||
def check_port_conflicts(targets: list[tuple[str, str, int]]) -> None:
|
||
result = subprocess.run(
|
||
["docker", "ps", "--format", "{{.Names}}\t{{.Ports}}"],
|
||
capture_output=True, text=True, timeout=CONNECT_TIMEOUT_SECONDS, check=False,
|
||
)
|
||
if result.returncode:
|
||
raise DatabaseReadinessError("无法读取 Docker 端口占用；请检查 daemon 和当前用户权限。")
|
||
for service, expected_container, port in targets:
|
||
own_binding = False
|
||
for line in result.stdout.splitlines():
|
||
name, _, bindings = line.partition("\t")
|
||
if not re.search(rf":{port}->", bindings):
|
||
continue
|
||
if name == expected_container:
|
||
own_binding = True
|
||
else:
|
||
raise DatabaseReadinessError(
|
||
f"PLANET_DB_PORT_OCCUPIED: {service} 端口 {port} 被 {name} 占用；"
|
||
"请调整 Planet 后端连接端口。"
|
||
)
|
||
if not own_binding:
|
||
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as listener:
|
||
try:
|
||
# Docker can reuse recently closed TCP connections; TIME_WAIT is not a listener.
|
||
listener.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
|
||
listener.bind(("127.0.0.1", port))
|
||
except OSError as error:
|
||
raise DatabaseReadinessError(
|
||
f"PLANET_DB_PORT_OCCUPIED: {service} 端口 {port} 被宿主机服务占用；"
|
||
"请检查监听进程或调整 Planet 后端连接端口。"
|
||
) from error
|
||
|
||
|
||
def check_published_port(url: URL) -> None:
|
||
if url.host not in LOOPBACK_HOSTS:
|
||
return # An explicitly configured external database has no local container mapping.
|
||
result = subprocess.run(
|
||
[
|
||
"docker",
|
||
"inspect",
|
||
"--format",
|
||
"{{.HostConfig.NetworkMode}}\n{{json .NetworkSettings.Ports}}",
|
||
POSTGRES_CONTAINER,
|
||
],
|
||
capture_output=True,
|
||
text=True,
|
||
timeout=CONNECT_TIMEOUT_SECONDS,
|
||
check=False,
|
||
)
|
||
if result.returncode:
|
||
raise DatabaseReadinessError(
|
||
"无法读取 planet_postgres；请检查当前 Docker context 和容器状态。"
|
||
)
|
||
network_mode, ports_json = result.stdout.strip().split("\n", 1)
|
||
if network_mode == "host":
|
||
return # Host networking intentionally has no published-port table.
|
||
ports = json.loads(ports_json) or {}
|
||
bindings = ports.get(POSTGRES_CONTAINER_PORT) or []
|
||
port = str(url.port or DEFAULT_POSTGRES_PORT)
|
||
if not any(binding.get("HostPort") == port for binding in bindings):
|
||
raise MissingPortBindingError(
|
||
f"planet_postgres 未向宿主机发布后端配置的端口 {port}；"
|
||
"请核对 Compose ports、DATABASE_URL 和 Docker context。"
|
||
"容器内部健康不代表宿主机可连接。"
|
||
)
|
||
|
||
|
||
async def check_connection(url: URL) -> None:
|
||
engine = create_async_engine(url, poolclass=NullPool, echo=False)
|
||
try:
|
||
async with asyncio.timeout(CONNECT_TIMEOUT_SECONDS):
|
||
async with engine.connect() as connection:
|
||
await connection.execute(text("SELECT 1"))
|
||
finally:
|
||
await engine.dispose()
|
||
|
||
|
||
def connection_diagnostic(error: BaseException) -> str:
|
||
pending = [error]
|
||
visited: set[int] = set()
|
||
while pending:
|
||
current = pending.pop()
|
||
if id(current) in visited:
|
||
continue
|
||
visited.add(id(current))
|
||
sqlstate = getattr(current, "sqlstate", None) or getattr(current, "pgcode", None)
|
||
if sqlstate in {"28P01", "28000"}:
|
||
return (
|
||
"PostgreSQL 认证失败：核对 DATABASE_URL 的账号密码及连接目标。"
|
||
"只改 POSTGRES_PASSWORD 不会更新 DATABASE_URL，也不会重设已有数据卷的密码。"
|
||
)
|
||
if sqlstate == "3D000":
|
||
return "目标数据库不存在：核对 DATABASE_URL 的库名与已有数据库。"
|
||
if isinstance(current, (TimeoutError, OSError)):
|
||
return "数据库连接被拒绝、超时或地址不可达：核对端口映射、监听服务和 Docker endpoint。"
|
||
for nested in (getattr(current, "orig", None), current.__cause__, current.__context__):
|
||
if isinstance(nested, BaseException):
|
||
pending.append(nested)
|
||
return "数据库连接检查失败：核对 backend/.env、进程环境变量和目标 PostgreSQL 服务日志。"
|
||
|
||
|
||
def main() -> int:
|
||
try:
|
||
if sys.argv[1:] == ["--ports"]:
|
||
print(*configured_ports())
|
||
return 0
|
||
url = make_url(backend_database_url())
|
||
print(
|
||
f"后端数据库目标: host={url.host!r} port={url.port or DEFAULT_POSTGRES_PORT} database={url.database!r}"
|
||
)
|
||
print(
|
||
"DATABASE_URL 来源: "
|
||
+ ("进程环境变量" if "DATABASE_URL" in os.environ else "backend 配置")
|
||
)
|
||
check_published_port(url)
|
||
asyncio.run(check_connection(url))
|
||
except MissingPortBindingError as error:
|
||
print(str(error), file=sys.stderr)
|
||
return 2
|
||
except DatabaseReadinessError as error:
|
||
print(str(error), file=sys.stderr)
|
||
return 1
|
||
except Exception as error:
|
||
# Driver/config exceptions can contain DSNs and passwords. Never echo their raw text.
|
||
print(connection_diagnostic(error), file=sys.stderr)
|
||
return 1
|
||
print("后端 PostgreSQL 连接与认证已通过")
|
||
return 0
|
||
|
||
|
||
if __name__ == "__main__":
|
||
raise SystemExit(main())
|