diff --git a/polaris-core/src/main/java/org/apache/polaris/core/persistence/AtomicOperationMetaStoreManager.java b/polaris-core/src/main/java/org/apache/polaris/core/persistence/AtomicOperationMetaStoreManager.java index e2c46c1515..c3841486a7 100644 --- a/polaris-core/src/main/java/org/apache/polaris/core/persistence/AtomicOperationMetaStoreManager.java +++ b/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 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 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 = findPrincipalById(callCtx, principalId); + if (principalEntity.isEmpty()) { return new PrincipalSecretsResult(BaseResult.ReturnStatus.ENTITY_NOT_FOUND, null); } diff --git a/polaris-core/src/main/java/org/apache/polaris/core/persistence/PolarisMetaStoreManager.java b/polaris-core/src/main/java/org/apache/polaris/core/persistence/PolarisMetaStoreManager.java index 566b10e644..cf3912fa95 100644 --- a/polaris-core/src/main/java/org/apache/polaris/core/persistence/PolarisMetaStoreManager.java +++ b/polaris-core/src/main/java/org/apache/polaris/core/persistence/PolarisMetaStoreManager.java @@ -470,6 +470,20 @@ default Optional findRootPrincipal(PolarisCallContext polarisCa return findPrincipalByName(polarisCallContext, PolarisEntityConstants.getRootPrincipalName()); } + default Optional 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 findPrincipalByName( PolarisCallContext polarisCallContext, String principalName) { EntityResult entityResult = diff --git a/polaris-core/src/testFixtures/java/org/apache/polaris/core/persistence/PolarisTestMetaStoreManager.java b/polaris-core/src/testFixtures/java/org/apache/polaris/core/persistence/PolarisTestMetaStoreManager.java index 514a14c30f..c8e486ee6e 100644 --- a/polaris-core/src/testFixtures/java/org/apache/polaris/core/persistence/PolarisTestMetaStoreManager.java +++ b/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); diff --git a/runtime/service/src/main/java/org/apache/polaris/service/admin/PolarisAdminService.java b/runtime/service/src/main/java/org/apache/polaris/service/admin/PolarisAdminService.java index 9c9074dba7..67ae7ad6e5 100644 --- a/runtime/service/src/main/java/org/apache/polaris/service/admin/PolarisAdminService.java +++ b/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 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())); } diff --git a/runtime/service/src/main/java/org/apache/polaris/service/auth/DefaultAuthenticator.java b/runtime/service/src/main/java/org/apache/polaris/service/auth/DefaultAuthenticator.java index 9abc75e38f..83a3a419d7 100644 --- a/runtime/service/src/main/java/org/apache/polaris/service/auth/DefaultAuthenticator.java +++ b/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 diff --git a/runtime/service/src/main/java/org/apache/polaris/service/auth/internal/broker/JWTBroker.java b/runtime/service/src/main/java/org/apache/polaris/service/auth/internal/broker/JWTBroker.java index 1cb13f5aac..71ea0d0548 100644 --- a/runtime/service/src/main/java/org/apache/polaris/service/auth/internal/broker/JWTBroker.java +++ b/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 principalLookup = + metaStoreManager.findPrincipalById(polarisCallContext, decodedToken.getPrincipalId()); + if (principalLookup.isEmpty()) { return TokenResponse.of(OAuthError.unauthorized_client); } String tokenString = @@ -191,15 +187,7 @@ private Optional 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()); } } diff --git a/runtime/service/src/test/java/org/apache/polaris/service/auth/internal/broker/JWTSymmetricKeyGeneratorTest.java b/runtime/service/src/test/java/org/apache/polaris/service/auth/internal/broker/JWTSymmetricKeyGeneratorTest.java index e3ff0518e3..651fc1a9d9 100644 --- a/runtime/service/src/test/java/org/apache/polaris/service/auth/internal/broker/JWTSymmetricKeyGeneratorTest.java +++ b/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( diff --git a/runtime/service/src/test/java/org/apache/polaris/service/auth/internal/broker/RSAKeyPairJWTBrokerTest.java b/runtime/service/src/test/java/org/apache/polaris/service/auth/internal/broker/RSAKeyPairJWTBrokerTest.java index 3bc4ff6a66..13bd7f3df2 100644 --- a/runtime/service/src/test/java/org/apache/polaris/service/auth/internal/broker/RSAKeyPairJWTBrokerTest.java +++ b/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,6 +48,7 @@ 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"; @@ -57,20 +56,13 @@ public void testSuccessfulTokenGeneration() throws Exception { 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 =