crypto_quant/core/db_manager.py

import pandas as pd
from sqlalchemy import create_engine, exc, text
import re, datetime
import logging
from core.utils import (
    transform_data_type,
    datetime_to_timestamp,
    check_date_time_format,
)

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


class DBData:
    def __init__(
        self, 
        db_url: str, 
        table_name: str = "crypto_market_data",
        columns: list = None
    ):
        self.db_url = db_url
        self.table_name = table_name
        self.temp_table_name = f"temp_{table_name}_{datetime.datetime.now().strftime('%Y%m%d_%H%M%S')}"
        self.columns = columns
        if self.columns is None:
            raise ValueError("columns不能为空")
        if len(self.columns) != len(set(self.columns)):
            raise ValueError("columns不能有重复")
    
    def create_insert_sql_text_by_temp_table(self, temp_table_name: str):
        """
        创建插入SQL语句（使用临时表）
        示例：
        INSERT INTO crypto_market_data
        (symbol, bar, timestamp, date_time, open, high, low, close, volume, volCcy, volCCyQuote, create_time)
        SELECT symbol, bar, timestamp, date_time, open, high, low, close, volume, volCcy, volCCyQuote, create_time
        FROM {temp_table_name}
        ON DUPLICATE KEY UPDATE
            open=VALUES(open), high=VALUES(high), low=VALUES(low), close=VALUES(close),
            volume=VALUES(volume), volCcy=VALUES(volCcy), volCCyQuote=VALUES(volCCyQuote), 
            date_time=VALUES(date_time), create_time=VALUES(create_time)
        """
        sql = f"""
        INSERT INTO {self.table_name}
        ({",".join(self.columns)})
        SELECT {",".join(self.columns)}
        FROM {temp_table_name}
        ON DUPLICATE KEY UPDATE
            {", ".join([f"{col}=VALUES({col})" for col in self.columns])}
        """
        return sql
    
    def create_insert_sql_text(self):
        """
        创建插入SQL语句（不使用临时表）
        示例：
        INSERT INTO crypto_market_data
        (symbol, bar, timestamp, date_time, open, high, low, close, volume, volCcy, volCCyQuote, create_time)
        VALUES (:symbol, :bar, :timestamp, :date_time, :open, :high, :low, :close, :volume, :volCcy, :volCCyQuote, :create_time)
        ON DUPLICATE KEY UPDATE
            open=VALUES(open), high=VALUES(high), low=VALUES(low), close=VALUES(close),
            volume=VALUES(volume), volCcy=VALUES(volCcy), volCCyQuote=VALUES(volCCyQuote), 
            date_time=VALUES(date_time), create_time=VALUES(create_time)
        """
        sql = f"""
            INSERT INTO {self.table_name}
            ({",".join(self.columns)})
            VALUES ({",".join([f":{col}" for col in self.columns])})
            ON DUPLICATE KEY UPDATE
                {", ".join([f"{col}=VALUES({col})" for col in self.columns])}
        """
        return sql

    def insert_data_to_mysql(self, df: pd.DataFrame):
        """
        将K线行情数据保存到MySQL的crypto_market_data表
        速度：⭐⭐⭐⭐⭐ 最快
        内存：⭐⭐⭐⭐ 中等
        适用场景：中小数据量（<10万条）
        :param df: K线数据DataFrame
        :param symbol: 交易对
        :param bar: K线周期
        :param db_url: 数据库连接URL
        """
        if df is None or df.empty:
            logging.warning("DataFrame为空，无需写入数据库。")
            return

        df = df[self.columns]
        # 建立数据库连接
        try:
            engine = create_engine(self.db_url)

            # 方案1：使用临时表 + 批量更新（推荐，速度最快）
            with engine.begin() as conn:
                # 将数据写入临时表
                df.to_sql(
                    name=self.temp_table_name,
                    con=engine,
                    if_exists="replace",
                    index=False,
                    method="multi",
                )

                # 使用INSERT ... ON DUPLICATE KEY UPDATE批量处理
                sql = text(self.create_insert_sql_text_by_temp_table(self.temp_table_name))
                conn.execute(sql)

                # 删除临时表
                conn.execute(text(f"DROP TABLE IF EXISTS {self.temp_table_name}"))

            logging.info("数据已成功写入数据库。")
        except Exception as e:
            logging.error(f"数据库连接或写入失败: {e}")
    
    def insert_data_to_mysql_fast(self, df: pd.DataFrame):
        """
        快速插入K线行情数据（方案2：使用executemany批量插入）
        速度：⭐⭐⭐⭐ 很快
        内存：⭐⭐⭐⭐⭐ 低
        适用场景：中等数据量
        """
        if df is None or df.empty:
            logging.warning("DataFrame为空，无需写入数据库。")
            return
        
        df = df[self.columns]
        try:
            engine = create_engine(self.db_url)
            with engine.begin() as conn:
                # 使用executemany批量插入
                sql = text(self.create_insert_sql_text())

                # 将DataFrame转换为字典列表
                data_dicts = [row.to_dict() for _, row in df.iterrows()]
                conn.execute(sql, data_dicts)

            logging.info("数据已成功写入数据库。")
        except Exception as e:
            logging.error(f"数据库连接或写入失败: {e}")
    
    def insert_data_to_mysql_chunk(self, df: pd.DataFrame, chunk_size: int = 1000):
        """
        分块插入K线行情数据（方案3：适合大数据量）
        速度：⭐⭐⭐ 中等
        内存：⭐⭐⭐⭐⭐ 最低
        适用场景：大数据量（>10万条）
        """
        if df is None or df.empty:
            logging.warning("DataFrame为空，无需写入数据库。")
            return
        
        df = df[self.columns]
        
        try:
            engine = create_engine(self.db_url)
            with engine.begin() as conn:
                # 分块处理
                total_rows = len(df)
                for i in range(0, total_rows, chunk_size):
                    chunk_df = df.iloc[i : i + chunk_size]
                    with engine.begin() as conn:
                        # 创建临时表
                        temp_table_name = f"{self.temp_table_name}_{i}"
                        
                        # 将数据写入临时表
                        chunk_df.to_sql(
                            name=temp_table_name,
                            con=engine,
                            if_exists="replace",
                            index=False,
                            method="multi",
                        )
                        
                        # 使用INSERT ... ON DUPLICATE KEY UPDATE批量处理
                        sql = text(self.create_insert_sql_text_by_temp_table(temp_table_name))
                        conn.execute(sql)
                        
                        # 删除临时表
                        conn.execute(text(f"DROP TABLE IF EXISTS {temp_table_name}"))
                        
                        logging.info(f"已处理 {min(i+chunk_size, total_rows)}/{total_rows} 条记录")
                logging.info("数据已成功写入数据库。")
        except Exception as e:
            logging.error(f"数据库连接或写入失败: {e}")
    
    def insert_data_to_mysql_simple(self, df: pd.DataFrame):
        """
        简单插入K线行情数据（方案4：直接使用to_sql，忽略重复）
        速度：⭐⭐⭐⭐⭐ 最快
        内存：⭐⭐⭐⭐ 中等
        注意：会抛出重复键错误，需要额外处理
        """
        if df is None or df.empty:
            logging.warning("DataFrame为空，无需写入数据库。")
            return
        
        df = df[self.columns]
        try:
            engine = create_engine(self.db_url)
            with engine.begin() as conn:
                df.to_sql(
                    name=self.table_name,
                    con=engine,
                    if_exists="append",
                    index=False,
                    method="multi",
                )
                logging.info("数据已成功写入数据库。")
        except Exception as e:
            logging.error(f"数据库连接或写入失败: {e}")

    def query_data(self, sql: str, condition_dict: dict, return_multi: bool = True):
        """
        查询数据
        :param sql: 查询SQL
        :param db_url: 数据库连接URL
        """
        try:
            engine = create_engine(self.db_url)
            with engine.connect() as conn:
                result = conn.execute(text(sql), condition_dict)
                if return_multi:
                    result = result.fetchall()
                    if result:
                        result_list = [
                            transform_data_type(dict(row._mapping)) for row in result
                        ]
                        return result_list
                    else:
                        return None
                else:
                    result = result.fetchone()
                    if result:
                        result_dict = transform_data_type(dict(result._mapping))
                        return result_dict
                    else:
                        return None
        except Exception as e:
            logging.error(f"查询数据出错: {e}")
            return None