diff --git a/data/database.py b/data/database.py index 9d6aed6..6fa4d57 100644 --- a/data/database.py +++ b/data/database.py @@ -1,22 +1,44 @@ import sqlite3 +import pymysql import pandas as pd from pathlib import Path from typing import Optional, List, Tuple from contextlib import contextmanager import os +from dotenv import load_dotenv + +load_dotenv() class DatabaseManager: - """Manages SQLite database for storing OHLCV data.""" + """Manages database for storing OHLCV data. + Supports SQLite (local) and MySQL (remote). + """ def __init__(self, db_path: str = "./data/market_data.db"): self.db_path = Path(db_path) - self.db_path.parent.mkdir(parents=True, exist_ok=True) + self.use_mysql = bool(os.getenv('MYSQL_HOST')) + + if not self.use_mysql: + self.db_path.parent.mkdir(parents=True, exist_ok=True) + self.initialize_db() @contextmanager def get_connection(self): """Context manager for database connection.""" - 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', + # pandas requires standard cursor or handling + ) + else: + conn = sqlite3.connect(self.db_path) + try: yield conn finally: @@ -27,27 +49,41 @@ class DatabaseManager: with self.get_connection() as conn: cursor = conn.cursor() - # Create OHLCV table - # Composite primary key ensures uniqueness for a specific candle - cursor.execute(""" - CREATE TABLE IF NOT EXISTS ohlcv ( - timestamp DATETIME, - symbol TEXT, - timeframe TEXT, - open REAL, - high REAL, - low REAL, - close REAL, - volume REAL, - PRIMARY KEY (symbol, timeframe, timestamp) - ) - """) - - # Create index for faster querying by time range - cursor.execute(""" - CREATE INDEX IF NOT EXISTS idx_ohlcv_time - ON ohlcv (symbol, timeframe, timestamp) - """) + if self.use_mysql: + # MySQL Schema + cursor.execute(""" + CREATE TABLE IF NOT EXISTS ohlcv ( + timestamp DATETIME, + symbol VARCHAR(20), + timeframe VARCHAR(10), + open DOUBLE, + high DOUBLE, + low DOUBLE, + close DOUBLE, + volume DOUBLE, + PRIMARY KEY (symbol, timeframe, timestamp), + INDEX idx_ohlcv_time (symbol, timeframe, timestamp) + ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4; + """) + else: + # SQLite Schema + cursor.execute(""" + CREATE TABLE IF NOT EXISTS ohlcv ( + timestamp DATETIME, + symbol TEXT, + timeframe TEXT, + open REAL, + high REAL, + low REAL, + close REAL, + volume REAL, + PRIMARY KEY (symbol, timeframe, timestamp) + ) + """) + cursor.execute(""" + CREATE INDEX IF NOT EXISTS idx_ohlcv_time + ON ohlcv (symbol, timeframe, timestamp) + """) conn.commit() @@ -61,7 +97,7 @@ class DatabaseManager: if not pd.api.types.is_datetime64_any_dtype(df['timestamp']): df['timestamp'] = pd.to_datetime(df['timestamp']) - # Convert to string for SQLite compatibility (ISO 8601) + # Convert to string for SQLite compatibility (ISO 8601) or MySQL df['timestamp'] = df['timestamp'].dt.strftime('%Y-%m-%d %H:%M:%S') # Add metadata columns if they don't exist @@ -71,38 +107,62 @@ class DatabaseManager: df['timeframe'] = timeframe # Prepare data for insertion (ensure order) - records = df[['timestamp', 'symbol', 'timeframe', 'open', 'high', 'low', 'close', 'volume']].to_records(index=False) + # records = df[['timestamp', 'symbol', 'timeframe', 'open', 'high', 'low', 'close', 'volume']].to_records(index=False) + # to_records returns numpy record array, which sqlite3 adapts. + # For MySQL/pymysql, we need list of tuples. + + cols = ['timestamp', 'symbol', 'timeframe', 'open', 'high', 'low', 'close', 'volume'] + data = [tuple(row) for row in df[cols].values] with self.get_connection() as conn: cursor = conn.cursor() - # efficiently insert or replace - # sqlite3 adapter handles datetime objects correctly - cursor.executemany(""" - INSERT OR REPLACE INTO ohlcv - (timestamp, symbol, timeframe, open, high, low, close, volume) - VALUES (?, ?, ?, ?, ?, ?, ?, ?) - """, records) + + if self.use_mysql: + # MySQL uses %s and REPLACE INTO + sql = """ + REPLACE INTO ohlcv + (timestamp, symbol, timeframe, open, high, low, close, volume) + VALUES (%s, %s, %s, %s, %s, %s, %s, %s) + """ + cursor.executemany(sql, data) + else: + # SQLite uses ? and INSERT OR REPLACE + sql = """ + INSERT OR REPLACE INTO ohlcv + (timestamp, symbol, timeframe, open, high, low, close, volume) + VALUES (?, ?, ?, ?, ?, ?, ?, ?) + """ + # sqlite3 accepts both list of tuples and numpy records, but list of tuples is safer for compatibility + cursor.executemany(sql, data) + conn.commit() def get_ohlcv(self, symbol: str, timeframe: str, start_date: Optional[str] = None, end_date: Optional[str] = None) -> pd.DataFrame: """Retrieve OHLCV data from database.""" - query = """ + + placeholder = '%s' if self.use_mysql else '?' + + query = f""" SELECT timestamp, open, high, low, close, volume FROM ohlcv - WHERE symbol = ? AND timeframe = ? + WHERE symbol = {placeholder} AND timeframe = {placeholder} """ params = [symbol, timeframe] if start_date: - query += " AND timestamp >= ?" + query += f" AND timestamp >= {placeholder}" params.append(start_date) if end_date: # Handle 'YYYY-MM-DD' vs full timestamp if len(str(end_date)) == 10: # YYYY-MM-DD - query += " AND timestamp < date(?, '+1 day')" + # SQLite has date(), MySQL uses DATE() or just string comparison + if self.use_mysql: + query += f" AND timestamp < DATE_ADD({placeholder}, INTERVAL 1 DAY)" + else: + query += f" AND timestamp < date({placeholder}, '+1 day')" else: - query += " AND timestamp <= ?" + query += f" AND timestamp <= {placeholder}" params.append(end_date) query += " ORDER BY timestamp ASC" @@ -114,26 +174,28 @@ class DatabaseManager: def get_latest_timestamp(self, symbol: str, timeframe: str) -> Optional[pd.Timestamp]: """Get the latest timestamp for a symbol/timeframe.""" + placeholder = '%s' if self.use_mysql else '?' with self.get_connection() as conn: cursor = conn.cursor() - cursor.execute(""" + cursor.execute(f""" SELECT MAX(timestamp) FROM ohlcv - WHERE symbol = ? AND timeframe = ? + WHERE symbol = {placeholder} AND timeframe = {placeholder} """, (symbol, timeframe)) result = cursor.fetchone() - return pd.to_datetime(result[0]) if result[0] else None + return pd.to_datetime(result[0]) if result and result[0] else None def get_earliest_timestamp(self, symbol: str, timeframe: str) -> Optional[pd.Timestamp]: """Get the earliest timestamp for a symbol/timeframe.""" + placeholder = '%s' if self.use_mysql else '?' with self.get_connection() as conn: cursor = conn.cursor() - cursor.execute(""" + cursor.execute(f""" SELECT MIN(timestamp) FROM ohlcv - WHERE symbol = ? AND timeframe = ? + WHERE symbol = {placeholder} AND timeframe = {placeholder} """, (symbol, timeframe)) result = cursor.fetchone() - return pd.to_datetime(result[0]) if result[0] else None + return pd.to_datetime(result[0]) if result and result[0] else None diff --git a/data/trading_db.py b/data/trading_db.py index a83438c..4f0180f 100644 --- a/data/trading_db.py +++ b/data/trading_db.py @@ -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}") diff --git a/requirements.txt b/requirements.txt index 161f195..0634808 100644 --- a/requirements.txt +++ b/requirements.txt @@ -38,3 +38,7 @@ flake8>=6.0.0 # Utilities tqdm>=4.66.0 + +# Database +pymysql>=1.1.0 +sqlalchemy>=2.0.0 diff --git a/scripts/migrate_sqlite_to_mysql.py b/scripts/migrate_sqlite_to_mysql.py new file mode 100644 index 0000000..05712cf --- /dev/null +++ b/scripts/migrate_sqlite_to_mysql.py @@ -0,0 +1,198 @@ +""" +SQLite to MySQL Migration Script + +This script migrates data from local SQLite databases (`trading_history.db` and `data/market_data.db`) +to a remote MySQL database defined in environment variables. + +Environment Variables required in .env: +- MYSQL_HOST +- MYSQL_PORT (default: 3306) +- MYSQL_USER +- MYSQL_PASSWORD +- MYSQL_DB + +Usage: + python scripts/migrate_sqlite_to_mysql.py +""" + +import sys +import os +import sqlite3 +import pymysql +from dotenv import load_dotenv +from tqdm import tqdm +import pandas as pd + +# Add project root to sys.path +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +load_dotenv() + +def get_mysql_connection(): + return 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 + ) + +def create_mysql_tables(conn): + with conn.cursor() as cursor: + # Create trades table + print("Creating 'trades' table in MySQL...") + 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; + """) + + # Create ohlcv table + print("Creating 'ohlcv' table in MySQL...") + cursor.execute(""" + CREATE TABLE IF NOT EXISTS ohlcv ( + timestamp DATETIME, + symbol VARCHAR(20), + timeframe VARCHAR(10), + open DOUBLE, + high DOUBLE, + low DOUBLE, + close DOUBLE, + volume DOUBLE, + PRIMARY KEY (symbol, timeframe, timestamp), + INDEX idx_time (timestamp) + ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4; + """) + conn.commit() + +def migrate_trades(mysql_conn): + sqlite_db_path = "trading_history.db" + if not os.path.exists(sqlite_db_path): + print(f"Skipping trades migration: {sqlite_db_path} not found.") + return + + print(f"Migrating trades from {sqlite_db_path}...") + + with sqlite3.connect(sqlite_db_path) as sqlite_conn: + sqlite_cursor = sqlite_conn.cursor() + + # Check if table exists + sqlite_cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='trades';") + if not sqlite_cursor.fetchone(): + print("Table 'trades' not found in SQLite.") + return + + sqlite_cursor.execute("SELECT * FROM trades") + rows = sqlite_cursor.fetchall() + columns = [description[0] for description in sqlite_cursor.description] + + if not rows: + print("No trades data to migrate.") + return + + print(f"Found {len(rows)} trade records.") + + # Insert into MySQL + with mysql_conn.cursor() as mysql_cursor: + # Prepare SQL statement + # Map columns. SQLite 'timestamp' is text, MySQL needs valid datetime format or string. + # We skip 'id' to let MySQL handle auto-increment or we can preserve it. + # Let's preserve ID if possible, but usually better to let destination handle it unless we need consistency. + # User wants migration, preserving ID is safer for history. + + placeholders = ', '.join(['%s'] * len(columns)) + cols_str = ', '.join(columns) + + sql = f"INSERT IGNORE INTO trades ({cols_str}) VALUES ({placeholders})" + + mysql_cursor.executemany(sql, rows) + mysql_conn.commit() + print(f"Migrated {mysql_cursor.rowcount} trade records.") + +def migrate_ohlcv(mysql_conn): + sqlite_db_path = "data/market_data.db" + if not os.path.exists(sqlite_db_path): + print(f"Skipping OHLCV migration: {sqlite_db_path} not found.") + return + + print(f"Migrating OHLCV from {sqlite_db_path}...") + + with sqlite3.connect(sqlite_db_path) as sqlite_conn: + # Use pandas for easier chunk reading + try: + # Check total count + total_rows = pd.read_sql_query("SELECT COUNT(*) FROM ohlcv", sqlite_conn).iloc[0, 0] + print(f"Found {total_rows} OHLCV records.") + + chunk_size = 10000 + offset = 0 + + with tqdm(total=total_rows, desc="Migrating OHLCV") as pbar: + while True: + df = pd.read_sql_query(f"SELECT * FROM ohlcv LIMIT {chunk_size} OFFSET {offset}", sqlite_conn) + if df.empty: + break + + # Convert DataFrame to list of tuples + data = [tuple(x) for x in df.to_numpy()] + + with mysql_conn.cursor() as mysql_cursor: + sql = """ + INSERT IGNORE INTO ohlcv + (timestamp, symbol, timeframe, open, high, low, close, volume) + VALUES (%s, %s, %s, %s, %s, %s, %s, %s) + """ + # Select columns explicitly to match SQL + cols = ['timestamp', 'symbol', 'timeframe', 'open', 'high', 'low', 'close', 'volume'] + batch_data = [tuple(row[col] for col in cols) for _, row in df.iterrows()] + + mysql_cursor.executemany(sql, batch_data) + mysql_conn.commit() + + offset += chunk_size + pbar.update(len(df)) + + except Exception as e: + print(f"Error reading/migrating OHLCV: {e}") + +def main(): + if not os.getenv('MYSQL_HOST'): + print("Error: MYSQL_HOST not set in .env") + return + + try: + print("Connecting to MySQL...") + mysql_conn = get_mysql_connection() + print("Connected.") + + create_mysql_tables(mysql_conn) + + migrate_trades(mysql_conn) + migrate_ohlcv(mysql_conn) + + mysql_conn.close() + print("\nMigration completed successfully.") + + except Exception as e: + print(f"\nMigration failed: {e}") + +if __name__ == "__main__": + main()