实战:python管理多个PostgreSQL数据库的连接

211 阅读4分钟

python管理多个PostgreSQL数据库的连接

主要实现了通过读取配置文件中的PostgreSQL服务信息,连接到相应的PostgreSQL数据库。提供了两个连接函数,postgresql_connectpostgresql_connect_encode,分别是普通的和带编码参数的连接函数。

操作步骤

导入相关库

python
from pyhive import hive
import cx_Oracle
import pymysql
import pymssql
from optparse import OptionParser
import logging
import sys
import traceback
import re
import warnings
with warnings.catch_warnings(record=True):
    import psycopg2
reload(sys)

2. 配置PostgreSQL连接信息

这里定义了一个包含多个PostgreSQL服务配置信息的列表pgconfigs。每个配置项由数据库名称、用户、密码、主机和端口组成。

python
pgconfigs = [
    {"asset_factory": {"database": "asset_factory", "user": "asset", "password": "Ro20ot16", "host": "10.251.80.202", "port": "65432"}},
    {"asset_register": {"database": "asset_register", "user": "asset", "password": "Ro20ot16", "host": "10.251.80.202", "port": "65432"}},
    {"creditdb": {"database": "creditdb", "user": "credit", "password": "hHJ98#pE40Y", "host": "10.251.101.175", "port": "18923"}},
    {"asset_portal": {"database": "asset_portal", "user": "asset", "password": "Ro20ot16", "host": "10.251.80.202", "port": "65432"}},
    {"tyjh": {"database": "tyjh", "user": "tyjh", "password": "k5y2dwoKcFm&^OsW", "host": "10.251.90.35", "port": "18922"}},
    {"cwrl": {"database": "cwrl", "user": "cwjh_yaxin", "password": "eKYy4R3&MbaaY3Zy", "host": "10.251.90.34", "port": "18921"}},
    {"dc_new_pg": {"database": "dc", "user": "dc", "password": "Ojjkcy@jVxKIeo5C", "host": "10.251.90.36", "port": "18923"}},
    {"credit_data": {"database": "credit_data", "user": "credit_data", "password": "@G*1f*3$3DzTH6%o", "host": "10.251.90.134", "port": "18923"}},
    {"new_jf_pg_cd": {"database": "accthuiju", "user": "datacenter", "password": "KridU593&%rj90", "host": "133.37.116.192", "port": "18921"}},
    {"new_accthuiju": {"database": "accthuiju", "user": "itf_ods", "password": "2dH~fZ^8", "host": "10.251.64.226", "port": "18921"}},
    {"crm3huijupg": {"database": "crm3huijupg", "user": "datacenter", "password": "lp*36^YD", "host": "10.251.65.93", "port": "18921"}},
    {"crm3hispg": {"database": "crm3hispg", "user": "datacenter", "password": "~4nD_jWG", "host": "10.251.64.196", "port": "18921"}},
    {"hana": {"database": "hana", "user": "hana", "password": "gvptHKXVNKEGw7pk", "host": "10.251.90.94", "port": "18924"}},
    {"dataos_71_pg_dev": {"database": "dacp", "user": "dacp", "password": "jxFgCKv9GJw2ohS3", "host": "10.251.110.104", "port": "18921"}}
]

3. 定义连接函数

普通连接函数

python
def postgresql_connect(servicename):
    global job_task_content_ori
    global job_resource_conn_info
    global job_hdfs_dir_info
    try:
        for subconfig in pgconfigs:
            if servicename in subconfig:
                pgconfig = subconfig[servicename]
                conn = psycopg2.connect(database=pgconfig["database"], user=pgconfig["user"],
                                        password=pgconfig["password"], host=pgconfig["host"], port=pgconfig["port"])
                return conn
        print("servicename %s not found in config dictionary" % servicename)
        sys.exit(6)
    except:
        traceback.print_exc()
        print("connect pg error " + pgconfig["database"])
        sys.exit(6)

上述代码中的postgresql_connect函数:

  • 遍历pgconfigs列表,查找与参数servicename匹配的配置。
  • 使用匹配的配置建立PostgreSQL连接。
  • 如果未找到匹配的服务名,打印错误信息并退出程序。

带编码参数的连接函数

python
def postgresql_connect_encode(servicename, encoding='UTF-8'):
    global job_task_content_ori
    global job_resource_conn_info
    global job_hdfs_dir_info
    try:
        for subconfig in pgconfigs:
            if servicename in subconfig:
                pgconfig = subconfig[servicename]
                conn = psycopg2.connect(database=pgconfig["database"], user=pgconfig["user"],
                                        password=pgconfig["password"], host=pgconfig["host"], port=pgconfig["port"],
                                        client_encoding=encoding)
                return conn
        print("servicename %s not found in config dictionary" % servicename)
        sys.exit(6)
    except:
        traceback.print_exc()
        print("connect pg error " + pgconfig["database"])
        sys.exit(6)

postgresql_connect_encode函数与前一个函数类似,唯一的区别是它接受一个可选的编码参数encoding,并将其传递给psycopg2.connect以设定客户端编码。

解释关键点

  1. 配置管理pgconfigs列表包含多个数据库配置,便于根据服务名快速查找匹配的配置项。
  2. 错误处理:在连接失败时,打印堆栈跟踪以帮助调试并退出程序。
  3. 编码管理:对于需要特定编码的连接,提供了额外的函数来设置编码。

使用示例

假设你要连接到名为 dataos_71_pg_dev 的PostgreSQL服务,可以像下面这样使用这些函数:

python
# 普通连接
conn = postgresql_connect("dataos_71_pg_dev")

# 带编码参数的连接
conn_with_encoding = postgresql_connect_encode("dataos_71_pg_dev", encoding='UTF-8')

完整python

#!/data/apps/python2715/bin/python
# -*- coding:utf-8 -*-

from pyhive import hive
import cx_Oracle
import pymysql
import pymssql
from optparse import OptionParser
import logging
import sys
import traceback
import re
import warnings
with warnings.catch_warnings(record=True):
    import psycopg2
reload(sys)

pgconfigs = [
    {"asset_factory": {"database": "asset_factory", "user": "asset", "password": "Ro20ot16", "host": "10.251.80.202", "port": "65432"}},
    {"asset_register": {"database": "asset_register", "user": "asset", "password": "Ro20ot16", "host": "10.251.80.202", "port": "65432"}},
    {"creditdb": {"database": "creditdb", "user": "credit", "password": "hHJ98#pE40Y", "host": "10.251.101.175", "port": "18923"}},
    {"asset_portal": {"database": "asset_portal", "user": "asset", "password": "Ro20ot16", "host": "10.251.80.202", "port": "65432"}},
    {"tyjh": {"database": "tyjh", "user": "tyjh", "password": "k5y2dwoKcFm&^OsW", "host": "10.251.90.35", "port": "18922"}},
    {"cwrl": {"database": "cwrl", "user": "cwjh_yaxin", "password": "eKYy4R3&MbaaY3Zy", "host": "10.251.90.34", "port": "18921"}},
    {"dc_new_pg": {"database": "dc", "user": "dc", "password": "Ojjkcy@jVxKIeo5C", "host": "10.251.90.36", "port": "18923"}},
    {"credit_data": {"database": "credit_data", "user": "credit_data", "password": "@G*1f*3$3DzTH6%o", "host": "10.251.90.134", "port": "18923"}},
    {"new_jf_pg_cd": {"database": "accthuiju", "user": "datacenter", "password": "KridU593&%rj90", "host": "133.37.116.192", "port": "18921"}},
    {"new_accthuiju": {"database": "accthuiju", "user": "itf_ods", "password": "2dH~fZ^8", "host": "10.251.64.226", "port": "18921"}},
    {"crm3huijupg": {"database": "crm3huijupg", "user": "datacenter", "password": "lp*36^YD", "host": "10.251.65.93", "port": "18921"}},
    {"crm3hispg": {"database": "crm3hispg", "user": "datacenter", "password": "~4nD_jWG", "host": "10.251.64.196", "port": "18921"}},
    {"hana": {"database": "hana", "user": "hana", "password": "gvptHKXVNKEGw7pk", "host": "10.251.90.94", "port": "18924"}},
    {"dataos_71_pg_dev": {"database": "dacp", "user": "dacp", "password": "jxFgCKv9GJw2ohS3", "host": "10.251.110.104", "port": "18921"}}
]


def postgresql_connect(servicename):
    global job_task_content_ori
    global job_resource_conn_info
    global job_hdfs_dir_info
    try:
        for subconfig in pgconfigs:
           if subconfig.has_key(servicename):
                pgconfig = subconfig[servicename]
                conn = psycopg2.connect(database=pgconfig["database"], user=pgconfig["user"],
                                        password=pgconfig["password"], host=pgconfig["host"], port=pgconfig["port"])
                return conn
        print("servinamename %s not found in config dictionary" % servicename)
        sys.exit(6)
    except:
        traceback.print_exc()
        print("connect pg error " + pgconfig["database"])
        sys.exit(6)


def postgresql_connect_encode(servicename, encoding='UTF-8'):
    global job_task_content_ori
    global job_resource_conn_info
    global job_hdfs_dir_info
    try:
        for subconfig in pgconfigs:
            if subconfig.has_key(servicename):
                pgconfig = subconfig[servicename]
                conn = psycopg2.connect(database=pgconfig["database"], user=pgconfig["user"],
                                        password=pgconfig["password"], host=pgconfig["host"], port=pgconfig["port"],
                                        client_encoding=encoding)
                return conn
        print("servinamename %s not found in config dictionary" % servicename)
        sys.exit(6)
    except:
        traceback.print_exc()
        print("connect pg error " + pgconfig["database"])
        sys.exit(6)
                   

改进点

  1. 不再使用 reload(sys) :

    • reload() 在 Python 3 中已经被移除,因为它不鼓励使用该功能。现在可以通过在程序开始时设置字符编码来避免这个问题。
  2. 替换 has_key() 函数:

    • has_key() 在 Python 3 中已经被移除,使用 in 语法更好:

      python
      if servicename in subconfig:
      
  3. 全局变量:

    • 如果全局变量在这个片段中并没有被显式使用,考虑移除该部分代码。如果需要使用这些变量,确保另行定义好。
  4. 异常处理:

    • 类似:

      python
      except Exception as e:
          traceback.print_exc()
          print(f"Connect pg error: {str(e)}")
      
    • 带有具体异常类和使用 f-string 获取更多详细信息。

  5. 将配置参数移出代码:

    • 将 pgconfigs 移到一个专用的配置文件或环境变量中,以便更好地管理数据库凭据并提高安全性。
  6. 使用 with 语句:

    • 可以使用 with 来管理数据库连接的上下文,以便于自动关闭连接。

改进后的代码

python
import psycopg2
import sys
import traceback

# 使用一个专用的配置文件或环境变量来存储敏感信息是最佳实践
pgconfigs = [
    {"asset_factory": {"database": "asset_factory", "user": "asset", "password": "Ro20ot16", "host": "10.251.80.202", "port": "65432"}},
    {"asset_register": {"database": "asset_register", "user": "asset", "password": "Ro20ot16", "host": "10.251.80.202", "port": "65432"}},
    # ...省略其他配置,按需增加...
]

def postgresql_connect(servicename):
    try:
        for subconfig in pgconfigs:
            if servicename in subconfig:
                pgconfig = subconfig[servicename]
                conn = psycopg2.connect(
                    database=pgconfig["database"],
                    user=pgconfig["user"],
                    password=pgconfig["password"],
                    host=pgconfig["host"],
                    port=pgconfig["port"]
                )
                return conn
        print(f"Service name {servicename} not found in config dictionary")
        sys.exit(6)
    except Exception as e:
        traceback.print_exc()
        print(f"Connect pg error: {str(e)}")
        sys.exit(6)

def postgresql_connect_encode(servicename, encoding='UTF-8'):
    try:
        for subconfig in pgconfigs:
            if servicename in subconfig:
                pgconfig = subconfig[servicename]
                conn = psycopg2.connect(
                    database=pgconfig["database"],
                    user=pgconfig["user"],
                    password=pgconfig["password"],
                    host=pgconfig["host"],
                    port=pgconfig["port"],
                    client_encoding=encoding
                )
                return conn
        print(f"Service name {servicename} not found in config dictionary")
        sys.exit(6)
    except Exception as e:
        traceback.print_exc()
        print(f"Connect pg error: {str(e)}")
        sys.exit(6)

总结

  • 使用 Python 3 的特性,如 in 语句替代已废弃的 has_key(),以及 f-string 增强对语句的简洁性和可读性。
  • 通过改进代码可以增强其兼容性、可读性和安全性。尤其对于数据库连接,尽量使用环境变量或配置管理工具来存储数据库凭据,以减少敏感数据泄露的风险。