#!/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())