psycopg2上下文管理器实现方案对比及commit调用疑问

我正在重构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()

我写了两种上下文管理器实现版本：

版本一 #

@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)

版本二 #

@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()，根据Psycopg2文档（2.5版本起，若连接在with语句中使用且块内无异常，该方法会自动调用），我是否可以移除这行代码？

问题1解答：两个版本的差异与选型 #

• 结构差异：
  版本一将连接和游标拆分为两个独立的上下文管理器，职责单一：get_connection仅负责连接的创建、异常回滚和关闭；get_cursor仅负责游标创建和关闭。版本二则把连接和游标绑定在同一个上下文管理器中，一次性返回两者。
• 灵活性差异：
  版本一更灵活——如果需要在同一个连接下创建多个游标，或者单独操作连接（比如手动调用conn.commit/rollback），拆分的方式更适配。版本二则绑定了连接和游标的生命周期，适合大多数单游标操作场景，但扩展性稍弱。
• 选型建议：
  如果业务场景以单游标操作为主，版本二写法更简洁，减少冗余代码；如果需要频繁在同一连接下切换游标，或有单独操作连接的需求，版本一的职责分离设计更实用，后期维护也更清晰。

问题2解答：是否可以移除conn.commit() #

可以移除，但需要满足几个前提：

• 确保提交操作完全处于连接的with块范围内。你的代码中insert_query在with块内执行，只要块内无异常抛出，Psycopg2会自动提交事务。
• 确认使用的Psycopg2版本为2.5及以上（当前主流环境基本都满足，老旧环境需额外验证）。
• 注意：如果insert_query存在在with块外被调用的可能，手动commit是必要的，但从你的代码逻辑来看，它仅在with块内执行，因此可以安全移除。

额外建议：你的random_query和insert_query直接依赖全局的cursor和conn变量，这不是健壮的设计。建议将游标和连接作为参数传入函数，比如：

def random_query(cursor, 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(conn, cursor, 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)

这样函数独立性更强，避免全局变量带来的潜在问题。

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