diff --git a/src/github.rs b/src/github.rs index 4ed0096b..8e0212a4 100644 --- a/src/github.rs +++ b/src/github.rs @@ -15,10 +15,12 @@ use std::{ }; use tracing as log; -#[derive(Debug, PartialEq, Eq, serde::Deserialize, Clone)] +pub type UserId = u64; + +#[derive(Debug, PartialEq, Eq, Hash, serde::Deserialize, Clone)] pub struct User { pub login: String, - pub id: u64, + pub id: UserId, } impl GithubClient { @@ -3048,7 +3050,7 @@ async fn project_items_by_status( } /// Retrieve all pull requests in status OPEN that are not drafts -pub async fn retrieve_pull_requests( +pub async fn retrieve_open_pull_requests( repo: &Repository, client: &GithubClient, ) -> anyhow::Result> { diff --git a/src/handlers.rs b/src/handlers.rs index f59b6295..6ec81752 100644 --- a/src/handlers.rs +++ b/src/handlers.rs @@ -1,5 +1,6 @@ use crate::config::{self, Config, ConfigurationError}; use crate::github::{Event, GithubClient, IssueCommentAction, IssuesAction, IssuesEvent}; +use crate::handlers::pr_tracking::ReviewerWorkqueue; use octocrab::Octocrab; use parser::command::{assign::AssignCommand, Command, Input}; use std::fmt; @@ -348,4 +349,7 @@ pub struct Context { pub db: crate::db::ClientPool, pub username: String, pub octocrab: Octocrab, + /// Represents the workqueue (assigned open PRs) of individual reviewers. + /// tokio's RwLock is used to avoid deadlocks, since we run on a single-threaded tokio runtime. + pub workqueue: Arc>, } diff --git a/src/handlers/assign.rs b/src/handlers/assign.rs index 9585a61b..7ef91540 100644 --- a/src/handlers/assign.rs +++ b/src/handlers/assign.rs @@ -23,7 +23,7 @@ use crate::{ config::{AssignConfig, WarnNonDefaultBranchException}, github::{self, Event, FileDiff, Issue, IssuesAction, Selection}, - handlers::{pr_tracking::has_user_capacity, Context, GithubClient, IssuesEvent}, + handlers::{Context, GithubClient, IssuesEvent}, interactions::EditIssueBody, }; use anyhow::{bail, Context as _}; @@ -39,7 +39,6 @@ use tracing as log; #[cfg(test)] mod tests { mod tests_candidates; - mod tests_db; mod tests_from_diff; } @@ -560,22 +559,22 @@ pub(super) async fn handle_command( } let db_client = ctx.db.get().await; if is_self_assign(&name, &event.user().login) { - let work_queue = has_user_capacity(&db_client, &name).await; - if work_queue.is_err() { - // NOTE: disabled for now, just log - log::warn!( - "[#{}] PR self-assign failed, DB reported that user {} has no review capacity. Ignoring.", - issue.number, - name - ); - // issue - // .post_comment( - // &ctx.github, - // &REVIEWER_HAS_NO_CAPACITY.replace("{username}", &name), - // ) - // .await?; - // return Ok(()); - } + // let work_queue = has_user_capacity(&db_client, &name).await; + // if work_queue.is_err() { + // // NOTE: disabled for now, just log + // log::warn!( + // "[#{}] PR self-assign failed, DB reported that user {} has no review capacity. Ignoring.", + // issue.number, + // name + // ); + // // issue + // // .post_comment( + // // &ctx.github, + // // &REVIEWER_HAS_NO_CAPACITY.replace("{username}", &name), + // // ) + // // .await?; + // // return Ok(()); + // } name.to_string() } else { @@ -802,7 +801,7 @@ impl fmt::Display for FindReviewerError { /// auto-assign groups, or rust-lang team names. It must have at least one /// entry. async fn find_reviewer_from_names( - db: &DbClient, + _db: &DbClient, teams: &Teams, config: &AssignConfig, issue: &Issue, @@ -841,24 +840,24 @@ async fn find_reviewer_from_names( } // filter out team members without capacity - let filtered_candidates = filter_by_capacity(db, &candidates) - .await - .expect("Error while filtering out team members"); - - if filtered_candidates.is_empty() { - // NOTE: disabled for now, just log - log::info!("[#{}] Filtered list of PR assignee is empty", issue.number); - // return Err(FindReviewerError::AllReviewersFiltered { - // initial: names.to_vec(), - // filtered: names.to_vec(), - // }); - } - - log::info!( - "[#{}] Filtered list of candidates: {:?}", - issue.number, - filtered_candidates - ); + // let filtered_candidates = filter_by_capacity(db, &candidates) + // .await + // .expect("Error while filtering out team members"); + // + // if filtered_candidates.is_empty() { + // // NOTE: disabled for now, just log + // log::info!("[#{}] Filtered list of PR assignee is empty", issue.number); + // // return Err(FindReviewerError::AllReviewersFiltered { + // // initial: names.to_vec(), + // // filtered: names.to_vec(), + // // }); + // } + // + // log::info!( + // "[#{}] Filtered list of candidates: {:?}", + // issue.number, + // filtered_candidates + // ); // Return unfiltered list of candidates Ok(candidates @@ -868,32 +867,6 @@ async fn find_reviewer_from_names( .to_string()) } -/// Filter out candidates not having review capacity -async fn filter_by_capacity( - db: &DbClient, - candidates: &HashSet<&str>, -) -> anyhow::Result> { - let usernames = candidates - .iter() - .map(|c| *c) - .collect::>() - .join(","); - - let q = format!( - " -SELECT username -FROM review_prefs r -JOIN users on users.user_id=r.user_id -AND username = ANY('{{ {} }}') -AND CARDINALITY(r.assigned_prs) < LEAST(COALESCE(r.max_assigned_prs,1000000))", - usernames - ); - let result = db.query(&q, &[]).await.context("Select DB error")?; - let candidates: HashSet = result.iter().map(|row| row.get("username")).collect(); - log::info!("DB returned these candidates: {:?}", candidates); - Ok(candidates) -} - /// Returns a list of candidate usernames (from relevant teams) to choose as a reviewer. fn candidate_reviewers_from_names<'a>( teams: &'a Teams, diff --git a/src/handlers/assign/tests/tests_db.rs b/src/handlers/assign/tests/tests_db.rs deleted file mode 100644 index 825e0945..00000000 --- a/src/handlers/assign/tests/tests_db.rs +++ /dev/null @@ -1,34 +0,0 @@ -#[cfg(test)] -mod tests { - use crate::handlers::assign::filter_by_capacity; - use crate::tests::run_test; - use std::collections::HashSet; - use tokio_postgres::GenericClient; - - #[tokio::test] - async fn find_reviewers_no_review_prefs() { - run_test(|ctx| async move { - ctx.add_user("usr1", 1).await; - ctx.add_user("usr2", 2).await; - let _users = filter_by_capacity( - ctx.db_client().await.client(), - &candidates(&["usr1", "usr2"]), - ) - .await?; - // FIXME: this test fails, because the query is wrong - // check_users(users, &["usr1", "usr2"]); - Ok(ctx) - }) - .await; - } - - fn candidates(users: &[&'static str]) -> HashSet<&'static str> { - users.into_iter().copied().collect() - } - - fn check_users(users: HashSet, expected: &[&'static str]) { - let mut users: Vec = users.into_iter().collect(); - users.sort(); - assert_eq!(users, expected); - } -} diff --git a/src/handlers/pr_tracking.rs b/src/handlers/pr_tracking.rs index f793fb12..6824fe7f 100644 --- a/src/handlers/pr_tracking.rs +++ b/src/handlers/pr_tracking.rs @@ -3,22 +3,36 @@ //! //! Purpose: //! -//! - Adds the PR to the workqueue of one team member (after the PR has been assigned) +//! - Adds the PR to the workqueue of one team member (after the PR has been assigned or reopened) //! - Removes the PR from the workqueue of one team member (after the PR has been unassigned or closed) -use super::assign::{FindReviewerError, REVIEWER_HAS_NO_CAPACITY, SELF_ASSIGN_HAS_NO_CAPACITY}; -use crate::db::users::record_username; -use crate::github::User; +use crate::github::{User, UserId}; use crate::{ config::ReviewPrefsConfig, github::{IssuesAction, IssuesEvent}, handlers::Context, - ReviewPrefs, }; -use anyhow::Context as _; -use tokio_postgres::Client as DbClient; +use std::collections::{HashMap, HashSet}; use tracing as log; +pub type PullRequestNumber = u64; + +/// Maps users to a set of currently assigned open non-draft pull requests. +/// We store this map in memory, rather than in the DB, because it can get desynced when webhooks +/// are missed. +/// It is thus reloaded when triagebot starts and also periodically, so it is not needed to store it +/// in the DB. +#[derive(Debug, Default)] +pub struct ReviewerWorkqueue { + reviewers: HashMap>, +} + +impl ReviewerWorkqueue { + pub fn new(reviewers: HashMap>) -> Self { + Self { reviewers } + } +} + pub(super) enum ReviewPrefsInput { Assigned { assignee: User }, Unassigned { assignee: User }, @@ -65,8 +79,6 @@ pub(super) async fn handle_input<'a>( event: &IssuesEvent, input: ReviewPrefsInput, ) -> anyhow::Result<()> { - let db_client = ctx.db.get().await; - log::info!("Handling event action {:?} in PR tracking", event.action); match input { @@ -74,41 +86,13 @@ pub(super) async fn handle_input<'a>( // (i.e. from the "Assignees" dropdown menu). // We need to also check assignee availability here. ReviewPrefsInput::Assigned { assignee } => { - 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 { - log::warn!( - "[#{}] DB reported that user {} has no review capacity. Ignoring.", - event.issue.number, - &assignee.login - ); - - // NOTE: disabled for now, just log - // 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?; - } - + let pr_number = event.issue.number; log::info!( - "Adding PR {} from workqueue of {} because they were assigned.", - event.issue.number, + "Adding PR {pr_number} from workqueue of {} because they were assigned.", assignee.login ); - upsert_pr_into_workqueue(&db_client, &assignee, event.issue.number) - .await - .context("Failed to add PR to work queue")?; + upsert_pr_into_workqueue(ctx, assignee.id, pr_number).await; } ReviewPrefsInput::Unassigned { assignee } => { let pr_number = event.issue.number; @@ -116,9 +100,7 @@ pub(super) async fn handle_input<'a>( "Removing PR {pr_number} from workqueue of {} because they were unassigned.", assignee.login ); - delete_pr_from_workqueue(&db_client, assignee.id, pr_number) - .await - .context("Failed to remove PR from work queue")?; + delete_pr_from_workqueue(ctx, assignee.id, pr_number).await; } ReviewPrefsInput::Closed => { for assignee in &event.issue.assignees { @@ -127,9 +109,7 @@ pub(super) async fn handle_input<'a>( "Removing PR {pr_number} from workqueue of {} because it was closed or merged.", assignee.login ); - delete_pr_from_workqueue(&db_client, assignee.id, pr_number) - .await - .context("Failed to to remove PR from work queue")?; + delete_pr_from_workqueue(ctx, assignee.id, pr_number).await; } } ReviewPrefsInput::Reopened => { @@ -139,9 +119,7 @@ pub(super) async fn handle_input<'a>( "Re-adding PR {pr_number} to workqueue of {} because it was (re)opened.", assignee.login ); - upsert_pr_into_workqueue(&db_client, &assignee, pr_number) - .await - .context("Failed to add PR to work queue")?; + upsert_pr_into_workqueue(ctx, assignee.id, pr_number).await; } } } @@ -149,72 +127,46 @@ pub(super) async fn handle_input<'a>( Ok(()) } -// Check user review capacity. -// Returns error if SQL query fails or user has no capacity -pub async fn has_user_capacity( - db: &crate::db::PooledClient, - assignee: &str, -) -> anyhow::Result { - 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()) +/// Get pull request assignments for a team member +pub async fn get_assigned_prs(ctx: &Context, user_id: UserId) -> HashSet { + ctx.workqueue + .read() + .await + .reviewers + .get(&user_id) + .cloned() + .unwrap_or_default() } /// Add a PR to the workqueue of a team member. /// Ensures no accidental PR duplicates. -async fn upsert_pr_into_workqueue( - db: &DbClient, - user: &User, - pr: u64, -) -> anyhow::Result { - // Ensure the user has entry in the `users` table - record_username(db, user.id, &user.login) - .await - .context("failed to record username")?; - - let q = " -INSERT INTO review_prefs -(user_id, assigned_prs) VALUES ($1, $2) -ON CONFLICT (user_id) -DO UPDATE SET assigned_prs = uniq(sort(array_append(review_prefs.assigned_prs, $3)));"; - db.execute(q, &[&(user.id as i64), &vec![pr as i32], &(pr as i32)]) +async fn upsert_pr_into_workqueue(ctx: &Context, user_id: UserId, pr: PullRequestNumber) { + ctx.workqueue + .write() .await - .context("Upsert DB error") + .reviewers + .entry(user_id) + .or_default() + .insert(pr); } /// Delete a PR from the workqueue of a team member -async fn delete_pr_from_workqueue( - db: &DbClient, - user_id: u64, - pr: u64, -) -> anyhow::Result { - let q = " -UPDATE review_prefs r -SET assigned_prs = array_remove(r.assigned_prs, $2) -WHERE r.user_id = $1;"; - db.execute(q, &[&(user_id as i64), &(pr as i32)]) - .await - .context("Update DB error") +async fn delete_pr_from_workqueue(ctx: &Context, user_id: UserId, pr: PullRequestNumber) { + let mut queue = ctx.workqueue.write().await; + if let Some(reviewer) = queue.reviewers.get_mut(&user_id) { + reviewer.remove(&pr); + } } #[cfg(test)] mod tests { use crate::config::Config; use crate::github::{Issue, IssuesAction, IssuesEvent, Repository, User}; - use crate::handlers::pr_tracking::{handle_input, parse_input, upsert_pr_into_workqueue}; + use crate::handlers::pr_tracking::{ + handle_input, parse_input, upsert_pr_into_workqueue, PullRequestNumber, + }; use crate::tests::github::{default_test_user, issue, pull_request, user}; use crate::tests::{run_test, TestContext}; - use tokio_postgres::GenericClient; #[tokio::test] async fn add_pr_to_workqueue_on_assign() { @@ -327,30 +279,29 @@ mod tests { .await; } - async fn check_assigned_prs(ctx: &TestContext, user: &User, expected_prs: &[i32]) { - let results = ctx - .db_client() - .await - .query( - "SELECT assigned_prs FROM review_prefs WHERE user_id = $1", - &[&(user.id as i64)], - ) + async fn check_assigned_prs( + ctx: &TestContext, + user: &User, + expected_prs: &[PullRequestNumber], + ) { + let mut assigned = ctx + .handler_ctx() + .workqueue + .read() .await - .unwrap(); - assert!(results.len() < 2); - let mut assigned = results - .get(0) - .map(|row| row.get::<_, Vec>(0)) - .unwrap_or_default(); + .reviewers + .get(&user.id) + .cloned() + .unwrap_or_default() + .into_iter() + .collect::>(); assigned.sort(); assert_eq!(assigned, expected_prs); } - async fn set_assigned_prs(ctx: &TestContext, user: &User, prs: &[i32]) { + async fn set_assigned_prs(ctx: &TestContext, user: &User, prs: &[PullRequestNumber]) { for &pr in prs { - upsert_pr_into_workqueue(&ctx.db_client().await.client(), user, pr as u64) - .await - .unwrap(); + upsert_pr_into_workqueue(ctx.handler_ctx(), user.id, pr).await; } check_assigned_prs(&ctx, user, prs).await; } diff --git a/src/handlers/pull_requests_assignment_update.rs b/src/handlers/pull_requests_assignment_update.rs index c5e6acef..48fb7285 100644 --- a/src/handlers/pull_requests_assignment_update.rs +++ b/src/handlers/pull_requests_assignment_update.rs @@ -1,12 +1,8 @@ -use std::collections::HashMap; - -use crate::db::users::record_username; -use crate::github::retrieve_pull_requests; +use crate::github::{retrieve_open_pull_requests, UserId}; +use crate::handlers::pr_tracking::{PullRequestNumber, ReviewerWorkqueue}; use crate::jobs::Job; -use crate::ReviewPrefs; -use anyhow::Context as _; use async_trait::async_trait; -use tokio_postgres::Client as DbClient; +use std::collections::{HashMap, HashSet}; pub struct PullRequestAssignmentUpdate; @@ -17,70 +13,23 @@ impl Job for PullRequestAssignmentUpdate { } async fn run(&self, ctx: &super::Context, _metadata: &serde_json::Value) -> anyhow::Result<()> { - let db = ctx.db.get().await; let gh = &ctx.github; tracing::trace!("starting pull_request_assignment_update"); let rust_repo = gh.repository("rust-lang/rust").await?; - let prs = retrieve_pull_requests(&rust_repo, &gh).await?; - - // delete all PR assignments before populating - init_table(&db).await?; - - // aggregate by user first - let aggregated = prs.into_iter().fold(HashMap::new(), |mut acc, (user, pr)| { - let (_, prs) = acc.entry(user.id).or_insert_with(|| (user, Vec::new())); - prs.push(pr); - acc - }); - - // populate the table - for (_user_id, (assignee, prs)) in &aggregated { - let assignee_id = assignee.id; - let _ = record_username(&db, assignee_id, &assignee.login).await; - create_team_member_workqueue(&db, assignee_id, &prs).await?; - } + let prs = retrieve_open_pull_requests(&rust_repo, &gh).await?; + + // Aggregate PRs by user + let aggregated: HashMap> = + prs.into_iter().fold(HashMap::new(), |mut acc, (user, pr)| { + let prs = acc.entry(user.id).or_default(); + prs.insert(pr as PullRequestNumber); + acc + }); + tracing::info!("PR assignments\n{aggregated:?}"); + *ctx.workqueue.write().await = ReviewerWorkqueue::new(aggregated); Ok(()) } } - -/// Truncate the review prefs table -async fn init_table(db: &DbClient) -> anyhow::Result { - let res = db - .execute("UPDATE review_prefs SET assigned_prs='{}';", &[]) - .await?; - Ok(res) -} - -/// Create a team member work queue -async fn create_team_member_workqueue( - db: &DbClient, - user_id: u64, - prs: &Vec, -) -> anyhow::Result { - let q = " -INSERT INTO review_prefs (user_id, assigned_prs) VALUES ($1, $2) -ON CONFLICT (user_id) -DO UPDATE SET assigned_prs = $2 -WHERE review_prefs.user_id=$1"; - db.execute(q, &[&(user_id as i64), prs]) - .await - .context("Insert DB error") -} - -/// Get pull request assignments for a team member -pub async fn get_review_prefs(db: &DbClient, user_id: u64) -> anyhow::Result { - let q = " -SELECT username,r.* -FROM review_prefs r -JOIN users on r.user_id=users.user_id -WHERE r.user_id = $1;"; - let row = db - .query_one(q, &[&(user_id as i64)]) - .await - .context("Error retrieving review preferences") - .unwrap(); - Ok(row.into()) -} diff --git a/src/lib.rs b/src/lib.rs index 1a674d60..35245c38 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,6 +1,6 @@ #![allow(clippy::new_without_default)] -use crate::github::PullRequestDetails; +use crate::github::{PullRequestDetails, UserId}; use anyhow::Context; use handlers::HandlerError; @@ -128,43 +128,37 @@ impl From for WebhookError { #[derive(Debug, Serialize)] pub struct ReviewPrefs { pub id: uuid::Uuid, - pub username: String, pub user_id: i64, - pub assigned_prs: Vec, pub max_assigned_prs: Option, } -impl ReviewPrefs { - fn to_string(&self) -> String { - let capacity = match self.max_assigned_prs { - Some(max) => format!("{}", max), - None => String::from("Not set (i.e. unlimited)"), - }; - let prs = self - .assigned_prs - .iter() - .map(|pr| format!("#{}", pr)) - .collect::>() - .join(", "); - format!( - "Username: {}\nAssigned PRs: {}\nReview capacity: {}", - self.username, prs, capacity - ) - } -} - impl From for ReviewPrefs { fn from(row: tokio_postgres::row::Row) -> Self { Self { id: row.get("id"), - username: row.get("username"), user_id: row.get("user_id"), - assigned_prs: row.get("assigned_prs"), 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> { + 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(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/src/main.rs b/src/main.rs index 5fa7a737..a35b9d97 100644 --- a/src/main.rs +++ b/src/main.rs @@ -6,6 +6,7 @@ use futures::StreamExt; use hyper::{header, Body, Request, Response, Server, StatusCode}; use route_recognizer::Router; use std::{env, net::SocketAddr, sync::Arc}; +use tokio::sync::RwLock; use tokio::{task, time}; use tower::{Service, ServiceExt}; use tracing as log; @@ -261,6 +262,7 @@ async fn run_server(addr: SocketAddr) -> anyhow::Result<()> { db: pool, github: gh, octocrab: oc, + workqueue: Arc::new(RwLock::new(Default::default())), }); // Run all jobs that don't have a schedule (one-off jobs) diff --git a/src/tests/mod.rs b/src/tests/mod.rs index e24a2125..8ded7a90 100644 --- a/src/tests/mod.rs +++ b/src/tests/mod.rs @@ -5,6 +5,8 @@ use crate::github::GithubClient; use crate::handlers::Context; use octocrab::Octocrab; use std::future::Future; +use std::sync::Arc; +use tokio::sync::RwLock; use tokio_postgres::config::Host; use tokio_postgres::{Config, GenericClient}; @@ -15,8 +17,8 @@ pub mod github; /// a database. pub struct TestContext { pool: ClientPool, + ctx: Context, db_name: String, - test_db_url: String, original_db_url: String, } @@ -54,18 +56,7 @@ impl TestContext { db::run_migrations(&mut *pool.get().await) .await .expect("Cannot run database migrations"); - Self { - pool, - db_name, - test_db_url, - original_db_url: db_url.to_string(), - } - } - /// Creates a fake handler context. - /// We currently do not mock outgoing nor incoming GitHub API calls, - /// so the API endpoints will not be actually working - pub fn handler_ctx(&self) -> Context { let octocrab = Octocrab::builder().build().unwrap(); let github = GithubClient::new( "gh-test-fake-token".to_string(), @@ -73,14 +64,29 @@ impl TestContext { "https://api.github.com/graphql".to_string(), "https://raw.githubusercontent.com".to_string(), ); - Context { + let ctx = Context { github, - db: self.create_pool(), + db: ClientPool::new(test_db_url), username: "triagebot-test".to_string(), octocrab, + workqueue: Arc::new(RwLock::new(Default::default())), + }; + + Self { + pool, + db_name, + original_db_url: db_url.to_string(), + ctx, } } + /// Returns a fake handler context. + /// We currently do not mock outgoing nor incoming GitHub API calls, + /// so the API endpoints will not be actually working. + pub fn handler_ctx(&self) -> &Context { + &self.ctx + } + pub async fn db_client(&self) -> PooledClient { self.pool.get().await } @@ -91,10 +97,6 @@ impl TestContext { .expect("Cannot create user"); } - fn create_pool(&self) -> ClientPool { - ClientPool::new(self.test_db_url.clone()) - } - async fn finish(self) { // Cleanup the test database // First, we need to stop using the database diff --git a/src/zulip.rs b/src/zulip.rs index 760ef905..91ed2654 100644 --- a/src/zulip.rs +++ b/src/zulip.rs @@ -1,9 +1,10 @@ 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::handlers::docs_update::docs_update; +use crate::handlers::pr_tracking::get_assigned_prs; use crate::handlers::project_goals::{self, ping_project_goals_owners}; -use crate::handlers::pull_requests_assignment_update::get_review_prefs; use crate::handlers::Context; use anyhow::{format_err, Context as _}; use std::env; @@ -163,7 +164,7 @@ fn handle_command<'a>( .map_err(|e| format_err!("Failed to parse movement, expected `move `: {e:?}.")), Some("meta") => add_meta_notification(&ctx, gh_id, words).await .map_err(|e| format_err!("Failed to parse `meta` command. Synopsis: meta : Add to your notification identified by (>0)\n\nError: {e:?}")), - Some("work") => query_pr_assignments(&ctx, gh_id, words).await + Some("work") => query_pr_assignments(ctx, gh_id, words).await .map_err(|e| format_err!("Failed to parse `work` command. Synopsis: work : shows your current PRs assignment\n\nError: {e:?}")), _ => { while let Some(word) = next { @@ -242,7 +243,7 @@ fn handle_command<'a>( } async fn query_pr_assignments( - ctx: &&Context, + ctx: &Context, gh_id: u64, mut words: impl Iterator, ) -> anyhow::Result> { @@ -253,18 +254,34 @@ async fn query_pr_assignments( let db_client = ctx.db.get().await; - let record = match subcommand { + let response = match subcommand { "show" => { - let rec = get_review_prefs(&db_client, gh_id).await; - if rec.is_err() { - anyhow::bail!("No preferences set.") - } - rec? + let mut assigned_prs = get_assigned_prs(ctx, gh_id) + .await + .into_iter() + .collect::>(); + assigned_prs.sort(); + + let prs = assigned_prs + .iter() + .map(|pr| format!("#{pr}")) + .collect::>() + .join(", "); + + let review_prefs = get_review_prefs(&db_client, gh_id).await?; + let capacity = match review_prefs.and_then(|p| p.max_assigned_prs) { + Some(max) => format!("{max}"), + None => String::from("Not set (i.e. unlimited)"), + }; + + let mut response = format!("Assigned PRs: {prs}\n"); + writeln!(response, "Review capacity: {capacity}")?; + response } _ => anyhow::bail!("Invalid subcommand."), }; - Ok(Some(record.to_string())) + Ok(Some(response)) } // This does two things: