python数据库连接封装成上下管理器的两种方法

方案1.自己写的上下文管理器

# coding=utf-8

import pandas as pd
import pymysql
import logging
from sqlalchemy import create_engine

logging.basicConfig(level=logging.INFO,
                    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)


class MysqlDB(object):
    """
    使用pymysql创建数据库连接,
    推荐使用with语句执行原生sql语句,建立并及时关闭数据库连接,
    比如: 建表,忽略插入等
    """

    def __init__(self, *args, **kwargs):
        self.schema = kwargs.get('schema', '')
        self.conf = kwargs.get('conf')
        self.host = self.conf.get('host')
        self.user = self.conf.get('user')
        self.pwd = self.conf.get('pwd')
        self.port = self.conf.get('port')
        self.conn = pymysql.connect(host=self.host, port=self.port,
                                    user=self.user, passwd=self.pwd,
                                    db=self.schema, charset='utf8')
        self.cursor = self.conn.cursor(pymysql.cursors.DictCursor)

    def __enter__(self):
        return self.cursor

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.conn.commit()
        self.cursor.close()
        self.conn.close()


class MysqlEngine(MysqlDB):
    """
    使用sqlalchemy创建的数据库连接引擎,
    适用于pandas相关的,读写操作,
    推荐使用with语句执行pandas相关的读写操作.
    比如: 批量查询与入库
    """

    def __init__(self, *args, **kwargs):
        MysqlDB.__init__(self, **kwargs)
        self.mysql_url = f'mysql+pymysql://{self.user}:{self.pwd}@{self.host}:{self.port}/{self.schema}?charset=utf8'
        self.engine = create_engine(self.mysql_url)

    def __enter__(self):
        return self.engine

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.engine.dispose()


def df_to_insert_sql(df, table_name, schema, ignore=False):
    """
    将DF转为insert sql语句
    返回sql语句
    """
    if isinstance(ignore, bool) and ignore == True:
        ignore = "ignore"
    else:
        ignore = ""
    df_dict = df.to_dict(orient='split')
    columns = ','.join(df_dict.get('columns'))
    values = ','.join(str(tuple(val)) for val in df_dict.get('data'))
    insert_sql = f'''
    insert {ignore} into {schema}.{table_name} ({columns})
    values {values}
    '''
    return insert_sql
# 应用示例:
db_conf = {
'user': 'test',
'pwd': '123456'
'port': 2206,
'host': 'test'
}
sql = 'show databases'
with MysqlDB(conf=db_conf) as db:
	a = db.execute(sql)
	data = db.fetchall()

with MysqlEngine(conf=db_conf) as engine:
	df = pd.read_sql_query(sql, con=engine)
    df.to_sql(name='temp_data',
              con=engine,
              index=False,
              schema='schema',
              if_exists="append",
              chunksize=10000,
              )

方案2:使用魔法模块装饰器创建(推荐使用)：

# 摘自框架databases.dbfactory.mysqldb.py
import contextlib   # 下文管理器工具模块
import logging

class MysqlDB_Rent(object):
    def __init__(self, *args, **kwargs):
        self.conf_name = kwargs.get("conf_name")   # 框架内容, 数据库配置名
        self.db_type = kwargs.get("db_type")     # 框架内容, 数据库业务类型,业务线等
        self.conf = mysql_conf.get(self.conf_name)  # 框架内容: 数据库配置名

    @contextlib.contextmanager
    def get_conn(self, *args, **kwargs):
        city = kwargs.get("city")                  # 框架内容, 城市相关
        self.conf_name = getConfigName(city=city, type=self.conf_name)    # 框架内容, 由于城市较多,分库后,导致城市不同,数据库配置不同
        conf = self.conf.get(kwargs.get("link_type", 'default'))
        # 获取主从连接
        host = conf.get("host")
        port = conf.get("port")
        user = conf.get("user")
        passwd = conf.get("passwd")
        db_name = get_db(type=self.conf_name, city=city)   # 框架内容, 数据库名字
        conn = pymysql.connect(host=host, port=port, user=user, password=passwd, db=db_name, charset="utf8")
        cursor = conn.cursor(cursor=pymysql.cursors.DictCursor)
        try:
            yield cursor    # yield关键词把上下文分割成两部分：yield之前就是__init__中的代码块；yield之后其实就是__exit__中的代码块
            logging.info(cursor._last_executed)
        except Exception as e:
            logging.error(e)
        finally:
            conn.commit()
            cursor.close()
            conn.close()

# 使用样例:
with MysqlDB(conf=conf_old).get_conn(city="") as cursor:
      sql = "select * from table"
      cursor.execute(sql)