如何在Postgres中用LANGUAGE SQL/BEGIN ATOMIC实现批量添用户到组的存储过程

需求 #

编写LANGUAGE SQL/BEGIN ATOMIC格式的存储过程，将用户身份批量分配到指定组，仅当用户尚未加入该组时执行插入。熟悉SQL Server T-SQL，但对PostgreSQL不熟悉，需适配现有表结构。

现有表结构（不可修改） #

CREATE TABLE IF NOT EXISTS user_identities (
    id BIGINT GENERATED ALWAYS AS IDENTITY PRIMARY KEY NOT NULL,
    -- more columns not relevant to this query
);

CREATE TABLE IF NOT EXISTS user_groups (
    id BIGINT GENERATED ALWAYS AS IDENTITY PRIMARY KEY NOT NULL,
    name TEXT NOT NULL
);

CREATE TABLE IF NOT EXISTS group_identities (
    user_id BIGINT REFERENCES user_identities(id) ON DELETE RESTRICT NOT NULL,
    group_id BIGINT REFERENCES user_groups(id) ON DELETE RESTRICT NOT NULL,
    PRIMARY KEY (user_id, group_id) -- 建议添加复合主键避免重复关联
);

SQL Server T-SQL 参考实现 #

CREATE OR ALTER PROC add_identities_to_group 
  @group_name varchar(50) NULL
AS BEGIN 
DECLARE @use_group_id INT 
IF @group_name IS NULL 
  SET @use_group_id = (SELECT TOP 1 id FROM user_groups WHERE id NOT IN (SELECT group_id FROM group_identities) ORDER BY id ASC) 
ELSE SET @use_group_id = (SELECT id FROM user_groups WHERE name = @group_name)

INSERT INTO group_identities (user_id, group_id) 
  SELECT id, @use_group_id FROM user_identities 
    WHERE id NOT IN (SELECT user_id FROM group_identities)
END 
GO 

遇到的问题 #

• 不清楚PostgreSQL中BEGIN ATOMIC原子存储过程如何结合变量与查询结果实现逻辑
• 不确定原子块中变量名与列名的区分规则
• 搜索"标准SQL"结果多指向SQL Server，难以找到适配PostgreSQL的具体方案

正确的PostgreSQL 原子存储过程实现 #

CREATE OR REPLACE PROCEDURE add_identities_to_group(
    group_name TEXT DEFAULT NULL
) LANGUAGE SQL
BEGIN ATOMIC
    DECLARE use_group_id BIGINT;

    -- 确定目标组ID
    IF group_name IS NULL THEN
        SELECT id INTO use_group_id
        FROM user_groups
        WHERE id NOT IN (SELECT group_id FROM group_identities)
        ORDER BY id ASC
        LIMIT 1; -- PostgreSQL使用LIMIT替代SQL Server的TOP 1
    ELSE
        SELECT id INTO use_group_id
        FROM user_groups
        WHERE name = group_name;
    END IF;

    -- 仅插入未加入目标组的用户关联记录
    INSERT INTO group_identities (user_id, group_id)
    SELECT ui.id, use_group_id
    FROM user_identities ui
    WHERE ui.id NOT IN (
        SELECT user_id
        FROM group_identities
        WHERE group_id = use_group_id
    );
END;

关键语法说明 #

1. 原子块与参数支持：BEGIN ATOMIC是PostgreSQL支持的标准SQL原子存储过程语法，完全支持输入参数及默认值设置
2. 变量声明与赋值：使用DECLARE声明变量，通过SELECT ... INTO完成变量赋值，这是SQL语言存储过程的标准方式
3. 语法差异适配：用LIMIT 1替代SQL Server的TOP 1；变量类型与表字段保持一致（BIGINT而非integer）
4. 逻辑修正：插入时检查用户是否未加入目标组（而非任何组），更贴合"仅当用户尚未加入该组时"的需求

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