project/mt_V2/db.py

import os
import pymysql
import threading
import time
import atexit
from collections import deque
from pymysql.cursors import DictCursor

# ============================================================
# 数据库配置
# ============================================================
DB_HOST = "120.24.26.108"
DB_PORT = 3307
DB_USER = "collect_user"
DB_PASSWORD = "collect123456"
DB_NAME = "drug_retrieve"      # 正式库；测试库改 drug_retrieve_test
DB_CHARSET = "utf8mb4"

# ============================================================
# 连接池
# ============================================================
MYSQL_POOL = None
POOL_LOCK = threading.Lock()


class PooledConnection:
    def __init__(self, conn, pool):
        self._conn = conn
        self._pool = pool
        self._closed = False

    def cursor(self, *args, **kwargs):
        return self._conn.cursor(*args, **kwargs)

    def commit(self):
        return self._conn.commit()

    def rollback(self):
        return self._conn.rollback()

    def ping(self, *args, **kwargs):
        return self._conn.ping(*args, **kwargs)

    def close(self):
        if self._closed:
            return
        self._closed = True
        try:
            self._pool.putconn(self._conn)
        except Exception:
            try:
                self._conn.close()
            except Exception:
                pass

    def __getattr__(self, item):
        return getattr(self._conn, item)


class MySQLConnectionPool:
    def __init__(self, minconn=1, maxconn=10, **connect_kwargs):
        self.minconn = int(minconn)
        self.maxconn = int(maxconn)
        self.connect_kwargs = dict(connect_kwargs)
        self._lock = threading.Lock()
        self._pool = deque()
        self._total = 0
        for _ in range(self.minconn):
            conn = pymysql.connect(**self.connect_kwargs)
            self._pool.append(conn)
            self._total += 1

    def getconn(self, timeout=None):
        start = time.time()
        while True:
            conn = None
            with self._lock:
                if self._pool:
                    conn = self._pool.popleft()
                elif self._total < self.maxconn:
                    conn = pymysql.connect(**self.connect_kwargs)
                    self._total += 1
            if conn is None:
                if timeout is not None and (time.time() - start) >= timeout:
                    raise TimeoutError("获取数据库连接超时")
                time.sleep(0.05)
                continue

            # 尝试 ping 确认连接有效，若无效则关闭并尝试重建
            try:
                conn.ping(reconnect=True)
                return PooledConnection(conn, self)
            except Exception:
                try:
                    conn.close()
                except Exception:
                    pass
                with self._lock:
                    self._total = max(0, self._total - 1)
                # 尝试新建连接并返回
                try:
                    new_conn = pymysql.connect(**self.connect_kwargs)
                    with self._lock:
                        self._total += 1
                    return PooledConnection(new_conn, self)
                except Exception:
                    # 若重建失败且超时则抛出，否则继续重试
                    if timeout is not None and (time.time() - start) >= timeout:
                        raise TimeoutError("获取数据库连接超时（重建失败）")
                    time.sleep(0.1)
                    continue

    def putconn(self, conn):
        try:
            conn.ping(reconnect=True)
            with self._lock:
                self._pool.append(conn)
        except Exception:
            try:
                conn.close()
            except Exception:
                pass
            with self._lock:
                self._total = max(0, self._total - 1)

    def closeall(self):
        with self._lock:
            while self._pool:
                c = self._pool.popleft()
                try:
                    c.close()
                except Exception:
                    pass
            self._total = 0


def _pool_kwargs():
    """连接池参数"""
    return {
        "host": DB_HOST,
        "port": DB_PORT,
        "user": DB_USER,
        "password": DB_PASSWORD,
        "database": DB_NAME,
        "charset": DB_CHARSET,
        "cursorclass": pymysql.cursors.Cursor,
    }


def get_mysql():
    """获取数据库连接（优先连接池，失败则直连）"""
    global MYSQL_POOL
    with POOL_LOCK:
        if MYSQL_POOL is None:
            try:
                MYSQL_POOL = MySQLConnectionPool(minconn=1, maxconn=20, **_pool_kwargs())
            except Exception:
                MYSQL_POOL = None
    if MYSQL_POOL is None:
        return pymysql.connect(**_pool_kwargs())
    return MYSQL_POOL.getconn()


def get_connection():
    """直连数据库（兼容旧脚本，DictCursor 模式）"""
    return pymysql.connect(
        host=DB_HOST,
        port=DB_PORT,
        user=DB_USER,
        password=DB_PASSWORD,
        database=DB_NAME,
        charset=DB_CHARSET,
        autocommit=False,
        cursorclass=DictCursor,
    )


# 在程序退出时确保池中所有物理连接被关闭
try:
    atexit.register(lambda: MYSQL_POOL.closeall() if MYSQL_POOL is not None else None)
except Exception:
    pass