Initial commit: 知识管理平台
主要功能: - 多用户管理系统 - 浏览器自动化(Playwright) - 任务编排和执行 - Docker容器化部署 - 数据持久化和日志管理 技术栈: - Flask 3.0.0 - Playwright 1.40.0 - SQLite with connection pooling - Docker + Docker Compose 部署说明详见README.md
This commit is contained in:
Yu Yon
250
db_pool.py
Executable file
250
db_pool.py
Executable file
@@ -0,0 +1,250 @@
|
||||
#!/usr/bin/env python3
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
数据库连接池模块
|
||||
使用queue实现固定大小的连接池,防止连接泄漏
|
||||
"""
|
||||
|
||||
import sqlite3
|
||||
import threading
|
||||
from queue import Queue, Empty
|
||||
import time
|
||||
|
||||
|
||||
class ConnectionPool:
|
||||
"""SQLite连接池"""
|
||||
|
||||
def __init__(self, database, pool_size=5, timeout=30):
|
||||
"""
|
||||
初始化连接池
|
||||
|
||||
Args:
|
||||
database: 数据库文件路径
|
||||
pool_size: 连接池大小(默认5)
|
||||
timeout: 获取连接超时时间(秒)
|
||||
"""
|
||||
self.database = database
|
||||
self.pool_size = pool_size
|
||||
self.timeout = timeout
|
||||
self._pool = Queue(maxsize=pool_size)
|
||||
self._lock = threading.Lock()
|
||||
self._created_connections = 0
|
||||
|
||||
# 预创建连接
|
||||
self._initialize_pool()
|
||||
|
||||
def _initialize_pool(self):
|
||||
"""预创建连接池中的连接"""
|
||||
for _ in range(self.pool_size):
|
||||
conn = self._create_connection()
|
||||
self._pool.put(conn)
|
||||
self._created_connections += 1
|
||||
|
||||
def _create_connection(self):
|
||||
"""创建新的数据库连接"""
|
||||
conn = sqlite3.connect(self.database, check_same_thread=False)
|
||||
conn.row_factory = sqlite3.Row
|
||||
# 设置WAL模式提高并发性能
|
||||
conn.execute('PRAGMA journal_mode=WAL')
|
||||
# 设置合理的超时时间
|
||||
conn.execute('PRAGMA busy_timeout=5000')
|
||||
return conn
|
||||
|
||||
def get_connection(self):
|
||||
"""
|
||||
从连接池获取连接
|
||||
|
||||
Returns:
|
||||
PooledConnection: 连接包装对象
|
||||
"""
|
||||
try:
|
||||
conn = self._pool.get(timeout=self.timeout)
|
||||
return PooledConnection(conn, self)
|
||||
except Empty:
|
||||
raise RuntimeError(f"无法在{self.timeout}秒内获取数据库连接")
|
||||
|
||||
def return_connection(self, conn):
|
||||
"""
|
||||
归还连接到连接池 [已修复Bug#7]
|
||||
|
||||
Args:
|
||||
conn: 要归还的连接
|
||||
"""
|
||||
import sqlite3
|
||||
from queue import Full
|
||||
|
||||
try:
|
||||
# 回滚任何未提交的事务
|
||||
conn.rollback()
|
||||
self._pool.put(conn, block=False)
|
||||
except sqlite3.Error as e:
|
||||
# 数据库相关错误，连接可能损坏
|
||||
print(f"归还连接失败(数据库错误): {e}")
|
||||
try:
|
||||
conn.close()
|
||||
except Exception:
|
||||
pass
|
||||
# 创建新连接补充
|
||||
with self._lock:
|
||||
try:
|
||||
new_conn = self._create_connection()
|
||||
self._pool.put(new_conn, block=False)
|
||||
except Exception as create_error:
|
||||
print(f"重建连接失败: {create_error}")
|
||||
except Full:
|
||||
# 队列已满（不应该发生）
|
||||
print(f"警告: 连接池已满，关闭多余连接")
|
||||
try:
|
||||
conn.close()
|
||||
except Exception:
|
||||
pass
|
||||
except Exception as e:
|
||||
print(f"归还连接失败(未知错误): {e}")
|
||||
try:
|
||||
conn.close()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def close_all(self):
|
||||
"""关闭所有连接"""
|
||||
while not self._pool.empty():
|
||||
try:
|
||||
conn = self._pool.get(block=False)
|
||||
conn.close()
|
||||
except Exception as e:
|
||||
print(f"关闭连接失败: {e}")
|
||||
|
||||
def get_stats(self):
|
||||
"""获取连接池统计信息"""
|
||||
return {
|
||||
'pool_size': self.pool_size,
|
||||
'available': self._pool.qsize(),
|
||||
'in_use': self.pool_size - self._pool.qsize(),
|
||||
'total_created': self._created_connections
|
||||
}
|
||||
|
||||
|
||||
class PooledConnection:
|
||||
"""连接池连接包装器,支持with语句自动归还"""
|
||||
|
||||
def __init__(self, conn, pool):
|
||||
"""
|
||||
初始化
|
||||
|
||||
Args:
|
||||
conn: 实际的数据库连接
|
||||
pool: 连接池对象
|
||||
"""
|
||||
self._conn = conn
|
||||
self._pool = pool
|
||||
self._cursor = None
|
||||
|
||||
def __enter__(self):
|
||||
"""支持with语句"""
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc_val, exc_tb):
|
||||
"""with语句结束时自动归还连接 [已修复Bug#3]"""
|
||||
try:
|
||||
if exc_type is not None:
|
||||
# 发生异常,回滚事务
|
||||
self._conn.rollback()
|
||||
print(f"数据库事务已回滚: {exc_type.__name__}")
|
||||
# 注意: 不自动commit，要求用户显式调用conn.commit()
|
||||
|
||||
if self._cursor:
|
||||
self._cursor.close()
|
||||
except Exception as e:
|
||||
print(f"关闭游标失败: {e}")
|
||||
finally:
|
||||
# 归还连接
|
||||
self._pool.return_connection(self._conn)
|
||||
|
||||
return False # 不抑制异常
|
||||
|
||||
def cursor(self):
|
||||
"""获取游标"""
|
||||
self._cursor = self._conn.cursor()
|
||||
return self._cursor
|
||||
|
||||
def commit(self):
|
||||
"""提交事务"""
|
||||
self._conn.commit()
|
||||
|
||||
def rollback(self):
|
||||
"""回滚事务"""
|
||||
self._conn.rollback()
|
||||
|
||||
def execute(self, sql, parameters=None):
|
||||
"""执行SQL"""
|
||||
cursor = self.cursor()
|
||||
if parameters:
|
||||
return cursor.execute(sql, parameters)
|
||||
return cursor.execute(sql)
|
||||
|
||||
def fetchone(self):
|
||||
"""获取一行"""
|
||||
if self._cursor:
|
||||
return self._cursor.fetchone()
|
||||
return None
|
||||
|
||||
def fetchall(self):
|
||||
"""获取所有行"""
|
||||
if self._cursor:
|
||||
return self._cursor.fetchall()
|
||||
return []
|
||||
|
||||
@property
|
||||
def lastrowid(self):
|
||||
"""最后插入的行ID"""
|
||||
if self._cursor:
|
||||
return self._cursor.lastrowid
|
||||
return None
|
||||
|
||||
@property
|
||||
def rowcount(self):
|
||||
"""影响的行数"""
|
||||
if self._cursor:
|
||||
return self._cursor.rowcount
|
||||
return 0
|
||||
|
||||
|
||||
# 全局连接池实例
|
||||
_pool = None
|
||||
_pool_lock = threading.Lock()
|
||||
|
||||
|
||||
def init_pool(database, pool_size=5):
|
||||
"""
|
||||
初始化全局连接池
|
||||
|
||||
Args:
|
||||
database: 数据库文件路径
|
||||
pool_size: 连接池大小
|
||||
"""
|
||||
global _pool
|
||||
with _pool_lock:
|
||||
if _pool is None:
|
||||
_pool = ConnectionPool(database, pool_size)
|
||||
print(f"✓ 数据库连接池已初始化 (大小: {pool_size})")
|
||||
|
||||
|
||||
def get_db():
|
||||
"""
|
||||
获取数据库连接(替代原有的get_db函数)
|
||||
|
||||
Returns:
|
||||
PooledConnection: 连接对象
|
||||
"""
|
||||
global _pool
|
||||
if _pool is None:
|
||||
raise RuntimeError("连接池未初始化,请先调用init_pool()")
|
||||
return _pool.get_connection()
|
||||
|
||||
|
||||
def get_pool_stats():
|
||||
"""获取连接池统计信息"""
|
||||
global _pool
|
||||
if _pool:
|
||||
return _pool.get_stats()
|
||||
return None
|
||||
Reference in New Issue
Block a user