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']
            )
            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