v2.0.0 done
This commit is contained in:
+262
@@ -0,0 +1,262 @@
|
||||
"""
|
||||
数据库连接与操作模块 - 性能优化版
|
||||
"""
|
||||
import pymysql
|
||||
import sqlalchemy
|
||||
from sqlalchemy import create_engine, text, event
|
||||
from sqlalchemy.pool import QueuePool
|
||||
from urllib.parse import quote
|
||||
from typing import Any, Dict, List, Optional, Tuple
|
||||
from contextlib import contextmanager
|
||||
|
||||
from app.config import AppConfig
|
||||
|
||||
|
||||
class DatabaseManager:
|
||||
"""数据库管理器 - 高性能版"""
|
||||
|
||||
def __init__(self, config: AppConfig):
|
||||
self.config = config
|
||||
self._engine: Optional[sqlalchemy.Engine] = None
|
||||
|
||||
@property
|
||||
def engine(self) -> sqlalchemy.Engine:
|
||||
"""获取 SQLAlchemy 引擎（带连接池）- 优化配置"""
|
||||
if self._engine is None:
|
||||
mysql = self.config.mysql
|
||||
self._engine = create_engine(
|
||||
f'mysql+pymysql://'
|
||||
f'{quote(mysql.user)}:'
|
||||
f'{quote(mysql.passwd)}@'
|
||||
f'{quote(mysql.host)}:'
|
||||
f'{mysql.port}/'
|
||||
f'{quote(mysql.dbname)}?charset=utf8mb4',
|
||||
poolclass=QueuePool,
|
||||
pool_size=10, # 增大连接池
|
||||
max_overflow=20, # 增大溢出连接
|
||||
pool_pre_ping=True,
|
||||
pool_recycle=3600,
|
||||
echo=False,
|
||||
# 性能优化参数
|
||||
connect_args={
|
||||
'local_infile': True, # 允许 LOAD DATA LOCAL
|
||||
'autocommit': False,
|
||||
}
|
||||
)
|
||||
return self._engine
|
||||
|
||||
@contextmanager
|
||||
def get_connection(self):
|
||||
"""获取 PyMySQL 连接（上下文管理器）"""
|
||||
mysql = self.config.mysql
|
||||
conn = pymysql.connect(
|
||||
host=mysql.host,
|
||||
port=mysql.port,
|
||||
user=mysql.user,
|
||||
password=mysql.passwd,
|
||||
database=mysql.dbname,
|
||||
charset='utf8mb4',
|
||||
cursorclass=pymysql.cursors.DictCursor,
|
||||
local_infile=True, # 允许 LOAD DATA LOCAL
|
||||
autocommit=False
|
||||
)
|
||||
try:
|
||||
yield conn
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
@contextmanager
|
||||
def get_fast_connection(self):
|
||||
"""获取高性能 PyMySQL 连接（用于批量插入）"""
|
||||
mysql = self.config.mysql
|
||||
conn = pymysql.connect(
|
||||
host=mysql.host,
|
||||
port=mysql.port,
|
||||
user=mysql.user,
|
||||
password=mysql.passwd,
|
||||
database=mysql.dbname,
|
||||
charset='utf8mb4',
|
||||
cursorclass=pymysql.cursors.Cursor, # 使用普通游标更快
|
||||
local_infile=True,
|
||||
autocommit=False,
|
||||
read_timeout=300,
|
||||
write_timeout=300
|
||||
)
|
||||
try:
|
||||
yield conn
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
def test_connection(self) -> Tuple[bool, str]:
|
||||
"""测试数据库连接"""
|
||||
try:
|
||||
with self.get_connection() as conn:
|
||||
with conn.cursor() as cursor:
|
||||
cursor.execute("SELECT 1")
|
||||
return True, "连接成功"
|
||||
except Exception as e:
|
||||
return False, str(e)
|
||||
|
||||
def get_tables(self) -> List[str]:
|
||||
"""获取所有表名"""
|
||||
with self.get_connection() as conn:
|
||||
with conn.cursor() as cursor:
|
||||
cursor.execute("SHOW TABLES")
|
||||
return [list(row.values())[0] for row in cursor.fetchall()]
|
||||
|
||||
def get_table_info(self, table_name: str) -> Dict[str, Any]:
|
||||
"""获取表信息"""
|
||||
with self.get_connection() as conn:
|
||||
with conn.cursor() as cursor:
|
||||
# 获取列信息
|
||||
cursor.execute(f"DESCRIBE `{table_name}`")
|
||||
columns = cursor.fetchall()
|
||||
|
||||
# 获取行数
|
||||
cursor.execute(f"SELECT COUNT(*) as count FROM `{table_name}`")
|
||||
count = cursor.fetchone()['count']
|
||||
|
||||
return {
|
||||
"name": table_name,
|
||||
"columns": columns,
|
||||
"row_count": count
|
||||
}
|
||||
|
||||
def query_table(
|
||||
self,
|
||||
table_name: str,
|
||||
page: int = 1,
|
||||
page_size: int = 50,
|
||||
filters: Optional[Dict[str, str]] = None,
|
||||
order_by: Optional[str] = None,
|
||||
order_dir: str = "ASC"
|
||||
) -> Dict[str, Any]:
|
||||
"""分页查询表数据"""
|
||||
offset = (page - 1) * page_size
|
||||
|
||||
with self.get_connection() as conn:
|
||||
with conn.cursor() as cursor:
|
||||
# 构建 WHERE 条件
|
||||
where_clause = ""
|
||||
params = []
|
||||
if filters:
|
||||
conditions = []
|
||||
for col, val in filters.items():
|
||||
if val:
|
||||
conditions.append(f"`{col}` LIKE %s")
|
||||
params.append(f"%{val}%")
|
||||
if conditions:
|
||||
where_clause = "WHERE " + " AND ".join(conditions)
|
||||
|
||||
# 获取总数
|
||||
count_sql = f"SELECT COUNT(*) as count FROM `{table_name}` {where_clause}"
|
||||
cursor.execute(count_sql, params)
|
||||
total = cursor.fetchone()['count']
|
||||
|
||||
# 构建排序
|
||||
order_clause = ""
|
||||
if order_by:
|
||||
order_clause = f"ORDER BY `{order_by}` {order_dir}"
|
||||
|
||||
# 查询数据
|
||||
query_sql = f"SELECT * FROM `{table_name}` {where_clause} {order_clause} LIMIT %s OFFSET %s"
|
||||
cursor.execute(query_sql, params + [page_size, offset])
|
||||
rows = cursor.fetchall()
|
||||
|
||||
return {
|
||||
"data": rows,
|
||||
"total": total,
|
||||
"page": page,
|
||||
"page_size": page_size,
|
||||
"total_pages": (total + page_size - 1) // page_size
|
||||
}
|
||||
|
||||
def delete_rows(self, table_name: str, condition: str, params: List[Any]) -> int:
|
||||
"""删除符合条件的行"""
|
||||
with self.get_connection() as conn:
|
||||
with conn.cursor() as cursor:
|
||||
sql = f"DELETE FROM `{table_name}` WHERE {condition}"
|
||||
cursor.execute(sql, params)
|
||||
conn.commit()
|
||||
return cursor.rowcount
|
||||
|
||||
def truncate_table(self, table_name: str) -> bool:
|
||||
"""清空表"""
|
||||
with self.get_connection() as conn:
|
||||
with conn.cursor() as cursor:
|
||||
cursor.execute(f"TRUNCATE TABLE `{table_name}`")
|
||||
conn.commit()
|
||||
return True
|
||||
|
||||
def drop_table(self, table_name: str) -> bool:
|
||||
"""删除表"""
|
||||
with self.get_connection() as conn:
|
||||
with conn.cursor() as cursor:
|
||||
cursor.execute(f"DROP TABLE IF EXISTS `{table_name}`")
|
||||
conn.commit()
|
||||
return True
|
||||
|
||||
def execute_sql(self, sql: str) -> Tuple[bool, Any]:
|
||||
"""执行自定义 SQL"""
|
||||
with self.get_connection() as conn:
|
||||
with conn.cursor() as cursor:
|
||||
try:
|
||||
cursor.execute(sql)
|
||||
if sql.strip().upper().startswith("SELECT"):
|
||||
return True, cursor.fetchall()
|
||||
else:
|
||||
conn.commit()
|
||||
return True, {"affected_rows": cursor.rowcount}
|
||||
except Exception as e:
|
||||
return False, str(e)
|
||||
|
||||
def bulk_insert(self, table_name: str, columns: List[str], data: List[Tuple],
|
||||
batch_size: int = 5000) -> int:
|
||||
"""
|
||||
高性能批量插入
|
||||
使用 executemany + 批量提交，比 to_sql 快 5-10 倍
|
||||
"""
|
||||
if not data:
|
||||
return 0
|
||||
|
||||
total_inserted = 0
|
||||
placeholders = ', '.join(['%s'] * len(columns))
|
||||
column_names = ', '.join([f'`{col}`' for col in columns])
|
||||
sql = f"INSERT INTO `{table_name}` ({column_names}) VALUES ({placeholders})"
|
||||
|
||||
with self.get_fast_connection() as conn:
|
||||
with conn.cursor() as cursor:
|
||||
# 优化插入性能的设置
|
||||
cursor.execute("SET autocommit=0")
|
||||
cursor.execute("SET unique_checks=0")
|
||||
cursor.execute("SET foreign_key_checks=0")
|
||||
|
||||
# 分批插入
|
||||
for i in range(0, len(data), batch_size):
|
||||
batch = data[i:i + batch_size]
|
||||
cursor.executemany(sql, batch)
|
||||
total_inserted += len(batch)
|
||||
|
||||
# 提交并恢复设置
|
||||
conn.commit()
|
||||
cursor.execute("SET unique_checks=1")
|
||||
cursor.execute("SET foreign_key_checks=1")
|
||||
cursor.execute("SET autocommit=1")
|
||||
|
||||
return total_inserted
|
||||
|
||||
def create_table_from_columns(self, table_name: str, columns: List[str]):
|
||||
"""根据列名创建表（所有列都是 VARCHAR(255)）"""
|
||||
column_defs = ', '.join([f'`{col}` VARCHAR(255)' for col in columns])
|
||||
sql = f"CREATE TABLE IF NOT EXISTS `{table_name}` ({column_defs}) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4"
|
||||
|
||||
with self.get_connection() as conn:
|
||||
with conn.cursor() as cursor:
|
||||
cursor.execute(sql)
|
||||
conn.commit()
|
||||
|
||||
def dispose(self):
|
||||
"""释放连接池"""
|
||||
if self._engine:
|
||||
self._engine.dispose()
|
||||
self._engine = None
|
||||
Reference in New Issue
Block a user