#!/usr/bin/env python # -*- coding: utf-8 -*- """ 配置加载器 统一从 config.yaml 加载配置 支持通过环境变量覆盖 """ import os import yaml from typing import Any, Dict, Optional # 项目根目录 BASE_DIR = os.path.dirname( os.path.dirname(os.path.abspath(__file__)) ) # 环境变量映射 ENV_MAPPING: Dict[str, str] = { # MySQL "MYSQL_HOST": "database.mysql.host", "MYSQL_PORT": "database.mysql.port", "MYSQL_USER": "database.mysql.user", "MYSQL_PASSWORD": "database.mysql.password", "MYSQL_DATABASE": "database.mysql.database", "MYSQL_POOL_SIZE": "database.mysql.pool_size", # Doris "DORIS_HOST": "database.doris.host", "DORIS_PORT": "database.doris.port", "DORIS_USER": "database.doris.user", "DORIS_PASSWORD": "database.doris.password", "DORIS_DATABASE": "database.doris.database", "DORIS_POOL_SIZE": "database.doris.pool_size", # PostgreSQL "PGSQL_HOST": "database.pgsql.host", "PGSQL_PORT": "database.pgsql.port", "PGSQL_USER": "database.pgsql.user", "PGSQL_PASSWORD": "database.pgsql.password", "PGSQL_DATABASE": "database.pgsql.database", "PGSQL_POOL_SIZE": "database.pgsql.pool_size", } # 类型转换 TYPE_CONVERTERS: Dict[str, type] = { "port": int, "pool_size": int, } def _get_base_dir() -> str: """获取项目根目录""" return BASE_DIR def _load_yaml() -> Dict[str, Any]: """加载 YAML 配置""" yaml_path = os.path.join(_get_base_dir(), "config.yaml") if not os.path.exists(yaml_path): return {} with open(yaml_path, "r", encoding="utf-8") as f: content = yaml.safe_load(f) return content if content else {} def _set_nested(config: Dict[str, Any], path: str, value: str) -> None: """设置嵌套配置""" keys = path.split(".") current = config for key in keys[:-1]: if key not in current: current[key] = {} current = current[key] last_key = keys[-1] # 类型转换 for type_key, converter in TYPE_CONVERTERS.items(): if type_key in last_key.lower() and isinstance(value, str): try: value = converter(value) except (ValueError, TypeError): pass current[last_key] = value def load_config() -> Dict[str, Any]: """ 加载配置 优先级: 环境变量 > config.yaml Returns: 配置字典 """ # 1. 加载 YAML 作为基础 config = _load_yaml() # 2. 环境变量覆盖 for env_key, config_path in ENV_MAPPING.items(): env_value = os.getenv(env_key) if env_value is not None: _set_nested(config, config_path, env_value) return config # 全局缓存 _config_cache: Optional[Dict[str, Any]] = None def get_config(key: Optional[str] = None, default: Any = None) -> Any: """ 获取配置值 Args: key: 配置路径,如 "database.mysql.host",None 返回全部 default: 默认值 Returns: 配置值 """ global _config_cache if _config_cache is None: _config_cache = load_config() if key is None: return _config_cache keys = key.split(".") current = _config_cache for k in keys: if isinstance(current, dict) and k in current: current = current[k] else: return default return current def reload_config() -> Dict[str, Any]: """重新加载配置""" global _config_cache _config_cache = None return load_config()