Python之使用上下文管理器實現連接MySQL

import pymysql
import traceback


class MySQLContextManager:
    """
    MySQL Context Manager
    """

    def __init__(self, host="localhost", port=3306, user="root", password="123456", database=""):
        self.host = host
        self.port = port
        self.user = user
        self.password = password
        self.database = database

    def __enter__(self):
        """
        上文管理器 返回的結果 會被賦給as 關鍵字之後的變量
        :return: cursor
        """
        try:
            self.db_conn = pymysql.connect(host=self.host, port=self.port, user=self.user, password=self.password,
                                           database=self.database, charset="utf8")
            self.cursor = self.db_conn.cursor()
            return self.cursor
        except Exception as e:
            traceback.print_exc()
            raise e

    def __exit__(self, exc_type, exc_val, exc_tb):
        """
        下文管理器 選擇性地處理包裝代碼塊中出現的異常
        :param exc_type: 異常類型 默認爲 None,發生異常,參數被填充
        :param exc_val: 異常實例 默認爲 None,發生異常,參數被填充
        :param exc_tb: 異常回溯 默認爲 None,發生異常,參數被填充
        :return:
        """
        try:
            if exc_type:
                self.db_conn.rollback()
                # 返回False 傳播異常
                # 返回True 終止異常 不要這麼做
                # 拋出不同的異常 代替 原有異常
                return False
            else:
                self.db_conn.commit()
        except Exception as e:
            raise e
        finally:
            self.cursor.close()
            self.db_conn.close()


if __name__ == '__main__':
    # 根據實際,傳入數據庫參數
    with MySQLContextManager() as db_cursor:
        select_sql = "select * from test"
        # 執行SQL語句 返回影響的行數
        db_cursor.execute(select_sql)
        # 返回執行結果
        result = db_cursor.fetchall()
        print(result)