110 lines
3.9 KiB
Python
110 lines
3.9 KiB
Python
import sqlite3
|
|
from datetime import datetime
|
|
import json
|
|
import os
|
|
from contextlib import contextmanager
|
|
|
|
class TradingDatabase:
|
|
"""交易数据库管理类 - 用于记录交易历史和预测结果"""
|
|
|
|
def __init__(self, db_path="trading_history.db"):
|
|
self.db_path = db_path
|
|
self._init_db()
|
|
|
|
@contextmanager
|
|
def get_connection(self):
|
|
"""数据库连接上下文管理器"""
|
|
conn = sqlite3.connect(self.db_path)
|
|
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 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()),
|
|
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('error_message'),
|
|
json.dumps(trade_data.get('raw_data', {}), ensure_ascii=False)
|
|
))
|
|
conn.commit()
|
|
except Exception as e:
|
|
print(f"写入数据库失败: {e}")
|
|
|
|
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]
|
|
return count > 0
|
|
except Exception as e:
|
|
print(f"检查已存在交易失败: {e}")
|
|
return False
|