feat: Add MySQL support and migration script

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
+            if self.use_mysql:
-            # Composite primary key ensures uniqueness for a specific candle
+                # MySQL Schema
-            cursor.execute("""
+                cursor.execute("""
-                CREATE TABLE IF NOT EXISTS ohlcv (
+                    CREATE TABLE IF NOT EXISTS ohlcv (
-                    timestamp DATETIME,
+                        timestamp DATETIME,
-                    symbol TEXT,
+                        symbol VARCHAR(20),
-                    timeframe TEXT,
+                        timeframe VARCHAR(10),
-                    open REAL,
+                        open DOUBLE,
-                    high REAL,
+                        high DOUBLE,
-                    low REAL,
+                        low DOUBLE,
-                    close REAL,
+                        close DOUBLE,
-                    volume REAL,
+                        volume DOUBLE,
-                    PRIMARY KEY (symbol, timeframe, timestamp)
+                        PRIMARY KEY (symbol, timeframe, timestamp),
-                )
+                        INDEX idx_ohlcv_time (symbol, timeframe, timestamp)
-            """)
+                    ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4;
-            
+                """)
-            # Create index for faster querying by time range
+            else:
-            cursor.execute("""
+                # SQLite Schema
-                CREATE INDEX IF NOT EXISTS idx_ohlcv_time 
+                cursor.execute("""
-                ON ohlcv (symbol, timeframe, timestamp)
+                    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
+            if self.use_mysql:
-            cursor.executemany("""
+                # MySQL uses %s and REPLACE INTO
-                INSERT OR REPLACE INTO ohlcv 
+                sql = """
-                (timestamp, symbol, timeframe, open, high, low, close, volume)
+                    REPLACE INTO ohlcv 
-                VALUES (?, ?, ?, ?, ?, ?, ?, ?)
+                    (timestamp, symbol, timeframe, open, high, low, close, volume)
-            """, records)
+                    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

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('''
+            if self.use_mysql:
-                CREATE TABLE IF NOT EXISTS trades (
+                # MySQL Schema
-                    id INTEGER PRIMARY KEY AUTOINCREMENT,
+                cursor.execute('''
-                    timestamp TEXT,
+                    CREATE TABLE IF NOT EXISTS trades (
-                    market_title TEXT,
+                        id INT AUTO_INCREMENT PRIMARY KEY,
-                    market_slug TEXT,
+                        timestamp DATETIME,
-                    model_name TEXT,
+                        market_title VARCHAR(255),
-                    prediction TEXT,
+                        market_slug VARCHAR(255),
-                    probability REAL,
+                        model_name VARCHAR(50),
-                    action TEXT,
+                        prediction VARCHAR(10),
-                    price REAL,
+                        probability DOUBLE,
-                    amount REAL,
+                        action VARCHAR(10),
-                    size REAL,
+                        price DOUBLE,
-                    status TEXT,
+                        amount DOUBLE,
-                    error_message TEXT,
+                        size DOUBLE,
-                    raw_data TEXT,
+                        status VARCHAR(20),
-                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
+                        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 (
+                # Prepare data
-                        timestamp, market_title, market_slug, model_name,
+                timestamp = trade_data.get('timestamp', datetime.now().isoformat())
-                        prediction, probability, action, price, amount, size,
+                raw_data_json = json.dumps(trade_data.get('raw_data', {}), ensure_ascii=False)
-                        status, error_message, raw_data
+                
-                    ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
+                if self.use_mysql:
-                ''', (
+                    sql = '''
-                    trade_data.get('timestamp', datetime.now().isoformat()),
+                        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 
+                if self.use_mysql:
-                    WHERE market_slug = ? AND status = ?
+                    sql = 'SELECT COUNT(*) FROM trades WHERE market_slug = %s AND status = %s'
-                ''', (market_slug, status))
+                else:
-                count = cursor.fetchone()[0]
+                    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}")

requirements.txt

@@ -38,3 +38,7 @@ flake8>=6.0.0
 
 # Utilities
 tqdm>=4.66.0
+
+# Database
+pymysql>=1.1.0
+sqlalchemy>=2.0.0

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()