SQL:通用关系型数据库连接简易抽象

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

 

posted @ 2021-01-06 17:08  士为知己  阅读(109)  评论(0)    收藏  举报