如何定义psycopg2连接PostgreSQL的上下文管理器类

要实现符合with语法规范的数据库连接类，核心是正确实现__enter__和__exit__两个魔术方法，注意你原有代码里的__exit__缺少必填的异常参数，同时需要补全事务提交/回滚、资源释放的逻辑。

完整实现代码如下：

import psycopg2
from psycopg2 import extras

class DBConn:
    def __init__(self, dns):
        self.dns = dns
        self.conn = None
        self.curs = None

    def __enter__(self):
        # 进入with上下文时初始化连接和游标
        self.conn = psycopg2.connect(**self.dns)
        # 如需返回字典格式的查询结果，指定RealDictCursor作为游标工厂
        self.curs = self.conn.cursor(cursor_factory=extras.RealDictCursor)
        # 返回值会被as关键字接收，可按需返回游标、连接或者类实例本身
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        # 退出上下文时统一处理事务和资源释放
        try:
            if exc_type is None:
                # 无异常则提交事务
                self.conn.commit()
            else:
                # 出现异常自动回滚
                self.conn.rollback()
        finally:
            # 无论执行是否出错，都确保关闭游标和连接
            if self.curs:
                self.curs.close()
            if self.conn:
                self.conn.close()
        # 返回False不拦截抛出的异常，方便上层捕获错误；不要随意返回True吞掉异常
        return False

使用方法 #

# 数据库连接配置
dns_config = {
    "dbname": "your_db_name",
    "user": "your_account",
    "password": "your_password",
    "host": "127.0.0.1",
    "port": 5432
}

# 直接通过with语句使用，退出块时自动完成提交/回滚、资源关闭
with DBConn(dns_config) as db:
    db.curs.execute("SELECT id, name FROM user_table LIMIT 10")
    result = db.curs.fetchall()
    print(result)

关键说明 #

• __exit__方法必须接收exc_type（异常类型）、exc_val（异常内容）、exc_tb（异常栈）三个位置参数，否则上下文管理器运行时会抛出参数不匹配错误
• 上述事务逻辑和psycopg2原生连接的上下文管理器行为完全一致，不需要手动调用commit/rollback
• 如果希望实例化类的时候就建立连接，而不是进入with块才创建连接，只需要把__enter__里创建连接、游标的逻辑移到__init__方法中即可，__exit__的清理逻辑不需要改动
• 如果希望as关键字直接拿到游标对象，只需要把__enter__方法的返回值改为self.curs，对应调整调用时的代码即可。

内容的提问来源于stack exchange，提问作者Ilya