157 lines
6.2 KiB
Python
157 lines
6.2 KiB
Python
import pymysql
|
|
from datetime import datetime
|
|
import json
|
|
import os
|
|
from contextlib import contextmanager
|
|
from dotenv import load_dotenv
|
|
|
|
load_dotenv()
|
|
|
|
|
|
class TradingDatabase:
|
|
"""交易数据库管理类 - 用于记录交易历史和预测结果 (MySQL)"""
|
|
|
|
def __init__(self, db_path=None):
|
|
"""初始化 (db_path 参数已废弃，仅为兼容性保留)"""
|
|
self._init_db()
|
|
|
|
@contextmanager
|
|
def get_connection(self):
|
|
"""获取 MySQL 数据库连接"""
|
|
conn = pymysql.connect(
|
|
host=os.getenv('MYSQL_HOST'),
|
|
port=int(os.getenv('MYSQL_PORT', 3306)),
|
|
user=os.getenv('MYSQL_USER'),
|
|
password=os.getenv('MYSQL_PASSWORD'),
|
|
database=os.getenv('MYSQL_DB'),
|
|
charset='utf8mb4',
|
|
cursorclass=pymysql.cursors.DictCursor
|
|
)
|
|
try:
|
|
yield conn
|
|
finally:
|
|
conn.close()
|
|
|
|
def _init_db(self):
|
|
"""初始化数据库表结构"""
|
|
with self.get_connection() as conn:
|
|
cursor = conn.cursor()
|
|
cursor.execute('''
|
|
CREATE TABLE IF NOT EXISTS trades (
|
|
id INT AUTO_INCREMENT PRIMARY KEY,
|
|
timestamp DATETIME,
|
|
market_title VARCHAR(255),
|
|
market_slug VARCHAR(255),
|
|
model_name VARCHAR(50),
|
|
prediction VARCHAR(10),
|
|
probability DOUBLE,
|
|
action VARCHAR(10),
|
|
price DOUBLE,
|
|
amount DOUBLE,
|
|
size DOUBLE,
|
|
status VARCHAR(20),
|
|
error_message TEXT,
|
|
raw_data JSON,
|
|
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
|
order_id VARCHAR(100),
|
|
INDEX idx_market_slug (market_slug),
|
|
INDEX idx_status (status)
|
|
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4;
|
|
''')
|
|
conn.commit()
|
|
|
|
# 检查并添加 order_id 列 (Migration)
|
|
self._ensure_column_exists(conn, 'trades', 'order_id', 'VARCHAR(100)')
|
|
|
|
def _ensure_column_exists(self, conn, table_name, column_name, column_type):
|
|
"""确保表中存在指定列"""
|
|
try:
|
|
cursor = conn.cursor()
|
|
cursor.execute(f"SHOW COLUMNS FROM {table_name} LIKE '{column_name}'")
|
|
if not cursor.fetchone():
|
|
print(f"正在添加列 {column_name} 到表 {table_name}...")
|
|
cursor.execute(f"ALTER TABLE {table_name} ADD COLUMN {column_name} {column_type}")
|
|
conn.commit()
|
|
except Exception as e:
|
|
print(f"检查/添加列失败: {e}")
|
|
|
|
def log_trade(self, trade_data: dict):
|
|
"""记录交易信息"""
|
|
try:
|
|
with self.get_connection() as conn:
|
|
cursor = conn.cursor()
|
|
|
|
timestamp = trade_data.get('timestamp', datetime.now().isoformat())
|
|
raw_data_json = json.dumps(trade_data.get('raw_data', {}), ensure_ascii=False)
|
|
|
|
sql = '''
|
|
INSERT INTO trades (
|
|
timestamp, market_title, market_slug, model_name,
|
|
prediction, probability, action, price, amount, size,
|
|
status, strategy, error_message, raw_data, order_id
|
|
) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
|
|
'''
|
|
|
|
cursor.execute(sql, (
|
|
timestamp,
|
|
trade_data.get('market_title'),
|
|
trade_data.get('market_slug'),
|
|
trade_data.get('model_name'),
|
|
trade_data.get('prediction'),
|
|
trade_data.get('probability'),
|
|
trade_data.get('action'),
|
|
trade_data.get('price'),
|
|
trade_data.get('amount'),
|
|
trade_data.get('size'),
|
|
trade_data.get('status'),
|
|
trade_data.get('strategy'),
|
|
trade_data.get('error_message'),
|
|
raw_data_json,
|
|
trade_data.get('order_id')
|
|
))
|
|
conn.commit()
|
|
except Exception as e:
|
|
print(f"写入数据库失败: {e}")
|
|
|
|
def check_existing_trade(self, market_slug: str, status: str = 'success') -> bool:
|
|
"""检查指定市场是否已有成功交易记录"""
|
|
try:
|
|
with self.get_connection() as conn:
|
|
cursor = conn.cursor()
|
|
sql = 'SELECT COUNT(*) as cnt FROM trades WHERE market_slug = %s AND status = %s'
|
|
cursor.execute(sql, (market_slug, status))
|
|
result = cursor.fetchone()
|
|
count = result['cnt'] if result else 0
|
|
return count > 0
|
|
except Exception as e:
|
|
print(f"检查已存在交易失败: {e}")
|
|
return False
|
|
|
|
def get_existing_trade(self, market_slug: str) -> dict:
|
|
"""获取指定市场的最近一次交易记录"""
|
|
try:
|
|
with self.get_connection() as conn:
|
|
cursor = conn.cursor()
|
|
# Get the latest trade for this market
|
|
sql = 'SELECT * FROM trades WHERE market_slug = %s ORDER BY id DESC LIMIT 1'
|
|
cursor.execute(sql, (market_slug,))
|
|
return cursor.fetchone()
|
|
except Exception as e:
|
|
print(f"获取交易记录失败: {e}")
|
|
return None
|
|
|
|
def update_trade_status(self, trade_id: int, status: str, error_message: str = None):
|
|
"""更新交易状态"""
|
|
try:
|
|
with self.get_connection() as conn:
|
|
cursor = conn.cursor()
|
|
if error_message:
|
|
sql = 'UPDATE trades SET status = %s, error_message = %s WHERE id = %s'
|
|
cursor.execute(sql, (status, error_message, trade_id))
|
|
else:
|
|
sql = 'UPDATE trades SET status = %s WHERE id = %s'
|
|
cursor.execute(sql, (status, trade_id))
|
|
conn.commit()
|
|
except Exception as e:
|
|
print(f"更新交易状态失败: {e}")
|