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
-            # 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