doris-sync/config/loader.py

#!/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()