# -*- coding:utf-8 -*- """ Author: qiaoxinjiu Create Data: 2020/11/6 17:34 """ import pymysql import pymongo # from DBUtils.PooledDB import PooledDB from dbutils.pooled_db import PooledDB import psycopg2 from psycopg2 import pool from psycopg2.extras import RealDictCursor from base_framework.public_tools.read_config import InitConfig from base_framework.public_tools.read_config import ReadConfig, get_current_config # from base_framework.public_tools import db_config as config from base_framework.base_config.current_pth import * """ @功能：创建数据库连接池 """ as_db = ['ZZYY'] class PgConnectionPool(InitConfig): pool = None pool_cache = dict() def __init__(self): try: super().__init__() self.DB_SSL = False except Exception as e: print(e) # 创建数据库连接conn和游标cursor def __enter__(self): self.conn = self.__getconn() self.cursor = self.conn.cursor() def __getconn(self, choose_db=None): # 如果未指定数据库，使用默认配置 if not choose_db: choose_db = 'default' try: self.pool = self.pool_cache[choose_db] except KeyError: rc = ReadConfig(config_file_path) # 根据choose_db值获取对应的PostgreSQL配置 if choose_db == 'default': db_host = rc.get_value(sections='PostgreSQL', options='db_test_host') db_port = rc.get_value(sections='PostgreSQL', options='db_test_port') db_name = rc.get_value(sections='PostgreSQL', options='db_test_dbname') db_user = rc.get_value(sections='PostgreSQL', options='db_test_user') db_password = rc.get_value(sections='PostgreSQL', options='db_test_password') db_min_cached = rc.get_value(sections='PostgreSQL', options='db_min_cached') db_max_cached = rc.get_value(sections='PostgreSQL', options='db_max_cached') db_max_shared = rc.get_value(sections='PostgreSQL', options='db_max_shared') db_max_connecyions = rc.get_value(sections='PostgreSQL', options='db_max_connecyions') db_max_usage = rc.get_value(sections='PostgreSQL', options='db_max_usage') else: # 可以扩展其他数据库连接 db_host = rc.get_value(sections='PostgreSQL', options=f'db_{choose_db}_host') db_port = rc.get_value(sections='PostgreSQL', options=f'db_{choose_db}_port') db_name = rc.get_value(sections='PostgreSQL', options=f'db_{choose_db}_name') db_user = rc.get_value(sections='PostgreSQL', options=f'db_{choose_db}_user') db_password = rc.get_value(sections='PostgreSQL', options=f'db_{choose_db}_password') db_min_cached = rc.get_value(sections='PostgreSQL', options=f'db_{choose_db}_min_cached') db_max_cached = rc.get_value(sections='PostgreSQL', options=f'db_{choose_db}_max_cached') db_max_shared = rc.get_value(sections='PostgreSQL', options=f'db_{choose_db}_max_shared') db_max_connecyions = rc.get_value(sections='PostgreSQL', options=f'db_{choose_db}_max_connecyions') db_max_usage = rc.get_value(sections='PostgreSQL', options=f'db_{choose_db}_max_usage') # PostgreSQL连接池配置 try: print("=" * 80) print("PostgreSQL连接池配置信息:") print(" 主机(Host): {}".format(db_host)) print(" 端口(Port): {}".format(db_port)) print(" 数据库名(Database): {}".format(db_name)) print(" 用户名(User): {}".format(db_user)) print(" 密码(Password): {} (已隐藏)".format('*' * len(db_password) if db_password else 'None')) print(" 最小缓存连接数(MinCached): {}".format(db_min_cached)) print(" 最大缓存连接数(MaxCached): {}".format(db_max_cached)) print(" 最大共享连接数(MaxShared): {}".format(db_max_shared)) print(" 最大连接数(MaxConnections): {}".format(db_max_connecyions)) print(" 最大使用次数(MaxUsage): {}".format(db_max_usage)) print(" SSL模式(SSLMode): {}".format('require' if self.DB_SSL else 'disable')) print(" 连接超时(ConnectTimeout): 30秒") print("=" * 80) self.pool = PooledDB( creator=psycopg2, host=db_host, port=int(db_port), user=db_user, password=db_password, database=db_name, mincached=int(db_min_cached), maxcached=int(db_max_cached), maxshared=int(db_max_shared), maxconnections=int(db_max_connecyions), blocking=True, maxusage=int(db_max_usage), setsession=None, # PostgreSQL特定参数 sslmode='require' if self.DB_SSL else 'disable', connect_timeout=30, keepalives=1, keepalives_idle=30, keepalives_interval=10, keepalives_count=5 ) self.pool_cache[choose_db] = self.pool print("PostgreSQL连接池创建成功") except Exception as e: error_msg = """ PostgreSQL连接池创建失败！ 连接配置信息: 主机(Host): {} 端口(Port): {} 数据库名(Database): {} 用户名(User): {} 密码(Password): {} (已隐藏) SSL模式(SSLMode): {} 连接超时(ConnectTimeout): 30秒 错误详情: {} """.format( db_host, db_port, db_name, db_user, '*' * len(db_password) if db_password else 'None', 'require' if self.DB_SSL else 'disable', str(e) ) print(error_msg) raise Exception(error_msg) try: return self.pool.connection() except Exception as e: # 尝试获取连接配置信息用于错误提示 try: rc = ReadConfig(config_file_path) db_host = rc.get_value(sections='PostgreSQL', options='db_test_host') db_port = rc.get_value(sections='PostgreSQL', options='db_test_port') db_name = rc.get_value(sections='PostgreSQL', options='db_test_dbname') db_user = rc.get_value(sections='PostgreSQL', options='db_test_user') except: db_host = db_port = db_name = db_user = 'unknown' error_msg = """ PostgreSQL连接获取失败！ 连接配置信息: 主机(Host): {} 端口(Port): {} 数据库名(Database): {} 用户名(User): {} 错误详情: {} """.format(db_host, db_port, db_name, db_user, str(e)) print(error_msg) raise Exception(error_msg) # 释放连接池资源 def __exit__(self, exc_type, exc_val, exc_tb): if self.cursor: self.cursor.close() if self.conn: self.conn.close() # 获取连接和游标（返回字典形式的结果） def getconn(self, choose_db=None): conn = self.__getconn(choose_db=choose_db) # 使用RealDictCursor返回字典形式的游标 cursor = conn.cursor(cursor_factory=RealDictCursor) return cursor, conn class MyConnectionPool(InitConfig): pool = None pool_cache = dict() def __init__(self): try: super().__init__() except Exception as e: print(e) self.current_business = get_current_config(section='run_evn_name', key='current_business') # 创建数据库连接conn和游标cursor def __enter__(self): self.conn = self.__getconn() self.cursor = self.conn.cursor() def __getconn(self, choose_db=None): current_team = ReadConfig(env_choose_path).get_value(sections='run_evn_name', options='current_team') if not choose_db: # 没有指定，则按小组默认设置 if current_team.upper() in as_db and choose_db is None: choose_db = 'as' elif current_team.upper() == "SE" and choose_db is None: choose_db = 'se' elif current_team.upper() == "XUEDAU" and choose_db is None: choose_db = 'xdu' elif current_team.upper() not in as_db and choose_db is None and self.current_business == 'hh': choose_db = 'hh' elif current_team.upper() not in as_db and choose_db is None and self.current_business == 'hhi': choose_db = 'hhi' try: self.pool = self.pool_cache[choose_db] except Exception as e: rc = ReadConfig(config_file_path) if choose_db == 'as': db_host = rc.get_value(sections='Mysql', options='db_as_svr') elif choose_db == 'se': db_host = rc.get_value(sections='Mysql', options='db_se_svr') elif choose_db == 'xdu': db_host = rc.get_value(sections='Mysql', options='db_xdu_svr') self.DB_TEST_USER = rc.get_value(sections='Mysql', options='db_xdu_user') self.DB_TEST_PASSWORD = rc.get_value(sections='Mysql', options='db_xdu_password') elif choose_db == 'hh' or choose_db == 'huohua': db_host = rc.get_value(sections='Mysql', options='db_hh_svr') elif choose_db == 'hhi': db_host = rc.get_value(sections='Mysql', options='db_hhi_svr') elif choose_db == 'hh.qa': # 自动化和信息化的数据都走huohua db_host = 'mysql.qa.huohua.cn' elif not choose_db: # 没有传入则默认走huohua choose_db = 'hh.qa' db_host = 'mysql.qa.huohua.cn' else: raise Exception("当前仅支持hh,hhi,as,se四个数据库服务器，而你选择的是：{}".format(choose_db)) default_db = 'sys' self.pool = PooledDB( creator=pymysql, host=db_host, port=int(self.DB_TEST_PORT), user=self.DB_TEST_USER, passwd=self.DB_TEST_PASSWORD, db=default_db, mincached=int(self.DB_MIN_CACHED), maxcached=int(self.DB_MAX_CACHED), maxshared=int(self.DB_MAX_SHARED), maxconnections=int(self.DB_MAX_CONNECYIONS), blocking=True, maxusage=int(self.DB_MAX_USAGE), setsession=None, use_unicode=True, charset=self.DB_CHARSET ) self.pool_cache[choose_db] = self.pool return self.pool.connection() # 释放连接池资源 def __exit__(self, exc_type, exc_val, exc_tb): self.cursor.close() self.conn.close() # 关闭连接归还给链接池 # def close(self): # self.cursor.close() # self.conn.close() # 从连接池中取出一个连接 def getconn(self, choose_db=None): conn = self.__getconn(choose_db=choose_db) # 字典形式返回 cursor = conn.cursor(pymysql.cursors.DictCursor) return cursor, conn # 获取连接池,实例化 def get_my_connection(): return MyConnectionPool() def get_pg_connection(): return PgConnectionPool() class MongoConnectionPool(InitConfig): def __init__(self): try: # super().__init__() super(MongoConnectionPool, self).__init__() except Exception as e: print(e) def mongo_connect(self): try: self.connect_ = pymongo.MongoClient(host=self.MONGO_HOST, port=int(self.MONGO_PORT), username=self.MONGO_USER, password=self.MONGO_PASSWORD, authSource="hulk_teach_marketing" ) except Exception as e: raise Exception("mongdb连接失败：{}".format(e)) return self.connect_ def get_my_mongo_connection(): return MongoConnectionPool()