video_detect/ds/db.py

import mysql.connector
from mysql.connector import Error
from .config import MYSQL_CONFIG

_connection_pool = None


class Database:
    """MySQL 连接池管理类"""
    pool_config = {
        "host": MYSQL_CONFIG.get("host", "localhost"),
        "port": int(MYSQL_CONFIG.get("port", 3306)),
        "user": MYSQL_CONFIG.get("user", "root"),
        "password": MYSQL_CONFIG.get("password", ""),
        "database": MYSQL_CONFIG.get("database", ""),
        "charset": MYSQL_CONFIG.get("charset", "utf8mb4"),
        "pool_name": "fastapi_pool",
        "pool_size": 5,
        "pool_reset_session": True
    }

    @classmethod
    def get_connection(cls):
        """获取数据库连接"""
        try:
            # 从连接池获取连接
            conn = mysql.connector.connect(**cls.pool_config)
            if conn.is_connected():
                return conn
        except Error as e:
            raise Exception(f"MySQL 连接失败: {str(e)}") from e

    @classmethod
    def close_connection(cls, conn, cursor=None):
        """关闭连接和游标"""
        try:
            if cursor:
                cursor.close()
            if conn and conn.is_connected():
                conn.close()
        except Error as e:
            raise Exception(f"MySQL 连接关闭失败: {str(e)}") from e

    @classmethod
    def close_all_connections(cls):
        """清理连接池（服务重启前调用）"""
        try:
            # 先检查属性是否存在、再判断是否有值
            if hasattr(cls, "_connection_pool") and cls._connection_pool:
                cls._connection_pool = None  # 重置连接池
                print("[Database] 连接池已重置、旧连接将被自动清理")
            else:
                print("[Database] 连接池未初始化或已重置、无需操作")
        except Exception as e:
            print(f"[Database] 重置连接池失败: {str(e)}")


# 暴露数据库操作工具
db = Database()