kg_agent/utils/mysql/db_config.py

import json
import os
from typing import Dict, Any
from dotenv import load_dotenv, find_dotenv


class DatabaseConfig:
    """数据库配置管理类"""

    def __init__(self, env_file: str = '.env'):
        self.env_file = env_file
        self._config = None

    def load_config(self) -> Dict[str, Any]:
        """从.env文件加载数据库配置"""
        if self._config is not None:
            return self._config

        # 使用 python-dotenv 加载环境变量
        # find_dotenv() 会自动查找 .env 文件
        env_path = find_dotenv(self.env_file)

        if not env_path:
            # 手动尝试多个可能的.env文件位置
            possible_paths = [
                # 当前工作目录
                os.path.join(os.getcwd(), self.env_file),
                # 项目根目录（从当前文件位置推算）
                os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(__file__))), self.env_file),
                # mysql目录的上级目录的上级目录
                os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))), self.env_file)
            ]

            for path in possible_paths:
                if os.path.exists(path):
                    env_path = path
                    break

            if not env_path:
                searched_paths = '\n'.join(possible_paths)
                raise FileNotFoundError(f"配置文件 {self.env_file} 不存在，已搜索以下路径:\n{searched_paths}")

        # 加载环境变量
        load_dotenv(env_path)

        # 从环境变量中读取 DB_INFO
        db_info_str = os.getenv('DB_INFO')
        if not db_info_str:
            raise ValueError("未找到DB_INFO配置")

        # 解析 JSON 格式的数据库配置
        try:
            self._config = json.loads(db_info_str)
        except json.JSONDecodeError as e:
            raise ValueError(f"DB_INFO配置格式错误: {e}")

        # 验证必要的配置项
        required_keys = ['host', 'database', 'user', 'passwd']
        for key in required_keys:
            if key not in self._config:
                raise ValueError(f"缺少必要的配置项: {key}")

        # 设置默认值
        self._config.setdefault('port', 3306)
        self._config.setdefault('charset', 'utf8')

        return self._config

    def get_connection_params(self) -> Dict[str, Any]:
        """获取数据库连接参数"""
        config = self.load_config()
        return {
            'host': config['host'],
            'port': config['port'],
            'user': config['user'],
            'password': config['passwd'],
            'database': config['database'],
            'charset': config['charset'],
            'autocommit': False,
            'cursorclass': None  # 将在连接池中设置
        }


# 全局配置实例
db_config = DatabaseConfig()