subuhui/crpto_predictor
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: Add MySQL support and migration script

Browse Source
This commit is contained in:
crypto_predictor committed 2026-02-06 16:35:08 +08:00
1 parent 006ccb2a1d
commit 8de6c54887
4 files changed
+427 -93
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

data/trading_db.py
+119 -49
View File
@@ -1,20 +1,50 @@
import sqlite3
import pymysql
from datetime import datetime
import json
import os
from contextlib import contextmanager
from dotenv import load_dotenv
load_dotenv()
class TradingDatabase:
"""交易数据库管理类 - 用于记录交易历史和预测结果"""
"""交易数据库管理类 - 用于记录交易历史和预测结果
支持 SQLite (本地) 和 MySQL (远程)。
如果 .env 中配置了 MYSQL_HOST，则优先使用 MySQL。
"""
def __init__(self, db_path="trading_history.db"):
self.db_path = db_path
self.use_mysql = bool(os.getenv('MYSQL_HOST'))
self._init_db()
@contextmanager
def get_connection(self):
"""数据库连接上下文管理器"""
conn = sqlite3.connect(self.db_path)
if self.use_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
)
else:
conn = sqlite3.connect(self.db_path)
# Make sqlite behave like dict cursor for consistency if needed,
# but existing code might expect tuple.
# Actually existing code uses cursor.execute directly.
# SQLite cursor returns tuples by default.
# PyMySQL DictCursor returns dicts.
# We need to be careful about compatibility.
# The existing log_trade uses INSERT which is fine.
# check_existing_trade uses SELECT count(*), fetchone()[0], which works for both if we don't use DictCursor for MySQL or handle it.
# Let's use standard cursor for MySQL for compatibility with existing fetchone()[0] usage.
try:
yield conn
finally:
@@ -24,49 +54,84 @@ class TradingDatabase:
"""初始化数据库表结构"""
with self.get_connection() as conn:
cursor = conn.cursor()
# 创建交易记录表
cursor.execute('''
CREATE TABLE IF NOT EXISTS trades (
id INTEGER PRIMARY KEY AUTOINCREMENT,
timestamp TEXT,
market_title TEXT,
market_slug TEXT,
model_name TEXT,
prediction TEXT,
probability REAL,
action TEXT,
price REAL,
amount REAL,
size REAL,
status TEXT,
error_message TEXT,
raw_data TEXT,
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
)
''')
if self.use_mysql:
# MySQL Schema
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,
INDEX idx_market_slug (market_slug),
INDEX idx_status (status)
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4;
''')
else:
# SQLite Schema
cursor.execute('''
CREATE TABLE IF NOT EXISTS trades (
id INTEGER PRIMARY KEY AUTOINCREMENT,
timestamp TEXT,
market_title TEXT,
market_slug TEXT,
model_name TEXT,
prediction TEXT,
probability REAL,
action TEXT,
price REAL,
amount REAL,
size REAL,
status TEXT,
error_message TEXT,
raw_data TEXT,
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
)
''')
conn.commit()
def log_trade(self, trade_data: dict):
"""
记录交易信息
Args:
trade_data: 包含交易信息的字典，支持字段:
timestamp, market_title, market_slug, model_name,
prediction, probability, action, price, amount, size,
status, error_message, raw_data
"""
try:
with self.get_connection() as conn:
cursor = conn.cursor()
cursor.execute('''
INSERT INTO trades (
timestamp, market_title, market_slug, model_name,
prediction, probability, action, price, amount, size,
status, error_message, raw_data
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
''', (
trade_data.get('timestamp', datetime.now().isoformat()),
# Prepare data
timestamp = trade_data.get('timestamp', datetime.now().isoformat())
raw_data_json = json.dumps(trade_data.get('raw_data', {}), ensure_ascii=False)
if self.use_mysql:
sql = '''
INSERT INTO trades (
timestamp, market_title, market_slug, model_name,
prediction, probability, action, price, amount, size,
status, error_message, raw_data
) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
'''
else:
sql = '''
INSERT INTO trades (
timestamp, market_title, market_slug, model_name,
prediction, probability, action, price, amount, size,
status, error_message, raw_data
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
'''
cursor.execute(sql, (
timestamp,
trade_data.get('market_title'),
trade_data.get('market_slug'),
trade_data.get('model_name'),
@@ -78,7 +143,7 @@ class TradingDatabase:
trade_data.get('size'),
trade_data.get('status'),
trade_data.get('error_message'),
json.dumps(trade_data.get('raw_data', {}), ensure_ascii=False)
raw_data_json
))
conn.commit()
except Exception as e:
@@ -87,22 +152,27 @@ class TradingDatabase:
def check_existing_trade(self, market_slug: str, status: str = 'success') -> bool:
"""
检查指定市场是否已有成功交易记录
Args:
market_slug: 市场slug
status: 交易状态 (默认 'success', 也可以是 'test')
Returns:
True 如果已有成功交易, 否则 False
"""
try:
with self.get_connection() as conn:
cursor = conn.cursor()
cursor.execute('''
SELECT COUNT(*) FROM trades
WHERE market_slug = ? AND status = ?
''', (market_slug, status))
count = cursor.fetchone()[0]
if self.use_mysql:
sql = 'SELECT COUNT(*) FROM trades WHERE market_slug = %s AND status = %s'
else:
sql = 'SELECT COUNT(*) FROM trades WHERE market_slug = ? AND status = ?'
cursor.execute(sql, (market_slug, status))
# Fetch result compatible with both
result = cursor.fetchone()
if isinstance(result, dict): # If DictCursor was used (not used here but for safety)
count = list(result.values())[0]
elif isinstance(result, (list, tuple)):
count = result[0]
else:
count = 0
return count > 0
except Exception as e:
print(f"检查已存在交易失败: {e}")