import pymysql
from pymysql.err import *


class DBContextManager:
    def __init__(self, db_config: dict):
        self.db_config = db_config
        self.connection = None
        self.cursor = None

    def __enter__(self):
        try:
            self.connection = pymysql.connect(
                host=self.db_config["host"],
                port=self.db_config["port"],
                user=self.db_config["user"],
                password=self.db_config["password"],
                db=self.db_config["db"],
                charset=self.db_config["charset"],
            )
            self.cursor = self.connection.cursor()
            return self.cursor
        except (OperationalError, KeyError) as err:
            print(err.args)
            return None

    def __exit__(self, exc_type, exc_val, exc_tb):
        if self.connection and self.cursor:
            if exc_type:
                print(exc_type, "\n", exc_val)
                self.connection.rollback()
            else:
                self.connection.commit()
            self.cursor.close()
            self.connection.close()
        return True