65 lines
1.9 KiB
Python
65 lines
1.9 KiB
Python
# -*- coding: utf-8 -*-
|
|
"""Quick utility for validating PostgreSQL connectivity (ASCII-only output)."""
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import os
|
|
import sys
|
|
|
|
PROJECT_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
|
|
if PROJECT_ROOT not in sys.path:
|
|
sys.path.insert(0, PROJECT_ROOT)
|
|
|
|
from database.connection import DatabaseConnection
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser(description="PostgreSQL connectivity smoke test")
|
|
parser.add_argument("--dsn", help="Override TEST_DB_DSN / env value")
|
|
parser.add_argument(
|
|
"--query",
|
|
default="SELECT 1 AS ok",
|
|
help="Custom SQL to run after connection (default: SELECT 1 AS ok)",
|
|
)
|
|
parser.add_argument(
|
|
"--timeout",
|
|
type=int,
|
|
default=10,
|
|
help="connect_timeout seconds passed to psycopg2 (capped at 20, default: 10)",
|
|
)
|
|
return parser.parse_args()
|
|
|
|
|
|
def main() -> int:
|
|
args = parse_args()
|
|
dsn = args.dsn or os.environ.get("TEST_DB_DSN")
|
|
if not dsn:
|
|
print("Missing DSN. Use --dsn or TEST_DB_DSN.", file=sys.stderr)
|
|
return 2
|
|
|
|
print(f"Trying connection: {dsn}")
|
|
try:
|
|
timeout = max(1, min(args.timeout, 20))
|
|
conn = DatabaseConnection(dsn, connect_timeout=timeout)
|
|
except Exception as exc: # pragma: no cover - diagnostic output
|
|
print("Connection failed:", exc, file=sys.stderr)
|
|
return 1
|
|
|
|
try:
|
|
result = conn.query(args.query)
|
|
print("Connection OK, query result:")
|
|
for row in result:
|
|
print(row)
|
|
conn.close()
|
|
return 0
|
|
except Exception as exc: # pragma: no cover - diagnostic output
|
|
print("Connection succeeded but query failed:", exc, file=sys.stderr)
|
|
try:
|
|
conn.close()
|
|
finally:
|
|
return 3
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|