81 lines
2.9 KiB
Python
81 lines
2.9 KiB
Python
# -*- 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
|