logistics/utils/gtools.py

import pymysql
from sqlalchemy import create_engine
from utils.config_manager import config


class MySQLconnect():
    def __init__(self, dbname: str = None):
        db_config = config.get_database_config()
        self.host = db_config.get("host", "192.168.100.33")
        self.port = db_config.get("port", 3306)
        self.user = db_config.get("username", "zhenggantian")
        self.password = db_config.get("password", "123456")
        self.dbname = dbname or db_config.get("database", "logistics")
        self.charset = db_config.get("charset", "utf8")
        self.pool_size = db_config.get("pool_size", 10)
        self.max_overflow = db_config.get("max_overflow", 5)
        self.pool_recycle = db_config.get("pool_recycle", 3600)

        if not isinstance(self.dbname, str):
            raise TypeError("dbname must be a string")

    def __enter__(self):
        self.eng = self.engine()
        self.con = self.connect()
        self.cur = self.con.cursor()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.eng.dispose()
        self.cur.close()
        self.con.close()
        if exc_val:
            raise

    def engine(self):
        return create_engine(
            f"mysql+pymysql://{self.user}:{self.password}@{self.host}:{self.port}/{self.dbname}",
            pool_size=self.pool_size,
            max_overflow=self.max_overflow,
            pool_recycle=self.pool_recycle
        )

    def connect(self):
        return pymysql.connect(
            host=self.host,
            port=self.port,
            database=self.dbname,
            user=self.user,
            password=self.password,
            charset=self.charset
        )