import pymysql
import psycopg2.extras
import pymssql
import logging
class DBMS(object):
""" The context manager has been implemented, Can be modified according to the actual situation"""
ENGINE_PORT_MAPS = {
3306: 'mysql', 5432: 'postgresql', 1433: 'sql-server'
}
def __init__(self, host: str, user: str, password: str, database: str, port: int = 3306, engine: str = 'mysql',
autocommit: bool = True, as_dict: bool = True):
"""
:param host: ip or domain
:param user:
:param password:
:param database:
:param port:
:param engine: engines = ('mysql', 'postgresql', 'sql-server')
:param autocommit: default True
:param as_dict: default True
Usage:
# >>> conn_kwargs = {
# 'host': 'host',
# 'user': 'user,
# 'password': 'password,
# 'database': 'database',
# 'port': 3306,
# }
# >>> try:
# >>> with DBMS(**conn_kwargs) as cursor:
# >>> cursor.execute("SELECT * FROM django_admin_log")
# >>> result = cursor.fetchall()
# >>> except Exception as e:
# >>> print(e)
"""
self._conn = None
self.autocommit = autocommit
port = int(port)
engine = self.ENGINE_PORT_MAPS.get(port, engine.lower())
engine_kwargs = dict(host=host, user=user, password=password, database=database, port=port)
if engine == 'mysql':
if as_dict:
engine_kwargs.update(cursorclass=pymysql.cursors.DictCursor)
self._conn = pymysql.connect(autocommit=autocommit, **engine_kwargs)
elif engine == 'postgresql':
if as_dict:
engine_kwargs.update(cursor_factory=psycopg2.extras.RealDictCursor)
self._conn = psycopg2.connect(**engine_kwargs)
self._conn.set_session(autocommit=autocommit)
elif engine == 'sql-server':
self._conn = pymssql.connect(autocommit=autocommit, as_dict=as_dict, **engine_kwargs)
else:
raise ValueError('engine must be one of {}'.format(self.ENGINE_PORT_MAPS.values()))
self._cursor = self._conn.cursor()
def __enter__(self):
return self._cursor
def __exit__(self, exc_type, exc_val, exc_tb):
self.close()
def __del__(self):
self.close()
def close(self):
try:
self._cursor.close()
self._conn.close()
except Exception as e:
logging.warning(e)
finally:
self._cursor = None
self._conn = None