PostgreSQL同时插入聚合行与更新聚合状态的实现问询

表结构与初始数据 #

创建表：

create table test (
    id serial, 
    contract varchar, 
    amount int, 
    aggregated int, 
    is_aggregate int
);

插入初始数据：

insert into test (contract, amount, aggregated) 
values ('abc', 100, 0), 
       ('abc', 200, 0), 
       ('xyz',  50, 0), 
       ('xyz',  60, 0);

初始表数据：

需求 #

需要按contract维度插入聚合行，同时将被聚合行的aggregated字段设为1，且操作需一次性完成以避免并发问题。期望结果如下：

当前尝试的SQL仅能插入聚合行，无法更新被聚合行的aggregated状态：

INSERT INTO test (contract, amount, aggregated, is_aggregate)
SELECT contract,
       SUM(amount) AS sum_amount,
       1,
       1
FROM test
WHERE aggregated IS NULL 
   OR aggregated = 0
GROUP BY contract
HAVING COUNT(*) > 1;

补充说明 #

多次聚合后的数据状态示例（省略time字段）：

最后一次聚合的行可通过aggregated IS NULL识别。

解决方案 #

使用PostgreSQL的事务+CTE锁定行的方式，确保操作原子性并避免并发问题：

BEGIN;

-- 锁定需要聚合的行，防止并发修改
WITH target_rows AS (
    SELECT id, contract, amount
    FROM test
    WHERE aggregated IS NULL OR aggregated = 0
    FOR UPDATE
),
-- 计算每个contract的聚合金额
aggregated_data AS (
    SELECT contract, SUM(amount) AS total_amount
    FROM target_rows
    GROUP BY contract
    HAVING COUNT(*) > 1
)
-- 插入聚合行，is_aggregate设为1，aggregated留空
INSERT INTO test (contract, amount, is_aggregate)
SELECT contract, total_amount, 1
FROM aggregated_data;

-- 更新被聚合行的aggregated字段为1
UPDATE test
SET aggregated = 1
WHERE id IN (SELECT id FROM target_rows);

COMMIT;

逻辑说明 #

1. 事务包裹：整个操作在一个事务中完成，确保原子性，要么全部成功要么全部回滚。
2. FOR UPDATE锁定：在target_rows中使用FOR UPDATE锁定筛选出的目标行，避免其他事务在当前操作期间修改这些行，解决并发问题。
3. 分步操作：先计算聚合数据并插入新行，再更新原行的aggregated状态，确保所有被聚合的行都被标记为已处理。
4. 兼容多次聚合：每次执行都会筛选出aggregated IS NULL或0的行，处理后标记为1，不会重复处理已聚合的行，适配后续新增数据或重聚合场景。

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