PL/pgSQL中实现逐行赋不同值的单条UPDATE操作方法

问题背景 #

我有一个存储元素与对应值映射关系的Python字典，示例如下：

db_rows_values = {
     <element_uuid_1>: 12,
     <element_uuid_2>: "abc",
     <element_uuid_3>: [123, 124, 125],
}

需要通过单条查询完成上述数据的更新。此前在Python中通过循环生成CASE语句拼接查询实现该需求，代码如下：

sql_query_elements_values_part = " ".join([f"WHEN '{element_row['element_id']}' "
                                           f"THEN '{ujson.dumps(element_row['value'])}'::JSONB "
                                           for element_row in db_row_values])

query_part_elements_values_update = f"""
    elements_value_update AS (
        UPDATE m2m_entries_n_elements 
        SET value = 
        CASE element_id 
            {sql_query_elements_values_part}
            ELSE NULL 
        END 
        WHERE element_id = ANY(%(elements_ids)s::UUID[])
        AND entry_id = ANY(%(entries_ids)s::UUID[])

        RETURNING element_id, entry_id, value

    ),
"""

现在需要将该逻辑改写为plpgsql实现，支持将db_rows_values以ROWTYPE数组或JSON格式传入，实现等价于原WHEN THEN分支的匹配赋值逻辑。

推荐实现：JSONB传参（无需自定义类型） #

这是最简便的方案，不需要提前在数据库创建自定义类型，逻辑和原CASE分支完全等价，性能优于拼接长CASE语句，同时彻底规避SQL注入风险。

函数定义 #

CREATE OR REPLACE FUNCTION batch_update_element_values(
    p_value_map JSONB, -- 直接传入Python端序列化后的字典，格式为{"<uuid1>": 12, "<uuid2>": "abc", ...}
    p_element_ids UUID[],
    p_entry_ids UUID[]
) RETURNS TABLE (
    element_id UUID,
    entry_id UUID,
    value JSONB
) AS $$
BEGIN
    RETURN QUERY
    WITH elements_value_update AS (
        UPDATE m2m_entries_n_elements t
        -- 等价于原CASE逻辑：键存在则取对应JSONB值，不存在则返回NULL，完全对齐ELSE NULL分支
        SET value = p_value_map -> t.element_id
        WHERE t.element_id = ANY(p_element_ids)
          AND t.entry_id = ANY(p_entry_ids)
        RETURNING t.element_id, t.entry_id, t.value
    )
    SELECT * FROM elements_value_update;
END;
$$ LANGUAGE plpgsql VOLATILE;

Python端调用示例 #

import ujson
import psycopg2

conn = psycopg2.connect("数据库连接串")
cur = conn.cursor()

cur.execute(
    "SELECT * FROM batch_update_element_values(%s::JSONB, %s::UUID[], %s::UUID[])",
    (
        ujson.dumps(db_rows_values), # 直接序列化原字典传入
        list(db_rows_values.keys()),
        entries_ids # 原逻辑中的entries_ids参数
    )
)
updated_records = cur.fetchall()
conn.commit()

可选实现：ROWTYPE数组传参 #

如果需要用复合类型数组传参，需要先在数据库创建对应的复合类型，再通过关联更新实现匹配逻辑。

第一步：创建复合类型 #

CREATE TYPE element_value_pair AS (
    element_id UUID,
    value JSONB
);

第二步：创建更新函数 #

CREATE OR REPLACE FUNCTION batch_update_element_values_rowtype(
    p_value_pairs element_value_pair[],
    p_element_ids UUID[],
    p_entry_ids UUID[]
) RETURNS TABLE (
    element_id UUID,
    entry_id UUID,
    value JSONB
) AS $$
BEGIN
    RETURN QUERY
    WITH value_map AS (
        SELECT (pair).element_id AS map_eid, (pair).value AS map_val
        FROM unnest(p_value_pairs) pair
    ),
    elements_value_update AS (
        UPDATE m2m_entries_n_elements t
        SET value = vm.map_val
        FROM value_map vm
        WHERE t.element_id = vm.map_eid
          AND t.element_id = ANY(p_element_ids)
          AND t.entry_id = ANY(p_entry_ids)
        RETURNING t.element_id, t.entry_id, t.value
    )
    SELECT * FROM elements_value_update;
END;
$$ LANGUAGE plpgsql VOLATILE;

注意：该实现默认仅更新数组中存在对应element_id的行，如果需要完全对齐原逻辑中「未匹配到的element_id对应value设为NULL」的规则，需要将关联改为左连接，覆盖所有符合WHERE条件的行。

方案对比 #

• JSONB方案无需提前定义数据库结构，跨语言调用适配成本低，JSONB按键查询时间复杂度为O(1)，更新条目量大时性能优势明显，优先选用。
• ROWTYPE数组方案适合数据库内部调用场景，跨语言传参时需要构造复合类型数组，适配成本更高。
• 两种参数化方案都彻底解决了原Python拼接SQL存在的注入风险。

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