Workflow for tracking PRs assignment

General overview at: rust-lang#1753

This patch implements the first part:
- a new DB table with just the fields to track how many PRs are
assigned to a contributor at any time
- Update this table everytime a PR is assigned or unassigned

No initial sync is planned at this time. Populating the table will happen over
time.

src/config.rs

@@ -35,6 +35,7 @@ pub(crate) struct Config {
     pub(crate) note: Option<NoteConfig>,
     pub(crate) mentions: Option<MentionsConfig>,
     pub(crate) no_merges: Option<NoMergesConfig>,
+    pub(crate) review_work_queue: Option<TeamMemberWorkQueueConfig>,
 }
 
 #[derive(PartialEq, Eq, Debug, serde::Deserialize)]
@@ -152,6 +153,13 @@ pub(crate) struct ShortcutConfig {
     _empty: (),
 }
 
+#[derive(PartialEq, Eq, Debug, serde::Deserialize)]
+#[serde(rename_all = "kebab-case")]
+pub(crate) struct TeamMemberWorkQueueConfig {
+    #[serde(default)]
+    pub(crate) enabled_for_teams: Vec<String>,
+}
+
 #[derive(PartialEq, Eq, Debug, serde::Deserialize)]
 pub(crate) struct PrioritizeConfig {
     pub(crate) label: String,
@@ -357,6 +365,9 @@ mod tests {
             [assign]
             users_on_vacation = ["jyn514"]
 
+            [review-work-queue]
+            enabled-for-teams = ["aaa", "bbb"]
+
             [note]
 
             [ping.compiler]
@@ -429,6 +440,9 @@ mod tests {
                 github_releases: None,
                 review_submitted: None,
                 review_requested: None,
+                review_work_queue: Some(TeamMemberWorkQueueConfig {
+                    enabled_for_teams: vec!["aaa".into(), "bbb".into()]
+                }),
                 mentions: None,
                 no_merges: None,
             }

src/db.rs

@@ -307,5 +307,14 @@ CREATE UNIQUE INDEX jobs_name_scheduled_at_unique_index
     ON jobs (
         name, scheduled_at
     );
+",
+    "
+CREATE table review_prefs (
+    id UUID DEFAULT gen_random_uuid() PRIMARY KEY,
+    user_id BIGINT REFERENCES users(user_id),
+    assigned_prs INT[] NOT NULL DEFAULT array[]::INT[],
+    num_assigned_prs INTEGER
+);
+CREATE UNIQUE INDEX review_prefs_user_id ON review_prefs(user_id);
 ",
 ];

src/handlers.rs

@@ -42,6 +42,7 @@ mod prioritize;
 mod relabel;
 mod review_requested;
 mod review_submitted;
+mod review_work_queue;
 mod rfc_helper;
 pub mod rustc_commits;
 mod shortcut;
@@ -166,6 +167,7 @@ issue_handlers! {
     no_merges,
     notify_zulip,
     review_requested,
+    review_work_queue,
 }
 
 macro_rules! command_handlers {

src/handlers/review_work_queue.rs

@@ -0,0 +1,169 @@
+use crate::{
+    config::TeamMemberWorkQueueConfig,
+    github::{IssuesAction, IssuesEvent},
+    handlers::Context,
+    TeamMemberWorkQueue,
+};
+use anyhow::Context as _;
+use tokio_postgres::Client as DbClient;
+use tracing as log;
+
+// This module updates the PR work queue of team members
+// - When a PR has been assigned, adds the PR to the work queue of team members
+// - When a PR is unassigned or closed, removes the PR from the work queue of all team members
+
+/// Get all assignees for a pull request
+async fn get_pr_assignees(
+    db: &DbClient,
+    issue_num: i32,
+) -> anyhow::Result<Vec<TeamMemberWorkQueue>> {
+    let q = "
+SELECT u.username, r.*
+FROM review_prefs r
+JOIN users u on u.user_id=r.user_id
+WHERE $1 = ANY (assigned_prs)";
+    let rows = db.query(q, &[&issue_num]).await?;
+    Ok(rows
+        .into_iter()
+        .filter_map(|row| Some(TeamMemberWorkQueue::from(row)))
+        .collect())
+}
+
+/// UPDATE a team member work queue
+async fn update_team_member_workqueue(
+    db: &DbClient,
+    assignee: &TeamMemberWorkQueue,
+) -> anyhow::Result<TeamMemberWorkQueue> {
+    let x = &assignee
+        .assigned_prs
+        .iter()
+        .map(|e| e.to_string())
+        .collect::<Vec<String>>()
+        .join(",");
+    let q = format!(
+        "
+UPDATE review_prefs r
+SET assigned_prs = '{{ {} }}', num_assigned_prs = $2
+FROM users u
+WHERE r.user_id=$1 AND u.user_id=r.user_id
+RETURNING u.username, r.*",
+        x
+    );
+    let num_assigned_prs = assignee.assigned_prs.len() as i32;
+    let rec = db
+        .query_one(&q, &[&assignee.user_id, &num_assigned_prs])
+        .await
+        .context("Update DB error")?;
+    Ok(rec.into())
+}
+
+/// Add a new user (if not existing)
+async fn maybe_create_team_member(
+    db: &DbClient,
+    user_id: i64,
+    username: &str,
+) -> anyhow::Result<u64> {
+    let q = "
+INSERT INTO users (user_id, username) VALUES ($1, $2)
+ON CONFLICT DO NOTHING";
+    let rec = db
+        .execute(q, &[&user_id, &username])
+        .await
+        .context("Insert user DB error")?;
+    Ok(rec)
+}
+
+/// INSERT or UPDATE (increasing by one) a team member work queue
+async fn upsert_team_member_workqueue(
+    db: &DbClient,
+    user_id: i64,
+    pr: i32,
+) -> anyhow::Result<u64, anyhow::Error> {
+    let q = format!(
+        "
+INSERT INTO review_prefs
+(user_id, assigned_prs, num_assigned_prs) VALUES ($1, '{{ {} }}', 1)
+ON CONFLICT (user_id)
+DO UPDATE SET assigned_prs = array_append(review_prefs.assigned_prs, $2), num_assigned_prs = review_prefs.num_assigned_prs + 1
+WHERE review_prefs.user_id=$1",
+        pr
+    );
+    db.execute(&q, &[&user_id, &pr])
+        .await
+        .context("Upsert DB error")
+}
+
+pub(super) struct ReviewPrefsInput {}
+
+pub(super) async fn parse_input(
+    _ctx: &Context,
+    event: &IssuesEvent,
+    config: Option<&TeamMemberWorkQueueConfig>,
+) -> Result<Option<ReviewPrefsInput>, String> {
+    // IMPORTANT: this config check MUST exist. Else, the triagebot will emit an error that this
+    // feature is not enabled
+    if config.is_none() {
+        return Ok(None);
+    }
+
+    // Do nothing if a) this is not a PR or b) it's not an action we need to handle
+    if !event.issue.is_pr()
+        || !matches!(
+            event.action,
+            IssuesAction::Assigned | IssuesAction::Unassigned | IssuesAction::Closed
+        )
+    {
+        return Ok(None);
+    }
+    Ok(Some(ReviewPrefsInput {}))
+}
+
+pub(super) async fn handle_input<'a>(
+    ctx: &Context,
+    _config: &TeamMemberWorkQueueConfig,
+    event: &IssuesEvent,
+    _inputs: ReviewPrefsInput,
+) -> anyhow::Result<()> {
+    let db_client = ctx.db.get().await;
+    let iss_num = event.issue.number as i32;
+
+    // Note: When changing assignees for a PR, we don't receive the assignee(s) removed, we receive
+    // an event `Unassigned` and the remaining assignees
+
+    // 1) Remove the PR from everyones' work queue
+    let mut current_assignees = get_pr_assignees(&db_client, iss_num).await?;
+    log::debug!("Removing assignment from user(s): {:?}", current_assignees);
+    for assignee in &mut current_assignees {
+        if let Some(index) = assignee
+            .assigned_prs
+            .iter()
+            .position(|value| *value == iss_num)
+        {
+            assignee.assigned_prs.swap_remove(index);
+        }
+        update_team_member_workqueue(&db_client, &assignee).await?;
+    }
+
+    // If closing a PR, nothing else to do
+    if event.action == IssuesAction::Closed {
+        return Ok(());
+    }
+
+    // 2) create or increase by one the team members work queue
+    // create team members if they don't exist
+    for u in event.issue.assignees.iter() {
+        let user_id = u.id.expect("Github user was expected! Please investigate.");
+
+        if let Err(err) = maybe_create_team_member(&db_client, user_id, &u.login).await {
+            log::error!("Failed to create user in DB: this PR assignment won't be tracked.");
+            return Err(err);
+        }
+
+        if let Err(err) = upsert_team_member_workqueue(&db_client, user_id, iss_num).await {
+            log::error!("Failed to track PR for user: this PR assignment won't be tracked.");
+            return Err(err);
+        }
+    }
+
+    Ok(())
+}

src/lib.rs

@@ -8,6 +8,7 @@ use crate::github::PullRequestDetails;
 use anyhow::Context;
 use handlers::HandlerError;
 use interactions::ErrorComment;
+use serde::Serialize;
 use std::fmt;
 use tracing as log;
 
@@ -122,6 +123,27 @@ impl From<anyhow::Error> for WebhookError {
     }
 }
 
+#[derive(Debug, Serialize, Clone)]
+pub struct TeamMemberWorkQueue {
+    pub username: String,
+    pub id: uuid::Uuid,
+    pub user_id: i64,
+    pub assigned_prs: Vec<i32>,
+    pub num_assigned_prs: Option<i32>,
+}
+
+impl From<tokio_postgres::row::Row> for TeamMemberWorkQueue {
+    fn from(row: tokio_postgres::row::Row) -> Self {
+        Self {
+            username: row.get("username"),
+            id: row.get("id"),
+            user_id: row.get("user_id"),
+            assigned_prs: row.get("assigned_prs"),
+            num_assigned_prs: row.get("num_assigned_prs"),
+        }
+    }
+}
+
 pub fn deserialize_payload<T: serde::de::DeserializeOwned>(v: &str) -> anyhow::Result<T> {
     let mut deserializer = serde_json::Deserializer::from_str(&v);
     let res: Result<T, _> = serde_path_to_error::deserialize(&mut deserializer);

triagebot.toml

@@ -3,6 +3,12 @@ allow-unauthenticated = ["bug", "invalid", "question", "enhancement"]
 
 [assign]
 
+# Enable this feature to keep track of the PR review assignment
+[review-work-queue]
+# A list of teams that will use the new PR review assignment workflow.
+# Team names match https://github.com/rust-lang/team/tree/master/teams
+enabled-for-teams=["compiler", "compiler-contributors"]
+
 [note]
 
 [ping.wg-meta]