Xiaoxia AI
83 lines
2.1 KiB
Python
83 lines
2.1 KiB
Python
"""
|
|
数据库连接池管理
|
|
"""
|
|
|
|
from typing import Optional
|
|
|
|
import psycopg2
|
|
from psycopg2 import pool
|
|
from psycopg2.extras import RealDictCursor
|
|
|
|
|
|
class DatabaseConnectionPool:
|
|
"""PostgreSQL 连接池"""
|
|
|
|
_instance: Optional["DatabaseConnectionPool"] = None
|
|
_pool: Optional[pool.ThreadedConnectionPool] = None
|
|
|
|
def __new__(cls):
|
|
if cls._instance is None:
|
|
cls._instance = super().__new__(cls)
|
|
return cls._instance
|
|
|
|
def initialize(
|
|
self,
|
|
connection_string: str,
|
|
minconn: int = 1,
|
|
maxconn: int = 10,
|
|
):
|
|
"""初始化连接池"""
|
|
if self._pool is None:
|
|
self._pool = pool.ThreadedConnectionPool(
|
|
minconn=minconn,
|
|
maxconn=maxconn,
|
|
dsn=connection_string,
|
|
)
|
|
|
|
def get_connection(self):
|
|
"""从连接池获取连接"""
|
|
if self._pool is None:
|
|
raise RuntimeError("Connection pool not initialized")
|
|
return self._pool.getconn()
|
|
|
|
def put_connection(self, conn):
|
|
"""将连接归还到连接池"""
|
|
if self._pool is not None:
|
|
self._pool.putconn(conn)
|
|
|
|
def close_all(self):
|
|
"""关闭所有连接"""
|
|
if self._pool is not None:
|
|
self._pool.closeall()
|
|
self._pool = None
|
|
|
|
|
|
# 全局连接池实例
|
|
db_pool = DatabaseConnectionPool()
|
|
|
|
|
|
class PooledConnection:
|
|
"""连接池上下文管理器"""
|
|
|
|
def __init__(self, cursor_factory=RealDictCursor):
|
|
self.cursor_factory = cursor_factory
|
|
self.conn = None
|
|
|
|
def __enter__(self):
|
|
self.conn = db_pool.get_connection()
|
|
if self.cursor_factory:
|
|
self.conn.cursor_factory = self.cursor_factory
|
|
return self.conn
|
|
|
|
def __exit__(self, exc_type, exc_val, exc_tb):
|
|
if self.conn:
|
|
if exc_type is not None:
|
|
self.conn.rollback()
|
|
db_pool.put_connection(self.conn)
|
|
return False
|
|
|
|
|
|
def get_db_connection():
|
|
"""获取数据库连接（用于依赖注入）"""
|
|
return PooledConnection()
|