jd-tb-yaofang-2/commons/conn_mysql.py

import logging
import os

import pymysql
from dbutils.pooled_db import PooledDB
from pymysql.cursors import DictCursor

from commons.config import (
    MYSQL_HOST,
    MYSQL_PORT,
    MYSQL_USER,
    MYSQL_PASSWORD,
    MYSQL_DATABASE,
    MYSQL_POOL_MAXCONNECTIONS,
    MYSQL_POOL_MIN_CACHED,
    MYSQL_POOL_MAX_CACHED,
)

logger = logging.getLogger(__name__)

_POOL_COMMON = {
    "maxconnections": MYSQL_POOL_MAXCONNECTIONS,
    "mincached": MYSQL_POOL_MIN_CACHED,
    "maxcached": MYSQL_POOL_MAX_CACHED,
    "blocking": True,
}


def _create_pool(
        host,
        port,
        user,
        password,
        database,
):
    return PooledDB(
        creator=pymysql,
        cursorclass=DictCursor,
        charset="utf8mb4",
        host=host,
        port=port,
        user=user,
        password=password,
        database=database,
        **_POOL_COMMON,
    )


class _MySQLPoolBase:
    """共用的取连接、写库、读库逻辑，避免两套类复制粘贴。"""

    def __init__(self, pool):
        self.pool = pool

    def get_conn(self):
        return self.pool.connection()

    def _close(self, conn, cursor) -> None:
        try:
            if cursor:
                cursor.close()
        finally:
            if conn:
                conn.close()

    def _write(self, fn, *args, **kwargs):
        conn = self.get_conn()
        cursor = conn.cursor()
        try:
            result = fn(cursor, *args, **kwargs)
            conn.commit()
            return result
        except Exception as e:
            logger.error("SQL 执行错误: %s", e)
            conn.rollback()
            return None
        finally:
            self._close(conn, cursor)

    def execute_many(self, sql, data):
        rows = list(data)
        if not rows:
            return

        def _do(c, s, d):
            c.executemany(s, d)

        self._write(_do, sql, rows)

    def execute_one(self, sql, data):
        def _do(c, s, d):
            c.execute(s, d)

        self._write(_do, sql, data)

    def update_data(self, sql, data):
        self.execute_one(sql, data)

    def select_data(self, sql, data=None):
        """只读查询，不应 commit。"""
        conn = self.get_conn()
        cursor = conn.cursor()
        try:
            if data is not None:
                cursor.execute(sql, data)
                return cursor.fetchall()
            else:
                cursor.execute(sql)
                return cursor.fetchall()
        except Exception as e:
            logger.error("SQL 查询错误: %s", e)
            return []
        finally:
            self._close(conn, cursor)

    def execute(self, sql, data=None):
        """通用写操作，返回受影响行数；失败返回 0。"""

        def _do(c, s, d):
            if d is not None:
                return c.execute(s, d)
            return c.execute(s)

        result = self._write(_do, sql, data)
        return result if isinstance(result, int) else 0


class MySQLPoolOn2(_MySQLPoolBase):
    """drug_retrieve 库。"""

    def __init__(self):
        pool = _create_pool(
            host=os.environ.get("MYSQL_DRUG_HOST", MYSQL_HOST),
            port=int(os.environ.get("MYSQL_DRUG_PORT", str(MYSQL_PORT))),
            user=os.environ.get("MYSQL_DRUG_USER", MYSQL_USER),
            password=os.environ.get("MYSQL_DRUG_PASSWORD", MYSQL_PASSWORD),
            database=os.environ.get("MYSQL_DRUG_DB", MYSQL_DATABASE),
        )
        super().__init__(pool)


if __name__ == "__main__":
    pool = MySQLPoolOn2()
    rows = pool.select_data("SELECT 1 AS one")
    print(rows)