AI_python_yoooger/ai2/sqlhelp.py

import yaml
import psycopg2

def read_sql_config(yaml_name):
    """
    读取 SQL 配置
    """
    yaml_path = f"{yaml_name}.yaml"
    with open(yaml_path, 'r', encoding='utf-8') as f:
        config = yaml.safe_load(f)
    sql_config = config.get('sql')
    if not sql_config:
        raise ValueError("未找到 'sql' 配置块")
    return sql_config

def get_user_power(user_id: str, yaml_name: str) -> int | None:
    """
    根据 user_id 查询数据库中对应的 power，找不到返回 None
    """
    conn = None
    try:
        sql_config = read_sql_config(yaml_name)
        conn = psycopg2.connect(
            dbname=sql_config['dbname'],
            user=sql_config['user'],
            password=sql_config['password'],
            host=sql_config['host'],
            port=sql_config['port']
        )
        
        conn.autocommit = True  # 只读查询推荐设置
        with conn.cursor() as cur:
            cur.execute(" WHERE user_id = %s", (user_id,))
            row = cur.fetchone()
            if row:
                return row[0]
            else:
                return None
    except Exception as e:
        print(f"数据库操作异常: {e}")
        return None
    finally:
        if conn:
            conn.close()