如何从PostgreSQL函数返回旧ID与重复行新ID的映射关系

问题说明 #

需要编写PostgreSQL函数，复制foo表中符合指定condition_col条件的行，将新行的condition_col设为目标值，同时返回原行ID与新行ID的对应关系。

示例数据 #

CREATE TABLE foo (
  id integer PRIMARY KEY GENERATED ALWAYS AS IDENTITY (START WITH 6),
  foo_content text NOT NULL,
  condition_col integer NOT NULL
);

INSERT INTO foo OVERRIDING SYSTEM VALUE VALUES
(  1 , 'Sing, O goddess', 99),
(  2 , 'the anger of Achilles', 99),
(  3 , 'son of Peleus', 99),
(  4 , 'that brought countless ills', 2),
(  5 , 'upon the Achaeans', 3);

修正后的函数实现 #

原写法无法关联新旧ID，核心原因是INSERT的RETURNING仅返回新生成ID，未保留原ID的关联。可以通过在INSERT的查询中携带原ID，再通过RETURNING同时输出新旧ID来解决：

CREATE OR REPLACE FUNCTION duplicate_entries(condition INT, new_condition INT)
RETURNS TABLE (
        old_id INT,
        new_id INT) AS
$$
BEGIN
    RETURN QUERY (
        WITH selected_rows AS (
            SELECT id, foo_content
            FROM foo
            WHERE condition_col = condition
        )
        INSERT INTO foo (foo_content, condition_col)
        SELECT foo_content, new_condition
        FROM selected_rows
        RETURNING selected_rows.id AS old_id, foo.id AS new_id
    );
END;
$$ LANGUAGE plpgsql;

验证结果 #

调用函数获取映射：

SELECT * FROM duplicate_entries(99, 100);

将得到预期输出：

查看更新后的表数据：

SELECT * FROM foo;

原理说明 #

1. selected_rows CTE先筛选出需要复制的行，同时保留原ID和内容字段；
2. INSERT操作从selected_rows中读取内容和新条件值插入表中；
3. RETURNING子句通过selected_rows.id获取原行ID，foo.id获取新生成的ID，直接建立两者的映射关系，无需临时表或额外关联操作。

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