From 595eb3c083001e105cd4c47c238153a7d939a78b Mon Sep 17 00:00:00 2001 From: apiraino Date: Wed, 13 Dec 2023 16:14:55 +0100 Subject: [PATCH] Workflow for tracking PRs assignment General overview at: https://github.com/rust-lang/triagebot/issues/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 | 14 +++ src/db.rs | 9 ++ src/handlers.rs | 2 + src/handlers/review_work_queue.rs | 162 ++++++++++++++++++++++++++++++ src/lib.rs | 22 ++++ triagebot.toml | 6 ++ 6 files changed, 215 insertions(+) create mode 100644 src/handlers/review_work_queue.rs diff --git a/src/config.rs b/src/config.rs index 8312aab35..5bd6b33f8 100644 --- a/src/config.rs +++ b/src/config.rs @@ -35,6 +35,7 @@ pub(crate) struct Config { pub(crate) note: Option, pub(crate) mentions: Option, pub(crate) no_merges: Option, + pub(crate) review_work_queue: Option, } #[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, +} + #[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, } diff --git a/src/db.rs b/src/db.rs index 272601b0b..59008b05a 100644 --- a/src/db.rs +++ b/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); ", ]; diff --git a/src/handlers.rs b/src/handlers.rs index d2bf47d87..10f9b250a 100644 --- a/src/handlers.rs +++ b/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 { diff --git a/src/handlers/review_work_queue.rs b/src/handlers/review_work_queue.rs new file mode 100644 index 000000000..051dea8ea --- /dev/null +++ b/src/handlers/review_work_queue.rs @@ -0,0 +1,162 @@ +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> { + 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 { + let q = " +UPDATE review_prefs r +SET assigned_prs = $2, num_assigned_prs = $3 +FROM users u +WHERE r.user_id=$1 AND u.user_id=r.user_id +RETURNING u.username, r.*"; + let num_assigned_prs = assignee.assigned_prs.len() as i32; + let rec = db + .query_one( + q, + &[&assignee.user_id, &assignee.assigned_prs, &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 { + 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) +} + +/// Create or increase by one a team member work queue +async fn upsert_team_member_workqueue( + db: &DbClient, + user_id: i64, + pr: i32, +) -> anyhow::Result { + let q = + " +INSERT INTO review_prefs +(user_id, assigned_prs, num_assigned_prs) VALUES ($1, $2, 1) +ON CONFLICT (user_id) +DO UPDATE SET assigned_prs = array_append(review_prefs.assigned_prs, $3), num_assigned_prs = review_prefs.num_assigned_prs + 1 +WHERE review_prefs.user_id=$1"; + let pr_v = vec![pr]; + db.execute(q, &[&user_id, &pr_v, &pr]) + .await + .context("Upsert DB error") +} + +pub(super) struct ReviewPrefsInput {} + +pub(super) async fn parse_input( + _ctx: &Context, + event: &IssuesEvent, + config: Option<&TeamMemberWorkQueueConfig>, +) -> Result, 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(()) +} diff --git a/src/lib.rs b/src/lib.rs index 629cf6466..d13b83c49 100644 --- a/src/lib.rs +++ b/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 for WebhookError { } } +#[derive(Debug, Serialize, Clone)] +pub struct TeamMemberWorkQueue { + pub username: String, + pub id: uuid::Uuid, + pub user_id: i64, + pub assigned_prs: Vec, + pub num_assigned_prs: Option, +} + +impl From 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(v: &str) -> anyhow::Result { let mut deserializer = serde_json::Deserializer::from_str(&v); let res: Result = serde_path_to_error::deserialize(&mut deserializer); diff --git a/triagebot.toml b/triagebot.toml index 4e45562e4..2e8a5a5f3 100644 --- a/triagebot.toml +++ b/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]