如何根据PostgreSQL插入语句返回结果决定事务提交或回滚

你在psql中看到的INSERT 0 1、INSERT 0 0返回值里，第二个数字就是当前语句实际插入/更新的行数，全程不需要执行额外SELECT查询，直接通过数据库内置能力就能拿到这个值，用来判断事务是提交还是回滚，常用实现分三类场景：

数据库端PL/pgSQL逻辑实现 #

如果事务逻辑直接在数据库侧通过匿名块、存储过程执行，用GET DIAGNOSTICS语法直接读取行统计值即可，无额外性能开销：

-- 测试表结构和你给出的一致
create table a(id int PRIMARY KEY, foo text, bar text);

DO $$
DECLARE
  v_affected_rows integer;
BEGIN
  -- 执行第一条UPSERT
  insert into a values(1, 'foo', 'bar') 
  on conflict (id) DO UPDATE SET foo = excluded.foo where a.bar = 'bar';
  -- 直接读取当前语句影响的行数
  GET DIAGNOSTICS v_affected_rows = ROW_COUNT;

  -- 无生效行直接回滚退出
  IF v_affected_rows = 0 THEN
    ROLLBACK;
    RETURN;
  END IF;

  -- 执行第二条UPSERT
  insert into a values(1, 'foo', 'bar') 
  on conflict (id) DO UPDATE SET foo = excluded.foo where a.bar = 'bara';
  GET DIAGNOSTICS v_affected_rows = ROW_COUNT;

  IF v_affected_rows = 0 THEN
    ROLLBACK;
    RETURN;
  END IF;

  -- 所有校验通过后提交事务
  COMMIT;
END $$;

这里拿到的v_affected_rows和psql客户端打印的返回值完全一致，统计逻辑是数据库执行语句时顺带完成的，没有额外IO开销。

应用层代码实现 #

如果逻辑是在业务应用代码中执行，所有主流PostgreSQL驱动都原生支持读取语句影响行数，不需要额外查询：

• Java JDBC：executeUpdate()方法的返回值就是影响行数
• Python psycopg2/psycopg3：游标执行完语句后，直接读取cursor.rowcount属性即可
• Go pgx驱动：Exec方法返回的结果对象调用RowsAffected()就能拿到对应数值

以Python为例的实现代码：

import psycopg2

conn = psycopg2.connect("dbname=foouser")
cur = conn.cursor()
try:
    cur.execute("""
        insert into a values(1, 'foo', 'bar')
        on conflict (id) DO UPDATE SET foo = excluded.foo where a.bar = %s
    """, ("bara",))
    # 直接读取影响行数，无额外查询
    affected = cur.rowcount
    if affected == 0:
        conn.rollback()
    else:
        # 执行后续业务逻辑
        conn.commit()
except Exception:
    conn.rollback()
finally:
    cur.close()
    conn.close()

psql脚本场景实现 #

如果是直接在psql客户端写脚本执行，可以用psql内置的ROW_COUNT变量做条件判断：

START TRANSACTION;
insert into a values(1, 'foo', 'bar') on conflict (id) DO UPDATE SET foo = excluded.foo where a.bar = 'bara';
-- 判断上一条语句的影响行数
\if :ROW_COUNT = 0
  \echo 无行被插入/更新，执行回滚
  ROLLBACK;
\else
  \echo 执行正常，提交事务
  COMMIT;
\endif

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