Merge pull request #1919 from Kobzol/setting-review-assignment-limit

Urgau · web-flow · commit da64ec5228cd · 2025-04-17T21:44:48.000Z
Implement setting review assignment limit
diff --git a/README.md b/README.md
@@ -99,6 +99,32 @@ Where the value in `--secret` is the secret value you place in `GITHUB_WEBHOOK_S
 
 If you are modifying code that sends message to Zulip and want to test your changes, you can register a [new free Zulip instance](https://zulip.com/new/). Before launching the triagebot locally, set the Zulip env vars to connect to your test instance (see example in `.env.sample`).
 
+You can also test Zulip webhooks locally with `curl`. For example, to test the Zulip hooks (commands sent to the
+Triagebot from the Rust lang Zulip), you start the triagebot on `localhost:8000` and then simulate a
+Zulip hook payload:
+``` sh
+curl http://localhost:8000/zulip-hook \
+    -H "Content-Type: application/json" \
+    -d '{
+        "data": "<CMD>",
+        "token": "<ZULIP_TOKEN>",
+        "message": {
+            "sender_id": <YOUR_ID>,
+            "recipient_id": <YOUR_ID>,
+            "sender_full_name": "Randolph Carter",
+            "sender_email": "r.carter@rust-lang.org",
+            "type": "stream"
+            }
+        }'
+```
+
+Where:
+- `CMD` is the exact command you would issue @triagebot on Zulip (ex. open a direct chat with the
+  bot and send "work show")
+- `ZULIP_TOKEN`: can be anything. Must correspond to the env var `$ZULIP_TOKEN` on your workstation
+- `YOUR_ID`: your GitHub user ID. Must be existing in your local triagebot database (table `users` and as
+  foreign key also in `review_prefs`)
+
 #### ngrok
 
 The following is an example of using <https://ngrok.com/> to provide webhook forwarding.
diff --git a/src/db.rs b/src/db.rs
@@ -10,6 +10,7 @@ use tokio_postgres::Client as DbClient;
 pub mod issue_data;
 pub mod jobs;
 pub mod notifications;
+pub mod review_prefs;
 pub mod rustc_commits;
 pub mod users;
 
diff --git a/src/db/review_prefs.rs b/src/db/review_prefs.rs
@@ -0,0 +1,129 @@
+use crate::db::users::record_username;
+use crate::github::{User, UserId};
+use anyhow::Context;
+use serde::Serialize;
+
+#[derive(Debug, Serialize)]
+pub struct ReviewPrefs {
+    pub id: uuid::Uuid,
+    pub user_id: i64,
+    pub max_assigned_prs: Option<i32>,
+}
+
+impl From<tokio_postgres::row::Row> for ReviewPrefs {
+    fn from(row: tokio_postgres::row::Row) -> Self {
+        Self {
+            id: row.get("id"),
+            user_id: row.get("user_id"),
+            max_assigned_prs: row.get("max_assigned_prs"),
+        }
+    }
+}
+
+/// Get team member review preferences.
+/// If they are missing, returns `Ok(None)`.
+pub async fn get_review_prefs(
+    db: &tokio_postgres::Client,
+    user_id: UserId,
+) -> anyhow::Result<Option<ReviewPrefs>> {
+    let query = "
+SELECT id, user_id, max_assigned_prs
+FROM review_prefs
+WHERE review_prefs.user_id = $1;";
+    let row = db
+        .query_opt(query, &[&(user_id as i64)])
+        .await
+        .context("Error retrieving review preferences")?;
+    Ok(row.map(|r| r.into()))
+}
+
+/// Updates review preferences of the specified user, or creates them
+/// if they do not exist yet.
+pub async fn upsert_review_prefs(
+    db: &tokio_postgres::Client,
+    user: User,
+    max_assigned_prs: Option<u32>,
+) -> anyhow::Result<u64, anyhow::Error> {
+    // We need to have the user stored in the DB to have a valid FK link in review_prefs
+    record_username(db, user.id, &user.login).await?;
+
+    let max_assigned_prs = max_assigned_prs.map(|v| v as i32);
+    let query = "
+INSERT INTO review_prefs(user_id, max_assigned_prs)
+VALUES ($1, $2)
+ON CONFLICT (user_id)
+DO UPDATE
+SET max_assigned_prs = excluded.max_assigned_prs";
+
+    let res = db
+        .execute(query, &[&(user.id as i64), &max_assigned_prs])
+        .await
+        .context("Error upserting review preferences")?;
+    Ok(res)
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::db::review_prefs::{get_review_prefs, upsert_review_prefs};
+    use crate::db::users::get_user;
+    use crate::tests::github::user;
+    use crate::tests::run_test;
+
+    #[tokio::test]
+    async fn insert_prefs_create_user() {
+        run_test(|ctx| async {
+            let db = ctx.db_client().await;
+
+            let user = user("Martin", 1);
+            upsert_review_prefs(&db, user.clone(), Some(1)).await?;
+
+            assert_eq!(get_user(&db, user.id).await?.unwrap(), user);
+
+            Ok(ctx)
+        })
+        .await;
+    }
+
+    #[tokio::test]
+    async fn insert_max_assigned_prs() {
+        run_test(|ctx| async {
+            let db = ctx.db_client().await;
+
+            upsert_review_prefs(&db, user("Martin", 1), Some(5)).await?;
+            assert_eq!(
+                get_review_prefs(&db, 1).await?.unwrap().max_assigned_prs,
+                Some(5)
+            );
+
+            Ok(ctx)
+        })
+        .await;
+    }
+
+    #[tokio::test]
+    async fn update_max_assigned_prs() {
+        run_test(|ctx| async {
+            let db = ctx.db_client().await;
+
+            upsert_review_prefs(&db, user("Martin", 1), Some(5)).await?;
+            assert_eq!(
+                get_review_prefs(&db, 1).await?.unwrap().max_assigned_prs,
+                Some(5)
+            );
+            upsert_review_prefs(&db, user("Martin", 1), Some(10)).await?;
+            assert_eq!(
+                get_review_prefs(&db, 1).await?.unwrap().max_assigned_prs,
+                Some(10)
+            );
+
+            upsert_review_prefs(&db, user("Martin", 1), None).await?;
+            assert_eq!(
+                get_review_prefs(&db, 1).await?.unwrap().max_assigned_prs,
+                None
+            );
+
+            Ok(ctx)
+        })
+        .await;
+    }
+}
diff --git a/src/lib.rs b/src/lib.rs
@@ -1,11 +1,10 @@
 #![allow(clippy::new_without_default)]
 
-use crate::github::{PullRequestDetails, UserId};
+use crate::github::PullRequestDetails;
 
 use anyhow::Context;
 use handlers::HandlerError;
 use interactions::ErrorComment;
-use serde::Serialize;
 use std::fmt;
 use tracing as log;
 
@@ -125,40 +124,6 @@ impl From<anyhow::Error> for WebhookError {
     }
 }
 
-#[derive(Debug, Serialize)]
-pub struct ReviewPrefs {
-    pub id: uuid::Uuid,
-    pub user_id: i64,
-    pub max_assigned_prs: Option<i32>,
-}
-
-impl From<tokio_postgres::row::Row> for ReviewPrefs {
-    fn from(row: tokio_postgres::row::Row) -> Self {
-        Self {
-            id: row.get("id"),
-            user_id: row.get("user_id"),
-            max_assigned_prs: row.get("max_assigned_prs"),
-        }
-    }
-}
-
-/// Get team member review preferences.
-/// If they are missing, returns `Ok(None)`.
-pub async fn get_review_prefs(
-    db: &tokio_postgres::Client,
-    user_id: UserId,
-) -> anyhow::Result<Option<ReviewPrefs>> {
-    let query = "
-SELECT id, user_id, max_assigned_prs
-FROM review_prefs
-WHERE review_prefs.user_id = $1;";
-    let row = db
-        .query_opt(query, &[&(user_id as i64)])
-        .await
-        .context("Error retrieving review preferences")?;
-    Ok(row.map(|r| r.into()))
-}
-
 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);
diff --git a/src/team_data.rs b/src/team_data.rs
@@ -1,6 +1,6 @@
 use crate::github::GithubClient;
 use anyhow::Context as _;
-use rust_team_data::v1::{Teams, ZulipMapping, BASE_URL};
+use rust_team_data::v1::{People, Teams, ZulipMapping, BASE_URL};
 use serde::de::DeserializeOwned;
 
 async fn by_url<T: DeserializeOwned>(client: &GithubClient, path: &str) -> anyhow::Result<T> {
@@ -36,3 +36,9 @@ pub async fn teams(client: &GithubClient) -> anyhow::Result<Teams> {
         .await
         .context("team-api: teams.json")
 }
+
+pub async fn people(client: &GithubClient) -> anyhow::Result<People> {
+    by_url(client, "/people.json")
+        .await
+        .context("team-api: people.json")
+}
diff --git a/src/zulip.rs b/src/zulip.rs
@@ -1,7 +1,7 @@
 use crate::db::notifications::add_metadata;
 use crate::db::notifications::{self, delete_ping, move_indices, record_ping, Identifier};
-use crate::get_review_prefs;
-use crate::github::{get_id_for_username, GithubClient};
+use crate::db::review_prefs::{get_review_prefs, upsert_review_prefs};
+use crate::github::{get_id_for_username, GithubClient, User};
 use crate::handlers::docs_update::docs_update;
 use crate::handlers::pr_tracking::get_assigned_prs;
 use crate::handlers::project_goals::{self, ping_project_goals_owners};
@@ -84,6 +84,19 @@ pub async fn to_github_id(client: &GithubClient, zulip_id: u64) -> anyhow::Resul
     Ok(map.users.get(&zulip_id).copied())
 }
 
+pub async fn username_from_gh_id(
+    client: &GithubClient,
+    gh_id: u64,
+) -> anyhow::Result<Option<String>> {
+    let people_map = crate::team_data::people(client).await?;
+    Ok(people_map
+        .people
+        .into_iter()
+        .filter(|(_, p)| p.github_id == gh_id)
+        .map(|p| p.0)
+        .next())
+}
+
 pub async fn to_zulip_id(client: &GithubClient, github_id: u64) -> anyhow::Result<Option<u64>> {
     let map = crate::team_data::zulip_map(client).await?;
     Ok(map
@@ -164,7 +177,7 @@ fn handle_command<'a>(
                 .map_err(|e| format_err!("Failed to parse movement, expected `move <from> <to>`: {e:?}.")),
             Some("meta") => add_meta_notification(&ctx, gh_id, words).await
                 .map_err(|e| format_err!("Failed to parse `meta` command. Synopsis: meta <num> <text>: Add <text> to your notification identified by <num> (>0)\n\nError: {e:?}")),
-            Some("work") => query_pr_assignments(ctx, gh_id, words).await
+            Some("work") => workqueue_commands(ctx, gh_id, words).await
                                                                     .map_err(|e| format_err!("Failed to parse `work` command. Synopsis: work <show>: shows your current PRs assignment\n\nError: {e:?}")),
             _ => {
                 while let Some(word) = next {
@@ -242,7 +255,9 @@ fn handle_command<'a>(
     })
 }
 
-async fn query_pr_assignments(
+/// Commands for working with the workqueue, e.g. showing how many PRs are assigned
+/// or modifying the PR review assignment limit.
+async fn workqueue_commands(
     ctx: &Context,
     gh_id: u64,
     mut words: impl Iterator<Item = &str>,
@@ -268,16 +283,59 @@ async fn query_pr_assignments(
                 .collect::<Vec<String>>()
                 .join(", ");
 
-            let review_prefs = get_review_prefs(&db_client, gh_id).await?;
+            let review_prefs = get_review_prefs(&db_client, gh_id)
+                .await
+                .context("cannot get review preferences")?;
             let capacity = match review_prefs.and_then(|p| p.max_assigned_prs) {
-                Some(max) => format!("{max}"),
+                Some(max) => max.to_string(),
                 None => String::from("Not set (i.e. unlimited)"),
             };
 
-            let mut response = format!("Assigned PRs: {prs}\n");
+            let mut response = format!("Assigned rust-lang/rust PRs: {prs}\n");
             writeln!(response, "Review capacity: {capacity}")?;
             response
         }
+        "set-pr-limit" => {
+            let max_assigned_prs = match words.next() {
+                Some(value) => {
+                    if words.next().is_some() {
+                        anyhow::bail!("Too many parameters.");
+                    }
+                    if value == "unlimited" {
+                        None
+                    } else {
+                        Some(value.parse::<u32>().context(
+                            "Wrong parameter format. Must be a positive integer or `unlimited` to unset the limit.",
+                        )?)
+                    }
+                }
+                None => anyhow::bail!("Missing parameter."),
+            };
+            let gh_username = username_from_gh_id(&ctx.github, gh_id)
+                .await?
+                .ok_or_else(|| {
+                    anyhow::anyhow!("Cannot find your GitHub username in the team database")
+                })?;
+
+            let user = User {
+                login: gh_username.clone(),
+                id: gh_id,
+            };
+            upsert_review_prefs(&db_client, user, max_assigned_prs)
+                .await
+                .context("Error occurred while setting review preferences.")?;
+            tracing::info!("Setting max assignment PRs of `{gh_username}` to {max_assigned_prs:?}");
+            format!(
+                "Review capacity set to {}",
+                match max_assigned_prs {
+                    Some(v) => v.to_string(),
+                    None => "unlimited".to_string(),
+                }
+            )
+        }
+        "help" => r"work show => show your assigned PRs
+work set-pr-limit <number>|unlimited => set the maximum number of PRs you can be assigned to"
+            .to_string(),
         _ => anyhow::bail!("Invalid subcommand."),
     };
 
