阶段性更新

etl_billiards/scripts/test_db_connection.py

@@ -0,0 +1,63 @@
+# -*- coding: utf-8 -*-
+"""Quick utility for validating PostgreSQL connectivity."""
+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=5,
+        help="connect_timeout seconds passed to psycopg2 (default: 5)",
+    )
+    return parser.parse_args()
+
+
+def main() -> int:
+    args = parse_args()
+    dsn = args.dsn or os.environ.get("TEST_DB_DSN")
+    if not dsn:
+        print("❌ 未提供 DSN，请通过 --dsn 或 TEST_DB_DSN 指定连接串", file=sys.stderr)
+        return 2
+
+    print(f"尝试连接: {dsn}")
+    try:
+        conn = DatabaseConnection(dsn, connect_timeout=args.timeout)
+    except Exception as exc:  # pragma: no cover - diagnostic output
+        print("❌ 连接失败:", exc, file=sys.stderr)
+        return 1
+
+    try:
+        result = conn.query(args.query)
+        print("✅ 连接成功，查询结果:")
+        for row in result:
+            print(row)
+        conn.close()
+        return 0
+    except Exception as exc:  # pragma: no cover - diagnostic output
+        print("⚠️ 连接成功但执行查询失败:", exc, file=sys.stderr)
+        try:
+            conn.close()
+        finally:
+            return 3
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())