AiDOPWarehouse/doc/db/mdp/contract_tests/_db.py

"""共享 MySQL 连接助手。连接参数取自环境变量，默认指向本地 aidopdev。

环境变量：
    MDP_DB_HOST (default 127.0.0.1)
    MDP_DB_PORT (default 3306)
    MDP_DB_USER (default root)
    MDP_DB_PASSWORD (default 空)
    MDP_DB_NAME (default aidopdev)
"""
import os

import pymysql


def get_conn():
    return pymysql.connect(
        host=os.environ.get("MDP_DB_HOST", "127.0.0.1"),
        port=int(os.environ.get("MDP_DB_PORT", "3306")),
        user=os.environ.get("MDP_DB_USER", "root"),
        password=os.environ.get("MDP_DB_PASSWORD", ""),
        database=os.environ.get("MDP_DB_NAME", "aidopdev"),
        charset="utf8mb4",
        cursorclass=pymysql.cursors.DictCursor,
        autocommit=True,
    )


def query(sql, params=None):
    conn = get_conn()
    try:
        with conn.cursor() as cur:
            cur.execute(sql, params or ())
            return cur.fetchall()
    finally:
        conn.close()