backtrader/database_manager.py

import os
import traceback

import pandas as pd
from sqlalchemy import create_engine, text, inspect

from config_manager import get_config_manager
from logger import get_logger


class DatabaseManager:
    """
    数据库管理类，负责数据的存储、查询和管理
    """

    def __init__(self):
        self.config = get_config_manager()
        self._engine = None
        self.logger = get_logger()

    def get_engine(self):
        """获取SQLite数据库引擎，如果不存在则创建"""
        if self._engine is not None:
            return self._engine

        db_path = self.config.get('sqlite.path')

        # 确保数据库目录存在
        os.makedirs(os.path.dirname(db_path), exist_ok=True)

        # 创建SQLite数据库引擎
        self._engine = create_engine(f'sqlite:///{db_path}', echo=False)
        return self._engine

    def table_exists(self, table_name):
        """
        检查表是否存在
        参数:
        table_name (str): 数据表名称
        返回:
        bool: 表是否存在
        """
        engine = self.get_engine()
        inspector = inspect(engine)
        return inspector.has_table(table_name)

    def get_existing_trade_dates(self, table_name):
        """
        从数据库中获取已有的交易日期
        参数:
        table_name (str): 数据表名称
        返回:
        set: 已存在于数据库中的交易日期集合
        """
        # 先检查表是否存在
        if not self.table_exists(table_name):
            self.logger.debug(f"表 '{table_name}' 不存在")
            return set()

        engine = self.get_engine()
        query = f"SELECT DISTINCT trade_date FROM {table_name}"
        try:
            with engine.connect() as connection:
                result = connection.execute(text(query))
                return {row[0] for row in result}
        except Exception as e:
            self.logger.error(f"获取已存在交易日期时出错: {e}")
            self.logger.debug(f"完整的错误追踪信息:\n{traceback.format_exc()}")
            return set()

    def load_df_from_db(self, table_name, conditions=None):
        """
        从数据库中加载数据
        参数:
        table_name (str): 数据表名称
        conditions (str): 过滤条件，如 "trade_date > '20230101'"
        返回:
        pandas.DataFrame: 查询结果
        """
        engine = self.get_engine()
        query = f"SELECT * FROM {table_name}"
        if conditions:
            query += f" WHERE {conditions}"
        try:
            return pd.read_sql(query, engine)
        except Exception as e:
            self.logger.error(f"从数据库加载数据时出错: {e}")
            self.logger.debug(f"完整的错误追踪信息:\n{traceback.format_exc()}")
            return pd.DataFrame()

    def save_df_to_db(self, df, table_name, if_exists='append'):
        """
        保存DataFrame到数据库
        参数:
        df (pandas.DataFrame): 要保存的数据
        table_name (str): 数据表名称
        if_exists (str): 如果表存在时的操作: 'fail', 'replace', 或 'append'
        返回:
        bool: 操作是否成功
        """
        if df.empty:
            self.logger.warning("警告: 尝试保存空的DataFrame到数据库")
            return False

        engine = self.get_engine()
        try:
            df.to_sql(table_name, engine, if_exists=if_exists, index=False)
            return True
        except Exception as e:
            self.logger.error(f"保存数据到数据库时出错: {e}")
            self.logger.debug(f"完整的错误追踪信息:\n{traceback.format_exc()}")
            return False

    def delete_existing_data_by_dates(self, table_name, trade_dates):
        """
        从数据库表中删除指定交易日期的数据

        参数:
        table_name (str): 数据表名称
        trade_dates (list): 需要删除的交易日期列表

        返回:
        bool: 操作是否成功
        """
        if not trade_dates:
            return True  # 如果没有日期需要删除，认为操作成功

        engine = self.get_engine()

        try:
            # 将日期列表转换为SQL安全的格式
            date_strings = [f"'{date}'" for date in trade_dates]
            dates_clause = ", ".join(date_strings)
            delete_query = f"DELETE FROM {table_name} WHERE trade_date IN ({dates_clause})"

            with engine.connect() as connection:
                connection.execute(text(delete_query))
                connection.commit()

            self.logger.info(f"成功从表 '{table_name}' 中删除 {len(trade_dates)} 个日期的数据")
            return True
        except Exception as e:
            self.logger.error(f"删除表 '{table_name}' 中的数据时出错: {e}")
            self.logger.debug(f"完整的错误追踪信息:\n{traceback.format_exc()}")
            return False