如何通过psycopg2实时获取PostgreSQL中Raise Info的输出？

问题背景 #

我有一个长期运行的PostgreSQL存储过程，通过Raise Info 'some status'向客户端反馈执行状态，示例函数如下：

Create or Replace Function Test_Raise() Returns void language plpgsql as $$
    DECLARE
        _X Text;
    Begin
        Raise Info 'Test 1 %', clock_timestamp();
        _x = pg_sleep(2);
        Raise Info 'Test 2 %', clock_timestamp();
        _x = pg_sleep(10);
        Raise Info 'Test 3 %', clock_timestamp();
    End;$$;

在psql或DataGrip中调用select Test_Raise();时，能实时显示每个Raise语句的输出，但使用psycopg2(v2.9.9)时，所有输出仅在查询结束后才批量显示。尝试设置conn.notices和conn.notifier后，只有notices能获取结果，但也是在查询完成后才追加到列表中。

复现代码：

import psycopg2


class Log(list):
    def append(self,item):
        print(item)
        super().append(item)


conn = psycopg2.connect(dbname='jade_replica', host='10.158.1.63', port=32222, user='postgres', password='moozle')
conn.notifies = Log()
conn.notices = Log()
cur = conn.cursor()
cur.execute("select * from test_raise()")

解决方案 #

要实现psycopg2实时捕获Raise Info输出，核心是打破默认阻塞式执行的限制，主动处理数据库的消息流。以下是可行的实现方法：

方法1：异步连接+手动轮询消息 #

启用psycopg2的异步连接，通过poll()轮询连接状态，同时实时提取notice消息：

import psycopg2
from psycopg2.extensions import POLL_OK, POLL_READ, POLL_WRITE


def process_notices(conn):
    # 循环提取所有待处理的notice
    while True:
        notice = conn.get_notice()
        if not notice:
            break
        print(f"实时状态: {notice.strip()}")


# 初始化异步连接
conn = psycopg2.connect(
    dbname='jade_replica', 
    host='10.158.1.63', 
    port=32222, 
    user='postgres', 
    password='moozle',
    async_=True
)

# 完成异步连接握手
while True:
    poll_state = conn.poll()
    if poll_state == POLL_OK:
        break
    # 等待连接就绪（无需额外操作，让系统调度）

# 执行存储过程
cur = conn.cursor()
cur.execute("select * from test_raise()")

# 轮询直到查询完成，同时处理实时notice
while True:
    poll_state = conn.poll()
    if poll_state == POLL_OK:
        break
    process_notices(conn)

# 处理最后剩余的notice
process_notices(conn)

cur.close()
conn.close()

关键说明 #

• 默认同步execute()会阻塞到查询结束，期间所有notice被缓存，查询完成后一次性推送到conn.notices。
• 异步连接模式下，poll()会分阶段返回连接状态，每次轮询时调用get_notice()就能及时获取到存储过程中输出的Raise Info消息。
• 如果需要在同步代码中实现类似效果，也可以通过多线程配合，但异步方式更轻量高效。

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