Agent/examples/auto_put_ad_mini/db/config.py

"""系统配置和白名单数据库操作

提供：
- 白名单账户查询
- 系统配置读写
- 配置缓存机制（避免频繁查询数据库）
"""

import json
import logging
from typing import Any, Dict, List, Optional
from datetime import datetime, timedelta

from .connection import get_connection

logger = logging.getLogger(__name__)

# 配置缓存（5分钟过期）
_config_cache: Dict[str, tuple[Any, datetime]] = {}
_cache_ttl = timedelta(minutes=5)


def get_whitelist_accounts() -> List[int]:
    """获取启用的白名单账户列表

    Returns:
        List[int]: 账户ID列表，如 [80769799, 71305011]

    Raises:
        Exception: 数据库查询失败
    """
    conn = None
    try:
        conn = get_connection()
        with conn.cursor() as cursor:
            sql = """
            SELECT account_id
            FROM account_whitelist
            WHERE enabled = TRUE
            ORDER BY id ASC
            """
            cursor.execute(sql)
            rows = cursor.fetchall()
            account_ids = [row["account_id"] for row in rows]
            logger.info(f"从数据库读取白名单账户：{len(account_ids)} 个")
            return account_ids
    except Exception as e:
        logger.error(f"查询白名单账户失败: {e}")
        raise
    finally:
        if conn:
            conn.close()


def get_system_config(key: str, default: Any = None, use_cache: bool = True) -> Any:
    """获取系统配置值

    Args:
        key: 配置键，如 'execution_enabled', 'cron_schedule'
        default: 默认值（配置不存在时返回）
        use_cache: 是否使用缓存（默认True）

    Returns:
        Any: 配置值，根据 value_type 自动转换类型

    Examples:
        >>> get_system_config('execution_enabled')  # 返回 False (boolean)
        >>> get_system_config('cron_schedule')  # 返回 '0 2 * * *' (string)
        >>> get_system_config('roi_low_factor')  # 返回 '0.75' (string)
    """
    # 检查缓存
    if use_cache and key in _config_cache:
        value, cached_at = _config_cache[key]
        if datetime.now() - cached_at < _cache_ttl:
            logger.debug(f"从缓存读取配置: {key} = {value}")
            return value

    conn = None
    try:
        conn = get_connection()
        with conn.cursor() as cursor:
            sql = """
            SELECT config_value, value_type
            FROM system_config
            WHERE config_key = %s AND enabled = TRUE
            """
            cursor.execute(sql, (key,))
            row = cursor.fetchone()

            if row is None:
                logger.warning(f"配置不存在: {key}，使用默认值: {default}")
                return default

            value = row["config_value"]
            value_type = row["value_type"]

            # 类型转换
            if value_type == "boolean":
                parsed_value = value.lower() in ("true", "1", "yes")
            elif value_type == "int":
                parsed_value = int(value)
            elif value_type == "json":
                parsed_value = json.loads(value)
            else:  # string
                parsed_value = value

            # 更新缓存
            _config_cache[key] = (parsed_value, datetime.now())
            logger.debug(f"从数据库读取配置: {key} = {parsed_value}")
            return parsed_value

    except Exception as e:
        logger.error(f"查询系统配置失败: {key}, {e}")
        return default
    finally:
        if conn:
            conn.close()


def update_system_config(
    key: str, value: Any, value_type: str = None, updated_by: str = None
) -> bool:
    """更新系统配置

    Args:
        key: 配置键
        value: 配置值
        value_type: 值类型（string, boolean, int, json），不指定则自动推断
        updated_by: 更新人

    Returns:
        bool: 更新成功返回 True

    Examples:
        >>> update_system_config('execution_enabled', False)
        >>> update_system_config('cron_schedule', '0 3 * * *')
    """
    # 自动推断类型
    if value_type is None:
        if isinstance(value, bool):
            value_type = "boolean"
        elif isinstance(value, int):
            value_type = "int"
        elif isinstance(value, (dict, list)):
            value_type = "json"
        else:
            value_type = "string"

    # 序列化值
    if value_type == "boolean":
        serialized_value = "true" if value else "false"
    elif value_type == "json":
        serialized_value = json.dumps(value, ensure_ascii=False)
    else:
        serialized_value = str(value)

    conn = None
    try:
        conn = get_connection()
        with conn.cursor() as cursor:
            sql = """
            UPDATE system_config
            SET config_value = %s,
                value_type = %s,
                updated_by = %s,
                updated_at = CURRENT_TIMESTAMP
            WHERE config_key = %s
            """
            affected_rows = cursor.execute(
                sql, (serialized_value, value_type, updated_by, key)
            )

            if affected_rows == 0:
                logger.warning(f"配置键不存在，尝试插入: {key}")
                insert_sql = """
                INSERT INTO system_config (config_key, config_value, value_type, updated_by)
                VALUES (%s, %s, %s, %s)
                """
                cursor.execute(insert_sql, (key, serialized_value, value_type, updated_by))

            # 清除缓存
            if key in _config_cache:
                del _config_cache[key]

            logger.info(f"更新系统配置: {key} = {value} (by {updated_by})")
            return True

    except Exception as e:
        logger.error(f"更新系统配置失败: {key}, {e}")
        return False
    finally:
        if conn:
            conn.close()


def get_all_system_configs() -> Dict[str, Any]:
    """获取所有启用的系统配置

    Returns:
        Dict[str, Any]: 配置字典，key为配置键，value为配置值

    Examples:
        >>> configs = get_all_system_configs()
        >>> print(configs)
        {'execution_enabled': False, 'cron_schedule': '0 2 * * *', ...}
    """
    conn = None
    try:
        conn = get_connection()
        with conn.cursor() as cursor:
            sql = """
            SELECT config_key, config_value, value_type
            FROM system_config
            WHERE enabled = TRUE
            ORDER BY id ASC
            """
            cursor.execute(sql)
            rows = cursor.fetchall()

            configs = {}
            for row in rows:
                key = row["config_key"]
                value = row["config_value"]
                value_type = row["value_type"]

                # 类型转换
                if value_type == "boolean":
                    configs[key] = value.lower() in ("true", "1", "yes")
                elif value_type == "int":
                    configs[key] = int(value)
                elif value_type == "json":
                    configs[key] = json.loads(value)
                else:
                    configs[key] = value

            logger.info(f"读取所有系统配置：{len(configs)} 项")
            return configs

    except Exception as e:
        logger.error(f"查询所有系统配置失败: {e}")
        return {}
    finally:
        if conn:
            conn.close()


def clear_config_cache():
    """清除配置缓存（用于测试或强制刷新）"""
    global _config_cache
    _config_cache = {}
    logger.info("配置缓存已清除")


if __name__ == "__main__":
    # 测试配置读取
    logging.basicConfig(level=logging.INFO)

    print("\n=== 测试白名单查询 ===")
    try:
        accounts = get_whitelist_accounts()
        print(f"白名单账户: {accounts}")
    except Exception as e:
        print(f"❌ 白名单查询失败: {e}")

    print("\n=== 测试系统配置查询 ===")
    try:
        execution_enabled = get_system_config("execution_enabled")
        cron_schedule = get_system_config("cron_schedule")
        roi_low_factor = get_system_config("roi_low_factor")
        print(f"执行开关: {execution_enabled} ({type(execution_enabled).__name__})")
        print(f"定时调度: {cron_schedule}")
        print(f"关停线系数: {roi_low_factor}")
    except Exception as e:
        print(f"❌ 系统配置查询失败: {e}")

    print("\n=== 测试所有配置查询 ===")
    try:
        all_configs = get_all_system_configs()
        for key, value in all_configs.items():
            print(f"  {key}: {value} ({type(value).__name__})")
    except Exception as e:
        print(f"❌ 所有配置查询失败: {e}")