41 lines
1.3 KiB
Python
41 lines
1.3 KiB
Python
"""Simple PostgreSQL connectivity smoke-checker."""
|
||
import os
|
||
import sys
|
||
import psycopg2
|
||
from psycopg2 import OperationalError
|
||
|
||
|
||
DEFAULT_DSN = os.environ.get(
|
||
"PG_DSN", "postgresql://local-Python:Neo-local-1991125@100.64.0.4:5432/LLZQ-test"
|
||
)
|
||
DEFAULT_TIMEOUT = max(1, min(int(os.environ.get("PG_CONNECT_TIMEOUT", 10)), 20))
|
||
|
||
|
||
def check_postgres_connection(dsn: str, timeout: int = DEFAULT_TIMEOUT) -> bool:
|
||
"""Return True if connection succeeds; print diagnostics otherwise."""
|
||
try:
|
||
conn = psycopg2.connect(dsn, connect_timeout=timeout)
|
||
with conn:
|
||
with conn.cursor() as cur:
|
||
cur.execute("SELECT 1;")
|
||
_ = cur.fetchone()
|
||
print(f"PostgreSQL 连接成功 (timeout={timeout}s)")
|
||
return True
|
||
except OperationalError as exc:
|
||
print("PostgreSQL 连接失败（OperationalError）：", exc)
|
||
except Exception as exc: # pragma: no cover - defensive
|
||
print("PostgreSQL 连接失败（其他异常）：", exc)
|
||
|
||
return False
|
||
|
||
|
||
if __name__ == "__main__":
|
||
dsn = sys.argv[1] if len(sys.argv) > 1 else DEFAULT_DSN
|
||
if not dsn:
|
||
print("缺少 DSN，请传入参数或设置 PG_DSN 环境变量。")
|
||
sys.exit(2)
|
||
|
||
ok = check_postgres_connection(dsn)
|
||
if not ok:
|
||
sys.exit(1)
|