589 lines
26 KiB
Python
589 lines
26 KiB
Python
from sqlalchemy import create_engine, MetaData, Table, select, update, delete, insert
|
||
from sqlalchemy.engine import URL
|
||
from sqlalchemy.exc import SQLAlchemyError
|
||
from typing import Optional, Dict, List, Any, Union
|
||
import pandas as pd
|
||
from log.LogUtil import LogUtil
|
||
import logging as d
|
||
|
||
#import os
|
||
#os.environ['IBM_DB_HOME'] = r"D:\win\IBM\SQLLIB\dsdriver"
|
||
#os.add_dll_directory(r"D:\win\IBM\SQLLIB\BIN")
|
||
#import ibm_db
|
||
|
||
"""
|
||
DBOperator - 通用数据库操作类
|
||
|
||
功能:
|
||
- 提供基础的数据库连接管理
|
||
- 支持SQL查询、更新、删除和插入操作
|
||
- 支持连接池配置
|
||
- 支持多种数据库类型
|
||
|
||
支持的数据库类型:
|
||
- MySQL
|
||
- PostgreSQL
|
||
- SQLite
|
||
- Oracle
|
||
- IBM DB2 (需额外配置)
|
||
|
||
依赖:
|
||
- sqlalchemy >= 1.4.0
|
||
- pandas
|
||
- ibm_db (仅用于IBM DB2)
|
||
- pymysql (用于MySQL)
|
||
- psycopg2 (用于PostgreSQL)
|
||
|
||
Author：
|
||
- Author : zoufuzhou
|
||
- Date : 2025-05-21 16:34:37
|
||
- Description : db operator tool
|
||
- LastEditTime : 2025-05-21 16:34:37
|
||
"""
|
||
|
||
class DBOperator:
|
||
"""通用数据库操作类"""
|
||
|
||
def __init__(self,
|
||
db_type: str,
|
||
host: str,
|
||
port: int,
|
||
database: str,
|
||
username: str,
|
||
password: str,
|
||
pool_size: int = 5):
|
||
"""
|
||
初始化数据库连接
|
||
:param db_type: 数据库类型(mysql/postgresql/sqlite/oracle等)
|
||
:param host: 主机地址
|
||
:param port: 端口号
|
||
:param database: 数据库名
|
||
:param username: 用户名
|
||
:param password: 密码
|
||
:param pool_size: 连接池大小
|
||
"""
|
||
self.db_type = db_type.lower()
|
||
self.host = host
|
||
self.port = port
|
||
self.database = database
|
||
self.username = username
|
||
self.password = password
|
||
|
||
# 创建连接引擎
|
||
self.engine = self._create_engine(pool_size)
|
||
self.metadata = MetaData()
|
||
|
||
def _create_engine(self, pool_size: int):
|
||
"""创建SQLAlchemy引擎"""
|
||
try:
|
||
# 生成连接URL
|
||
url = self._generate_connection_url()
|
||
d.info(f"创建{self.db_type}数据库引擎")
|
||
|
||
return create_engine(
|
||
url,
|
||
pool_size=pool_size,
|
||
max_overflow=10,
|
||
pool_pre_ping=True,
|
||
pool_recycle=3600
|
||
)
|
||
except Exception as e:
|
||
d.error(f"创建数据库引擎失败: {e}")
|
||
raise
|
||
|
||
def _generate_connection_url(self) -> URL:
|
||
"""生成数据库连接URL"""
|
||
if self.db_type == "mysql":
|
||
return URL.create(
|
||
"mysql+pymysql",
|
||
username=self.username,
|
||
password=self.password,
|
||
host=self.host,
|
||
port=self.port,
|
||
database=self.database
|
||
)
|
||
elif self.db_type == "postgresql":
|
||
return URL.create(
|
||
"postgresql+psycopg2",
|
||
username=self.username,
|
||
password=self.password,
|
||
host=self.host,
|
||
port=self.port,
|
||
database=self.database
|
||
)
|
||
elif self.db_type == "sqlite":
|
||
return f"sqlite:///{self.database}"
|
||
elif self.db_type == "oracle":
|
||
return URL.create(
|
||
"oracle+cx_oracle",
|
||
username=self.username,
|
||
password=self.password,
|
||
host=self.host,
|
||
port=self.port,
|
||
database=self.database
|
||
)
|
||
elif self.db_type == "db2":
|
||
return URL.create(
|
||
"db2+ibm_db",
|
||
username=self.username,
|
||
password=self.password,
|
||
host=self.host,
|
||
port=self.port,
|
||
database=self.database
|
||
)
|
||
else:
|
||
raise ValueError(f"不支持的数据库类型: {self.db_type}")
|
||
|
||
def execute_query(self, table_name: str,
|
||
filters: Optional[Dict[str, Any]] = None,
|
||
columns: Optional[List[str]] = None,
|
||
as_dataframe: bool = False) -> Union[List[Dict[str, Any]], pd.DataFrame]:
|
||
"""
|
||
执行查询
|
||
:param table_name: 表名
|
||
:param filters: 过滤条件字典 {列名: 值}
|
||
:param columns: 要查询的列名列表
|
||
:param as_dataframe: 是否返回DataFrame
|
||
:return: 结果字典列表或DataFrame
|
||
"""
|
||
try:
|
||
table = Table(table_name, self.metadata, autoload_with=self.engine)
|
||
|
||
# 构建查询
|
||
query = select(table)
|
||
|
||
# 添加过滤条件(不区分大小写)
|
||
if filters:
|
||
# 查找匹配的列名(不区分大小写)
|
||
conditions = []
|
||
for k, v in filters.items():
|
||
matching_cols = [col for col in table.columns if col.name.lower() == k.lower()]
|
||
if matching_cols:
|
||
col_name = matching_cols[0].name
|
||
col_obj = getattr(table.c, col_name)
|
||
|
||
# 处理特殊操作符
|
||
if isinstance(v, dict) and len(v) == 1:
|
||
op, val = next(iter(v.items()))
|
||
if op == '$like':
|
||
conditions.append(col_obj.like(val))
|
||
elif op == '$gt':
|
||
conditions.append(col_obj > val)
|
||
elif op == '$gte':
|
||
conditions.append(col_obj >= val)
|
||
elif op == '$lt':
|
||
conditions.append(col_obj < val)
|
||
elif op == '$lte':
|
||
conditions.append(col_obj <= val)
|
||
elif op == '$between':
|
||
if isinstance(val, (list, tuple)) and len(val) == 2:
|
||
conditions.append(col_obj.between(val[0], val[1]))
|
||
else:
|
||
# 默认等于操作
|
||
conditions.append(col_obj == v)
|
||
if conditions:
|
||
query = query.where(*conditions)
|
||
|
||
# 指定列(不区分大小写)
|
||
if columns:
|
||
# 查找匹配的列名(不区分大小写)
|
||
selected_cols = []
|
||
for col in columns:
|
||
matching_cols = [table_col for table_col in table.columns
|
||
if table_col.name.lower() == col.lower()]
|
||
if matching_cols:
|
||
selected_cols.append(getattr(table.c, matching_cols[0].name))
|
||
if selected_cols:
|
||
query = query.with_only_columns(*selected_cols)
|
||
|
||
# 执行查询
|
||
with self.engine.connect() as conn:
|
||
result = conn.execute(query)
|
||
# 处理DB2结果集，确保列名和值正确映射
|
||
columns = result.keys()
|
||
result_data = [dict(zip(columns, row)) for row in result]
|
||
|
||
if as_dataframe:
|
||
return pd.DataFrame(result_data)
|
||
return result_data
|
||
|
||
except SQLAlchemyError as e:
|
||
d.error(f"查询执行失败: {e}")
|
||
return pd.DataFrame() if as_dataframe else []
|
||
|
||
def execute_update(self, table_name: str,
|
||
data: Dict[str, Any],
|
||
filters: Optional[Dict[str, Any]] = None) -> int:
|
||
"""
|
||
执行更新
|
||
:param table_name: 表名
|
||
:param data: 更新数据字典 {列名: 新值}
|
||
:param filters: 过滤条件字典 {列名: 值}
|
||
:return: 影响的行数
|
||
"""
|
||
try:
|
||
table = Table(table_name, self.metadata, autoload_with=self.engine)
|
||
|
||
# 构建更新语句
|
||
stmt = update(table).values(**data)
|
||
|
||
# 添加过滤条件(不区分大小写)
|
||
if filters:
|
||
# 查找匹配的列名(不区分大小写)
|
||
conditions = []
|
||
for k, v in filters.items():
|
||
matching_cols = [col for col in table.columns if col.name.lower() == k.lower()]
|
||
if matching_cols:
|
||
conditions.append(getattr(table.c, matching_cols[0].name) == v)
|
||
if conditions:
|
||
stmt = stmt.where(*conditions)
|
||
|
||
# 执行更新
|
||
with self.engine.begin() as conn:
|
||
result = conn.execute(stmt)
|
||
return result.rowcount
|
||
|
||
except SQLAlchemyError as e:
|
||
d.error(f"更新执行失败: {e}")
|
||
return 0
|
||
|
||
def execute_insert(self, table_name: str, data: Dict[str, Any]) -> int:
|
||
"""
|
||
执行插入
|
||
:param table_name: 表名
|
||
:param data: 插入数据字典 {列名: 值}
|
||
:return: 插入的行数(通常为1)
|
||
"""
|
||
try:
|
||
table = Table(table_name, self.metadata, autoload_with=self.engine)
|
||
|
||
# 执行插入(不区分大小写)
|
||
with self.engine.begin() as conn:
|
||
# 转换数据键名为表列名(不区分大小写)
|
||
insert_data = {}
|
||
for k, v in data.items():
|
||
matching_cols = [col for col in table.columns if col.name.lower() == k.lower()]
|
||
if matching_cols:
|
||
insert_data[matching_cols[0].name] = v
|
||
if insert_data:
|
||
result = conn.execute(insert(table).values(**insert_data))
|
||
return result.rowcount
|
||
|
||
except SQLAlchemyError as e:
|
||
d.error(f"插入执行失败: {e}")
|
||
return 0
|
||
|
||
def execute_delete(self, table_name: str,
|
||
filters: Dict[str, Any]) -> int:
|
||
"""
|
||
执行删除
|
||
:param table_name: 表名
|
||
:param filters: 过滤条件字典 {列名: 值}
|
||
:return: 影响的行数
|
||
"""
|
||
try:
|
||
table = Table(table_name, self.metadata, autoload_with=self.engine)
|
||
|
||
# 执行删除(不区分大小写)
|
||
with self.engine.begin() as conn:
|
||
# 查找匹配的列名(不区分大小写)
|
||
conditions = []
|
||
for k, v in filters.items():
|
||
matching_cols = [col for col in table.columns if col.name.lower() == k.lower()]
|
||
if matching_cols:
|
||
conditions.append(getattr(table.c, matching_cols[0].name) == v)
|
||
if conditions:
|
||
result = conn.execute(delete(table).where(*conditions))
|
||
return result.rowcount
|
||
|
||
except SQLAlchemyError as e:
|
||
d.error(f"删除执行失败: {e}")
|
||
return 0
|
||
|
||
def write_dataframe(self,
|
||
df: pd.DataFrame,
|
||
table_name: str,
|
||
upsert_keys: Optional[List[str]] = None,
|
||
batch_size: int = 1000) -> int:
|
||
"""
|
||
将DataFrame写入数据库表(支持存在更新，不存在插入)
|
||
:param df: 要写入的DataFrame
|
||
:param table_name: 目标表名
|
||
:param upsert_keys: 用于判断记录是否存在的键列名列表
|
||
:param batch_size: 批量处理大小
|
||
:return: 处理的行数
|
||
"""
|
||
if df.empty:
|
||
d.warning("空DataFrame，跳过写入")
|
||
return 0
|
||
|
||
try:
|
||
# 反射表结构
|
||
table = Table(table_name, self.metadata, autoload_with=self.engine)
|
||
|
||
# 检查DataFrame列是否匹配表结构(不区分大小写)
|
||
df_columns_lower = [col.lower() for col in df.columns]
|
||
missing_cols = [col.name for col in table.columns
|
||
if not col.nullable and col.name.lower() not in df_columns_lower]
|
||
if missing_cols:
|
||
raise ValueError(f"DataFrame缺少必填列: {missing_cols}")
|
||
|
||
# 记录表结构详情
|
||
table_info = {col.name: str(col.type) for col in table.columns}
|
||
# d.info(f"表{table_name}结构: {table_info}")
|
||
|
||
# 转换数据类型以匹配表定义(不区分大小写)
|
||
for col in table.columns:
|
||
# 查找匹配的DataFrame列名(不区分大小写)
|
||
matching_cols = [df_col for df_col in df.columns
|
||
if df_col.lower() == col.name.lower()]
|
||
if matching_cols:
|
||
df_col = matching_cols[0] # 取第一个匹配的列名
|
||
if 'INT' in str(col.type):
|
||
df[df_col] = df[df_col].astype('int64')
|
||
elif 'FLOAT' in str(col.type) or 'DECIMAL' in str(col.type):
|
||
df[df_col] = df[df_col].astype('float64')
|
||
elif 'DATE' in str(col.type) or 'TIME' in str(col.type):
|
||
df[df_col] = pd.to_datetime(df[df_col])
|
||
|
||
total_rows = len(df)
|
||
processed_rows = 0
|
||
|
||
# 分批处理
|
||
for i in range(0, total_rows, batch_size):
|
||
batch = df.iloc[i:i + batch_size]
|
||
data_list = batch.to_dict('records')
|
||
|
||
with self.engine.begin() as conn:
|
||
if upsert_keys is None:
|
||
# 纯插入模式，失败则转为更新
|
||
for data in data_list:
|
||
try:
|
||
# 转换数据键名为表列名(不区分大小写)
|
||
insert_data = {}
|
||
for k, v in data.items():
|
||
matching_cols = [col for col in table.columns if col.name.lower() == k.lower()]
|
||
if matching_cols:
|
||
insert_data[matching_cols[0].name] = v
|
||
if insert_data:
|
||
conn.execute(insert(table).values(**insert_data))
|
||
except SQLAlchemyError as e:
|
||
# 如果是唯一约束冲突，则转为更新
|
||
if "unique constraint" in str(e).lower() or "duplicate key" in str(e).lower():
|
||
d.warning(f"插入冲突，转为更新: {e}")
|
||
|
||
# 尝试提取冲突键
|
||
conflict_keys = []
|
||
if self.db_type == 'postgresql':
|
||
# PostgreSQL错误信息通常包含约束名
|
||
conflict_keys = [k for k in data.keys()
|
||
if k in table.primary_key.columns]
|
||
elif self.db_type == 'mysql':
|
||
# MySQL错误信息通常包含键名
|
||
conflict_keys = [k for k in data.keys()
|
||
if k in table.primary_key.columns]
|
||
elif self.db_type == 'db2':
|
||
# DB2错误信息通常包含键名
|
||
conflict_keys = [k for k in data.keys()
|
||
if k in table.primary_key.columns]
|
||
|
||
if not conflict_keys:
|
||
conflict_keys = list(data.keys())
|
||
|
||
# 执行更新
|
||
update_stmt = update(table).values(**data)
|
||
for key in conflict_keys:
|
||
update_stmt = update_stmt.where(
|
||
getattr(table.c, key) == data[key]
|
||
)
|
||
conn.execute(update_stmt)
|
||
else:
|
||
# 其他错误仍然抛出
|
||
# d.error(f"插入失败 - 表结构: {table_info}")
|
||
d.error(f"尝试插入的数据: {data}")
|
||
if "NOT NULL" in str(e):
|
||
missing = [col.name for col in table.columns
|
||
if not col.nullable and col.name not in data]
|
||
raise ValueError(f"缺少必填字段: {missing}") from e
|
||
raise
|
||
else:
|
||
# Upsert模式
|
||
if self.db_type == 'postgresql':
|
||
# PostgreSQL使用ON CONFLICT语法
|
||
for data in data_list:
|
||
# 转换数据键名为表列名(不区分大小写)
|
||
insert_data = {}
|
||
for k, v in data.items():
|
||
matching_cols = [col for col in table.columns if col.name.lower() == k.lower()]
|
||
if matching_cols:
|
||
insert_data[matching_cols[0].name] = v
|
||
|
||
# 转换upsert_keys为表列名(不区分大小写)
|
||
actual_upsert_keys = []
|
||
for key in upsert_keys:
|
||
matching_cols = [col for col in table.columns if col.name.lower() == key.lower()]
|
||
if matching_cols:
|
||
actual_upsert_keys.append(matching_cols[0].name)
|
||
|
||
if insert_data and actual_upsert_keys:
|
||
update_dict = {k: v for k, v in insert_data.items()
|
||
if k not in actual_upsert_keys}
|
||
stmt = insert(table).values(**insert_data)
|
||
stmt = stmt.on_conflict_do_update(
|
||
index_elements=actual_upsert_keys,
|
||
set_=update_dict
|
||
)
|
||
conn.execute(stmt)
|
||
elif self.db_type == 'mysql':
|
||
# MySQL使用ON DUPLICATE KEY UPDATE语法
|
||
for data in data_list:
|
||
# 转换数据键名为表列名(不区分大小写)
|
||
insert_data = {}
|
||
for k, v in data.items():
|
||
matching_cols = [col for col in table.columns if col.name.lower() == k.lower()]
|
||
if matching_cols:
|
||
insert_data[matching_cols[0].name] = v
|
||
|
||
# 转换upsert_keys为表列名(不区分大小写)
|
||
actual_upsert_keys = []
|
||
for key in upsert_keys:
|
||
matching_cols = [col for col in table.columns if col.name.lower() == key.lower()]
|
||
if matching_cols:
|
||
actual_upsert_keys.append(matching_cols[0].name)
|
||
|
||
if insert_data and actual_upsert_keys:
|
||
update_dict = {k: v for k, v in insert_data.items()
|
||
if k not in actual_upsert_keys}
|
||
stmt = insert(table).values(**insert_data)
|
||
stmt = stmt.on_duplicate_key_update(**update_dict)
|
||
conn.execute(stmt)
|
||
elif self.db_type == 'db2':
|
||
# DB2使用MERGE语法
|
||
for data in data_list:
|
||
# 转换数据键名为表列名(不区分大小写)
|
||
insert_data = {}
|
||
for k, v in data.items():
|
||
matching_cols = [col for col in table.columns if col.name.lower() == k.lower()]
|
||
if matching_cols:
|
||
insert_data[matching_cols[0].name] = v
|
||
|
||
# 转换upsert_keys为表列名(不区分大小写)
|
||
actual_upsert_keys = []
|
||
for key in upsert_keys:
|
||
matching_cols = [col for col in table.columns if col.name.lower() == key.lower()]
|
||
if matching_cols:
|
||
actual_upsert_keys.append(matching_cols[0].name)
|
||
|
||
if insert_data and actual_upsert_keys:
|
||
# 构建MERGE语句
|
||
update_cols = [col for col in insert_data.keys()
|
||
if col not in actual_upsert_keys]
|
||
merge_set = {col: insert_data[col] for col in update_cols}
|
||
merge_condition = {key: insert_data[key] for key in actual_upsert_keys}
|
||
|
||
# 先尝试更新
|
||
update_stmt = update(table).values(**merge_set)
|
||
for key in actual_upsert_keys:
|
||
update_stmt = update_stmt.where(
|
||
getattr(table.c, key) == merge_condition[key]
|
||
)
|
||
result = conn.execute(update_stmt)
|
||
|
||
# 如果没有更新到行，则插入
|
||
if result.rowcount == 0:
|
||
conn.execute(insert(table).values(**insert_data))
|
||
else:
|
||
raise ValueError(f"不支持的数据库类型: {self.db_type}")
|
||
|
||
processed_rows += len(data_list)
|
||
d.info(f"已处理{processed_rows}/{total_rows}行到{table_name}")
|
||
|
||
return processed_rows
|
||
|
||
except SQLAlchemyError as e:
|
||
d.error(f"DataFrame写入失败: {e}")
|
||
raise
|
||
|
||
def close(self):
|
||
"""关闭所有连接"""
|
||
self.engine.dispose()
|
||
d.info("数据库连接已关闭")
|
||
|
||
# 使用示例
|
||
if __name__ == "__main__":
|
||
# DB2示例
|
||
db2_db = DBOperator(
|
||
db_type="db2",
|
||
host="192.168.137.100",
|
||
port=50000,
|
||
database="appdb",
|
||
username="som",
|
||
password="dscdsc1"
|
||
)
|
||
LogUtil.init("app")
|
||
|
||
# # 查询示例
|
||
|
||
# # db2_results = db2_db.execute_query("T_TCM_DEC", {"entId": 'test7'}, as_dataframe=True)
|
||
db2_results = db2_db.execute_query("T_TCM_DEC", {"entId": {"$like":'test%'}}, as_dataframe=True)
|
||
# # db2_results = db2_db.execute_query("T_TCM_DEC",as_dataframe=True)
|
||
print("DB2经理员工:", db2_results)
|
||
#
|
||
# # # 插入示例
|
||
# # db2_insert_count = db2_db.execute_insert("T_TCM_DEC", {
|
||
# # "entId": "test3",
|
||
# # "f1e_ten_deviate": 1
|
||
# # })
|
||
# # print(f"DB2插入了{db2_insert_count}条记录")
|
||
#
|
||
# # DataFrame写入示例
|
||
# try:
|
||
# # 创建测试DataFrame - 包含所有必填字段并确保类型匹配
|
||
# test_data = {
|
||
# 'entid': ['test7', 'test8', 'test9'],
|
||
# 'F1e_ten_deviate': [1, 0, 1]
|
||
# }
|
||
#
|
||
#
|
||
# test_df = pd.DataFrame(test_data)
|
||
#
|
||
# # # 打印测试数据详情
|
||
# # print("\n测试数据详情:")
|
||
# # print(test_df.dtypes)
|
||
#
|
||
# # print("\n测试DataFrame写入DB2:")
|
||
# # print(test_df)
|
||
#
|
||
# # 1. 纯插入模式
|
||
# # rows_inserted = db2_db.write_dataframe(
|
||
# # df=test_df,
|
||
# # table_name="T_TCM_DEC",
|
||
# # batch_size=2
|
||
# # )
|
||
# # print(f"成功插入{rows_inserted}行到T_TCM_DEC表")
|
||
#
|
||
# # # 2. Upsert模式
|
||
# test_data_update = {
|
||
# 'entid': ['test7', 'test8', 'test10'], # test10是新记录
|
||
# 'f1e_ten_deviate': [2, 8, 10]
|
||
# }
|
||
# # 确保包含所有必填字段
|
||
# test_df_update = pd.DataFrame(test_data_update)
|
||
#
|
||
# print("\n测试Upsert操作:")
|
||
# print(test_df_update)
|
||
#
|
||
# rows_processed = db2_db.write_dataframe(
|
||
# df=test_df_update,
|
||
# table_name="T_TCM_DEC",
|
||
# upsert_keys=['entid'], # 使用entid作为唯一键
|
||
# batch_size=2
|
||
# )
|
||
# print(f"成功处理{rows_processed}行(更新2行+插入1行)")
|
||
#
|
||
# except Exception as e:
|
||
# print(f"DataFrame写入失败: {e}")
|
||
#
|
||
# # 关闭连接
|
||
# db2_db.close()
|