35 lines
1.1 KiB
Python
35 lines
1.1 KiB
Python
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
|
|
|