82 lines
2.4 KiB
Python
82 lines
2.4 KiB
Python
"""
|
|
数据库连接管理
|
|
"""
|
|
import pymysql
|
|
from contextlib import contextmanager
|
|
import os
|
|
import logging
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class Database:
|
|
"""数据库连接类"""
|
|
|
|
def __init__(self):
|
|
self.host = os.getenv('DB_HOST', 'localhost')
|
|
self.port = int(os.getenv('DB_PORT', 3306))
|
|
self.user = os.getenv('DB_USER', 'root')
|
|
self.password = os.getenv('DB_PASSWORD', '')
|
|
self.database = os.getenv('DB_NAME', 'auto_trade_sys')
|
|
|
|
@contextmanager
|
|
def get_connection(self):
|
|
"""获取数据库连接（上下文管理器）"""
|
|
conn = None
|
|
try:
|
|
conn = pymysql.connect(
|
|
host=self.host,
|
|
port=self.port,
|
|
user=self.user,
|
|
password=self.password,
|
|
database=self.database,
|
|
charset='utf8mb4',
|
|
cursorclass=pymysql.cursors.DictCursor,
|
|
autocommit=False
|
|
)
|
|
yield conn
|
|
except Exception as e:
|
|
if conn:
|
|
conn.rollback()
|
|
logger.error(f"数据库连接错误: {e}")
|
|
raise
|
|
finally:
|
|
if conn:
|
|
conn.close()
|
|
|
|
def execute_query(self, query, params=None):
|
|
"""执行查询，返回所有结果"""
|
|
with self.get_connection() as conn:
|
|
with conn.cursor() as cursor:
|
|
cursor.execute(query, params)
|
|
conn.commit()
|
|
return cursor.fetchall()
|
|
|
|
def execute_one(self, query, params=None):
|
|
"""执行查询，返回单条结果"""
|
|
with self.get_connection() as conn:
|
|
with conn.cursor() as cursor:
|
|
cursor.execute(query, params)
|
|
conn.commit()
|
|
return cursor.fetchone()
|
|
|
|
def execute_update(self, query, params=None):
|
|
"""执行更新，返回影响行数"""
|
|
with self.get_connection() as conn:
|
|
with conn.cursor() as cursor:
|
|
affected = cursor.execute(query, params)
|
|
conn.commit()
|
|
return affected
|
|
|
|
def execute_many(self, query, params_list):
|
|
"""批量执行"""
|
|
with self.get_connection() as conn:
|
|
with conn.cursor() as cursor:
|
|
affected = cursor.executemany(query, params_list)
|
|
conn.commit()
|
|
return affected
|
|
|
|
|
|
# 全局数据库实例
|
|
db = Database()
|