426 lines
15 KiB
Python
426 lines
15 KiB
Python
#import os
|
||
#os.environ['IBM_DB_HOME'] = r"D:\win\IBM\SQLLIB\dsdriver"
|
||
#os.add_dll_directory(r"D:\win\IBM\SQLLIB\BIN")
|
||
from typing import Optional, List, Dict, Any, TypeAlias, Union
|
||
import pandas as pd
|
||
import ibm_db
|
||
from log.LogUtil import LogUtil
|
||
import logging as d
|
||
|
||
# 定义DB2连接类型别名
|
||
DB2Connection: TypeAlias = Any
|
||
"""
|
||
DB2数据库操作工具类
|
||
|
||
功能:
|
||
- 连接池管理
|
||
- 自动重连机制
|
||
- 安全参数处理
|
||
- 基础查询执行
|
||
|
||
支持的数据库:
|
||
- IBM DB2 (Windows/Linux)
|
||
|
||
依赖库:
|
||
- ibm_db (IBM DB2 Python驱动)
|
||
- pandas (数据处理)
|
||
- logging (日志记录)
|
||
|
||
作者: zoufuzhou
|
||
版本: 1.0
|
||
"""
|
||
|
||
|
||
class DB2Tool:
|
||
"""DB2数据库操作工具类"""
|
||
|
||
def __init__(self,
|
||
host: str,
|
||
port: int,
|
||
database: str,
|
||
username: str,
|
||
password: str,
|
||
conn_pool_size: int = 5):
|
||
"""
|
||
初始化DB2连接参数
|
||
:param host: 主机地址
|
||
:param port: 端口号
|
||
:param database: 数据库名
|
||
:param username: 用户名
|
||
:param password: 密码
|
||
:param conn_pool_size: 连接池大小
|
||
"""
|
||
|
||
# 再初始化连接参数
|
||
self.host = host
|
||
self.port = port
|
||
self.database = database
|
||
self.username = username
|
||
self.password = password
|
||
self.conn_pool_size = conn_pool_size
|
||
self.connection_pool = []
|
||
|
||
try:
|
||
self._init_connection_pool()
|
||
except Exception as e:
|
||
d.error(f"初始化连接池失败: {e}")
|
||
raise
|
||
|
||
def _init_connection_pool(self):
|
||
"""初始化连接池"""
|
||
for _ in range(self.conn_pool_size):
|
||
conn = self._create_connection()
|
||
if conn:
|
||
self.connection_pool.append(conn)
|
||
|
||
def _create_connection(self, retries: int = 3) -> Optional[DB2Connection]:
|
||
"""创建单个连接
|
||
:param retries: 连接重试次数，默认3次
|
||
"""
|
||
attempt = 0
|
||
last_error = None
|
||
|
||
while attempt < retries:
|
||
attempt += 1
|
||
try:
|
||
# 安全处理连接参数
|
||
def safe_str(s: str) -> str:
|
||
if not isinstance(s, str):
|
||
s = str(s)
|
||
try:
|
||
# 先尝试UTF-8编码
|
||
s.encode('utf-8')
|
||
# 替换可能引起问题的特殊字符
|
||
return ''.join(c if c.isalnum() or c in ('_', '-', '.') else '_' for c in s)
|
||
except UnicodeError:
|
||
# 对于无法UTF-8编码的字符串，使用更安全的处理方式
|
||
return ''.join(c if ord(c) < 128 and c.isalnum() else '_' for c in s)
|
||
|
||
# 构建安全连接字符串(添加编码参数)
|
||
conn_params = [
|
||
f"DATABASE={safe_str(self.database)}",
|
||
f"HOSTNAME={safe_str(self.host)}",
|
||
f"PORT={self.port}",
|
||
"PROTOCOL=TCPIP",
|
||
f"UID={safe_str(self.username)}",
|
||
f"PWD={safe_str(self.password)}",
|
||
"CONNECTTIMEOUT=10",
|
||
"CLIENT_LOCALE=en_US.UTF-8", # 强制客户端使用UTF-8
|
||
"DB_LOCALE=en_US.UTF-8" # 强制数据库使用UTF-8
|
||
]
|
||
conn_str = ";".join(conn_params) + ";"
|
||
d.debug(f"完整连接字符串: {conn_str}")
|
||
|
||
d.debug(f"安全处理后的连接字符串: {conn_str}")
|
||
|
||
d.debug(f"尝试连接(第{attempt}次): {conn_str}")
|
||
conn = ibm_db.connect(conn_str, "", "")
|
||
|
||
# 验证连接是否有效
|
||
if ibm_db.active(conn):
|
||
d.info(f"DB2连接成功(第{attempt}次尝试)")
|
||
return conn
|
||
ibm_db.close(conn)
|
||
|
||
except Exception as e: # ibm_db通用异常捕获
|
||
if hasattr(ibm_db, 'stmt_errormsg'):
|
||
last_error = ibm_db.stmt_errormsg()
|
||
else:
|
||
last_error = str(e)
|
||
d.warning(f"DB2连接失败(第{attempt}次): {last_error}")
|
||
if attempt < retries:
|
||
import time
|
||
time.sleep(1) # 等待1秒后重试
|
||
|
||
except Exception as e:
|
||
last_error = str(e)
|
||
d.error(f"连接过程中发生意外错误(第{attempt}次): {last_error}")
|
||
break
|
||
|
||
d.error(f"所有连接尝试失败，最后错误: {last_error}")
|
||
return None
|
||
|
||
def get_connection(self) -> Optional[DB2Connection]:
|
||
"""从连接池获取连接"""
|
||
if not self.connection_pool:
|
||
d.warning("连接池为空，创建新连接")
|
||
return self._create_connection()
|
||
return self.connection_pool.pop()
|
||
|
||
def release_connection(self, conn: DB2Connection):
|
||
"""释放连接回连接池"""
|
||
if conn and len(self.connection_pool) < self.conn_pool_size:
|
||
self.connection_pool.append(conn)
|
||
|
||
def execute_query(self,
|
||
sql: str,
|
||
params: Optional[tuple] = None,
|
||
as_dataframe: bool = False) -> Union[List[Dict[str, Any]], pd.DataFrame]:
|
||
"""
|
||
执行查询语句
|
||
:param sql: SQL语句
|
||
:param params: 参数
|
||
:param as_dataframe: 是否返回DataFrame
|
||
:return: 结果列表或DataFrame
|
||
"""
|
||
conn = self.get_connection()
|
||
if not conn:
|
||
return []
|
||
|
||
try:
|
||
stmt = ibm_db.prepare(conn, sql)
|
||
if params:
|
||
ibm_db.execute(stmt, params)
|
||
else:
|
||
ibm_db.execute(stmt)
|
||
|
||
result = []
|
||
row = ibm_db.fetch_assoc(stmt)
|
||
while row:
|
||
result.append(row)
|
||
row = ibm_db.fetch_assoc(stmt)
|
||
|
||
if as_dataframe:
|
||
return pd.DataFrame(result) if result else pd.DataFrame()
|
||
return result
|
||
except Exception as e:
|
||
d.error(f"查询执行失败: {e}")
|
||
return pd.DataFrame() if as_dataframe else []
|
||
finally:
|
||
ibm_db.free_stmt(stmt)
|
||
self.release_connection(conn)
|
||
|
||
def execute_update(self, sql: str, params: Optional[tuple] = None) -> int:
|
||
"""
|
||
执行更新语句
|
||
:param sql: SQL语句
|
||
:param params: 参数
|
||
:return: 影响的行数
|
||
"""
|
||
conn = self.get_connection()
|
||
if not conn:
|
||
return 0
|
||
|
||
try:
|
||
stmt = ibm_db.prepare(conn, sql)
|
||
if params:
|
||
ibm_db.execute(stmt, params)
|
||
else:
|
||
ibm_db.execute(stmt)
|
||
return ibm_db.num_rows(stmt)
|
||
except Exception as e:
|
||
d.error(f"更新执行失败: {e}")
|
||
return 0
|
||
finally:
|
||
ibm_db.free_stmt(stmt)
|
||
self.release_connection(conn)
|
||
|
||
def execute_transaction(self, sql_list: List[str], params_list: Optional[List[tuple]] = None) -> bool:
|
||
"""
|
||
执行事务
|
||
:param sql_list: SQL语句列表
|
||
:param params_list: 参数列表
|
||
:return: 是否成功
|
||
"""
|
||
conn = self.get_connection()
|
||
if not conn:
|
||
return False
|
||
|
||
try:
|
||
ibm_db.autocommit(conn, ibm_db.SQL_AUTOCOMMIT_OFF)
|
||
|
||
def safe_encode(value):
|
||
"""安全编码参数值"""
|
||
if isinstance(value, str):
|
||
try:
|
||
return value.encode('utf-8').decode('utf-8')
|
||
except UnicodeError:
|
||
return value.encode('ascii', 'replace').decode('ascii')
|
||
return value
|
||
|
||
for i, sql in enumerate(sql_list):
|
||
try:
|
||
d.debug(f"准备执行SQL: {sql}")
|
||
stmt = ibm_db.prepare(conn, sql)
|
||
|
||
if params_list and i < len(params_list):
|
||
params = params_list[i]
|
||
d.debug(f"参数列表: {params}")
|
||
|
||
# 使用原生参数绑定方式
|
||
for param_idx, param_val in enumerate(params, start=1):
|
||
if isinstance(param_val, str):
|
||
# 字符串参数特殊处理
|
||
try:
|
||
ibm_db.bind_param(stmt, param_idx, param_val.encode('utf-8'), ibm_db.SQL_PARAM_INPUT)
|
||
except Exception as e:
|
||
d.error(f"参数绑定失败(位置{param_idx}): {param_val}")
|
||
raise
|
||
else:
|
||
# 非字符串参数直接绑定
|
||
ibm_db.bind_param(stmt, param_idx, param_val, ibm_db.SQL_PARAM_INPUT)
|
||
|
||
d.debug("执行带绑定参数的SQL")
|
||
ibm_db.execute(stmt)
|
||
else:
|
||
d.debug("执行无参数SQL")
|
||
ibm_db.execute(stmt)
|
||
|
||
ibm_db.free_stmt(stmt)
|
||
except Exception as e:
|
||
d.error(f"执行SQL语句失败: {sql}\n错误: {e}")
|
||
raise
|
||
|
||
ibm_db.commit(conn)
|
||
return True
|
||
except Exception as e:
|
||
d.error(f"事务执行失败: {e}", exc_info=True)
|
||
try:
|
||
ibm_db.rollback(conn)
|
||
except Exception as rollback_err:
|
||
d.error(f"回滚失败: {rollback_err}")
|
||
return False
|
||
finally:
|
||
try:
|
||
ibm_db.autocommit(conn, ibm_db.SQL_AUTOCOMMIT_ON)
|
||
self.release_connection(conn)
|
||
except Exception as e:
|
||
d.error(f"释放连接失败: {e}")
|
||
|
||
def close_all_connections(self):
|
||
"""关闭所有连接"""
|
||
for conn in self.connection_pool:
|
||
ibm_db.close(conn)
|
||
self.connection_pool.clear()
|
||
d.info("所有DB2连接已关闭")
|
||
|
||
def write_dataframe(self,
|
||
df: pd.DataFrame,
|
||
table_name: str,
|
||
batch_size: int = 1000,
|
||
use_transaction: bool = True,
|
||
upsert_keys: Optional[List[str]] = None,
|
||
update_columns: Optional[List[str]] = None) -> int:
|
||
"""
|
||
将DataFrame写入数据库表(支持存在更新，不存在插入)
|
||
:param df: 要写入的DataFrame
|
||
:param table_name: 目标表名
|
||
:param batch_size: 批量大小
|
||
:param use_transaction: 是否使用事务
|
||
:param upsert_keys: 用于判断记录是否存在的键列名列表
|
||
:param update_columns: 需要更新的列名列表(为None时更新所有非键列)
|
||
:return: 处理的行数
|
||
"""
|
||
if df.empty:
|
||
d.warning("空DataFrame，跳过写入")
|
||
return 0
|
||
|
||
# 获取列名
|
||
columns = df.columns.tolist()
|
||
total_rows = len(df)
|
||
processed_rows = 0
|
||
|
||
# 判断操作模式
|
||
if upsert_keys is None:
|
||
# 纯插入模式
|
||
placeholders = ",".join(["?"] * len(columns))
|
||
sql = f"INSERT INTO {table_name} ({','.join(columns)}) VALUES ({placeholders})"
|
||
else:
|
||
# Upsert模式
|
||
if update_columns is None:
|
||
update_columns = [col for col in columns if col not in upsert_keys]
|
||
|
||
# 构建MERGE语句(DB2语法)
|
||
merge_set = ",".join([f"{col}=EXCLUDED.{col}" for col in update_columns])
|
||
merge_condition = " AND ".join([f"TARGET.{key}=EXCLUDED.{key}" for key in upsert_keys])
|
||
sql = f"""
|
||
MERGE INTO {table_name} AS TARGET
|
||
USING (VALUES ({','.join(['?']*len(columns))})) AS EXCLUDED ({','.join(columns)})
|
||
ON ({merge_condition})
|
||
WHEN MATCHED THEN UPDATE SET {merge_set}
|
||
WHEN NOT MATCHED THEN INSERT ({','.join(columns)}) VALUES ({','.join([f'EXCLUDED.{col}' for col in columns])})
|
||
"""
|
||
|
||
# 分批处理
|
||
for i in range(0, total_rows, batch_size):
|
||
batch = df.iloc[i:i + batch_size]
|
||
params = [tuple(row) for row in batch.itertuples(index=False)]
|
||
|
||
if use_transaction:
|
||
# 使用事务写入
|
||
success = self.execute_transaction([sql] * len(batch), params)
|
||
if not success:
|
||
raise RuntimeError(f"写入表{table_name}失败")
|
||
else:
|
||
# 不使用事务，逐条写入
|
||
for param in params:
|
||
affected = self.execute_update(sql, param)
|
||
if affected <= 0 and upsert_keys is None: # 纯插入模式才检查影响行数
|
||
raise RuntimeError(f"写入表{table_name}失败")
|
||
|
||
processed_rows += len(batch)
|
||
d.info(f"已处理{processed_rows}/{total_rows}行到{table_name}")
|
||
|
||
return processed_rows
|
||
|
||
def __del__(self):
|
||
"""析构函数"""
|
||
self.close_all_connections()
|
||
|
||
# 使用示例
|
||
if __name__ == "__main__":
|
||
# 配置数据库连接参数
|
||
db2 = DB2Tool(
|
||
host="192.168.137.100",
|
||
port=50000,
|
||
database="appdb",
|
||
username="som",
|
||
password="dscdsc1"
|
||
)
|
||
LogUtil.init("app")
|
||
|
||
# 查询示例
|
||
results = db2.execute_query("SELECT * FROM T_SOM_STAT WHERE STEEL_GRADE = ? ", ("steel",), as_dataframe=True)
|
||
print("查询结果:", results)
|
||
|
||
# 更新示例
|
||
affected_rows = db2.execute_update("UPDATE T_TCM_DEC SET USE_TEN_DEVIATE = ? WHERE entid = ?", (1, "test1"))
|
||
print("影响行数:", affected_rows)
|
||
|
||
# 事务示例
|
||
transaction_sqls = [
|
||
"INSERT INTO T_TCM_DEC (entid, USE_TEN_DEVIATE) VALUES ('test1', 1)",
|
||
"UPDATE T_TCM_DEC SET USE_TEN_DEVIATE = 2 WHERE entid = 'test1'"
|
||
]
|
||
success = db2.execute_transaction(transaction_sqls)
|
||
print("事务执行结果:", success)
|
||
|
||
# DataFrame写入数据库示例
|
||
try:
|
||
# 创建测试DataFrame
|
||
test_data = {
|
||
'entid': ['test4', 'test5', 'test6'],
|
||
'USE_TEN_DEVIATE': [1, 0, 1]
|
||
}
|
||
test_df = pd.DataFrame(test_data)
|
||
|
||
# 测试Upsert模式
|
||
test_data_update = {
|
||
'entid': ['test4', 'test5', 'test9'], # test7是新记录
|
||
'USE_TEN_DEVIATE': [2, 9, 1]
|
||
}
|
||
test_df_update = pd.DataFrame(test_data_update)
|
||
print(test_df_update)
|
||
|
||
rows_processed = db2.write_dataframe(
|
||
df=test_df_update,
|
||
table_name="T_TCM_DEC",
|
||
batch_size=2,
|
||
upsert_keys=['entid'], # 使用entid作为唯一键
|
||
update_columns=['USE_TEN_DEVIATE'] # 更新这些列
|
||
)
|
||
print(f"成功处理{rows_processed}行(更新2行,插入1行)")
|
||
|
||
except Exception as e:
|
||
print(f"DataFrame写入失败: {e}")
|