Agent/examples/auto_put_ad_mini/db/connection.py

"""数据库连接管理

提供统一的数据库连接接口，支持环境变量配置
参考 content_finder/db/connection.py
"""

import os
import logging
from typing import Optional

try:
    import pymysql
    import pymysql.cursors
except ImportError:
    pymysql = None

logger = logging.getLogger(__name__)


def get_connection():
    """获取数据库连接

    从环境变量读取配置：
    - DB_HOST: 数据库主机地址
    - DB_PORT: 数据库端口（默认3306）
    - DB_USER: 数据库用户名
    - DB_PASSWORD: 数据库密码
    - DB_NAME: 数据库名称

    Returns:
        pymysql.Connection: 数据库连接对象

    Raises:
        ImportError: pymysql 未安装
        ValueError: 数据库配置缺失
        Exception: 连接失败
    """
    if pymysql is None:
        raise ImportError(
            "pymysql 未安装，请运行: pip install pymysql\n"
            "或在 requirements.txt 中添加 pymysql>=1.0.0"
        )

    # 读取环境变量
    host = os.getenv("DB_HOST", "").strip()
    port = int(os.getenv("DB_PORT", "3306"))
    user = os.getenv("DB_USER", "").strip()
    password = os.getenv("DB_PASSWORD", "")
    database = os.getenv("DB_NAME", "").strip()

    # 验证必需配置
    if not all([host, user, database]):
        raise ValueError(
            "数据库配置缺失！请在 .env 文件或环境变量中设置：\n"
            "  DB_HOST=数据库主机地址\n"
            "  DB_USER=数据库用户名\n"
            "  DB_PASSWORD=数据库密码\n"
            "  DB_NAME=数据库名称\n"
            "  DB_PORT=3306  # 可选，默认3306"
        )

    try:
        conn = pymysql.connect(
            host=host,
            port=port,
            user=user,
            password=password,
            database=database,
            charset="utf8mb4",
            cursorclass=pymysql.cursors.DictCursor,  # 返回字典格式
            autocommit=True,  # 自动提交（简化事务管理）
        )
        logger.debug(f"数据库连接成功: {user}@{host}:{port}/{database}")
        return conn
    except Exception as e:
        logger.error(f"数据库连接失败: {e}")
        raise


def test_connection() -> bool:
    """测试数据库连接

    Returns:
        bool: 连接成功返回 True，失败返回 False
    """
    try:
        conn = get_connection()
        with conn.cursor() as cursor:
            cursor.execute("SELECT 1")
            result = cursor.fetchone()
        conn.close()
        logger.info("数据库连接测试成功")
        return True
    except Exception as e:
        logger.error(f"数据库连接测试失败: {e}")
        return False


if __name__ == "__main__":
    # 测试数据库连接
    logging.basicConfig(level=logging.INFO)
    if test_connection():
        print("✅ 数据库连接正常")
    else:
        print("❌ 数据库连接失败，请检查配置")