mysql.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
# @Version     : Python 3.11.4
# @Software    : Sublime Text 4
# @Author      : StudentCWZ
# @Email       : StudentCWZ@outlook.com
# @Date        : 2023/11/12 11:08
# @File        : mysql.py
# @Description :
"""

from typing import Optional

import pymysql
from dbutils.pooled_db import PooledDB


class MySQLHelper:

    def __init__(self, **kwargs):
        """
        Initialize the MySQLHelper class, create a database connection pool
        """
        self.kwargs = kwargs if kwargs else {}
        self.conn = None
        self.db_pool = PooledDB(pymysql, **self.kwargs)

    def __enter__(self):
        """
        Implement the __enter__ magic method, support the with statement
        """
        self.conn = self.db_pool.connection()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """
        Implement the __exit__ magic method, support the with statement
        """
        if self.conn:
            self.conn.close()
        if exc_type:
            print(f"Error occurred: {exc_val}")

    def execute(self, sql: str) -> Optional[tuple]:
        """
        Execute the SQL statement and retrieve the results
        """
        if not self.conn:
            print("No connection available.")
            return None

        try:
            # Use the cursor object to execute SQL statement
            with self.conn.cursor() as cursor:
                cursor.execute(sql)
                result = cursor.fetchall()
                return result

        except Exception as e:
            print(f"Error executing SQL: {e}")
            return None