feat(rollups): add rollup_contents mapping

.gitattributes

@@ -0,0 +1 @@
+/.sqlx/*.json binary linguist-generated=true

migrations/20251221160635_rollup_contents.down.sql

@@ -0,0 +1 @@
+DROP TABLE IF EXISTS rollup_contents;

migrations/20251221160635_rollup_contents.up.sql

@@ -0,0 +1,11 @@
+CREATE TABLE IF NOT EXISTS rollup_contents
+(
+    -- We have to store the PR through their numbers rather than a FK to the pull_request table. This is due to the rollup PR
+    -- being created at the time of record insertion in this table, so we won't have an entry in pull_request for it
+    rollup_pr_number  BIGINT NOT NULL,
+    member_pr_number  BIGINT NOT NULL,
+    repository TEXT   NOT NULL,
+
+    PRIMARY KEY (repository, rollup_pr_number, member_pr_number)
+);
+CREATE INDEX rollup_contents_member_pr_number_idx ON rollup_contents (repository, member_pr_number);

src/database/client.rs

@@ -2,7 +2,10 @@ use sqlx::PgPool;
 
 use crate::bors::comment::CommentTag;
 use crate::bors::{BuildKind, PullRequestStatus, RollupMode};
-use crate::database::operations::{get_merge_queue_prs, update_pr_auto_build_id};
+use crate::database::operations::{
+    create_rollup_pr_content, get_merge_queue_prs, get_rollup_pr_contents, get_rollups_for_pr,
+    update_pr_auto_build_id,
+};
 use crate::database::{
     BuildModel, BuildStatus, CommentModel, PullRequestModel, RepoModel, TreeState, WorkflowModel,
     WorkflowStatus, WorkflowType,
@@ -344,4 +347,39 @@ impl PgDbClient {
     pub async fn delete_tagged_bot_comment(&self, comment: &CommentModel) -> anyhow::Result<()> {
         delete_tagged_bot_comment(&self.pool, comment.id).await
     }
+
+    /// Register the contents of a rollup in the DB.
+    /// The contents are stored as rollup PR number - member PR number records, scoped under a specific repository.
+    /// All records are inserted in a single transaction.
+    pub async fn create_rollup(
+        &self,
+        repo: &GithubRepoName,
+        rollup_pr_number: &PullRequestNumber,
+        member_pr_numbers: &[PullRequestNumber],
+    ) -> anyhow::Result<()> {
+        let mut tx = self.pool.begin().await?;
+        for member_pr_number in member_pr_numbers {
+            create_rollup_pr_content(&mut *tx, repo, rollup_pr_number, member_pr_number).await?;
+        }
+        tx.commit().await?;
+        Ok(())
+    }
+
+    /// Returns the numbers of all the PRs associated with the requested rollup.
+    pub async fn get_rollup_pr_contents(
+        &self,
+        repo: &GithubRepoName,
+        rollup_pr_number: &PullRequestNumber,
+    ) -> anyhow::Result<Vec<PullRequestNumber>> {
+        get_rollup_pr_contents(&self.pool, repo, rollup_pr_number).await
+    }
+
+    /// Returns all the rollups associated with the requested PR.
+    pub async fn get_rollups_for_pr(
+        &self,
+        repo: &GithubRepoName,
+        pr_number: &PullRequestNumber,
+    ) -> anyhow::Result<Vec<PullRequestNumber>> {
+        get_rollups_for_pr(&self.pool, repo, pr_number).await
+    }
 }

src/database/operations.rs

@@ -1094,3 +1094,80 @@ pub(crate) async fn clear_auto_build(
     })
     .await
 }
+
+pub(crate) async fn create_rollup_pr_content(
+    executor: impl PgExecutor<'_>,
+    repo: &GithubRepoName,
+    rollup_pr_number: &PullRequestNumber,
+    member_pr_number: &PullRequestNumber,
+) -> anyhow::Result<()> {
+    measure_db_query("create_rollup_pr", || async {
+        sqlx::query!(
+            r#"
+        INSERT INTO rollup_contents (repository, rollup_pr_number, member_pr_number)
+        VALUES ($1, $2, $3)
+        "#,
+            repo as &GithubRepoName,
+            rollup_pr_number.0 as i32,
+            member_pr_number.0 as i32
+        )
+        .execute(executor)
+        .await?;
+        Ok(())
+    })
+    .await
+}
+
+pub(crate) async fn get_rollup_pr_contents(
+    executor: impl PgExecutor<'_>,
+    repo: &GithubRepoName,
+    rollup_pr_number: &PullRequestNumber,
+) -> anyhow::Result<Vec<PullRequestNumber>> {
+    measure_db_query("get_rollup_pr_contents", || async {
+        let pr_numbers = sqlx::query_scalar!(
+            r#"
+            SELECT
+                member_pr_number as "member_pr_number: i64"
+            FROM rollup_contents
+            WHERE repository = $1
+              AND rollup_pr_number = $2
+            "#,
+            repo as &GithubRepoName,
+            rollup_pr_number.0 as i32,
+        )
+        .fetch_all(executor)
+        .await?
+        .into_iter()
+        .map(|num| PullRequestNumber(num as u64))
+        .collect();
+        Ok(pr_numbers)
+    })
+    .await
+}
+
+pub(crate) async fn get_rollups_for_pr(
+    executor: impl PgExecutor<'_>,
+    repo: &GithubRepoName,
+    pr_number: &PullRequestNumber,
+) -> anyhow::Result<Vec<PullRequestNumber>> {
+    measure_db_query("get_rollups_for_pr", || async {
+        let rollup_pr_numbers = sqlx::query_scalar!(
+            r#"
+            SELECT
+                rollup_pr_number as "rollup_pr_number: i64"
+            FROM rollup_contents
+            WHERE repository = $1
+              AND member_pr_number = $2
+            "#,
+            repo as &GithubRepoName,
+            pr_number.0 as i32,
+        )
+        .fetch_all(executor)
+        .await?
+        .into_iter()
+        .map(|num| PullRequestNumber(num as u64))
+        .collect();
+        Ok(rollup_pr_numbers)
+    })
+    .await
+}

src/github/rollup.rs

@@ -153,6 +153,7 @@ async fn create_rollup(
         repo_owner,
         pr_nums,
     } = rollup_state;
+    let repo = GithubRepoName::new(&repo_owner, &repo_name);
 
     let username = user_client.username;
 
@@ -175,13 +176,7 @@ async fn create_rollup(
     // Validate PRs
     let mut rollup_prs = Vec::new();
     for &num in &pr_nums {
-        match db
-            .get_pull_request(
-                &GithubRepoName::new(&repo_owner, &repo_name),
-                (num as u64).into(),
-            )
-            .await?
-        {
+        match db.get_pull_request(&repo, (num as u64).into()).await? {
             Some(pr) => {
                 if !pr.is_rollupable() {
                     return Err(RollupError::PullRequestNotRollupable {
@@ -323,6 +318,14 @@ async fn create_rollup(
         .await
         .context("Cannot create PR")?;
 
+    db.create_rollup(
+        &repo,
+        &pr.number,
+        &successes.iter().map(|pr| pr.number).collect::<Vec<_>>(),
+    )
+    .await
+    .context("Cannot create rollup contents record - ensure there are no duplicates in the pr_nums input")?;
+
     Ok(pr)
 }
 
@@ -427,6 +430,19 @@ mod tests {
             make_rollup(ctx, &[&pr2, &pr3, &pr4, &pr5])
                 .await?
                 .assert_status(StatusCode::TEMPORARY_REDIRECT);
+            assert_eq!(
+                ctx.db()
+                    .get_rollup_pr_contents(&default_repo_name(), &PullRequestNumber(6))
+                    .await?,
+                vec![pr2.number, pr3.number, pr5.number]
+            );
+            // Failed merges never make it to the rollup, so this PR shouldn't be present
+            assert_eq!(
+                ctx.db()
+                    .get_rollups_for_pr(&default_repo_name(), &PullRequestNumber(4))
+                    .await?,
+                vec![]
+            );
             Ok(())
         })
         .await;
@@ -497,6 +513,16 @@ mod tests {
                 .assert_status(StatusCode::TEMPORARY_REDIRECT)
                 .get_header("location");
             insta::assert_snapshot!(pr_url, @"https://github.com/rust-lang/borstest/pull/4");
+            assert_eq!(
+                ctx.db()
+                    .get_rollup_pr_contents(&pr2.repo, &PullRequestNumber(4))
+                    .await?,
+                vec![pr2.number, pr3.number]
+            );
+            assert_eq!(
+                ctx.db().get_rollups_for_pr(&pr2.repo, &pr2.number).await?,
+                vec![PullRequestNumber(4)]
+            );
             Ok(())
         })
         .await;
@@ -512,6 +538,60 @@ mod tests {
         ");
     }
 
+    /// Ensure that a PR can be associated to multiple rollups in case of CI failure in the initial rollup
+    #[sqlx::test]
+    async fn multiple_rollups_same_pr(pool: sqlx::PgPool) {
+        let gh = run_test((pool, rollup_state()), async |ctx: &mut BorsTester| {
+            let pr2 = ctx.open_pr((), |_| {}).await?;
+            let pr3 = ctx.open_pr((), |_| {}).await?;
+            ctx.approve(pr2.id()).await?;
+            ctx.approve(pr3.id()).await?;
+
+            make_rollup(ctx, &[&pr2, &pr3]).await?;
+            // Simulate a maintainer closing the rollup due to a CI failure requiring popping a PR from the contents
+            ctx.set_pr_status_closed(4).await?;
+            make_rollup(ctx, &[&pr3]).await?;
+            // Ensure both rollups have the requested PRs
+            assert_eq!(
+                ctx.db()
+                    .get_rollup_pr_contents(&pr2.repo, &PullRequestNumber(4))
+                    .await?,
+                vec![pr2.number, pr3.number]
+            );
+            assert_eq!(
+                ctx.db()
+                    .get_rollup_pr_contents(&pr2.repo, &PullRequestNumber(5))
+                    .await?,
+                vec![pr3.number]
+            );
+            // And ensure PR 3 is associated to both rollups
+            assert_eq!(
+                ctx.db().get_rollups_for_pr(&pr3.repo, &pr3.number).await?,
+                vec![PullRequestNumber(4), PullRequestNumber(5)]
+            );
+            Ok(())
+        })
+        .await;
+        let repo = gh.get_repo(());
+        insta::assert_snapshot!(repo.lock().get_pr(4).description, @r"
+        Successful merges:
+
+         - #2 (Title of PR 2)
+         - #3 (Title of PR 3)
+
+        r? @ghost
+        @rustbot modify labels: rollup
+        ");
+        insta::assert_snapshot!(repo.lock().get_pr(5).description, @r"
+        Successful merges:
+
+         - #3 (Title of PR 3)
+
+        r? @ghost
+        @rustbot modify labels: rollup
+        ");
+    }
+
     #[sqlx::test]
     async fn rollup_order_by_priority(pool: sqlx::PgPool) {
         let gh = run_test((pool, rollup_state()), async |ctx: &mut BorsTester| {