Add CSV groups output, and optional expanding

Author: illicitonion

src/bin/trainee-tracker.rs

@@ -88,6 +88,10 @@ async fn main() {
             "/groups/google",
             get(trainee_tracker::frontend::list_google_groups),
         )
+        .route(
+            "/groups/google.csv",
+            get(trainee_tracker::frontend::list_google_groups_csv),
+        )
         .layer(session_layer)
         .with_state(server_state);
 

src/frontend.rs

@@ -1,21 +1,13 @@
-use std::{
-    collections::{BTreeMap, BTreeSet},
-    str::FromStr,
-};
+use std::collections::{BTreeMap, BTreeSet};
 
 use anyhow::Context;
 use askama::Template;
 use axum::{
     extract::{OriginalUri, Path, Query, State},
-    response::Html,
+    response::{Html, IntoResponse, Response},
 };
-use email_address::EmailAddress;
 use futures::future::join_all;
-use gsuite_api::{
-    types::{Group, Member},
-    Response,
-};
-use http::Uri;
+use http::{header::CONTENT_TYPE, StatusCode, Uri};
 use serde::Deserialize;
 use tower_sessions::Session;
 
@@ -25,7 +17,7 @@ use crate::{
     course::{
         fetch_batch_metadata, get_batch, Attendance, Batch, BatchMetadata, Course, Submission,
     },
-    google_groups::{groups_client, GoogleGroup},
+    google_groups::{get_groups, groups_client, GoogleGroup},
     octocrab::octocrab,
     prs::{MaybeReviewerStaffOnlyDetails, PrState, ReviewerInfo},
     reviewer_staff_info::get_reviewer_staff_info,
@@ -248,68 +240,85 @@ pub(crate) struct Redirect {
 #[derive(Template)]
 #[template(path = "google-groups.html")]
 struct GoogleGroups {
-    pub groups: Vec<GoogleGroup>,
+    pub groups: BTreeSet<GoogleGroup>,
+}
+
+#[derive(Deserialize)]
+pub struct GroupListParams {
+    #[serde(default)]
+    expand: bool,
 }
 
 pub async fn list_google_groups(
     session: Session,
     State(server_state): State<ServerState>,
     OriginalUri(original_uri): OriginalUri,
+    Query(params): Query<GroupListParams>,
 ) -> Result<Html<String>, Error> {
     let client = groups_client(&session, server_state, original_uri).await?;
-    let groups_response = client
-        .groups()
-        .list_all(
-            "my_customer",
-            "codeyourfuture.io",
-            gsuite_api::types::DirectoryGroupsListOrderBy::Email,
-            "",
-            gsuite_api::types::SortOrder::Ascending,
-            "",
-        )
-        .await
-        .context("Failed to list Google groups")?;
-    let groups = error_for_status(groups_response)?;
-    let group_member_futures = groups
+    let mut groups = get_groups(&client).await?;
+    if params.expand {
+        groups
+            .expand_recursively()
+            .context("Failed to expand groups recursively")?;
+    }
+    Ok(Html(
+        GoogleGroups {
+            groups: groups.groups,
+        }
+        .render()
+        .unwrap(),
+    ))
+}
+
+pub async fn list_google_groups_csv(
+    session: Session,
+    State(server_state): State<ServerState>,
+    OriginalUri(original_uri): OriginalUri,
+    Query(params): Query<GroupListParams>,
+) -> Result<Csv, Error> {
+    let client = groups_client(&session, server_state, original_uri).await?;
+    let mut groups = get_groups(&client).await?;
+    if params.expand {
+        groups
+            .expand_recursively()
+            .context("Failed to expand groups recursively")?;
+    }
+
+    let member_count = groups
+        .groups
         .iter()
-        .map(|Group { id, .. }| async { client.members().list_all(id, false, "").await })
-        .collect::<Vec<_>>();
-    let group_members = join_all(group_member_futures).await;
+        .map(|group| group.members.len())
+        .max()
+        .unwrap_or(0);
 
-    let result = groups
-        .into_iter()
-        .zip(group_members.into_iter())
-        .map(|(group, members)| {
-            let members =
-                error_for_status(members.context("Failed to list Google group members")?)?;
-            Ok(GoogleGroup {
-                email: EmailAddress::from_str(&group.email).with_context(|| {
-                    format!("Failed to parse group email address {}", group.email)
-                })?,
-                members: members
-                    .into_iter()
-                    .map(|Member { email, .. }| {
-                        Ok(EmailAddress::from_str(&email).with_context(|| {
-                            format!(
-                                "Failed to parse group member email address {} (member of {})",
-                                email, group.email
-                            )
-                        })?)
-                    })
-                    .collect::<Result<_, anyhow::Error>>()?,
-            })
-        })
-        .collect::<Result<_, Error>>()?;
-    Ok(Html(GoogleGroups { groups: result }.render().unwrap()))
+    // Manually writing a CSV because the CSV crate doesn't like different numbers of fields per record.
+    let mut out = String::new();
+    out += "group";
+    for i in 0..member_count {
+        out += &format!(",member{}", i + 1);
+    }
+    out += "\n";
+
+    for group in groups.groups {
+        out += group.email.as_str();
+        for member in group.members {
+            out += ",";
+            out += member.as_str();
+        }
+        out += "\n"
+    }
+    Ok(Csv(out))
 }
 
-fn error_for_status<T: std::fmt::Debug>(response: Response<T>) -> Result<T, Error> {
-    if !response.status.is_success() {
-        Err(Error::Fatal(anyhow::anyhow!(
-            "Got bad response from Google Groups API: {:?}",
-            response
-        )))
-    } else {
-        Ok(response.body)
+pub struct Csv(String);
+
+impl IntoResponse for Csv {
+    fn into_response(self) -> axum::response::Response {
+        Response::builder()
+            .header(CONTENT_TYPE, "text/csv")
+            .status(StatusCode::OK)
+            .body(axum::body::Body::from(self.0))
+            .expect("Failed to build response")
     }
 }

src/google_groups.rs

@@ -1,8 +1,15 @@
-use std::collections::BTreeSet;
+use std::{
+    collections::{BTreeMap, BTreeSet},
+    str::FromStr,
+};
 
 use anyhow::Context;
 use email_address::EmailAddress;
-use gsuite_api::Client;
+use futures::future::join_all;
+use gsuite_api::{
+    types::{Group, Member},
+    Client, Response,
+};
 use http::Uri;
 use tower_sessions::Session;
 
@@ -45,6 +52,7 @@ pub async fn groups_client(
     }
 }
 
+#[derive(Debug, PartialEq, Eq, PartialOrd, Ord)]
 pub(crate) struct GoogleGroup {
     pub email: EmailAddress,
     pub members: BTreeSet<EmailAddress>,
@@ -57,3 +65,170 @@ impl GoogleGroup {
         format!("https://groups.google.com/a/{domain}/g/{user}")
     }
 }
+
+#[derive(Debug, PartialEq, Eq)]
+pub(crate) struct GoogleGroups {
+    pub groups: BTreeSet<GoogleGroup>,
+}
+
+pub(crate) async fn get_groups(client: &Client) -> Result<GoogleGroups, Error> {
+    let groups_response = client
+        .groups()
+        .list_all(
+            "my_customer",
+            "codeyourfuture.io",
+            gsuite_api::types::DirectoryGroupsListOrderBy::Email,
+            "",
+            gsuite_api::types::SortOrder::Ascending,
+            "",
+        )
+        .await
+        .context("Failed to list Google groups")?;
+    let groups = error_for_status(groups_response)?;
+    let group_member_futures = groups
+        .iter()
+        .map(|Group { id, .. }| async { client.members().list_all(id, false, "").await })
+        .collect::<Vec<_>>();
+    let group_members = join_all(group_member_futures).await;
+
+    let groups = groups
+        .into_iter()
+        .zip(group_members.into_iter())
+        .map(|(group, members)| {
+            let members =
+                error_for_status(members.context("Failed to list Google group members")?)?;
+            Ok(GoogleGroup {
+                email: EmailAddress::from_str(&group.email).with_context(|| {
+                    format!("Failed to parse group email address {}", group.email)
+                })?,
+                members: members
+                    .into_iter()
+                    .map(|Member { email, .. }| {
+                        Ok(EmailAddress::from_str(&email).with_context(|| {
+                            format!(
+                                "Failed to parse group member email address {} (member of {})",
+                                email, group.email
+                            )
+                        })?)
+                    })
+                    .collect::<Result<_, anyhow::Error>>()?,
+            })
+        })
+        .collect::<Result<_, Error>>()?;
+    Ok(GoogleGroups { groups })
+}
+
+impl GoogleGroups {
+    pub(crate) fn expand_recursively(&mut self) -> Result<(), anyhow::Error> {
+        let mut index = BTreeMap::new();
+        let groups = self
+            .groups
+            .iter()
+            .map(|GoogleGroup { email, .. }| email.clone())
+            .collect::<BTreeSet<_>>();
+        for group in &self.groups {
+            index.insert(group.email.clone(), group.members.clone());
+        }
+        let mut iteration = 0;
+        loop {
+            let mut changed = false;
+            if iteration > 15 {
+                return Err(anyhow::anyhow!("Reached recursion limit expanding groups"));
+            }
+            let mut to_replace: BTreeMap<
+                EmailAddress,
+                BTreeMap<EmailAddress, BTreeSet<EmailAddress>>,
+            > = BTreeMap::new();
+            for (group, members) in index.iter() {
+                for member in members.iter() {
+                    if groups.contains(member) {
+                        to_replace
+                            .entry(group.clone())
+                            .or_default()
+                            .insert(member.clone(), index.get(member).unwrap().clone());
+                    }
+                }
+            }
+            for (group, replacements) in to_replace {
+                for (to_replace, replacements) in replacements {
+                    index.get_mut(&group).unwrap().remove(&to_replace);
+                    index.get_mut(&group).unwrap().extend(replacements);
+                    changed = true;
+                }
+            }
+            if !changed {
+                break;
+            }
+            iteration += 1;
+        }
+        self.groups = index
+            .into_iter()
+            .map(|(email, members)| GoogleGroup { email, members })
+            .collect();
+        Ok(())
+    }
+}
+
+fn error_for_status<T: std::fmt::Debug>(response: Response<T>) -> Result<T, Error> {
+    if !response.status.is_success() {
+        Err(Error::Fatal(anyhow::anyhow!(
+            "Got bad response from Google Groups API: {:?}",
+            response
+        )))
+    } else {
+        Ok(response.body)
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use email_address::EmailAddress;
+    use maplit::btreeset;
+
+    use crate::google_groups::{GoogleGroup, GoogleGroups};
+
+    #[test]
+    fn test_expand_recursively() {
+        let outer_group = EmailAddress::new_unchecked("container@example.com");
+        let inner_group = EmailAddress::new_unchecked("inner@example.com");
+        let inner_members = btreeset![
+            EmailAddress::new_unchecked("someone@example.com"),
+            EmailAddress::new_unchecked("someone-else@example.com")
+        ];
+        let other_member = EmailAddress::new_unchecked("external@example.com");
+        let all_members = btreeset![
+            other_member.clone(),
+            EmailAddress::new_unchecked("someone@example.com"),
+            EmailAddress::new_unchecked("someone-else@example.com")
+        ];
+
+        let mut input = GoogleGroups {
+            groups: btreeset![
+                GoogleGroup {
+                    email: outer_group.clone(),
+                    members: btreeset![inner_group.clone(), other_member.clone()],
+                },
+                GoogleGroup {
+                    email: inner_group.clone(),
+                    members: inner_members.clone(),
+                }
+            ],
+        };
+
+        let want = GoogleGroups {
+            groups: btreeset![
+                GoogleGroup {
+                    email: outer_group,
+                    members: all_members,
+                },
+                GoogleGroup {
+                    email: inner_group,
+                    members: inner_members,
+                }
+            ],
+        };
+
+        input.expand_recursively().unwrap();
+        assert_eq!(input, want);
+    }
+}