数据库、redis连接池

import contextlib
import redis
import pymysql

IS_DEV = False
if IS_DEV:
    # 正式版redis
    REDIS = {
        'host': '**********.redis.rds.aliyuncs.com',
        'port': 6379,
        'db': 11,
        'password': '********'
    }
    # 正式版数据库
    DB_INFO = {
        'host': '************.mysql.rds.aliyuncs.com',
        'port': 3306,
        'user': 'mysqlsa',
        'password': '********',
        'db': '*****agent'
    }
    ADS_INFO = {
        'host': '************.ads.aliyuncs.com',
        'port': 3306,
        'user': 'agent_sa',
        'password': '********',
        'db': '*****agent'
    }
else:
    # 测试版 线下 redis
    REDIS = {
        'host': '192.168.0.8',
        'port': 6379,
        'db': 10,
        'password': '******'
    }
    # 测试版数据库
    DB_INFO = {
        'host': '***********.mysql.rds.aliyuncs.com',
        'port': 3306,
        'user': 'mysqlsa',
        'password': '*******',
        'db': '******agent'
    }
    ADS_INFO = {
        'host': '***********.ads.aliyuncs.com',
        'port': 3306,
        'user': 'mysqlsa',
        'password': '********',
        'db': '******agent'
    }


# 创建数据库连接池池
class DatabaseConnectPool:
    '''
    数据库连接池类  支持多库多池
    '''

    def __init__(self, module, credentials, size=10):
        '''
        :param module: 数据库驱动 例如 pymysql
        :param credentials: 数据库连接信息
        credentials={
                'DB_INFO' = {
                'host': 'rm-uf6t4r3u8vea8u340971.mysql.rds.aliyuncs.com',
                'port': 3306,
                'user': 'res_sql',
                'password': 'Lansi321',
                'db': 'lansi_agent'
            }
        }
        :param size: 每个连接池最大连接数量
        '''
        self.__module = module
        self.__credentials = credentials
        self.__databases = {}
        self.__size = size

    def get(self, pool_name):
        '''
        根据连接池名称获取连接
        :param pool_name:
        :return:
        '''
        if pool_name not in self.__credentials:
            raise Exception('数据库信息不存在')
        elif pool_name not in self.__databases:
            self.__databases[pool_name] = []
            self.create_connection(pool_name)
        elif not self.__databases[pool_name]:
            self.create_connection(pool_name)
        con = self.__databases[pool_name].pop()
        try:
            con.ping()
        except Exception as e:
            con = e
        return con

    def create_connection(self, pool_name):
        '''
        创建数据库连接
        :param pool_name:
        :return:
        '''
        if pool_name not in self.__credentials:
            raise Exception('数据库信息不存在')
        db_info = self.__credentials[pool_name]
        con = self.__module.connect(**db_info)
        self.__databases[pool_name].append(con)

    def put(self, pool_name, con):
        '''
        把数据库连接放回连接池
        :param pool_name:
        :param con:
        :return:
        '''
        if isinstance(con, self.__module.connections.Connection):
            if len(self.__databases[pool_name]) > self.__size:
                con.close()
            self.__databases[pool_name].insert(0, con)
        else:
            try:
                con.close()
            except:
                pass

    def __del__(self):
        for con_list in self.__databases.values():
            for con in con_list:
                try:
                    con.close()
                except:
                    pass


db_info = {
    # 连接池名称: 数据库信息
    'rds': DB_INFO,
    'ads': ADS_INFO
}
db_pool = DatabaseConnectPool(pymysql, db_info)


def search_sql_pool(total_sql, pool_name='rds', result_type='dict'):
    '''
    多库查询
    :param total_sql: 被执行的sql
    :param pool_name: 连接池名称
    :param db_name: 数据库名
    :param result_type: 返回结果类型  默认字典
    :return:
    '''
    # 从数据库连接池池中获取连接池
    db = db_pool.get(pool_name)
    # 如果没有取出连接  直到取到连接
    while isinstance(db, str):
        db = db_pool.get(pool_name)
    if result_type == 'dict':
        try:
            with db.cursor(pymysql.cursors.DictCursor) as cursor:
                # 使⽤execute⽅法执⾏SQL语句
                cursor.execute(total_sql)
                # 获取数据
                data = cursor.fetchall()
            # 提交
            db.commit()
        except Exception as e:
            data = e
        finally:
            db_pool.put(pool_name, db)
    # 返回列表类型查询结果
    else:
        try:
            with db.cursor() as cursor:
                # 使⽤execute⽅法执⾏SQL语句
                cursor.execute(total_sql)
                # 获取数据
                data = cursor.fetchall()
            # 提交
            db.commit()
        except Exception as e:
            data = e
        finally:
            db_pool.put(pool_name, db)
    return data


def rds_sql(total_sql):
    '''
    RDS查询  返回字典
    :param total_sql:
    :return:
    '''
    return search_sql_pool(total_sql, pool_name='rds')


def rds_list_sql(total_sql):
    '''
    RDS查询  返回列表
    :param total_sql:
    :return:
    '''
    return search_sql_pool(total_sql, pool_name='rds', result_type='list')


def ads_sql(total_sql):
    '''
    ADS查询  返回字典
    :param total_sql:
    :return:
    '''
    return search_sql_pool(total_sql, pool_name='ads')


def ads_list_sql(total_sql):
    '''
    ADS查询  返回列表
    :param total_sql:
    :return:
    '''
    return search_sql_pool(total_sql, pool_name='ads', result_type='list')


# 创建redis连接池
class RedisConnectPool:
    '''
    redis连接池 使用的是对redis模块封装了一层
    '''

    def __init__(self, redis_info, size=10):
        '''
        初始化连接池
        :param redis_info: {
            'host': '192.168.0.8',
            'port': 6379,
            'db': 10,
            'password': 'Lansi.123'
        }
        :param size: 每个连接池最大连接数量
        '''
        self.__redis_info = redis_info
        self.__connections = []
        self.__size = size

    @contextlib.contextmanager
    def get(self):
        '''
        获取一个连接
        :return:
        '''
        if not self.__connections:
            self.create_connection()
        con = self.__connections.pop()
        con.ping()
        yield con
        self.put(con)

    def create_connection(self):
        '''
        创建连接
        :return:
        '''
        con = redis.Redis(decode_responses=True, **self.__redis_info)
        self.__connections.append(con)

    def put(self, con):
        '''
        把连接放回连接池
        :param con:
        :return:
        '''
        if isinstance(con, str):
            return
        if self.__size < len(self.__connections):
            con.close()
        elif isinstance(con, redis.client.Redis):
            self.__connections.insert(0, con)
        else:
            try:
                con.close()
            except:
                pass

    def __del__(self):
        for con in self.__connections:
            try:
                con.close()
            except:
                pass


redis_pool = RedisConnectPool(REDIS)

if __name__ == '__main__':
    date = rds_sql('select now()')
    if isinstance(date, str):
        print('ERROR: {}'.format(date))
    else:
        print(date)

    with redis_pool.get() as con:
        print(con.keys('*'))
        print(con.get('test_key'))