[SCIM 4.2/4] List groups and users in a single query

nexus/db-queries/src/db/datastore/scim_provider_store.rs

@@ -922,26 +922,36 @@ impl<'a> CrdbScimProviderStore<'a> {
         Ok(convert_to_scim_group(new_group, members))
     }
 
-    async fn list_groups_in_txn(
+    async fn list_groups_with_members(
         &self,
         conn: &async_bb8_diesel::Connection<DbConnection>,
         err: OptionalError<ProviderStoreError>,
         filter: Option<FilterOp>,
     ) -> Result<Vec<StoredParts<Group>>, diesel::result::Error> {
-        use nexus_db_schema::schema::silo_group::dsl;
+        use nexus_db_schema::schema::silo_group::dsl as group_dsl;
+        use nexus_db_schema::schema::silo_group_membership::dsl as membership_dsl;
+        use std::collections::HashMap;
 
-        let mut query = dsl::silo_group
-            .filter(dsl::silo_id.eq(self.authz_silo.id()))
-            .filter(dsl::user_provision_type.eq(model::UserProvisionType::Scim))
-            .filter(dsl::time_deleted.is_null())
+        let mut query = group_dsl::silo_group
+            .left_join(
+                membership_dsl::silo_group_membership
+                    .on(group_dsl::id.eq(membership_dsl::silo_group_id)),
+            )
+            .filter(group_dsl::silo_id.eq(self.authz_silo.id()))
+            .filter(
+                group_dsl::user_provision_type
+                    .eq(model::UserProvisionType::Scim),
+            )
+            .filter(group_dsl::time_deleted.is_null())
             .into_boxed();
 
         match filter {
             Some(FilterOp::DisplayNameEq(display_name)) => {
                 // displayName is defined as `"caseExact" : false` in RFC 7643,
                 // section 8.7.1
                 query = query.filter(
-                    lower(dsl::display_name).eq(lower(display_name.clone())),
+                    lower(group_dsl::display_name)
+                        .eq(lower(display_name.clone())),
                 );
             }
 
@@ -959,20 +969,60 @@ impl<'a> CrdbScimProviderStore<'a> {
             }
         }
 
-        let groups = query
-            .select(model::SiloGroup::as_returning())
+        // Select group fields and optional member user_id
+        type GroupRow = (model::SiloGroup, Option<Uuid>);
+        type GroupsMap =
+            HashMap<Uuid, (Option<model::SiloGroup>, Vec<SiloUserUuid>)>;
+
+        let rows: Vec<GroupRow> = query
+            .select((
+                model::SiloGroup::as_select(),
+                membership_dsl::silo_user_id.nullable(),
+            ))
             .load_async(conn)
             .await?;
 
-        let mut returned_groups = Vec::with_capacity(groups.len());
+        // Group the results by group_id
+        let mut groups_map: GroupsMap = HashMap::new();
 
-        for group in groups {
-            let members = self
-                .get_group_members_for_group_in_txn(
-                    conn,
-                    group.identity.id.into(),
-                )
-                .await?;
+        for (group, maybe_user_id) in rows {
+            let group_id = group.identity.id.into_untyped_uuid();
+
+            let entry = groups_map
+                .entry(group_id)
+                .or_insert_with(|| (None, Vec::new()));
+
+            // Store the group on first occurrence
+            if entry.0.is_none() {
+                entry.0 = Some(group);
+            }
+
+            // If this row has a member, add it to the group's members
+            if let Some(user_id) = maybe_user_id {
+                entry.1.push(SiloUserUuid::from_untyped_uuid(user_id));
+            }
+        }
+
+        // Convert to the expected return type
+        let mut returned_groups = Vec::with_capacity(groups_map.len());
+
+        for (_group_id, (maybe_group, members)) in groups_map {
+            let group = maybe_group.expect("group should always be present");
+
+            let members = if members.is_empty() {
+                None
+            } else {
+                let mut id_ord_map = IdOrdMap::with_capacity(members.len());
+                for user_id in members {
+                    id_ord_map
+                        .insert_unique(GroupMember {
+                            resource_type: Some(ResourceType::User.to_string()),
+                            value: Some(user_id.to_string()),
+                        })
+                        .expect("user_id should be unique");
+                }
+                Some(id_ord_map)
+            };
 
             let SiloGroup::Scim(group) = group.into() else {
                 // With the user provision type filter, this should never be
@@ -1720,14 +1770,7 @@ impl<'a> ProviderStore for CrdbScimProviderStore<'a> {
         let err: OptionalError<ProviderStoreError> = OptionalError::new();
 
         let groups = self
-            .datastore
-            .transaction_retry_wrapper("scim_list_groups")
-            .transaction(&conn, |conn| {
-                let err = err.clone();
-                let filter = filter.clone();
-
-                async move { self.list_groups_in_txn(&conn, err, filter).await }
-            })
+            .list_groups_with_members(&conn, err.clone(), filter)
             .await
             .map_err(|e| {
                 if let Some(e) = err.take() {

nexus/tests/integration_tests/scim.rs

@@ -1995,7 +1995,7 @@ async fn test_scim_user_admin_group_priv_conflict(
 }
 
 #[nexus_test]
-async fn test_scim_list_users_with_groups(cptestctx: &ControlPlaneTestContext) {
+async fn test_scim_list_users_and_groups(cptestctx: &ControlPlaneTestContext) {
     let client = &cptestctx.external_client;
     let nexus = &cptestctx.server.server_context().nexus;
     let opctx = OpContext::for_tests(
@@ -2101,7 +2101,7 @@ async fn test_scim_list_users_with_groups(cptestctx: &ControlPlaneTestContext) {
     .execute_and_parse_unwrap()
     .await;
 
-    let _group3: scim2_rs::Group = NexusRequest::new(
+    let group3: scim2_rs::Group = NexusRequest::new(
         RequestBuilder::new(client, Method::POST, "/scim/v2/Groups")
             .header(http::header::CONTENT_TYPE, "application/scim+json")
             .header(
@@ -2184,4 +2184,60 @@ async fn test_scim_list_users_with_groups(cptestctx: &ControlPlaneTestContext) {
     // user5 should have no groups
     let user5 = find_user(&users[4].id);
     assert!(user5.groups.is_none());
+
+    // List all groups and verify members
+    let response: scim2_rs::ListResponse = NexusRequest::new(
+        RequestBuilder::new(client, Method::GET, "/scim/v2/Groups")
+            .header(http::header::CONTENT_TYPE, "application/scim+json")
+            .header(
+                http::header::AUTHORIZATION,
+                format!("Bearer oxide-scim-{}", created_token.bearer_token),
+            )
+            .allow_non_dropshot_errors()
+            .expect_status(Some(StatusCode::OK)),
+    )
+    .execute_and_parse_unwrap()
+    .await;
+
+    let returned_groups: Vec<scim2_rs::Group> = serde_json::from_value(
+        serde_json::to_value(&response.resources).unwrap(),
+    )
+    .unwrap();
+
+    // Find our created groups in the response
+    let find_group = |group_id: &str| {
+        returned_groups
+            .iter()
+            .find(|g| g.id == group_id)
+            .expect("group should be in list")
+    };
+
+    // group1 should have 3 members
+    let returned_group1 = find_group(&group1.id);
+    assert!(returned_group1.members.is_some());
+    let group1_members = returned_group1.members.as_ref().unwrap();
+    assert_eq!(group1_members.len(), 3);
+    let group1_member_ids: std::collections::HashSet<_> = group1_members
+        .iter()
+        .map(|m| m.value.as_ref().unwrap().as_str())
+        .collect();
+    assert!(group1_member_ids.contains(users[0].id.as_str()));
+    assert!(group1_member_ids.contains(users[1].id.as_str()));
+    assert!(group1_member_ids.contains(users[2].id.as_str()));
+
+    // group2 should have 2 members
+    let returned_group2 = find_group(&group2.id);
+    assert!(returned_group2.members.is_some());
+    let group2_members = returned_group2.members.as_ref().unwrap();
+    assert_eq!(group2_members.len(), 2);
+    let group2_member_ids: std::collections::HashSet<_> = group2_members
+        .iter()
+        .map(|m| m.value.as_ref().unwrap().as_str())
+        .collect();
+    assert!(group2_member_ids.contains(users[0].id.as_str()));
+    assert!(group2_member_ids.contains(users[3].id.as_str()));
+
+    // group3 should have no members
+    let returned_group3 = find_group(&group3.id);
+    assert!(returned_group3.members.is_none());
 }