crypto_quant/core/db_trade_data.py

import pandas as pd
import logging
from typing import Optional, List, Dict, Any, Union
from core.db_manager import DBData
from core.utils import check_date_time_format, datetime_to_timestamp

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s: %(message)s")


class DBTradeData:
    def __init__(
        self, 
        db_url: str
    ):
        self.db_url = db_url
        self.table_name = "crypto_trade_data"
        self.columns = [
            "symbol",
            "ts",
            "date_time",
            "tradeId",
            "side",
            "sz",
            "px",
            "create_time",
        ]
        self.db_manager = DBData(db_url, self.table_name, self.columns)

    def _process_time_parameter(self, time_param: Optional[Union[str, int]]) -> Optional[int]:
        """
        处理时间参数，统一转换为时间戳
        :param time_param: 时间参数（字符串或整数）
        :return: 时间戳或None
        """
        if time_param is None:
            return None
            
        if isinstance(time_param, int):
            return time_param
            
        if isinstance(time_param, str):
            if time_param.isdigit():
                return int(time_param)
            else:
                parsed_time = check_date_time_format(time_param)
                if parsed_time is None:
                    logging.warning(f"日期时间格式错误: {time_param}")
                    return None
                return datetime_to_timestamp(parsed_time)
        
        return None

    def _build_query_conditions(
        self, 
        symbol: Optional[str] = None, 
        side: Optional[str] = None, 
        start: Optional[Union[str, int]] = None, 
        end: Optional[Union[str, int]] = None,
        additional_conditions: Optional[List[str]] = None
    ) -> tuple[List[str], Dict[str, Any]]:
        """
        构建查询条件
        :param symbol: 交易对
        :param side: 交易方向
        :param start: 开始时间
        :param end: 结束时间
        :param additional_conditions: 额外的查询条件
        :return: (条件列表, 参数字典)
        """
        conditions = additional_conditions or []
        condition_dict = {}
        
        if symbol:
            conditions.append("symbol = :symbol")
            condition_dict["symbol"] = symbol
        if side:
            conditions.append("side = :side")
            condition_dict["side"] = side
            
        # 处理时间参数
        start_timestamp = self._process_time_parameter(start)
        end_timestamp = self._process_time_parameter(end)
        
        if start_timestamp is not None:
            conditions.append("ts >= :start")
            condition_dict["start"] = start_timestamp
        if end_timestamp is not None:
            conditions.append("ts <= :end")
            condition_dict["end"] = end_timestamp
            
        return conditions, condition_dict

    def insert_data_to_mysql(self, df: pd.DataFrame) -> None:
        """
        将交易数据保存到MySQL的crypto_trade_data表
        速度：⭐⭐⭐⭐⭐ 最快
        内存：⭐⭐⭐⭐ 中等
        适用场景：中小数据量（<10万条）
        :param df: 交易数据DataFrame
        """
        if df is None or df.empty:
            logging.warning("DataFrame为空，无需写入数据库。")
            return

        self.db_manager.insert_data_to_mysql(df)
    
    def insert_data_to_mysql_fast(self, df: pd.DataFrame) -> None:
        """
        快速插入交易数据（方案2：使用executemany批量插入）
        速度：⭐⭐⭐⭐ 很快
        内存：⭐⭐⭐⭐⭐ 低
        适用场景：中等数据量
        :param df: 交易数据DataFrame
        """
        if df is None or df.empty:
            logging.warning("DataFrame为空，无需写入数据库。")
            return
        
        self.db_manager.insert_data_to_mysql_fast(df)
    
    def insert_data_to_mysql_chunk(self, df: pd.DataFrame, chunk_size: int = 1000) -> None:
        """
        分块插入交易数据（方案3：适合大数据量）
        速度：⭐⭐⭐ 中等
        内存：⭐⭐⭐⭐⭐ 最低
        适用场景：大数据量（>10万条）
        :param df: 交易数据DataFrame
        :param chunk_size: 分块大小
        """
        if df is None or df.empty:
            logging.warning("DataFrame为空，无需写入数据库。")
            return
        
        self.db_manager.insert_data_to_mysql_chunk(df, chunk_size)
    
    def insert_data_to_mysql_simple(self, df: pd.DataFrame) -> None:
        """
        简单插入交易数据（方案4：直接使用to_sql，忽略重复）
        速度：⭐⭐⭐⭐⭐ 最快
        内存：⭐⭐⭐⭐ 中等
        注意：会抛出重复键错误，需要额外处理
        """
        if df is None or df.empty:
            logging.warning("DataFrame为空，无需写入数据库。")
            return
        
        self.db_manager.insert_data_to_mysql_simple(df)
    
    def query_latest_data(self, symbol: str) -> Optional[Dict[str, Any]]:
        """
        查询最新交易数据
        :param symbol: 交易对
        :return: 最新交易记录或None
        """
        sql = """
        SELECT * FROM crypto_trade_data
        WHERE symbol = :symbol
        ORDER BY ts DESC
        LIMIT 1
        """
        condition_dict = {"symbol": symbol}
        return self.db_manager.query_data(sql, condition_dict, return_multi=False)
    
    def query_earliest_data(self, symbol: str) -> Optional[Dict[str, Any]]:
        """
        查询最早交易数据
        :param symbol: 交易对
        :return: 最早交易记录或None
        """
        sql = """
        SELECT * FROM crypto_trade_data
        WHERE symbol = :symbol
        ORDER BY ts ASC
        LIMIT 1
        """
        condition_dict = {"symbol": symbol}
        return self.db_manager.query_data(sql, condition_dict, return_multi=False)
    
    def query_data_by_tradeId(self, tradeId: str) -> Optional[Dict[str, Any]]:
        """
        根据交易ID查询交易数据
        :param tradeId: 交易ID
        :return: 交易记录或None
        """
        sql = """
        SELECT * FROM crypto_trade_data
        WHERE tradeId = :tradeId
        """
        condition_dict = {"tradeId": tradeId}
        return self.db_manager.query_data(sql, condition_dict, return_multi=False)
    
    def query_trade_data_by_symbol(self, symbol: str, start: Optional[Union[str, int]] = None, end: Optional[Union[str, int]] = None) -> Optional[List[Dict[str, Any]]]:
        """
        根据交易对查询交易数据
        :param symbol: 交易对
        :param start: 开始时间
        :param end: 结束时间
        :return: 交易记录列表或None
        """
        conditions, condition_dict = self._build_query_conditions(symbol=symbol, start=start, end=end)
        
        where_clause = " AND ".join(conditions) if conditions else "1=1"
        sql = f"""
            SELECT * FROM crypto_trade_data
            WHERE {where_clause}
            ORDER BY ts ASC
        """
        
        return self.db_manager.query_data(sql, condition_dict, return_multi=True)
    
    def query_trade_data_by_side(self, side: str, symbol: Optional[str] = None, start: Optional[Union[str, int]] = None, end: Optional[Union[str, int]] = None) -> Optional[List[Dict[str, Any]]]:
        """
        根据交易方向查询交易数据
        :param side: 交易方向 (buy/sell)
        :param symbol: 交易对（可选）
        :param start: 开始时间
        :param end: 结束时间
        :return: 交易记录列表或None
        """
        conditions, condition_dict = self._build_query_conditions(symbol=symbol, side=side, start=start, end=end)
        
        where_clause = " AND ".join(conditions)
        sql = f"""
            SELECT * FROM crypto_trade_data
            WHERE {where_clause}
            ORDER BY ts ASC
        """
        
        return self.db_manager.query_data(sql, condition_dict, return_multi=True)
    
    def query_buy_trades(self, symbol: Optional[str] = None, start: Optional[Union[str, int]] = None, end: Optional[Union[str, int]] = None) -> Optional[List[Dict[str, Any]]]:
        """
        查询买入交易记录
        :param symbol: 交易对（可选）
        :param start: 开始时间
        :param end: 结束时间
        :return: 买入交易记录列表或None
        """
        return self.query_trade_data_by_side("buy", symbol, start, end)
    
    def query_sell_trades(self, symbol: Optional[str] = None, start: Optional[Union[str, int]] = None, end: Optional[Union[str, int]] = None) -> Optional[List[Dict[str, Any]]]:
        """
        查询卖出交易记录
        :param symbol: 交易对（可选）
        :param start: 开始时间
        :param end: 结束时间
        :return: 卖出交易记录列表或None
        """
        return self.query_trade_data_by_side("sell", symbol, start, end)
    
    def get_trade_statistics(self, symbol: Optional[str] = None, start: Optional[Union[str, int]] = None, end: Optional[Union[str, int]] = None) -> Optional[Dict[str, Any]]:
        """
        获取交易统计信息
        :param symbol: 交易对（可选）
        :param start: 开始时间
        :param end: 结束时间
        :return: 统计信息或None
        """
        conditions, condition_dict = self._build_query_conditions(symbol=symbol, start=start, end=end)
        
        where_clause = " AND ".join(conditions) if conditions else "1=1"
        sql = f"""
            SELECT 
                COUNT(*) as total_trades,
                SUM(CASE WHEN side = 'buy' THEN 1 ELSE 0 END) as buy_count,
                SUM(CASE WHEN side = 'sell' THEN 1 ELSE 0 END) as sell_count,
                SUM(CASE WHEN side = 'buy' THEN sz ELSE 0 END) as total_buy_volume,
                SUM(CASE WHEN side = 'sell' THEN sz ELSE 0 END) as total_sell_volume,
                SUM(CASE WHEN side = 'buy' THEN sz * px ELSE 0 END) as total_buy_value,
                SUM(CASE WHEN side = 'sell' THEN sz * px ELSE 0 END) as total_sell_value,
                AVG(px) as avg_price,
                MIN(px) as min_price,
                MAX(px) as max_price,
                MIN(ts) as first_trade_time,
                MAX(ts) as last_trade_time
            FROM crypto_trade_data
            WHERE {where_clause}
        """
        
        return self.db_manager.query_data(sql, condition_dict, return_multi=False)
    
    def get_volume_price_analysis(self, symbol: Optional[str] = None, start: Optional[Union[str, int]] = None, end: Optional[Union[str, int]] = None) -> Optional[List[Dict[str, Any]]]:
        """
        获取成交量价格分析
        :param symbol: 交易对（可选）
        :param start: 开始时间
        :param end: 结束时间
        :return: 分析结果列表或None
        """
        conditions, condition_dict = self._build_query_conditions(symbol=symbol, start=start, end=end)
        
        where_clause = " AND ".join(conditions) if conditions else "1=1"
        sql = f"""
            SELECT 
                side,
                COUNT(*) as trade_count,
                SUM(sz) as total_volume,
                SUM(sz * px) as total_value,
                AVG(px) as avg_price,
                MIN(px) as min_price,
                MAX(px) as max_price,
                AVG(sz) as avg_volume
            FROM crypto_trade_data
            WHERE {where_clause}
            GROUP BY side
            ORDER BY side
        """
        
        return self.db_manager.query_data(sql, condition_dict, return_multi=True)
    
    def get_recent_trades(self, symbol: Optional[str] = None, limit: int = 100) -> Optional[List[Dict[str, Any]]]:
        """
        获取最近的交易记录
        :param symbol: 交易对（可选）
        :param limit: 返回记录数量
        :return: 最近交易记录列表或None
        """
        conditions = []
        condition_dict = {}
        
        if symbol:
            conditions.append("symbol = :symbol")
            condition_dict["symbol"] = symbol
        
        where_clause = " AND ".join(conditions) if conditions else "1=1"
        sql = f"""
            SELECT * FROM crypto_trade_data
            WHERE {where_clause}
            ORDER BY ts DESC
            LIMIT :limit
        """
        condition_dict["limit"] = limit
        
        return self.db_manager.query_data(sql, condition_dict, return_multi=True)
    
    def get_trades_by_price_range(self, min_price: float, max_price: float, symbol: Optional[str] = None, start: Optional[Union[str, int]] = None, end: Optional[Union[str, int]] = None) -> Optional[List[Dict[str, Any]]]:
        """
        根据价格范围查询交易记录
        :param min_price: 最低价格
        :param max_price: 最高价格
        :param symbol: 交易对（可选）
        :param start: 开始时间
        :param end: 结束时间
        :return: 交易记录列表或None
        """
        conditions, condition_dict = self._build_query_conditions(symbol=symbol, start=start, end=end)
        
        conditions.append("px BETWEEN :min_price AND :max_price")
        condition_dict["min_price"] = min_price
        condition_dict["max_price"] = max_price
        
        where_clause = " AND ".join(conditions)
        sql = f"""
            SELECT * FROM crypto_trade_data
            WHERE {where_clause}
            ORDER BY ts ASC
        """
        
        return self.db_manager.query_data(sql, condition_dict, return_multi=True)
    
    def get_trades_by_volume_range(self, min_volume: float, max_volume: float, symbol: Optional[str] = None, start: Optional[Union[str, int]] = None, end: Optional[Union[str, int]] = None) -> Optional[List[Dict[str, Any]]]:
        """
        根据成交量范围查询交易记录
        :param min_volume: 最低成交量
        :param max_volume: 最高成交量
        :param symbol: 交易对（可选）
        :param start: 开始时间
        :param end: 结束时间
        :return: 交易记录列表或None
        """
        conditions, condition_dict = self._build_query_conditions(symbol=symbol, start=start, end=end)
        
        conditions.append("sz BETWEEN :min_volume AND :max_volume")
        condition_dict["min_volume"] = min_volume
        condition_dict["max_volume"] = max_volume
        
        where_clause = " AND ".join(conditions)
        sql = f"""
            SELECT * FROM crypto_trade_data
            WHERE {where_clause}
            ORDER BY ts ASC
        """
        
        return self.db_manager.query_data(sql, condition_dict, return_multi=True)