Bound the catalog-role count taken when dropping a catalog

polaris-core/src/main/java/org/apache/polaris/core/persistence/AtomicOperationMetaStoreManager.java

@@ -1158,15 +1158,13 @@ public void deletePrincipalSecrets(
       }
 
       // get the list of catalog roles, at most 2
-      List<PolarisBaseEntity> catalogRoles =
-          ms.listFullEntities(
+      List<EntityNameLookupRecord> catalogRoles =
+          ms.listEntities(
                   callCtx,
                   catalogId,
                   catalogId,
                   PolarisEntityType.CATALOG_ROLE,
                   PolarisEntitySubType.ANY_SUBTYPE,
-                  entity -> true,
-                  Function.identity(),
                   PageToken.fromLimit(2))
               .items();
 
@@ -1178,7 +1176,17 @@ public void deletePrincipalSecrets(
       // if 1, drop the last catalog role. Should be the catalog admin role but don't validate this
       if (!catalogRoles.isEmpty()) {
         // drop the last catalog role in that catalog, should be the admin catalog role
-        this.dropEntity(callCtx, ms, catalogRoles.get(0));
+        EntityNameLookupRecord lastCatalogRole = catalogRoles.get(0);
+        PolarisBaseEntity catalogRoleToDrop =
+            ms.lookupEntity(
+                callCtx,
+                lastCatalogRole.getCatalogId(),
+                lastCatalogRole.getId(),
+                lastCatalogRole.getTypeCode());
+        // null means it was dropped concurrently, which leaves nothing to do
+        if (catalogRoleToDrop != null) {
+          this.dropEntity(callCtx, ms, catalogRoleToDrop);
+        }
       }
     } else if (refreshEntityToDrop.getType() == PolarisEntityType.NAMESPACE) {
       if (ms.hasChildren(

polaris-core/src/test/java/org/apache/polaris/core/persistence/PolarisTreeMapAtomicOperationMetaStoreManagerTest.java

@@ -19,14 +19,23 @@
 package org.apache.polaris.core.persistence;
 
 import static org.apache.polaris.core.persistence.PrincipalSecretsGenerator.RANDOM_SECRETS;
+import static org.assertj.core.api.Assertions.assertThat;
 
+import java.util.List;
+import java.util.Map;
 import org.apache.polaris.core.PolarisCallContext;
 import org.apache.polaris.core.PolarisDefaultDiagServiceImpl;
 import org.apache.polaris.core.PolarisDiagnostics;
+import org.apache.polaris.core.entity.PolarisBaseEntity;
+import org.apache.polaris.core.entity.PolarisEntityConstants;
+import org.apache.polaris.core.entity.PolarisEntitySubType;
+import org.apache.polaris.core.entity.PolarisEntityType;
+import org.apache.polaris.core.persistence.pagination.PageToken;
 import org.apache.polaris.core.persistence.transactional.TreeMapMetaStore;
 import org.apache.polaris.core.persistence.transactional.TreeMapTransactionalPersistenceImpl;
 import org.junit.jupiter.api.Disabled;
 import org.junit.jupiter.api.Test;
+import org.mockito.ArgumentCaptor;
 import org.mockito.Mockito;
 
 public class PolarisTreeMapAtomicOperationMetaStoreManagerTest
@@ -50,4 +59,57 @@ public PolarisTestMetaStoreManager createPolarisTestMetaStoreManager() {
       "AtomicOperationMetaStoreManager calls storePrincipalSecrets outside a transaction, which is incompatible with "
           + "TreeMap's transactional slice reads. Collision detection is covered by JDBC and NoSQL backend tests.")
   protected void testResetCredentialsClientIdCollision() {}
+
+  @Test
+  void dropCatalogListsAtMostTwoCatalogRoleIdentities() {
+    PolarisDiagnostics diagServices = new PolarisDefaultDiagServiceImpl();
+    TreeMapMetaStore store = new TreeMapMetaStore(diagServices);
+    TreeMapTransactionalPersistenceImpl spiedStore =
+        Mockito.spy(
+            new TreeMapTransactionalPersistenceImpl(
+                diagServices, store, Mockito.mock(), RANDOM_SECRETS));
+    PolarisTestMetaStoreManager tm =
+        new PolarisTestMetaStoreManager(
+            new AtomicOperationMetaStoreManager(clock, diagServices),
+            new PolarisCallContext(() -> "testRealm", spiedStore));
+
+    PolarisBaseEntity catalog =
+        new PolarisBaseEntity(
+            PolarisEntityConstants.getNullId(),
+            tm.polarisMetaStoreManager().generateNewEntityId(tm.polarisCallContext()).getId(),
+            PolarisEntityType.CATALOG,
+            PolarisEntitySubType.NULL_SUBTYPE,
+            PolarisEntityConstants.getRootEntityId(),
+            "counted");
+    catalog =
+        tm.polarisMetaStoreManager()
+            .createCatalog(tm.polarisCallContext(), catalog, List.of())
+            .getCatalog();
+    Mockito.clearInvocations(spiedStore);
+
+    tm.polarisMetaStoreManager()
+        .dropEntityIfExists(tm.polarisCallContext(), null, catalog, Map.of(), false);
+
+    Mockito.verify(spiedStore, Mockito.never())
+        .listFullEntities(
+            Mockito.any(),
+            Mockito.anyLong(),
+            Mockito.anyLong(),
+            Mockito.eq(PolarisEntityType.CATALOG_ROLE),
+            Mockito.any(),
+            Mockito.any(),
+            Mockito.any(),
+            Mockito.any());
+
+    ArgumentCaptor<PageToken> token = ArgumentCaptor.forClass(PageToken.class);
+    Mockito.verify(spiedStore)
+        .listEntities(
+            Mockito.any(),
+            Mockito.anyLong(),
+            Mockito.anyLong(),
+            Mockito.eq(PolarisEntityType.CATALOG_ROLE),
+            Mockito.any(),
+            token.capture());
+    assertThat(token.getValue().pageSize()).hasValue(2);
+  }
 }