CORE-131 adminGetUsersByIDs searches more kinds of ids (#1576)

broadinstitute · Oct 31, 2024 · 0f47b00 · 0f47b00
1 parent 58be30b
commit 0f47b00
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 4 deletions.
diff --git a/src/main/resources/swagger/api-docs.yaml b/src/main/resources/swagger/api-docs.yaml
@@ -199,7 +199,7 @@ paths:
     post:
       tags:
         - Admin
-      summary: Gets a list of users for a list of Sam User IDs
+      summary: Gets a list of users for a list of Sam, Azure B2C, or Google Subject IDs
       operationId: adminGetUsersByIDs
       requestBody:
         content:

diff --git a/src/main/scala/org/broadinstitute/dsde/workbench/sam/dataAccess/PostgresDirectoryDAO.scala b/src/main/scala/org/broadinstitute/dsde/workbench/sam/dataAccess/PostgresDirectoryDAO.scala
@@ -430,7 +430,19 @@ class PostgresDirectoryDAO(protected val writeDbRef: DbReference, protected val
     } else {
       readOnlyTransaction("batchLoadUsers", samRequestContext) { implicit session =>
         val userTable = UserTable.syntax
-        val loadUserQuery = samsql"select ${userTable.resultAll} from ${UserTable as userTable} where ${userTable.id} in (${samUserIds})"
+        // the with clause is to keep the query size down, we only send the samUserIds once and reuse it in each unioned query
+        val loadUserQuery =
+          samsql"""
+                  with sam_user_ids (user_id) as (values ${samUserIds.map(id => samsqls"($id)")})
+                  select ${userTable.resultAll} from ${UserTable as userTable}
+                  join sam_user_ids ids on ids.user_id = ${userTable.id}
+                  union
+                  select ${userTable.resultAll} from ${UserTable as userTable}
+                  join sam_user_ids ids on ids.user_id = ${userTable.azureB2cId}
+                  union
+                  select ${userTable.resultAll} from ${UserTable as userTable}
+                  join sam_user_ids ids on ids.user_id = ${userTable.googleSubjectId}
+                  """
 
         loadUserQuery
           .map(UserTable(userTable))

diff --git a/...est/scala/org/broadinstitute/dsde/workbench/sam/dataAccess/PostgresDirectoryDAOSpec.scala b/...est/scala/org/broadinstitute/dsde/workbench/sam/dataAccess/PostgresDirectoryDAOSpec.scala
@@ -630,8 +630,19 @@ class PostgresDirectoryDAOSpec extends RetryableAnyFreeSpec with Matchers with B
         assume(databaseEnabled, databaseEnabledClue)
         val users = Seq.range(0, 10).map(_ => Generator.genWorkbenchUserBoth.sample.get)
         users.foreach(user => dao.createUser(user, samRequestContext).unsafeRunSync())
-        val loadedUsers = dao.batchLoadUsers(users.map(_.id).toSet, samRequestContext).unsafeRunSync()
-        loadedUsers should contain theSameElementsAs users
+        val loadedUsersBySamId = dao.batchLoadUsers(users.map(_.id).toSet, samRequestContext).unsafeRunSync()
+        loadedUsersBySamId should contain theSameElementsAs users
+
+        val b2cIds = users.flatMap(_.azureB2CId.map(id => WorkbenchUserId(id.value))).toSet
+        val loadedUsersByAzureB2cId = dao.batchLoadUsers(b2cIds, samRequestContext).unsafeRunSync()
+        loadedUsersByAzureB2cId should contain theSameElementsAs users
+
+        val googleSubjectIds = users.flatMap(_.googleSubjectId.map(id => WorkbenchUserId(id.value))).toSet
+        val loadedUsersByGoogleSubjectId = dao.batchLoadUsers(googleSubjectIds, samRequestContext).unsafeRunSync()
+        loadedUsersByGoogleSubjectId should contain theSameElementsAs users
+
+        val loadedBy2Ids = dao.batchLoadUsers(googleSubjectIds ++ b2cIds, samRequestContext).unsafeRunSync()
+        loadedBy2Ids should contain theSameElementsAs users
       }
     }