Add webhook handler to update PR workload queues

src/config.rs

@@ -39,6 +39,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 +318,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)| {

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);
  ",
 ];

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,
 }
 

src/handlers/pr_tracking.rs

@@ -0,0 +1,137 @@
+//! This module updates the PR work queue of the Rust project contributors
+//!
+//! Purpose:
+//!
+//! - Adds the PR to the work queue of one or more team member (when the PR has been assigned)
+//! - Removes the PR from the work queue of one or more team members (when the PR is unassigned or closed)
+//!
+//! Notes:
+//!
+//! - When assigning or unassigning a PR, we only receive the list of current assignee(s)
+//! - In case of a complete PR unassignment the list will be empty (i.e. does not have assignee(s) removed)
+//! - The code in this handler MUST be idempotent because GH emits a webhook trigger for every assignee
+
+use crate::{
+    config::ReviewPrefsConfig,
+    db::notifications::record_username,
+    github::{IssuesAction, IssuesEvent},
+    handlers::Context,
+};
+use anyhow::Context as _;
+use itertools::Itertools;
+use tokio_postgres::Client as DbClient;
+use tracing as log;
+
+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
+    // and if the action is an assignment or unassignment
+    if !event.issue.is_pr()
+        || !matches!(
+            event.action,
+            IssuesAction::Assigned | IssuesAction::Unassigned
+        )
+    {
+        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;
+
+    // ensure all team members involved in this action are existing in the `users` table
+    for user in event.issue.assignees.iter() {
+        if let Err(err) = record_username(&db_client, user.id.unwrap(), &user.login)
+            .await
+            .context("failed to record username")
+        {
+            log::error!("record username: {:?}", err);
+        }
+    }
+
+    let user_ids: Vec<i64> = event
+        .issue
+        .assignees
+        .iter()
+        .map(|u| u.id.unwrap() as i64)
+        .collect();
+
+    // unassign PR from team members
+    let _ = delete_pr_from_workqueue(&db_client, &user_ids, event.issue.number)
+        .await
+        .context("Failed to remove PR from workqueue");
+
+    // assign PR to team members in the list
+    if user_ids.len() > 0 {
+        let _ = upsert_pr_to_workqueue(&db_client, &user_ids, event.issue.number)
+            .await
+            .context("Failed to assign PR");
+    }
+
+    Ok(())
+}
+
+/// Add a PR to the work queue of a number of team members
+async fn upsert_pr_to_workqueue(
+    db: &DbClient,
+    user_ids: &Vec<i64>,
+    pr: u64,
+) -> anyhow::Result<u64, anyhow::Error> {
+    let values = user_ids
+        .iter()
+        .map(|uid| format!("({}, '{{ {} }}')", uid, pr))
+        .join(", ");
+    let q = format!(
+        "
+INSERT INTO review_prefs
+(user_id, assigned_prs) VALUES {}
+ON CONFLICT (user_id)
+DO UPDATE SET assigned_prs = uniq(sort(array_append(review_prefs.assigned_prs, $1)))",
+        values
+    );
+    db.execute(&q, &[&(pr as i32)])
+        .await
+        .context("Upsert DB error")
+}
+
+/// Delete a PR from the workqueue of teams members.
+/// If `exclude_user_ids` is empty, delete the PR from everyone's workqueue
+/// Else, delete the PR from everyone BUT those in the list
+async fn delete_pr_from_workqueue(
+    db: &DbClient,
+    exclude_user_ids: &Vec<i64>,
+    pr: u64,
+) -> anyhow::Result<u64, anyhow::Error> {
+    let mut q = String::from(
+        "
+UPDATE review_prefs r
+SET assigned_prs = array_remove(r.assigned_prs, $1)",
+    );
+
+    if exclude_user_ids.len() > 0 {
+        q.push_str(&format!(
+            " WHERE r.user_id NOT IN ('{{ {} }}')",
+            exclude_user_ids.iter().join(","),
+        ));
+    }
+    db.execute(&q, &[&(pr as i32)])
+        .await
+        .context("Update DB error")
+}