如何创建使用上下文管理器的 postgres 连接？

Question

我试图创建一个包装器 class ，它使用输入/退出来允许上下文管理器关闭 postgres 数据库连接。

class ClosingConnection:
def __init__(self, schema_name: str) -> None:
    """
    :param schema_name: the db schema (i.e. the tenant)
    """
    super().__init__()
    self.schema_name = schema_name

def __enter__(self):
    try:
        self.conn = psycopg2.connect(
                host=os.environ["DB_HOST"],
                port=os.environ["DB_PORT"],
                database=os.environ["DB_NAME"],
                user=os.environ["DB_USERNAME"],
                password=password,  
                options=f"-c search_path={self.schema_name}",
                cursor_factory=psycopg2.extras.DictCursor,
        )
        return self.conn
    except psycopg2.OperationalError:
        pass

def __exit__(self, exc_type, exc_val, exc_tb):
    self.conn.close()

用法

def get_db(schema_name: str):
    return ClosingConnection(schema_name)

with get_db("test") as db:
    cursor = db.cursor()
    cursor.execute("some sql")
    rv = cursor.fetchall()
    cursor.close()
    

这可以很好地检索数据，但是如果我执行 upsert，则只会执行第一个 upsert。 其余的 rv 正确返回，但实际数据库从未更新。

相反，如果我摆脱关闭 class，这工作正常

用法

def get_db(schema_name: str):
    try:
        self.conn = psycopg2.connect(
                host=os.environ["DB_HOST"],
                port=os.environ["DB_PORT"],
                database=os.environ["DB_NAME"],
                user=os.environ["DB_USERNAME"],
                password=password,  
                options=f"-c search_path={self.schema_name}",
                cursor_factory=psycopg2.extras.DictCursor,
        )
        return self.conn
    except psycopg2.OperationalError:
        pass

conn = get_db("test") 
try:      
    cursor = db.cursor()
    cursor.execute("some sql")
    rv = cursor.fetchall()
    cursor.close()
finally:
    conn.close()

出于好奇，这里是我正在使用的 upsert 命令（我也尝试过带有 upsert 版本的旧版本）

INSERT INTO settings_account (setting_name, setting_value)
    VALUES (%(setting_name)s, %(setting_value)s)
    ON CONFLICT (setting_name)
    DO
        UPDATE SET setting_value=EXCLUDED.setting_value
    RETURNING *

为什么包装器 class 导致数据库写入失败？ 包装器 class 的想法是否有缺陷？

Answer 1

感谢上面的评论，我能够得到这个工作：

    class ClosingConnection:
def __init__(self, schema_name: str) -> None:
    """
    :param schema_name: the db schema (i.e. the tenant)
    """
    super().__init__()
    self.schema_name = schema_name

def __enter__(self):
    try:
        self.conn = psycopg2.connect(
                host=os.environ["DB_HOST"],
                port=os.environ["DB_PORT"],
                database=os.environ["DB_NAME"],
                user=os.environ["DB_USERNAME"],
                password=password,  
                options=f"-c search_path={self.schema_name}",
                cursor_factory=psycopg2.extras.DictCursor,
        )
        return self.conn
    except psycopg2.OperationalError:
        pass

def __exit__(self, exc_type, exc_val, exc_tb):
    self.conn.commit()
    self.conn.close()

def get_db(schema_name: str):
    return ClosingConnection(schema_name)

with get_db("test") as db:
    with db.cursor() as cursor:
        cursor.execute("some sql")
        rv = cursor.fetchall()