为何PostgreSQL代码块需返回1？动态执行预编译语句行计数异常

问题现象 #

我们在实现通用存储过程动态执行预编译语句时遇到了奇怪的行为：

• 当预编译的UPDATE语句带有RETURNING 1时，存储过程能正确输出影响行数（比如3行）；
• 但去掉RETURNING 1后，存储过程输出的行计数变成0，和实际影响行数不符。

简化示例代码：

drop table if exists test;
create table test (str text);
insert into test(str) values ('row1'), ('row2'), ('row3');

-- 通用存储过程
drop procedure if exists run_sql(text);
create procedure run_sql(named_statement text)
as $$
declare
    _row_count bigint = -1;
begin
    execute 'execute ' || named_statement;
    get diagnostics _row_count := row_count;
    raise notice 'Rows updated: %', _row_count;
end
$$ language plpgsql;

-- 测试用预编译语句（去掉RETURNING 1后行计数异常）
prepare update_test as
        update test
        set str = str
        where true;
call run_sql('update_test');
deallocate update_test;

直接在PL/pgSQL块中执行预编译语句时，即使没有RETURNING也能拿到正确行计数，但通过动态EXECUTE 'execute <prepared_stmt>'调用时就不行。

原因分析 #

这是PostgreSQL PL/pgSQL的上下文特性导致的：

• GET DIAGNOSTICS ... ROW_COUNT只能捕获当前PL/pgSQL上下文直接执行的SQL语句的影响行数；
• 当使用EXECUTE 'execute <prepared_stmt>'这种嵌套执行时，内层预编译语句的行计数不会传递到外层的动态EXECUTE上下文里；
• 只有当预编译语句带有RETURNING时，内层执行会返回结果集，外层的EXECUTE会处理这个结果集，此时ROW_COUNT会被设置为结果集的行数（也就是实际影响行数）。

解决方案 #

方案一：利用事务级统计函数（推荐） #

通过查询PostgreSQL提供的事务级统计函数，在执行预编译语句前后获取计数差值，从而得到准确的影响行数。同时修复原代码的SQL注入风险：

drop procedure if exists run_sql(text);
create procedure run_sql(named_statement text)
as $$
declare
    _row_count bigint = 0;
    _prev_ins bigint;
    _prev_upd bigint;
    _prev_del bigint;
    _curr_ins bigint;
    _curr_upd bigint;
    _curr_del bigint;
begin
    -- 获取执行前的事务内操作计数
    select pg_stat_get_xact_tuples_inserted(current_database()),
           pg_stat_get_xact_tuples_updated(current_database()),
           pg_stat_get_xact_tuples_deleted(current_database())
    into _prev_ins, _prev_upd, _prev_del;

    -- 安全执行预编译语句（用format避免SQL注入）
    execute format('EXECUTE %I', named_statement);

    -- 获取执行后的事务内操作计数
    select pg_stat_get_xact_tuples_inserted(current_database()),
           pg_stat_get_xact_tuples_updated(current_database()),
           pg_stat_get_xact_tuples_deleted(current_database())
    into _curr_ins, _curr_upd, _curr_del;

    -- 计算总影响行数
    _row_count := (_curr_ins - _prev_ins) + (_curr_upd - _prev_upd) + (_curr_del - _prev_del);

    raise notice 'Rows affected: %', _row_count;

    -- 后续业务逻辑
end
$$ language plpgsql;

注意：如果存储过程在执行预编译语句前后还有其他数据修改操作，需要调整逻辑单独统计对应操作的计数（比如只统计更新行数）。

方案二：动态创建PL/pgSQL块捕获行计数 #

通过动态创建临时PL/pgSQL块执行预编译语句，利用PostgreSQL的通知机制传递行计数：

drop procedure if exists run_sql(text);
create procedure run_sql(named_statement text)
as $$
declare
    _row_count bigint = 0;
    _notify_msg text;
begin
    -- 动态执行内部块，获取行计数并发送通知
    execute format('
        do $$
        declare
            cnt bigint;
        begin
            execute %L;
            get diagnostics cnt := row_count;
            perform pg_notify(''row_count_notify'', cnt::text);
        end $$;
    ', named_statement);

    -- 监听并获取通知内容
    listen row_count_notify;
    fetch notification into _notify_msg;
    _row_count := _notify_msg::bigint;
    unlisten row_count_notify;

    raise notice 'Rows affected: %', _row_count;

    -- 后续业务逻辑
end
$$ language plpgsql;

这个方法适合复杂场景，但要注意并发执行时的通知冲突问题。

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