Add findPrincipalById helper (#2810)

* Add findPrincipalById helper

this simplifies frequent usage of the lower level `loadEntity` api (similar to the
existing `findPrincipalByName` helper)

Author: XN137

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

@@ -844,14 +844,12 @@ private void revokeGrantRecord(
     BasePersistence ms = callCtx.getMetaStore();
 
     // if not found, the principal must have been dropped
-    EntityResult loadEntityResult =
-        loadEntity(
-            callCtx, PolarisEntityConstants.getNullId(), principalId, PolarisEntityType.PRINCIPAL);
-    if (loadEntityResult.getReturnStatus() != BaseResult.ReturnStatus.SUCCESS) {
+    Optional<PrincipalEntity> principalLookup = findPrincipalById(callCtx, principalId);
+    if (principalLookup.isEmpty()) {
       return new PrincipalSecretsResult(BaseResult.ReturnStatus.ENTITY_NOT_FOUND, null);
     }
 
-    PolarisBaseEntity principal = loadEntityResult.getEntity();
+    PrincipalEntity principal = principalLookup.get();
     Map<String, String> internalProps = principal.getInternalPropertiesAsMap();
 
     boolean doReset =
@@ -895,11 +893,10 @@ private void revokeGrantRecord(
       String customClientSecret) {
     // get metastore we should be using
     BasePersistence ms = callCtx.getMetaStore();
+
     // if not found, the principal must have been dropped
-    EntityResult loadEntityResult =
-        loadEntity(
-            callCtx, PolarisEntityConstants.getNullId(), principalId, PolarisEntityType.PRINCIPAL);
-    if (loadEntityResult.getReturnStatus() != BaseResult.ReturnStatus.SUCCESS) {
+    Optional<PrincipalEntity> principalEntity = findPrincipalById(callCtx, principalId);
+    if (principalEntity.isEmpty()) {
       return new PrincipalSecretsResult(BaseResult.ReturnStatus.ENTITY_NOT_FOUND, null);
     }
 

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

@@ -470,6 +470,20 @@ default Optional<PrincipalEntity> findRootPrincipal(PolarisCallContext polarisCa
     return findPrincipalByName(polarisCallContext, PolarisEntityConstants.getRootPrincipalName());
   }
 
+  default Optional<PrincipalEntity> findPrincipalById(
+      PolarisCallContext polarisCallContext, long principalId) {
+    EntityResult loadResult =
+        loadEntity(
+            polarisCallContext,
+            PolarisEntityConstants.getNullId(),
+            principalId,
+            PolarisEntityType.PRINCIPAL);
+    if (!loadResult.isSuccess()) {
+      return Optional.empty();
+    }
+    return Optional.of(loadResult.getEntity()).map(PrincipalEntity::of);
+  }
+
   default Optional<PrincipalEntity> findPrincipalByName(
       PolarisCallContext polarisCallContext, String principalName) {
     EntityResult entityResult =

polaris-core/src/testFixtures/java/org/apache/polaris/core/persistence/PolarisTestMetaStoreManager.java

@@ -404,7 +404,7 @@ void ensureGrantRecordRemoved(
   }
 
   /** Create a principal */
-  PolarisBaseEntity createPrincipal(String name) {
+  PrincipalEntity createPrincipal(String name) {
     // create new principal identity
     PrincipalEntity principalEntity =
         new PrincipalEntity.Builder()
@@ -490,14 +490,11 @@ PolarisBaseEntity createPrincipal(String name) {
             .getPrincipalSecrets();
     Assertions.assertThat(secrets.getMainSecret()).isNotEqualTo(reloadSecrets.getMainSecret());
 
-    PolarisBaseEntity reloadPrincipal =
+    PrincipalEntity reloadPrincipal =
         polarisMetaStoreManager
-            .loadEntity(
-                this.polarisCallContext,
-                0L,
-                createPrincipalResult.getPrincipal().getId(),
-                createPrincipalResult.getPrincipal().getType())
-            .getEntity();
+            .findPrincipalById(
+                this.polarisCallContext, createPrincipalResult.getPrincipal().getId())
+            .orElseThrow();
     internalProperties = reloadPrincipal.getInternalPropertiesAsMap();
     Assertions.assertThat(
             internalProperties.get(
@@ -549,11 +546,10 @@ PolarisBaseEntity createPrincipal(String name) {
     Assertions.assertThat(reloadSecrets.getMainSecretHash()).isNotEqualTo(newMainSecretHash);
     Assertions.assertThat(reloadSecrets.getSecondarySecretHash()).isNotEqualTo(newMainSecretHash);
 
-    PolarisBaseEntity newPrincipal =
+    PrincipalEntity newPrincipal =
         polarisMetaStoreManager
-            .loadEntity(
-                this.polarisCallContext, 0L, principalEntity.getId(), principalEntity.getType())
-            .getEntity();
+            .findPrincipalById(this.polarisCallContext, principalEntity.getId())
+            .orElseThrow();
     internalProperties = newPrincipal.getInternalPropertiesAsMap();
     Assertions.assertThat(
             internalProperties.get(
@@ -582,11 +578,10 @@ PolarisBaseEntity createPrincipal(String name) {
     Assertions.assertThat(postResetCredentials.getSecondarySecretHash())
         .isNotEqualTo(reloadSecrets.getSecondarySecretHash());
 
-    PolarisBaseEntity finalPrincipal =
+    PrincipalEntity finalPrincipal =
         polarisMetaStoreManager
-            .loadEntity(
-                this.polarisCallContext, 0L, principalEntity.getId(), principalEntity.getType())
-            .getEntity();
+            .findPrincipalById(this.polarisCallContext, principalEntity.getId())
+            .orElseThrow();
     internalProperties = finalPrincipal.getInternalPropertiesAsMap();
     Assertions.assertThat(
             internalProperties.get(
@@ -1342,8 +1337,8 @@ PolarisBaseEntity createTestCatalog(String catalogName) {
     grantToGrantee(catalog, R2, PR2, PolarisPrivilege.CATALOG_ROLE_USAGE);
 
     // also create two new principals
-    PolarisBaseEntity P1 = this.createPrincipal("P1");
-    PolarisBaseEntity P2 = this.createPrincipal("P2");
+    PrincipalEntity P1 = this.createPrincipal("P1");
+    PrincipalEntity P2 = this.createPrincipal("P2");
 
     // assign PR1 and PR2 to this principal
     grantToGrantee(null, PR1, P1, PolarisPrivilege.PRINCIPAL_ROLE_USAGE);

runtime/service/src/main/java/org/apache/polaris/service/admin/PolarisAdminService.java

@@ -1182,15 +1182,17 @@ public void deletePrincipal(String name) {
               "Failed to %s secrets for principal '%s'",
               shouldReset ? "reset" : "rotate", principalName));
     }
-    PolarisEntity newPrincipal =
-        PolarisEntity.of(
-            metaStoreManager.loadEntity(
-                getCurrentPolarisContext(),
-                0L,
-                currentPrincipalEntity.getId(),
-                currentPrincipalEntity.getType()));
+    Optional<PrincipalEntity> updatedPrincipalEntity =
+        metaStoreManager.findPrincipalById(
+            getCurrentPolarisContext(), currentPrincipalEntity.getId());
+    if (updatedPrincipalEntity.isEmpty()) {
+      throw new IllegalStateException(
+          String.format(
+              "Failed to reload principal '%s' by id: %s",
+              principalName, currentPrincipalEntity.getId()));
+    }
     return new PrincipalWithCredentials(
-        PrincipalEntity.of(newPrincipal).asPrincipal(),
+        updatedPrincipalEntity.get().asPrincipal(),
         new PrincipalWithCredentialsCredentials(
             newSecrets.getPrincipalClientId(), newSecrets.getMainSecret()));
   }

runtime/service/src/main/java/org/apache/polaris/service/auth/DefaultAuthenticator.java

@@ -115,14 +115,10 @@ protected PrincipalEntity resolvePrincipalEntity(PolarisCredential credentials)
       // otherwise, use the principal name to load the entity.
       if (credentials.getPrincipalId() != null && credentials.getPrincipalId() > 0) {
         principal =
-            PrincipalEntity.of(
-                metaStoreManager
-                    .loadEntity(
-                        callContext.getPolarisCallContext(),
-                        0L,
-                        credentials.getPrincipalId(),
-                        PolarisEntityType.PRINCIPAL)
-                    .getEntity());
+            metaStoreManager
+                .findPrincipalById(
+                    callContext.getPolarisCallContext(), credentials.getPrincipalId())
+                .orElse(null);
       } else if (credentials.getPrincipalName() != null) {
         principal =
             metaStoreManager

runtime/service/src/main/java/org/apache/polaris/service/auth/internal/broker/JWTBroker.java

@@ -28,10 +28,8 @@
 import java.util.UUID;
 import org.apache.iceberg.exceptions.NotAuthorizedException;
 import org.apache.polaris.core.PolarisCallContext;
-import org.apache.polaris.core.entity.PolarisEntityType;
 import org.apache.polaris.core.entity.PrincipalEntity;
 import org.apache.polaris.core.persistence.PolarisMetaStoreManager;
-import org.apache.polaris.core.persistence.dao.entity.EntityResult;
 import org.apache.polaris.core.persistence.dao.entity.PrincipalSecretsResult;
 import org.apache.polaris.service.auth.DefaultAuthenticator;
 import org.apache.polaris.service.auth.PolarisCredential;
@@ -106,11 +104,9 @@ public TokenResponse generateFromToken(
       LOGGER.error("Failed to verify the token", e.getCause());
       return TokenResponse.of(OAuthError.invalid_client);
     }
-    EntityResult principalLookup =
-        metaStoreManager.loadEntity(
-            polarisCallContext, 0L, decodedToken.getPrincipalId(), PolarisEntityType.PRINCIPAL);
-    if (!principalLookup.isSuccess()
-        || principalLookup.getEntity().getType() != PolarisEntityType.PRINCIPAL) {
+    Optional<PrincipalEntity> principalLookup =
+        metaStoreManager.findPrincipalById(polarisCallContext, decodedToken.getPrincipalId());
+    if (principalLookup.isEmpty()) {
       return TokenResponse.of(OAuthError.unauthorized_client);
     }
     String tokenString =
@@ -191,15 +187,7 @@ private Optional<PrincipalEntity> findPrincipalEntity(
     if (!principalSecrets.getPrincipalSecrets().matchesSecret(clientSecret)) {
       return Optional.empty();
     }
-    EntityResult result =
-        metaStoreManager.loadEntity(
-            polarisCallContext,
-            0L,
-            principalSecrets.getPrincipalSecrets().getPrincipalId(),
-            PolarisEntityType.PRINCIPAL);
-    if (!result.isSuccess() || result.getEntity().getType() != PolarisEntityType.PRINCIPAL) {
-      return Optional.empty();
-    }
-    return Optional.of(PrincipalEntity.of(result.getEntity()));
+    return metaStoreManager.findPrincipalById(
+        polarisCallContext, principalSecrets.getPrincipalSecrets().getPrincipalId());
   }
 }

runtime/service/src/test/java/org/apache/polaris/service/auth/internal/broker/JWTSymmetricKeyGeneratorTest.java

@@ -24,13 +24,11 @@
 import com.auth0.jwt.JWTVerifier;
 import com.auth0.jwt.algorithms.Algorithm;
 import com.auth0.jwt.interfaces.DecodedJWT;
+import java.util.Optional;
 import org.apache.polaris.core.PolarisCallContext;
-import org.apache.polaris.core.entity.PolarisBaseEntity;
-import org.apache.polaris.core.entity.PolarisEntitySubType;
-import org.apache.polaris.core.entity.PolarisEntityType;
 import org.apache.polaris.core.entity.PolarisPrincipalSecrets;
+import org.apache.polaris.core.entity.PrincipalEntity;
 import org.apache.polaris.core.persistence.PolarisMetaStoreManager;
-import org.apache.polaris.core.persistence.dao.entity.EntityResult;
 import org.apache.polaris.core.persistence.dao.entity.PrincipalSecretsResult;
 import org.apache.polaris.service.types.TokenType;
 import org.junit.jupiter.api.Test;
@@ -43,23 +41,17 @@ public class JWTSymmetricKeyGeneratorTest {
   public void testJWTSymmetricKeyGenerator() {
     PolarisCallContext polarisCallContext = new PolarisCallContext(null, null, null);
     PolarisMetaStoreManager metastoreManager = Mockito.mock(PolarisMetaStoreManager.class);
+    long principalId = 123L;
     String mainSecret = "test_secret";
     String clientId = "test_client_id";
     PolarisPrincipalSecrets principalSecrets =
-        new PolarisPrincipalSecrets(1L, clientId, mainSecret, "otherSecret");
+        new PolarisPrincipalSecrets(principalId, clientId, mainSecret, "otherSecret");
     Mockito.when(metastoreManager.loadPrincipalSecrets(polarisCallContext, clientId))
         .thenReturn(new PrincipalSecretsResult(principalSecrets));
-    PolarisBaseEntity principal =
-        new PolarisBaseEntity(
-            0L,
-            1L,
-            PolarisEntityType.PRINCIPAL,
-            PolarisEntitySubType.NULL_SUBTYPE,
-            0L,
-            "principal");
-    Mockito.when(
-            metastoreManager.loadEntity(polarisCallContext, 0L, 1L, PolarisEntityType.PRINCIPAL))
-        .thenReturn(new EntityResult(principal));
+    PrincipalEntity principal =
+        new PrincipalEntity.Builder().setId(principalId).setName("principal").build();
+    Mockito.when(metastoreManager.findPrincipalById(polarisCallContext, principalId))
+        .thenReturn(Optional.of(principal));
     TokenBroker generator = new SymmetricKeyJWTBroker(metastoreManager, 666, () -> "polaris");
     TokenResponse token =
         generator.generateFromClientSecrets(

runtime/service/src/test/java/org/apache/polaris/service/auth/internal/broker/RSAKeyPairJWTBrokerTest.java

@@ -28,14 +28,12 @@
 import jakarta.inject.Inject;
 import java.security.interfaces.RSAPrivateKey;
 import java.security.interfaces.RSAPublicKey;
+import java.util.Optional;
 import org.apache.polaris.core.PolarisCallContext;
 import org.apache.polaris.core.config.PolarisConfigurationStore;
-import org.apache.polaris.core.entity.PolarisBaseEntity;
-import org.apache.polaris.core.entity.PolarisEntitySubType;
-import org.apache.polaris.core.entity.PolarisEntityType;
 import org.apache.polaris.core.entity.PolarisPrincipalSecrets;
+import org.apache.polaris.core.entity.PrincipalEntity;
 import org.apache.polaris.core.persistence.PolarisMetaStoreManager;
-import org.apache.polaris.core.persistence.dao.entity.EntityResult;
 import org.apache.polaris.core.persistence.dao.entity.PrincipalSecretsResult;
 import org.apache.polaris.service.types.TokenType;
 import org.junit.jupiter.api.Test;
@@ -50,27 +48,21 @@ public class RSAKeyPairJWTBrokerTest {
   public void testSuccessfulTokenGeneration() throws Exception {
     var keyPair = PemUtils.generateKeyPair();
 
+    final long principalId = 123L;
     final String clientId = "test-client-id";
     final String scope = "PRINCIPAL_ROLE:TEST";
 
     PolarisCallContext polarisCallContext = new PolarisCallContext(null, null, configurationStore);
     PolarisMetaStoreManager metastoreManager = Mockito.mock(PolarisMetaStoreManager.class);
     String mainSecret = "client-secret";
     PolarisPrincipalSecrets principalSecrets =
-        new PolarisPrincipalSecrets(1L, clientId, mainSecret, "otherSecret");
+        new PolarisPrincipalSecrets(principalId, clientId, mainSecret, "otherSecret");
     Mockito.when(metastoreManager.loadPrincipalSecrets(polarisCallContext, clientId))
         .thenReturn(new PrincipalSecretsResult(principalSecrets));
-    PolarisBaseEntity principal =
-        new PolarisBaseEntity(
-            0L,
-            1L,
-            PolarisEntityType.PRINCIPAL,
-            PolarisEntitySubType.NULL_SUBTYPE,
-            0L,
-            "principal");
-    Mockito.when(
-            metastoreManager.loadEntity(polarisCallContext, 0L, 1L, PolarisEntityType.PRINCIPAL))
-        .thenReturn(new EntityResult(principal));
+    PrincipalEntity principal =
+        new PrincipalEntity.Builder().setId(principalId).setName("principal").build();
+    Mockito.when(metastoreManager.findPrincipalById(polarisCallContext, principalId))
+        .thenReturn(Optional.of(principal));
     KeyProvider provider = new LocalRSAKeyProvider(keyPair);
     TokenBroker tokenBroker = new RSAKeyPairJWTBroker(metastoreManager, 420, provider);
     TokenResponse token =