初始提交：飞球 ETL 系统全量代码

database/connection.py

@@ -0,0 +1,80 @@
+# -*- coding: utf-8 -*-
+"""数据库连接管理器（限制最大连接超时时间）。"""
+
+import psycopg2
+import psycopg2.extras
+
+
+class DatabaseConnection:
+    """封装 psycopg2 连接，支持会话参数和超时保护。"""
+
+    def __init__(self, dsn: str, session: dict = None, connect_timeout: int = None):
+        self._dsn = dsn
+        self._session = session or {}
+        self._connect_timeout = connect_timeout
+        self.conn = self._open_connection()
+
+    def _open_connection(self):
+        """创建并初始化连接（包含会话参数）。"""
+        timeout_val = self._connect_timeout if self._connect_timeout is not None else 5
+        # 生产环境要求：数据库连接超时不得超过 20 秒。
+        timeout_val = max(1, min(int(timeout_val), 20))
+
+        conn = psycopg2.connect(self._dsn, connect_timeout=timeout_val)
+        conn.autocommit = False
+
+        # 会话参数（时区、语句超时等）
+        if self._session:
+            with conn.cursor() as c:
+                if self._session.get("timezone"):
+                    c.execute("SET TIME ZONE %s", (self._session["timezone"],))
+                if self._session.get("statement_timeout_ms") is not None:
+                    c.execute(
+                        "SET statement_timeout = %s",
+                        (int(self._session["statement_timeout_ms"]),),
+                    )
+                if self._session.get("lock_timeout_ms") is not None:
+                    c.execute(
+                        "SET lock_timeout = %s", (int(self._session["lock_timeout_ms"]),)
+                    )
+                if self._session.get("idle_in_tx_timeout_ms") is not None:
+                    c.execute(
+                        "SET idle_in_transaction_session_timeout = %s",
+                        (int(self._session["idle_in_tx_timeout_ms"]),),
+                    )
+        return conn
+
+    def query(self, sql: str, args=None):
+        """Execute a query and fetch all rows."""
+        with self.conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as c:
+            c.execute(sql, args)
+            return c.fetchall()
+
+    def execute(self, sql: str, args=None):
+        """Execute a SQL statement without returning rows."""
+        with self.conn.cursor() as c:
+            c.execute(sql, args)
+
+    def commit(self):
+        """Commit current transaction."""
+        self.conn.commit()
+
+    def rollback(self):
+        """Rollback current transaction."""
+        self.conn.rollback()
+
+    def close(self):
+        """Safely close the connection."""
+        try:
+            self.conn.close()
+        except Exception:
+            pass
+
+    def ensure_open(self) -> bool:
+        """确保连接可用，若已关闭则尝试重连。"""
+        try:
+            if getattr(self.conn, "closed", 0):
+                self.conn = self._open_connection()
+            return True
+        except Exception:
+            return False