Import pull request assignment into triagebot

General overview at: rust-lang#1753

- Added a new DB table with the fields to track how many PRs are
  assigned to a contributor
- Initial DB table population with a one-off job, manually run.

github-graphql/PullRequestsOpen.gql

@@ -0,0 +1,26 @@
+query PullRequestsOpen ($repo_owner: String!, $repo_name: String!, $after: String) {
+  repository(owner: $repo_owner, name: $repo_name) {
+    pullRequests(first: 100, after: $after, states:OPEN, labels: ["S-waiting-on-review","T-compiler"]) {
+      pageInfo {
+        hasNextPage
+        endCursor
+      }
+      nodes {
+        number
+        updatedAt
+        createdAt
+        assignees(first: 10) {
+          nodes {
+            login
+            id
+          }
+        }
+        labels(first:5, orderBy:{field:NAME, direction:DESC}) {
+          nodes {
+            name
+          }
+        }
+      }
+    }
+  }
+}

github-graphql/src/lib.rs

@@ -89,6 +89,7 @@ pub mod queries {
     #[derive(cynic::QueryFragment, Debug)]
     pub struct User {
         pub login: String,
+        pub id: cynic::Id,
     }
 
     #[derive(cynic::QueryFragment, Debug)]
@@ -385,3 +386,44 @@ pub mod project_items {
         pub date: Option<Date>,
     }
 }
+
+/// Retrieve all pull requests waiting on review from T-compiler
+/// GraphQL query: see file github-graphql/PullRequestsOpen.gql
+pub mod pull_requests_open {
+    use crate::queries::{LabelConnection, PullRequestConnection, UserConnection};
+
+    use super::queries::DateTime;
+    use super::schema;
+
+    #[derive(cynic::QueryVariables, Clone, Debug)]
+    pub struct PullRequestsOpenVariables<'a> {
+        pub repo_owner: &'a str,
+        pub repo_name: &'a str,
+        pub after: Option<String>,
+    }
+
+    #[derive(cynic::QueryFragment, Debug)]
+    #[cynic(graphql_type = "Query", variables = "PullRequestsOpenVariables")]
+    pub struct PullRequestsOpen {
+        #[arguments(owner: $repo_owner, name: $repo_name)]
+        pub repository: Option<Repository>,
+    }
+
+    #[derive(cynic::QueryFragment, Debug)]
+    #[cynic(variables = "PullRequestsOpenVariables")]
+    pub struct Repository {
+        #[arguments(first: 100, after: $after, states: "OPEN", labels: ["S-waiting-on-review","T-compiler"])]
+        pub pull_requests: PullRequestConnection,
+    }
+
+    #[derive(cynic::QueryFragment, Debug)]
+    pub struct PullRequest {
+        pub number: i32,
+        pub updated_at: DateTime,
+        pub created_at: DateTime,
+        #[arguments(first: 10)]
+        pub assignees: UserConnection,
+        #[arguments(first: 5, orderBy: { direction: "DESC", field: "NAME" })]
+        pub labels: Option<LabelConnection>,
+    }
+}

src/db.rs

@@ -308,4 +308,12 @@ CREATE UNIQUE INDEX jobs_name_scheduled_at_unique_index
         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[]
+);
+CREATE UNIQUE INDEX review_prefs_user_id ON review_prefs(user_id);
+ ",
 ];

src/github.rs

@@ -2592,6 +2592,81 @@ async fn project_items_by_status(
     Ok(all_items)
 }
 
+/// Retrieve all pull requests in status OPEN, not DRAFT, with label: S-waiting-on-review
+pub async fn retrieve_pull_requests(
+    repo: &Repository,
+    client: &GithubClient,
+) -> anyhow::Result<Vec<(Vec<User>, i32)>> {
+    use cynic::QueryBuilder;
+    use github_graphql::pull_requests_open::{PullRequestsOpen, PullRequestsOpenVariables};
+
+    let repo_owner = repo.owner();
+    let repo_name = repo.name();
+
+    let mut prs = vec![];
+
+    let mut vars = PullRequestsOpenVariables {
+        repo_owner,
+        repo_name,
+        after: None,
+    };
+    loop {
+        let query = PullRequestsOpen::build(vars.clone());
+        let req = client.post(&client.graphql_url);
+        let req = req.json(&query);
+
+        let data: cynic::GraphQlResponse<PullRequestsOpen> = client.json(req).await?;
+        if let Some(errors) = data.errors {
+            anyhow::bail!("There were graphql errors. {:?}", errors);
+        }
+        let repository = data
+            .data
+            .ok_or_else(|| anyhow::anyhow!("No data returned."))?
+            .repository
+            .ok_or_else(|| anyhow::anyhow!("No repository."))?;
+        prs.extend(repository.pull_requests.nodes);
+
+        let page_info = repository.pull_requests.page_info;
+        if !page_info.has_next_page || page_info.end_cursor.is_none() {
+            break;
+        }
+        vars.after = page_info.end_cursor;
+    }
+
+    log::trace!("Retrieved {} PRs", prs.len());
+
+    let prs: Vec<_> = prs
+        .into_iter()
+        .filter_map(|pr| {
+            if pr.is_draft {
+                return None;
+            }
+
+            let assignees = pr
+                .assignees
+                .nodes
+                .iter()
+                .map(|user| {
+                    let x = user.id.clone();
+                    let user_id_p = x
+                        .into_inner()
+                        .parse::<i64>()
+                        .expect("Failed to parse a cynic::Id into a i64");
+                    let user_id = Some(user_id_p);
+                    User {
+                        login: user.login.clone(),
+                        id: user_id,
+                    }
+                })
+                .collect();
+
+            Some((assignees, pr.number))
+        })
+        .collect();
+
+    Ok(prs)
+}
+
 pub enum DesignMeetingStatus {
     Proposed,
     Scheduled,

src/handlers.rs

@@ -39,6 +39,7 @@ mod notification;
 mod notify_zulip;
 mod ping;
 mod prioritize;
+pub mod pull_requests_assignment_update;
 mod relabel;
 mod review_requested;
 mod review_submitted;

src/handlers/pull_requests_assignment_update.rs

@@ -0,0 +1,68 @@
+use crate::github::retrieve_pull_requests;
+use crate::jobs::Job;
+use anyhow::Context as _;
+use async_trait::async_trait;
+use tokio_postgres::Client as DbClient;
+
+pub struct PullRequestAssignmentUpdate;
+
+#[async_trait]
+impl Job for PullRequestAssignmentUpdate {
+    fn name(&self) -> &'static str {
+        "pull_request_assignment_update"
+    }
+
+    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");
+
+        // delete everything before populating
+        init_table(&db).await?;
+
+        let rust_repo = gh.repository("rust-lang/rust").await?;
+        let prs = retrieve_pull_requests(&rust_repo, &gh).await?;
+
+        // populate the table
+        for (assignees, pr_num) in prs {
+            for assignee in assignees {
+                let assignee_id = assignee.id.expect("checked");
+                ensure_team_member(&db, assignee_id, &assignee.login).await?;
+                create_team_member_workqueue(&db, assignee_id, pr_num).await?;
+            }
+        }
+
+        Ok(())
+    }
+}
+
+/// Truncate the review prefs table
+async fn init_table(db: &DbClient) -> anyhow::Result<u64> {
+    let res = db.execute("TRUNCATE review_prefs;", &[]).await?;
+    Ok(res)
+}
+
+/// Add a new user (if not existing)
+async fn ensure_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)
+}
+
+/// Create a team member work queue
+async fn create_team_member_workqueue(
+    db: &DbClient,
+    user_id: i64,
+    pr: i32,
+) -> anyhow::Result<u64, anyhow::Error> {
+    let q = "INSERT INTO review_prefs (user_id, assigned_prs) VALUES ($1, $2);";
+    db.execute(q, &[&user_id, &vec![pr], &pr])
+        .await
+        .context("Insert DB error")
+}

src/jobs.rs

@@ -50,6 +50,7 @@ use crate::{
     db::jobs::JobSchedule,
     handlers::{
         docs_update::DocsUpdateJob,
+        pull_requests_assignment_update::PullRequestAssignmentUpdate,
         rustc_commits::RustcCommitsJob,
         types_planning_updates::{
             TypesPlanningMeetingThreadOpenJob, TypesPlanningMeetingUpdatesPing,
@@ -73,6 +74,7 @@ pub fn jobs() -> Vec<Box<dyn Job + Send + Sync>> {
         Box::new(RustcCommitsJob),
         Box::new(TypesPlanningMeetingThreadOpenJob),
         Box::new(TypesPlanningMeetingUpdatesPing),
+        Box::new(PullRequestAssignmentUpdate),
     ]
 }
 
@@ -119,7 +121,7 @@ fn jobs_defined() {
     unique_all_job_names.dedup();
     assert_eq!(all_job_names, unique_all_job_names);
 
-    // Also ensure that our defalt jobs are release jobs
+    // Also ensure that our default jobs are release jobs
     let default_jobs = default_jobs();
     default_jobs
         .iter()