如何用Diesel、Rocket与PostgreSQL实现批量Upsert并返回结果

核心思路 #

把原来逐行循环插入的逻辑改成两次批量操作：先批量处理所有帖子的Upsert并获取对应ID，再批量处理所有评论的Upsert，将数据库往返次数从O(N+M)（N为帖子数，M为评论数）压缩到2次。

步骤1：补全数据库约束（关键） #

评论表当前没有明确的唯一约束，Upsert需要指定冲突判断条件，建议补充：

• 若评论ID是唯一标识：给comment.id添加主键约束
• 若同一帖子下不允许重复内容：添加联合唯一约束(post_id, content)

修改后的Schema示例：

create table post 
(
    id   integer not null primary key,
    name text    not null
);

create table comment
(
    id      integer primary key,
    post_id integer not null references post(id),
    content text not null,
    -- 可选：按帖子+内容去重的约束
    unique(post_id, content)
);

步骤2：批量Upsert帖子并返回ID #

利用Diesel的批量插入API，结合ON CONFLICT实现Upsert，同时返回所有帖子的ID（无论插入还是更新），为后续评论关联做准备。

步骤3：预处理评论数据，关联帖子ID #

将每个帖子的评论与对应的帖子ID绑定，生成统一的评论批量插入集合。

步骤4：批量Upsert评论 #

用批量插入API处理所有评论，根据之前确定的唯一约束配置ON CONFLICT逻辑。

优化后完整代码 #

#[database("example_db")]
pub struct ExampleDB(diesel::PgConnection);

use diesel::prelude::*;
use crate::schema::{post, comment}; // 假设已通过diesel_cli生成schema

// 适配Diesel的插入结构体
#[derive(Insertable)]
#[diesel(table_name = post)]
struct InsertablePost {
    id: i32,
    name: String,
}

#[derive(Insertable)]
#[diesel(table_name = comment)]
struct InsertableComment {
    id: Option<i32>, // 若评论ID由客户端提供则改为i32，数据库自增则留空
    post_id: i32,
    content: String,
}

// PostJSON转InsertablePost的转换逻辑
impl From<PostJSON> for InsertablePost {
    fn from(json: PostJSON) -> Self {
        InsertablePost {
            id: /* 从PostJSON获取或生成ID，匹配原逻辑 */,
            name: json.name,
        }
    }
}

async fn insert_posts(posts_to_insert: Vec<PostJSON>, db: ExampleDB) -> Result<(), diesel::result::Error> {
    db.run(|conn| {
        // 1. 转换为批量帖子插入集合
        let batch_posts: Vec<InsertablePost> = posts_to_insert.iter().map(|p| p.clone().into()).collect();

        // 2. 批量Upsert帖子，返回所有帖子ID
        let post_ids: Vec<i32> = diesel::insert_into(post::table)
            .values(&batch_posts)
            .on_conflict(post::id)
            .do_update()
            .set(post::name.eq(diesel::excluded(post::name))) // 冲突时更新name字段
            .returning(post::id)
            .get_results(conn)?;

        // 3. 预处理所有评论，绑定对应post_id
        let mut batch_comments = Vec::new();
        for (idx, post_json) in posts_to_insert.iter().enumerate() {
            let current_post_id = post_ids[idx];
            for comment_json in &post_json.comments {
                batch_comments.push(InsertableComment {
                    id: /* 评论ID，根据实际场景赋值 */,
                    post_id: current_post_id,
                    content: comment_json.content.clone(),
                });
            }
        }

        // 4. 批量Upsert评论
        diesel::insert_into(comment::table)
            .values(&batch_comments)
            .on_conflict(comment::id) // 若用ID作为唯一约束
            // 若用联合约束则替换为：.on_conflict((comment::post_id, comment::content))
            .do_update()
            .set(comment::content.eq(diesel::excluded(comment::content))) // 冲突时更新内容
            .execute(conn)?;

        Ok(())
    }).await
}

关键注意事项 #

• 异步适配：确保使用Diesel的异步版本（diesel_async），若用同步Diesel则调整db.run内的逻辑为同步执行。
• ID生成逻辑：若帖子/评论ID由数据库自增，需移除插入结构体中的ID字段，通过returning获取数据库生成的ID。
• 事务原子性：可将整个批量操作包裹在事务中，确保帖子和评论的操作要么全成功要么全失败，只需在db.run内添加conn.transaction(|| { /* 批量操作逻辑 */ })。

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