From 9c88e4bc637e762a5b57ab16f2a93cba97d16af1 Mon Sep 17 00:00:00 2001
From: apiraino <apiraino@users.noreply.github.com>
Date: Thu, 29 Feb 2024 10:38:02 +0100
Subject: [PATCH] Add webhook handler to update PR workload queues

---
 src/config.rs               |  12 ++++
 src/db.rs                   |   1 +
 src/github.rs               |  13 ++++-
 src/handlers.rs             |   2 +
 src/handlers/pr_tracking.rs | 111 ++++++++++++++++++++++++++++++++++++
 5 files changed, 137 insertions(+), 2 deletions(-)
 create mode 100644 src/handlers/pr_tracking.rs

diff --git a/src/config.rs b/src/config.rs
index 244c35ef..e1ecce2d 100644
--- a/src/config.rs
+++ b/src/config.rs
@@ -15,6 +15,10 @@ lazy_static::lazy_static! {
         RwLock::new(HashMap::new());
 }
 
+// This struct maps each possible option of the triagebot.toml.
+// See documentation of options at: https://forge.rust-lang.org/triagebot/pr-assignment.html#configuration
+// When adding a new config option to the triagebot.toml, it must be also mapped here
+// Will be used by the `issue_handlers!()` or `command_handlers!()` macros.
 #[derive(PartialEq, Eq, Debug, serde::Deserialize)]
 #[serde(rename_all = "kebab-case")]
 #[serde(deny_unknown_fields)]
@@ -39,6 +43,7 @@ pub(crate) struct Config {
     // We want this validation to run even without the entry in the config file
     #[serde(default = "ValidateConfig::default")]
     pub(crate) validate_config: Option<ValidateConfig>,
+    pub(crate) pr_tracking: Option<ReviewPrefsConfig>,
 }
 
 #[derive(PartialEq, Eq, Debug, serde::Deserialize)]
@@ -317,6 +322,12 @@ pub(crate) struct GitHubReleasesConfig {
     pub(crate) changelog_branch: String,
 }
 
+#[derive(PartialEq, Eq, Debug, serde::Deserialize)]
+pub(crate) struct ReviewPrefsConfig {
+    #[serde(default)]
+    _empty: (),
+}
+
 fn get_cached_config(repo: &str) -> Option<Result<Arc<Config>, ConfigurationError>> {
     let cache = CONFIG_CACHE.read().unwrap();
     cache.get(repo).and_then(|(config, fetch_time)| {
@@ -463,6 +474,7 @@ mod tests {
                 mentions: None,
                 no_merges: None,
                 validate_config: Some(ValidateConfig {}),
+                pr_tracking: None,
             }
         );
     }
diff --git a/src/db.rs b/src/db.rs
index 6d5b5e8c..1da96372 100644
--- a/src/db.rs
+++ b/src/db.rs
@@ -332,6 +332,7 @@ CREATE table review_prefs (
     assigned_prs INT[] NOT NULL DEFAULT array[]::INT[]
 );",
     "
+CREATE EXTENSION intarray;
 CREATE UNIQUE INDEX review_prefs_user_id ON review_prefs(user_id);
  ",
 ];
diff --git a/src/github.rs b/src/github.rs
index 3b8d7ae8..c078fe56 100644
--- a/src/github.rs
+++ b/src/github.rs
@@ -286,8 +286,11 @@ pub struct Issue {
     ///
     /// Example: `https://github.com/octocat/Hello-World/pull/1347`
     pub html_url: String,
+    // User performing an `action`
     pub user: User,
     pub labels: Vec<Label>,
+    // Users assigned to the issue/pr after `action` has been performed
+    // These are NOT the same as `IssueEvent.assignee`
     pub assignees: Vec<User>,
     /// Indicator if this is a pull request.
     ///
@@ -954,8 +957,14 @@ pub enum IssuesAction {
     Unpinned,
     Closed,
     Reopened,
-    Assigned,
-    Unassigned,
+    Assigned {
+        /// Github users assigned to the issue / pull request
+        assignee: User,
+    },
+    Unassigned {
+        /// Github users removed from the issue / pull request
+        assignee: User,
+    },
     Labeled {
         /// The label added from the issue
         label: Label,
diff --git a/src/handlers.rs b/src/handlers.rs
index 855be5aa..f3045a55 100644
--- a/src/handlers.rs
+++ b/src/handlers.rs
@@ -38,6 +38,7 @@ mod note;
 mod notification;
 mod notify_zulip;
 mod ping;
+pub mod pr_tracking;
 mod prioritize;
 pub mod pull_requests_assignment_update;
 mod relabel;
@@ -168,6 +169,7 @@ issue_handlers! {
     no_merges,
     notify_zulip,
     review_requested,
+    pr_tracking,
     validate_config,
 }
 
diff --git a/src/handlers/pr_tracking.rs b/src/handlers/pr_tracking.rs
new file mode 100644
index 00000000..bdbb86fd
--- /dev/null
+++ b/src/handlers/pr_tracking.rs
@@ -0,0 +1,111 @@
+//! This module updates the PR workqueue of the Rust project contributors
+//!
+//! Purpose:
+//!
+//! - Adds the PR to the workqueue of one team member (when the PR has been assigned)
+//! - Removes the PR from the workqueue of one team member (when the PR is unassigned or closed)
+
+use crate::{
+    config::ReviewPrefsConfig,
+    db::notifications::record_username,
+    github::{IssuesAction, IssuesEvent},
+    handlers::Context,
+};
+use anyhow::Context as _;
+use tokio_postgres::Client as DbClient;
+
+pub(super) struct ReviewPrefsInput {}
+
+pub(super) async fn parse_input(
+    _ctx: &Context,
+    event: &IssuesEvent,
+    config: Option<&ReviewPrefsConfig>,
+) -> Result<Option<ReviewPrefsInput>, String> {
+    // NOTE: this config check MUST exist. Else, the triagebot will emit an error
+    // about this feature not being enabled
+    if config.is_none() {
+        return Ok(None);
+    };
+
+    // Execute this handler only if this is a PR ...
+    if !event.issue.is_pr() {
+        return Ok(None);
+    }
+
+    // ... and if the action is an assignment or unassignment with an assignee
+    let (IssuesAction::Assigned { assignee: _ } | IssuesAction::Unassigned { assignee: _ }) =
+        &event.action
+    else {
+        return Ok(None);
+    };
+    Ok(Some(ReviewPrefsInput {}))
+}
+
+pub(super) async fn handle_input<'a>(
+    ctx: &Context,
+    _config: &ReviewPrefsConfig,
+    event: &IssuesEvent,
+    _inputs: ReviewPrefsInput,
+) -> anyhow::Result<()> {
+    let db_client = ctx.db.get().await;
+
+    // extract the assignee matching the assignment or unassignment enum variants or return and ignore this handler
+    let IssuesEvent {
+        action: IssuesAction::Assigned { assignee } | IssuesAction::Unassigned { assignee },
+        ..
+    } = event
+    else {
+        return Ok(());
+    };
+
+    // ensure the team member object of this action exists in the `users` table
+    record_username(&db_client, assignee.id.unwrap(), &assignee.login)
+        .await
+        .context("failed to record username")?;
+
+    if let IssuesAction::Unassigned { assignee: _ } = &event.action {
+        delete_pr_from_workqueue(&db_client, assignee.id.unwrap(), event.issue.number)
+            .await
+            .context("Failed to remove PR from workqueue")?;
+    }
+
+    if let IssuesAction::Assigned { assignee: _ } = &event.action {
+        upsert_pr_into_workqueue(&db_client, assignee.id.unwrap(), event.issue.number)
+            .await
+            .context("Failed to add PR to workqueue")?;
+    }
+
+    Ok(())
+}
+
+/// Add a PR to the workqueue of a team member.
+/// Ensures no accidental PR duplicates.
+async fn upsert_pr_into_workqueue(
+    db: &DbClient,
+    user_id: u64,
+    pr: u64,
+) -> anyhow::Result<u64, anyhow::Error> {
+    let q = "
+INSERT INTO review_prefs
+(user_id, assigned_prs) VALUES ($1, $2)
+ON CONFLICT (user_id)
+DO UPDATE SET assigned_prs = uniq(sort(array_append(review_prefs.assigned_prs, $3)));";
+    db.execute(q, &[&(user_id as i64), &vec![pr as i32], &(pr as i32)])
+        .await
+        .context("Upsert DB error")
+}
+
+/// Delete a PR from the workqueue of a team member
+async fn delete_pr_from_workqueue(
+    db: &DbClient,
+    user_id: u64,
+    pr: u64,
+) -> anyhow::Result<u64, anyhow::Error> {
+    let q = "
+UPDATE review_prefs r
+SET assigned_prs = array_remove(r.assigned_prs, $2)
+WHERE r.user_id = $1;";
+    db.execute(q, &[&(user_id as i64), &(pr as i32)])
+        .await
+        .context("Update DB error")
+}