接入数据库

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

data/fetcher.py

@@ -10,6 +10,8 @@ from binance.client import Client
 import pytz
 from dotenv import load_dotenv
 
+from data.database import DatabaseManager
+
 # 加载环境变量
 load_dotenv()
 
@@ -23,10 +25,10 @@ class BinanceDataFetcher:
     }
 
     def __init__(self, cache_dir: str = "./data/cache"):
-        """初始化Binance客户端和缓存目录
-
+        """初始化Binance客户端和数据库管理器
+        
         Args:
-            cache_dir: 缓存下载数据的目录
+            cache_dir: 缓存/数据库目录
         """
         # 配置代理
         proxies = None
@@ -58,8 +60,8 @@ class BinanceDataFetcher:
             print("  3. 或者使用已缓存的数据")
             self.client = None
 
-        self.cache_dir = Path(cache_dir)
-        self.cache_dir.mkdir(parents=True, exist_ok=True)
+        # 初始化数据库管理器
+        self.db_manager = DatabaseManager(os.path.join(cache_dir, "market_data.db"))
 
     def fetch_ohlcv(
         self,
@@ -81,21 +83,7 @@ class BinanceDataFetcher:
         Returns:
             包含以下列的DataFrame: timestamp, open, high, low, close, volume
         """
-        # 生成缓存文件名
-        cache_file = self._get_cache_filename(symbol, timeframe, start_date, end_date)
-
-        # 检查缓存
-        if use_cache and cache_file.exists():
-            print(f"从缓存加载数据：{cache_file}")
-            return pd.read_csv(cache_file, parse_dates=['timestamp'])
-
-        # 检查客户端是否可用
-        if self.client is None:
-            raise ConnectionError(
-                "Binance客户端不可用。请使用VPN或使用已缓存的数据。"
-            )
-
-        # 设置默认日期
+        # 设置日期范围
         if end_date is None:
             end_dt = datetime.now(pytz.UTC)
         else:
@@ -106,6 +94,34 @@ class BinanceDataFetcher:
         else:
             start_dt = datetime.strptime(start_date, '%Y-%m-%d').replace(tzinfo=pytz.UTC)
 
+        start_str = start_dt.strftime('%Y-%m-%d')
+        end_str = end_dt.strftime('%Y-%m-%d')
+
+        # 1. 尝试从数据库获取数据
+        df = pd.DataFrame()
+        if use_cache:
+            df = self.db_manager.get_ohlcv(symbol, timeframe, start_str, end_str)
+            
+            # 检查数据是否足够
+            if not df.empty:
+                min_ts = df['timestamp'].min().replace(tzinfo=pytz.UTC)
+                max_ts = df['timestamp'].max().replace(tzinfo=pytz.UTC)
+                
+                # 简单判断覆盖范围 (容忍少量缺失)
+                if min_ts <= start_dt + timedelta(days=2) and max_ts >= end_dt - timedelta(days=2):
+                    print(f"从数据库加载数据：{symbol} {timeframe} ({len(df)} rows)")
+                    return df
+
+        # 2. 从API获取数据
+        if self.client is None:
+            if not df.empty:
+                 print(f"⚠ API不可用，返回数据库现有数据 ({len(df)} rows)")
+                 return df
+            
+            raise ConnectionError(
+                "Binance客户端不可用。请使用VPN或确保数据库有数据。"
+            )
+
         # Binance使用USDT交易对
         pair = f"{symbol}USDT"
         interval = self.INTERVALS[timeframe]
@@ -113,53 +129,47 @@ class BinanceDataFetcher:
         print(f"获取 {symbol} {timeframe} 数据，时间范围：{start_dt.date()} 至 {end_dt.date()}...")
 
         # 获取K线数据
-        klines = self.client.get_historical_klines(
-            symbol=pair,
-            interval=interval,
-            start_str=str(int(start_dt.timestamp() * 1000)),
-            end_str=str(int(end_dt.timestamp() * 1000))
-        )
+        try:
+            klines = self.client.get_historical_klines(
+                symbol=pair,
+                interval=interval,
+                start_str=str(int(start_dt.timestamp() * 1000)),
+                end_str=str(int(end_dt.timestamp() * 1000))
+            )
+        except Exception as e:
+            if not df.empty:
+                print(f"⚠ API请求失败 ({e})，返回数据库现有数据")
+                return df
+            raise e
 
         if not klines:
-            raise ValueError(f"未获取到 {symbol} {timeframe} 的数据")
+             if not df.empty:
+                 return df
+             raise ValueError(f"未获取到 {symbol} {timeframe} 的数据")
 
         # 转换为DataFrame
-        df = pd.DataFrame(klines, columns=[
+        new_df = pd.DataFrame(klines, columns=[
             'timestamp', 'open', 'high', 'low', 'close', 'volume',
             'close_time', 'quote_volume', 'trades', 'taker_buy_base',
             'taker_buy_quote', 'ignore'
         ])
 
         # 选择并转换列
-        df = df[['timestamp', 'open', 'high', 'low', 'close', 'volume']].copy()
-        df['timestamp'] = pd.to_datetime(df['timestamp'], unit='ms')
+        new_df = new_df[['timestamp', 'open', 'high', 'low', 'close', 'volume']].copy()
+        new_df['timestamp'] = pd.to_datetime(new_df['timestamp'], unit='ms')
 
         for col in ['open', 'high', 'low', 'close', 'volume']:
-            df[col] = df[col].astype(float)
+            new_df[col] = new_df[col].astype(float)
 
         # 按时间排序
-        df = df.sort_values('timestamp').reset_index(drop=True)
+        new_df = new_df.sort_values('timestamp').reset_index(drop=True)
 
-        print(f"已获取 {len(df)} 行数据，时间范围：{df['timestamp'].min()} 至 {df['timestamp'].max()}")
+        print(f"已获取 {len(new_df)} 行数据")
 
-        # 保存到缓存
-        df.to_csv(cache_file, index=False)
-        print(f"已缓存数据至：{cache_file}")
-
-        return df
-
-    def _get_cache_filename(
-        self,
-        symbol: str,
-        timeframe: str,
-        start_date: Optional[str],
-        end_date: Optional[str]
-    ) -> Path:
-        """根据参数生成缓存文件名"""
-        start = start_date or "1y"
-        end = end_date or "now"
-        filename = f"{symbol}_{timeframe}_{start}_{end}.csv"
-        return self.cache_dir / filename
+        # 保存到数据库
+        self.db_manager.save_ohlcv(new_df, symbol, timeframe)
+        
+        return new_df
 
     def get_latest_price(self, symbol: Literal['BTC', 'ETH', 'SOL']) -> float:
         """获取符号的当前价格
@@ -187,5 +197,8 @@ if __name__ == "__main__":
     print(f"日期范围：{df['timestamp'].min()} 至 {df['timestamp'].max()}")
 
     # 获取当前价格
-    current_price = fetcher.get_latest_price('BTC')
-    print(f"\n当前BTC价格：${current_price:,.2f}")
+    try:
+        current_price = fetcher.get_latest_price('BTC')
+        print(f"\n当前BTC价格：${current_price:,.2f}")
+    except:
+        print("\n无法获取当前价格")