refactor(csa-362): move IamClient into SDK as IdentityEndpoint

- Add IdentityEndpoint.java extending HeraclesEndpoint in core SDK
- Add IamGroupRef.java model class for IAM group references
- Register client.identity in AtlanClient alongside other endpoints
- Remove IamClient.kt from admin-export package
- Update Users.kt and AdminExporter.kt to use ctx.client.identity
- Update IamClientTest.kt to test IdentityEndpoint via forTesting()
- Add branch trigger to merge.yml for CI build

Author: bladata1990

.github/workflows/merge.yml

@@ -6,7 +6,7 @@ permissions:
 
 on:
   push:
-    branches: [main]
+    branches: [main, fix/csa-362-admin-export-scale]
   workflow_dispatch:
     inputs:
       branch:

samples/packages/admin-export/src/main/kotlin/com/atlan/pkg/ae/AdminExporter.kt

@@ -44,13 +44,6 @@ object AdminExporter {
             val glossaryMap = preloadGlossaryNameMap(ctx)
             val connectionMap = preloadConnectionMap(ctx)
 
-            // Build IAM client for bulk user/group fetches via Redis-backed identity API.
-            // Replaces per-user Keycloak calls that fail at high user volumes (60K+).
-            // Auth: client_credentials grant with client_id=atlan-backend.
-            // CLIENT_SECRET is injected from argo-client-creds secret into the main container.
-            val identityBase = "http://heracles-service.heracles.svc.cluster.local"
-            val bearerToken = ctx.client.impersonate.escalate()
-            val iamClient = IamClient(baseUrl = identityBase, bearerToken = bearerToken, logger = logger)
             val xlsxOutput = ctx.config.fileFormat == "XLSX"
 
             val outputDirectory = validatePathIsSafe(od)
@@ -89,7 +82,7 @@ object AdminExporter {
                 ExcelWriter(xlsxFileActual.toString()).use { xlsx ->
                     ctx.config.objectsToInclude.forEach { objectName ->
                         when (objectName) {
-                            "users" -> Users(ctx, xlsx.createSheet("Users"), logger, iamClient).export()
+                            "users" -> Users(ctx, xlsx.createSheet("Users"), logger).export()
                             "groups" -> Groups(ctx, xlsx.createSheet("Groups"), logger).export()
                             "personas" -> Personas(ctx, xlsx.createSheet("Personas"), glossaryMap, connectionMap, logger).export()
                             "purposes" -> Purposes(ctx, xlsx.createSheet("Purposes"), logger).export()
@@ -101,7 +94,7 @@ object AdminExporter {
             } else {
                 ctx.config.objectsToInclude.forEach { objectName ->
                     when (objectName) {
-                        "users" -> CSVWriter(usersFileActual.toString()).use { csv -> Users(ctx, csv, logger, iamClient).export() }
+                        "users" -> CSVWriter(usersFileActual.toString()).use { csv -> Users(ctx, csv, logger).export() }
                         "groups" -> CSVWriter(groupsFileActual.toString()).use { csv -> Groups(ctx, csv, logger).export() }
                         "personas" -> CSVWriter(personasFileActual.toString()).use { csv -> Personas(ctx, csv, glossaryMap, connectionMap, logger).export() }
                         "purposes" -> CSVWriter(purposesFileActual.toString()).use { csv -> Purposes(ctx, csv, logger).export() }

samples/packages/admin-export/src/main/kotlin/com/atlan/pkg/ae/IamClient.kt

@@ -4,9 +4,9 @@ package com.atlan.pkg.ae.exports
 
 import AdminExportCfg
 import com.atlan.api.UsersEndpoint
+import com.atlan.model.admin.IamGroupRef
 import com.atlan.model.admin.UserRequest
 import com.atlan.pkg.PackageContext
-import com.atlan.pkg.ae.IamClient
 import com.atlan.pkg.serde.TabularWriter
 import com.atlan.pkg.serde.cell.TimestampXformer
 import mu.KLogger
@@ -16,7 +16,6 @@ class Users(
     private val ctx: PackageContext<AdminExportCfg>,
     private val writer: TabularWriter,
     private val logger: KLogger,
-    private val iamClient: IamClient,
 ) {
     fun export() {
         logger.info { "Exporting all users..." }
@@ -56,20 +55,20 @@ class Users(
         logger.info { "Fetched ${allUsers.size} users. Resolving group memberships via IAM API..." }
 
         // Bulk-fetch group memberships — ~120 Redis calls instead of 60K Keycloak calls
-        val userGroupsMap = iamClient.getUserGroups(allUsers.mapNotNull { it.id })
+        val userGroupsMap = ctx.client.identity.getUserGroups(allUsers.mapNotNull { it.id })
 
         // Collect all unique group IDs seen across all users, then resolve aliases in bulk
         val allGroupIds =
             userGroupsMap.values
                 .flatten()
                 .map { it.id }
                 .distinct()
-        val groupAliasMap = iamClient.getGroupAliases(allGroupIds)
+        val groupAliasMap = ctx.client.identity.getGroupAliases(allGroupIds)
 
         val ts = Instant.now().toString()
         allUsers.forEach { user ->
             val personas = user.personas?.joinToString("\n") { it.displayName ?: "" } ?: ""
-            val groups = userGroupsMap[user.id] ?: emptyList<IamClient.GroupRef>()
+            val groups = userGroupsMap[user.id] ?: emptyList<IamGroupRef>()
             val technicalNames = groups.joinToString("\n") { it.name }
             val nontechnicalNames = groups.joinToString("\n") { groupAliasMap[it.id] ?: it.name }
             val designation =

samples/packages/admin-export/src/main/kotlin/com/atlan/pkg/ae/exports/Users.kt

@@ -1,8 +1,7 @@
 /* SPDX-License-Identifier: Apache-2.0
    Copyright 2024 Atlan Pte. Ltd. */
-import com.atlan.pkg.ae.IamClient
+import com.atlan.api.IdentityEndpoint
 import com.sun.net.httpserver.HttpServer
-import mu.KotlinLogging
 import org.testng.Assert.assertEquals
 import org.testng.Assert.assertTrue
 import org.testng.annotations.AfterClass
@@ -13,14 +12,13 @@ import java.net.URLDecoder
 import java.nio.charset.StandardCharsets
 
 /**
- * Unit tests for IamClient using a local HTTP stub server.
+ * Unit tests for IdentityEndpoint using a local HTTP stub server.
  * Verifies bulk user-group and group-alias fetching, batching behaviour,
  * and graceful handling of empty inputs and missing optional fields.
  */
 class IamClientTest {
-    private val logger = KotlinLogging.logger {}
     private lateinit var server: HttpServer
-    private lateinit var client: IamClient
+    private lateinit var endpoint: IdentityEndpoint
 
     @BeforeClass
     fun setup() {
@@ -34,20 +32,16 @@ class IamClientTest {
                     query.split("&").firstOrNull { it.startsWith("filter=") }?.removePrefix("filter=") ?: "{}",
                     StandardCharsets.UTF_8,
                 )
-            // Extract IDs from {"id":{"$in":["id1","id2",...]}}
             val ids =
                 Regex("\"([^\"]+)\"")
-                    .findAll(
-                        filter.substringAfter("\$in\":[").substringBefore("]"),
-                    ).map { it.groupValues[1] }
+                    .findAll(filter.substringAfter("\$in\":[").substringBefore("]"))
+                    .map { it.groupValues[1] }
                     .toList()
-
             val body =
                 ids
                     .joinToString(",", "[", "]") { id ->
                         """{"id":"$id","groups":[{"id":"grp-$id","name":"group-$id"}]}"""
                     }.toByteArray()
-
             exchange.sendResponseHeaders(200, body.size.toLong())
             exchange.responseBody.use { it.write(body) }
         }
@@ -62,34 +56,25 @@ class IamClientTest {
                 )
             val ids =
                 Regex("\"([^\"]+)\"")
-                    .findAll(
-                        filter.substringAfter("\$in\":[").substringBefore("]"),
-                    ).map { it.groupValues[1] }
+                    .findAll(filter.substringAfter("\$in\":[").substringBefore("]"))
+                    .map { it.groupValues[1] }
                     .toList()
-
             val body =
                 ids
                     .joinToString(",", "[", "]") { id ->
-                        when {
-                            id.contains("no-alias") -> """{"id":"$id"}"""
-
-                            // missing attributes entirely
-                            else -> """{"id":"$id","attributes":{"alias":"Alias for $id"}}"""
+                        if (id.contains("no-alias")) {
+                            """{"id":"$id"}"""
+                        } else {
+                            """{"id":"$id","attributes":{"alias":"Alias for $id"}}"""
                         }
                     }.toByteArray()
-
             exchange.sendResponseHeaders(200, body.size.toLong())
             exchange.responseBody.use { it.write(body) }
         }
 
         server.start()
         val port = server.address.port
-        client =
-            IamClient(
-                baseUrl = "http://localhost:$port",
-                bearerToken = "test-token",
-                logger = logger,
-            )
+        endpoint = IdentityEndpoint.forTesting("http://localhost:$port", "test-token")
     }
 
     @AfterClass
@@ -99,21 +84,20 @@ class IamClientTest {
 
     @Test
     fun getUserGroupsEmptyInputReturnsEmptyMap() {
-        val result = client.getUserGroups(emptyList())
+        val result = endpoint.getUserGroups(emptyList())
         assertTrue(result.isEmpty(), "Expected empty map for empty user ID list")
     }
 
     @Test
     fun getGroupAliasesEmptyInputReturnsEmptyMap() {
-        val result = client.getGroupAliases(emptyList())
+        val result = endpoint.getGroupAliases(emptyList())
         assertTrue(result.isEmpty(), "Expected empty map for empty group ID list")
     }
 
     @Test
     fun getUserGroupsReturnsMappedGroups() {
         val userIds = listOf("user-1", "user-2", "user-3")
-        val result = client.getUserGroups(userIds)
-
+        val result = endpoint.getUserGroups(userIds)
         assertEquals(result.size, 3)
         userIds.forEach { userId ->
             val groups = result[userId]
@@ -126,8 +110,7 @@ class IamClientTest {
     @Test
     fun getGroupAliasesReturnsMappedAliases() {
         val groupIds = listOf("grp-user-1", "grp-user-2")
-        val result = client.getGroupAliases(groupIds)
-
+        val result = endpoint.getGroupAliases(groupIds)
         assertEquals(result.size, 2)
         groupIds.forEach { groupId ->
             assertEquals(result[groupId], "Alias for $groupId")
@@ -136,35 +119,23 @@ class IamClientTest {
 
     @Test
     fun getGroupAliasesMissingAttributesReturnsEmptyString() {
-        val result = client.getGroupAliases(listOf("no-alias-grp"))
+        val result = endpoint.getGroupAliases(listOf("no-alias-grp"))
         assertEquals(result["no-alias-grp"], "", "Expected empty string when attributes are missing")
     }
 
     @Test
     fun getUserGroupsBatchesLargeInputs() {
-        // 550 users should produce 2 batches (500 + 50) with batchSize=500
-        val userIds = (1..550).map { "batch-user-$it" }
-        val requestCount =
-            java.util.concurrent.atomic
-                .AtomicInteger(0)
-
-        // Wrap client with a smaller batchSize to verify chunking without needing 550 ids
-        val smallBatchClient =
-            IamClient(
-                baseUrl = "http://localhost:${server.address.port}",
-                bearerToken = "test-token",
-                logger = logger,
-            )
-        val result = smallBatchClient.getUserGroups(userIds.take(7), batchSize = 3)
-        // 7 users, batchSize=3 → 3 batches: [3, 3, 1]
+        // 7 users with batchSize=3 → 3 batches: [3, 3, 1]
+        val userIds = (1..7).map { "user-$it" }
+        val result = endpoint.getUserGroups(userIds, 3)
         assertEquals(result.size, 7, "All 7 users should be present in result")
     }
 
     @Test
     fun getGroupAliasesBatchesLargeInputs() {
+        // 7 groups with batchSize=3 → 3 batches: [3, 3, 1]
         val groupIds = (1..7).map { "grp-user-$it" }
-        val result = client.getGroupAliases(groupIds, batchSize = 3)
-        // 7 groups, batchSize=3 → 3 batches: [3, 3, 1]
+        val result = endpoint.getGroupAliases(groupIds, 3)
         assertEquals(result.size, 7, "All 7 groups should be present in result")
     }
 }