project/mt_V2/db.py

import os
import pymysql
import threading
import time
import atexit
from collections import deque


# simple direct connector (used by some scripts)
from pymysql.cursors import DictCursor

def get_connection():
    return pymysql.connect(
        host="120.24.26.108",
        port=3307,
        user="root",
        password="zhijiayun123456",
        database="drug_retrieve",
        charset="utf8mb4",
        autocommit=False,
        cursorclass=DictCursor,
    )


# connection pool used by a_mt.py
MYSQL_POOL = None


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 get_mysql():
    global MYSQL_POOL
    if 'MYSQL_POOL' not in globals() or MYSQL_POOL is None:
        MYSQL_POOL = None
        try:
            MYSQL_POOL = MySQLConnectionPool(
                minconn=1,
                maxconn=20,
                host="120.24.26.108",
                port=3307,
                user="root",
                password="zhijiayun123456",
                database="drug_retrieve",
                charset="utf8mb4",
                cursorclass=pymysql.cursors.Cursor,
            )
        except Exception:
            MYSQL_POOL = None
    if MYSQL_POOL is None:
        return pymysql.connect(
            host="120.24.26.108",
            port=3307,
            user="root",
            password="zhijiayun123456",
            database="drug_retrieve",
            charset="utf8mb4",
            cursorclass=pymysql.cursors.Cursor,
        )
    return MYSQL_POOL.getconn()


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