backtrader/database_manager.py

import os
import traceback

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

from logger import get_logger
from utils import load_config


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

    def __init__(self):
        self.config = load_config()
        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['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 get_table_name(self, key):
        """根据表键名获取实际表名"""
        return self.config['sqlite']['table_name'].get(key, key)

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

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

        table_name = self.get_table_name(table_key)
        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_key, conditions=None):
        """
        从数据库中加载数据
        参数:
        table_key (str): 表键名
        conditions (str): 过滤条件，如 "trade_date > '20230101'"
        返回:
        pandas.DataFrame: 查询结果
        """
        table_name = self.get_table_name(table_key)
        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_key, if_exists='append'):
        """
        保存DataFrame到数据库
        参数:
        df (pandas.DataFrame): 要保存的数据
        table_key (str): 表键名
        if_exists (str): 如果表存在时的操作: 'fail', 'replace', 或 'append'
        返回:
        bool: 操作是否成功
        """
        if df.empty:
            self.logger.warning("警告: 尝试保存空的DataFrame到数据库")
            return False

        table_name = self.get_table_name(table_key)
        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