planet/scripts/check_database_connection.py

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