如何改进psycopg2异步游标执行包装函数，实现灵活结果获取并确保连接正确关闭

嘿，我来帮你梳理下代码里的问题，然后给出可行的改进方案～

首先你的代码有两个核心坑：

• 你写了async函数，但用的是同步版的psycopg2 API（psycopg2.connect），这会直接阻塞异步事件循环，完全浪费了异步的优势；
• 你直接返回fetchone()的结果，没法灵活调用fetchall()/fetchmany()，而且原代码里finally块直接关闭连接，就算你想返回游标，外面调用时连接已经关了，游标也会直接失效。

要实现你想要的「调用函数后能自由获取结果」的需求，同时确保连接正确关闭，我们需要用psycopg2的异步接口（或者更现代的psycopg3异步版），再配合异步上下文管理器来管理连接和游标的生命周期——这样连接会在你操作结果的期间保持打开，用完自动关闭，还能自动处理异常回滚。

方案一：基于psycopg2的异步实现 #

首先确保安装了带异步支持的psycopg2：

pip install psycopg2-binary

然后改写你的代码为异步上下文管理器，这样能灵活操作游标，同时自动管理连接：

import asyncio
import psycopg2
from psycopg2.extensions import register_adapter, AsIs
from psycopg2.extras import RealDictCursor  # 可选：返回字典格式的结果，更易读

# 注册适配器，确保异步操作能正确处理元组参数
register_adapter(tuple, AsIs)

class QueryExecutor:
    def __init__(self, query, *args):
        self.query = query
        self.args = args
        self.connection = None
        self.cursor = None

    # 异步上下文进入逻辑：建立连接、执行查询
    async def __aenter__(self):
        self.connection = await psycopg2.extensions.async_connect(**db_params)
        self.cursor = await self.connection.cursor(cursor_factory=RealDictCursor)
        await self.cursor.execute(self.query, self.args)
        return self.cursor

    # 异步上下文退出逻辑：提交/回滚、关闭游标和连接
    async def __aexit__(self, exc_type, exc, tb):
        if exc_type is not None:
            # 有异常就回滚
            await self.connection.rollback()
        else:
            # 无异常就提交
            await self.connection.commit()
        # 确保资源释放
        await self.cursor.close()
        await self.connection.close()

# 用法示例
async def main():
    # 用async with包裹，操作期间连接保持打开
    async with QueryExecutor(
        "SELECT * FROM users WHERE id = %s AND email = %s", 
        "test", "test"
    ) as cursor:
        result_one = await cursor.fetchone()  # 获取单条结果
        result_all = await cursor.fetchall()  # 获取所有结果
        result_many = await cursor.fetchmany(5)  # 获取指定数量结果
        print("单条结果:", result_one)

asyncio.run(main())

方案二：更现代的psycopg3异步版（推荐） #

如果你愿意升级到psycopg3（现在官方更推荐的版本），它的异步API更简洁直观：
先安装依赖：

pip install "psycopg[async]"

然后实现代码：

import asyncio
import psycopg
from psycopg.rows import dict_row  # 同样返回字典格式结果

async def get_query_cursor(query, *args):
    # 建立异步连接
    conn = await psycopg.AsyncConnection.connect(**db_params)
    try:
        cur = conn.cursor(row_factory=dict_row)
        await cur.execute(query, args)
        yield cur  # 返回游标供外部操作
        await conn.commit()  # 无异常则提交
    except Exception as e:
        await conn.rollback()  # 有异常则回滚
        raise e
    finally:
        # 不管成功失败，都关闭游标和连接
        await cur.close()
        await conn.close()

# 用法示例
async def main():
    async for cursor in get_query_cursor(
        "SELECT * FROM users WHERE id = %s AND email = %s",
        "test", "test"
    ):
        result_one = await cursor.fetchone()
        result_all = await cursor.fetchall()
        print("所有结果:", result_all)

asyncio.run(main())

关键改进点总结 #

1. 必须用异步API：同步的psycopg2不能直接在async函数里用，会阻塞事件循环，一定要用async_connect（psycopg2）或AsyncConnection（psycopg3）；
2. 用上下文管理器管理资源：连接和游标必须在你操作结果的期间保持打开，用完自动关闭，避免资源泄漏；
3. 异常处理：一定要在异常时回滚，保证数据一致性；
4. 灵活获取结果：返回游标而不是直接返回结果，让你可以根据需求调用fetchone()/fetchall()/fetchmany()。

备注：内容来源于stack exchange，提问作者ILOVEANGELDUST