Ensure user capacity is respected on PR assignment

This check is specifically used when assigning from the Github web UI

Signed-off-by: apiraino <apiraino@users.noreply.github.com>

Author: apiraino

.env.sample

@@ -20,4 +20,8 @@ GITHUB_WEBHOOK_SECRET=MUST_BE_CONFIGURED
 # ZULIP_API_TOKEN=yyy
 
 # Authenticates inbound webhooks from Github
-# ZULIP_TOKEN=xxx
+# ZULIP_TOKEN=xxx
+
+# Use another endpoint to retrieve teams of the Rust project (useful for local testing)
+# default: https://team-api.infra.rust-lang.org/v1
+# TEAMS_API_URL=http://localhost:8080

src/handlers.rs

@@ -207,7 +207,7 @@ macro_rules! issue_handlers {
 
 // Handle events that happened on issues
 //
-// This is for events that happen only on issues (e.g. label changes).
+// This is for events that happen only on issues or pull requests (e.g. label changes or assignments).
 // Each module in the list must contain the functions `parse_input` and `handle_input`.
 issue_handlers! {
     assign,
@@ -328,7 +328,7 @@ macro_rules! command_handlers {
 //
 // This is for handlers for commands parsed by the `parser` crate.
 // Each variant of `parser::command::Command` must be in this list,
-// preceded by the module containing the coresponding `handle_command` function
+// preceded by the module containing the corresponding `handle_command` function
 command_handlers! {
     assign: Assign,
     glacier: Glacier,

src/handlers/assign.rs

@@ -728,7 +728,6 @@ impl fmt::Display for FindReviewerError {
                     f,
                     "Could not assign reviewer from: `{}`.\n\
                      User(s) `{}` are either the PR author, already assigned, or on vacation. \
-                     If it's a team, we could not find any candidates.\n\
                      Please use `r?` to specify someone else to assign.",
                     initial.join(","),
                     filtered.join(","),
@@ -820,14 +819,11 @@ SELECT username
 FROM review_prefs r
 JOIN users on users.user_id=r.user_id
 AND username = ANY('{{ {} }}')
-AND ARRAY_LENGTH(r.assigned_prs, 1) < max_assigned_prs",
+AND CARDINALITY(r.assigned_prs) < LEAST(COALESCE(r.max_assigned_prs,1000000))",
         usernames
     );
     let result = db.query(&q, &[]).await.context("Select DB error")?;
-    let mut candidates: HashSet<String> = HashSet::new();
-    for row in result {
-        candidates.insert(row.get("username"));
-    }
+    let candidates: HashSet<String> = result.iter().map(|row| row.get("username")).collect();
     Ok(candidates)
 }
 

src/handlers/pr_tracking.rs

@@ -1,19 +1,23 @@
 //! This module updates the PR workqueue of the Rust project contributors
+//! Runs after a PR has been assigned or unassigned
 //!
 //! 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)
+//! - Adds the PR to the workqueue of one team member (after the PR has been assigned)
+//! - Removes the PR from the workqueue of one team member (after the PR has been unassigned or closed)
 
 use crate::{
     config::ReviewPrefsConfig,
     db::notifications::record_username,
     github::{IssuesAction, IssuesEvent},
     handlers::Context,
+    ReviewPrefs,
 };
 use anyhow::Context as _;
 use tokio_postgres::Client as DbClient;
 
+use super::assign::{FindReviewerError, REVIEWER_HAS_NO_CAPACITY, SELF_ASSIGN_HAS_NO_CAPACITY};
+
 pub(super) struct ReviewPrefsInput {}
 
 pub(super) async fn parse_input(
@@ -49,7 +53,7 @@ pub(super) async fn handle_input<'a>(
 ) -> 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
+    // extract the assignee or ignore this handler and return
     let IssuesEvent {
         action: IssuesAction::Assigned { assignee } | IssuesAction::Unassigned { assignee },
         ..
@@ -66,18 +70,61 @@ pub(super) async fn handle_input<'a>(
     if matches!(event.action, IssuesAction::Unassigned { .. }) {
         delete_pr_from_workqueue(&db_client, assignee.id, event.issue.number)
             .await
-            .context("Failed to remove PR from workqueue")?;
+            .context("Failed to remove PR from work ueue")?;
     }
 
+    // This handler is reached also when assigning a PR using the Github UI
+    // (i.e. from the "Assignees" dropdown menu).
+    // We need to also check assignee availability here.
     if matches!(event.action, IssuesAction::Assigned { .. }) {
+        let work_queue = has_user_capacity(&db_client, &assignee.login)
+            .await
+            .context("Failed to retrieve user work queue");
+
+        // if user has no capacity, revert the PR assignment (GitHub has already assigned it)
+        // and post a comment suggesting what to do
+        if let Err(_) = work_queue {
+            event
+                .issue
+                .remove_assignees(&ctx.github, crate::github::Selection::One(&assignee.login))
+                .await?;
+
+            let msg = if assignee.login.to_lowercase() == event.issue.user.login.to_lowercase() {
+                SELF_ASSIGN_HAS_NO_CAPACITY.replace("{username}", &assignee.login)
+            } else {
+                REVIEWER_HAS_NO_CAPACITY.replace("{username}", &assignee.login)
+            };
+            event.issue.post_comment(&ctx.github, &msg).await?;
+        }
+
         upsert_pr_into_workqueue(&db_client, assignee.id, event.issue.number)
             .await
-            .context("Failed to add PR to workqueue")?;
+            .context("Failed to add PR to work queue")?;
     }
 
     Ok(())
 }
 
+// TODO: we should just fetch the number of assigned prs and max assigned prs. The caller should do the check.
+pub async fn has_user_capacity(
+    db: &crate::db::PooledClient,
+    assignee: &str,
+) -> anyhow::Result<ReviewPrefs, FindReviewerError> {
+    let q = "
+SELECT username, r.*
+FROM review_prefs r
+JOIN users ON users.user_id = r.user_id
+WHERE username = $1
+AND CARDINALITY(r.assigned_prs) < LEAST(COALESCE(r.max_assigned_prs,1000000));";
+    let rec = db.query_one(q, &[&assignee]).await;
+    if let Err(_) = rec {
+        return Err(FindReviewerError::ReviewerHasNoCapacity {
+            username: assignee.to_string(),
+        });
+    }
+    Ok(rec.unwrap().into())
+}
+
 /// Add a PR to the workqueue of a team member.
 /// Ensures no accidental PR duplicates.
 async fn upsert_pr_into_workqueue(