Files
ylt_diy/lib/db.py

112 lines
4.2 KiB
Python

"""
数据库连接配置
"""
import os
import pymysql
from contextlib import contextmanager
from lib.logger import log_info, log_error, log_warning
def get_db_config():
"""从环境变量或默认值获取数据库配置"""
return {
'host': os.environ.get('DB_HOST', 'haoslm2.xicp.net'),
'port': int(os.environ.get('DB_PORT', 10216)),
'user': os.environ.get('DB_USER', 'root'),
'password': os.environ.get('DB_PASSWORD', 'DsideaL147258369'),
'database': os.environ.get('DB_NAME', 'yltcharge'),
'charset': 'utf8mb4',
'cursorclass': pymysql.cursors.DictCursor,
'autocommit': True,
'connect_timeout': 30,
'read_timeout': 120,
'write_timeout': 120
}
DB_CONFIG = get_db_config()
@contextmanager
def get_connection():
"""获取数据库连接(上下文管理器)"""
try:
conn = pymysql.connect(**DB_CONFIG)
yield conn
except pymysql.Error as e:
log_error(f'数据库连接失败: {e}', 'db', exc_info=True)
raise
finally:
try:
conn.close()
except:
pass
def execute_query(sql, params=None, retry=2):
"""执行查询并返回结果"""
for attempt in range(retry + 1):
try:
with get_connection() as conn:
with conn.cursor() as cursor:
cursor.execute(sql, params)
result = cursor.fetchall()
if attempt > 0:
log_info(f'[数据库] 查询重试成功,第{attempt+1}次尝试', 'db')
return result
except pymysql.err.OperationalError as e:
if attempt < retry and (e.args[0] == 2013 or e.args[0] == 2006):
log_warning(f'[数据库] 查询连接断开,正在重试(第{attempt+1}次): {e}', 'db')
import time
time.sleep(1)
continue
log_error(f'查询失败: {e}\nSQL: {sql}\nParams: {params}', 'db', exc_info=True)
raise
except Exception as e:
log_error(f'查询失败: {e}\nSQL: {sql}\nParams: {params}', 'db', exc_info=True)
raise
def execute_update(sql, params=None, retry=2):
"""执行更新操作"""
for attempt in range(retry + 1):
try:
with get_connection() as conn:
with conn.cursor() as cursor:
affected = cursor.execute(sql, params)
if attempt > 0:
log_info(f'[数据库] 更新重试成功,第{attempt+1}次尝试', 'db')
return cursor.rowcount
except pymysql.err.OperationalError as e:
if attempt < retry and (e.args[0] == 2013 or e.args[0] == 2006):
log_warning(f'[数据库] 更新连接断开,正在重试(第{attempt+1}次): {e}', 'db')
import time
time.sleep(1)
continue
log_error(f'更新失败: {e}\nSQL: {sql}\nParams: {params}', 'db', exc_info=True)
raise
except Exception as e:
log_error(f'更新失败: {e}\nSQL: {sql}\nParams: {params}', 'db', exc_info=True)
raise
def execute_insert(sql, params=None, retry=2):
"""执行插入操作并返回插入ID"""
for attempt in range(retry + 1):
try:
with get_connection() as conn:
with conn.cursor() as cursor:
cursor.execute(sql, params)
last_id = cursor.lastrowid
if attempt > 0:
log_info(f'[数据库] 插入重试成功,第{attempt+1}次尝试', 'db')
return last_id
except pymysql.err.OperationalError as e:
if attempt < retry and (e.args[0] == 2013 or e.args[0] == 2006):
log_warning(f'[数据库] 插入连接断开,正在重试(第{attempt+1}次): {e}', 'db')
import time
time.sleep(1)
continue
log_error(f'插入失败: {e}\nSQL: {sql}\nParams: {params}', 'db', exc_info=True)
raise
except Exception as e:
log_error(f'插入失败: {e}\nSQL: {sql}\nParams: {params}', 'db', exc_info=True)
raise