如何用可复用异步函数/生成器精简psycopg连接池样板代码？

我用psycopg通过连接池连接PostgreSQL，运行正常，但每个需要执行事务SQL的函数都有三层嵌套：

原代码实现 #

/app/db.py #

from os import getenv
from psycopg_pool import AsyncConnectionPool

pool = AsyncConnectionPool(getenv('POSTGRES_URL'))

/app/foo.py #

from db import pool
from psycopg.rows import dict_row


async def create_foo(**kwargs):
    foo = {}
    async with pool.connection() as conn:
        async with conn.transaction():
            async with conn.cursor(row_factory=dict_row) as cur:
                # 执行SQL查询
    return foo


async def update_foo(foo_id, **kwargs):
    foo = {}
    async with pool.connection() as conn:
        async with conn.transaction():
            async with conn.cursor(row_factory=dict_row) as cur:
                # 执行SQL查询
    return foo

重构尝试及报错 #

我想把这部分逻辑抽象成辅助函数，于是做了如下重构：

重构后的/app/db.py #

from contextlib import asynccontextmanager
from os import getenv
from psycopg_pool import AsyncConnectionPool

pool = AsyncConnectionPool(getenv('POSTGRES_URL'))


@asynccontextmanager
async def get_tx_cursor(**kwargs):
    async with pool.connection() as conn:
        conn.transaction()
        cur = conn.cursor(**kwargs)
        yield cur

调用代码（/app/foo.py） #

from db import get_tx_cursor
from psycopg.rows import dict_row


async def create_foo(**kwargs):
    foo = {}
    async with get_tx_cursor(row_factory=dict_row) as cur:
        # 执行SQL查询
    return foo

但运行时报错：

TypeError: '_AsyncGeneratorContextManager' object does not support the context manager protocol

我还尝试了另一种变体：

async def get_tx_cursor(**kwargs):
    async with pool.connection() as conn:
        async with conn.transaction():
            async with conn.cursor(**kwargs) as cur:
                yield cur

结果类似，似乎无法直接使用生成器。

需求 #

有没有无需额外依赖库的简洁方法，能将cursor暴露给调用函数？

使用版本 #

• python: 3.11
• psycopg: 3.1.8
• psycopg-pool: 3.1.6

解决方案 #

正确的异步上下文管理器实现 #

问题出在两个点：

1. 第一个重构中，conn.transaction()没有用async with管理事务上下文，conn.cursor()也未通过async with处理，导致资源无法正确释放；
2. 第二个变体缺少@asynccontextmanager装饰器，直接返回的异步生成器无法被async with识别。

正确的辅助函数需要完整嵌套所有异步上下文，并通过装饰器标记：

# /app/db.py
from contextlib import asynccontextmanager
from os import getenv
from psycopg_pool import AsyncConnectionPool

pool = AsyncConnectionPool(getenv('POSTGRES_URL'))


@asynccontextmanager
async def get_tx_cursor(**kwargs):
    # 从连接池获取并管理连接
    async with pool.connection() as conn:
        # 开启并管理事务上下文
        async with conn.transaction():
            # 创建并管理游标上下文
            async with conn.cursor(**kwargs) as cur:
                # 将游标暴露给调用方
                yield cur

调用方式保持不变 #

# /app/foo.py
from db import get_tx_cursor
from psycopg.rows import dict_row


async def create_foo(**kwargs):
    foo = {}
    async with get_tx_cursor(row_factory=dict_row) as cur:
        # 示例：执行插入并返回结果
        await cur.execute(
            "INSERT INTO foo (name, value) VALUES (%s, %s) RETURNING *",
            (kwargs['name'], kwargs['value'])
        )
        foo = await cur.fetchone()
    return foo


async def update_foo(foo_id, **kwargs):
    foo = {}
    async with get_tx_cursor(row_factory=dict_row) as cur:
        # 示例：执行更新并返回结果
        await cur.execute(
            "UPDATE foo SET value = %s WHERE id = %s RETURNING *",
            (kwargs['value'], foo_id)
        )
        foo = await cur.fetchone()
    return foo

说明 #

• 该实现会自动处理连接的获取/释放、事务的提交/回滚、游标的创建/关闭；
• 如果async with块内抛出异常，事务会自动回滚，连接也会正确归还到连接池；
• 调用方只需专注编写SQL逻辑，无需重复编写三层嵌套代码。

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