接入数据库

data/database.py

@@ -0,0 +1,139 @@
+import sqlite3
+import pandas as pd
+from pathlib import Path
+from typing import Optional, List, Tuple
+from contextlib import contextmanager
+import os
+
+class DatabaseManager:
+    """Manages SQLite database for storing OHLCV data."""
+
+    def __init__(self, db_path: str = "./data/cache/market_data.db"):
+        self.db_path = Path(db_path)
+        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)
+        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()
+            
+            # 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)
+            """)
+            
+            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)
+        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)
+        
+        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)
+            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 = """
+            SELECT timestamp, open, high, low, close, volume 
+            FROM ohlcv 
+            WHERE symbol = ? AND timeframe = ?
+        """
+        params = [symbol, timeframe]
+
+        if start_date:
+            query += " AND timestamp >= ?"
+            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')"
+            else:
+                 query += " AND timestamp <= ?"
+            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."""
+        with self.get_connection() as conn:
+            cursor = conn.cursor()
+            cursor.execute("""
+                SELECT MAX(timestamp) 
+                FROM ohlcv 
+                WHERE symbol = ? AND timeframe = ?
+            """, (symbol, timeframe))
+            result = cursor.fetchone()
+            
+        return pd.to_datetime(result[0]) if result[0] else None
+
+    def get_earliest_timestamp(self, symbol: str, timeframe: str) -> Optional[pd.Timestamp]:
+        """Get the earliest timestamp for a symbol/timeframe."""
+        with self.get_connection() as conn:
+            cursor = conn.cursor()
+            cursor.execute("""
+                SELECT MIN(timestamp) 
+                FROM ohlcv 
+                WHERE symbol = ? AND timeframe = ?
+            """, (symbol, timeframe))
+            result = cursor.fetchone()
+            
+        return pd.to_datetime(result[0]) if result[0] else None