在psycopg2中实现上下文管理器的正确方法是什么？

Question

我一直在使用函数重构我的 psycopg2 代码，之前我将所有代码都放在 try- except-finally 块上，但是我不太确定如何实现上下文管理器来处理连接和光标。我的 SQL 查询工作如下：

def random_query(schema, table, username, number_of_files):
    random_query = sql.SQL("SELECT * FROM {schema}.{table} WHERE username = {username} ORDER BY RANDOM() LIMIT {limit}").format(
        schema=sql.Identifier(schema),
        table=sql.Identifier(table),
        username=sql.Literal(username),
        limit=sql.Literal(number_of_files)
        )
    cursor.execute(random_query)
    return cursor.fetchone()

def insert_query(schema, table, values):
    insert_query = sql.SQL("INSERT INTO {schema}.{table}(shortcode, username, filename, extension) VALUES ({shortcode}, {username}, {filename}, {extension})").format(
        schema=sql.Identifier(schema),
        table=sql.Identifier(table),
        shortcode=sql.Literal(values[0]),
        username=sql.Literal(values[1]),
        filename=sql.Literal(values[2]),
        extension=sql.Literal(values[3])
        )
    cursor.execute(insert_query)
    conn.commit()

1. 第一版：

@contextmanager
def get_connection():
    connection = psycopg2.connect(**DB_CONNECTION)
    try:
        yield connection
    except Exception as err:
        connection.rollback()
        print('Error: ', err)
        raise
    finally:
        if (connection):
            connection.close()
            print("Connection is closed.")

@contextmanager
def get_cursor(connection):
    cursor = connection.cursor()
    try:
        yield cursor
    finally:
        cursor.close()

with get_connection() as conn, get_cursor(conn) as cursor:
    random_record = random_query('test_schema', 'test_table', 'username', 1)
    insert_query('test_schema', 'test_table2', random_record)

2. 第二版：

@contextmanager
def sql_connection():
    connection = psycopg2.connect(**DB_CONNECTION)
    cursor = connection.cursor()
    try:
        yield connection,cursor
    except Exception as err:
        connection.rollback()
        print('Error : ', err)
        raise
    finally:
        if (connection):
            cursor.close()
            connection.close()
            print("Connection is closed")

with sql_connection() as (conn, cursor):
    random_record = random_query('test_schema', 'test_table', 'username', 1)
    insert_query('test_schema', 'test_table2', random_record)

我的问题是：

1. 第一版和第二版有什么区别吗？哪一个更可取？
2. 正如您在

   insert_query

   中看到的，有一行调用

   conn.commit()

   从文档中，我知道如果我们使用上下文管理器，则没有必要这样做。我可以删除它们吗？

2.5版本更改：如果连接用在with语句中， 如果没有引发异常，则自动调用该方法 带块。

Answer 1

这两个版本都不是更好的，你仍然通过重复行为使事情变得过于复杂。根据此处的示例连接：

import psycopg2

connection = psycopg2.connect(**DB_CONNECTION)

with connection:
    with connection.cursor() as cur:
        cur.execute(<sql>)

with connection:
    with connection.cursor() as cur:
        cur.execute(<other_sql>)

提交、回滚连接和关闭游标都已为您完成。当您不想再使用该连接时，您所要做的就是

connection.close()

。

更新

问题和答案适用于

psycopg2

，如果您使用

psycopg(3)

，则该版本中的连接上下文管理器行为已更改。在

psycopg(3)

中，

with connection

将在完成时关闭连接，而在

psycopg2

中，它只是关闭事务。