crpto_predictor/data/database.py

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 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.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."""
        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:
            conn.close()

    def initialize_db(self):
        """Create tables if they don't exist."""
        with self.get_connection() as conn:
            cursor = conn.cursor()
            
            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()

    def save_ohlcv(self, df: pd.DataFrame, symbol: str, timeframe: str):
        """Save OHLCV data to database (Upsert)."""
        if df.empty:
            return

        # Ensure timestamp is datetime
        df = df.copy()
        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) or MySQL
        df['timestamp'] = df['timestamp'].dt.strftime('%Y-%m-%d %H:%M:%S')
        
        # Add metadata columns if they don't exist
        if 'symbol' not in df.columns:
            df['symbol'] = symbol
        if 'timeframe' not in df.columns:
            df['timeframe'] = timeframe
        
        # Prepare data for insertion (ensure order)
        # 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()
            
            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."""
        
        placeholder = '%s' if self.use_mysql else '?'
        
        query = f"""
            SELECT timestamp, open, high, low, close, volume 
            FROM ohlcv 
            WHERE symbol = {placeholder} AND timeframe = {placeholder}
        """
        params = [symbol, timeframe]

        if start_date:
            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
                # 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 += f" AND timestamp <= {placeholder}"
            params.append(end_date)
            
        query += " ORDER BY timestamp ASC"

        with self.get_connection() as conn:
            df = pd.read_sql_query(query, conn, params=params, parse_dates=['timestamp'])
            
        return df

    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(f"""
                SELECT MAX(timestamp) 
                FROM ohlcv 
                WHERE symbol = {placeholder} AND timeframe = {placeholder}
            """, (symbol, timeframe))
            result = cursor.fetchone()
            
        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(f"""
                SELECT MIN(timestamp) 
                FROM ohlcv 
                WHERE symbol = {placeholder} AND timeframe = {placeholder}
            """, (symbol, timeframe))
            result = cursor.fetchone()
            
        return pd.to_datetime(result[0]) if result and result[0] else None