Skip to content

Add findPrincipalById helper #2810

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Oct 17, 2025
Merged

Add findPrincipalById helper #2810

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
c790f1b
Add findPrincipalById helper
XN137 Oct 14, 2025
a6710d4
review: improve error message in rotateOrResetCredentialsHelper
XN137 Oct 17, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 6 additions & 9 deletions ...re/src/main/java/org/apache/polaris/core/persistence/AtomicOperationMetaStoreManager.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 =
Expand Down Expand Up @@ -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);
}

Expand Down
14 changes: 14 additions & 0 deletions polaris-core/src/main/java/org/apache/polaris/core/persistence/PolarisMetaStoreManager.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -470,6 +470,20 @@ default Optional<PrincipalEntity> findRootPrincipal(PolarisCallContext polarisCa
return findPrincipalByName(polarisCallContext, PolarisEntityConstants.getRootPrincipalName());
}

default Optional<PrincipalEntity> findPrincipalById(
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Having this helper method in principle SGTM. However, given that PolarisMetaStoreManager is one of the main plugin / integrations points, I'm not sure about expanding its API surface with one more method.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

why do you think its expanding the API surface?
its implementation is just reusing existing methods but encodes more information about the proper way to use them and about the type of entity they will return.
also simplifying code around "exception" handling.

these defaults methods probably never need to be re-implemented by actual classes.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It gives callers two options: 1) use the old method 2) use the helper methods... In my mind that expands the API.

I'm thinking that adding methods like this is probably an indication that we may need to rework the interface as a whole.

This is more of a point for thinking about, not an objection, really.

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 =
Expand Down
31 changes: 13 additions & 18 deletions ...rc/testFixtures/java/org/apache/polaris/core/persistence/PolarisTestMetaStoreManager.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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);
Expand Down
18 changes: 10 additions & 8 deletions runtime/service/src/main/java/org/apache/polaris/service/admin/PolarisAdminService.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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()));
}
Expand Down
12 changes: 4 additions & 8 deletions runtime/service/src/main/java/org/apache/polaris/service/auth/DefaultAuthenticator.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
22 changes: 5 additions & 17 deletions runtime/service/src/main/java/org/apache/polaris/service/auth/internal/broker/JWTBroker.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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 =
Expand Down Expand Up @@ -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());
}
}
24 changes: 8 additions & 16 deletions ...st/java/org/apache/polaris/service/auth/internal/broker/JWTSymmetricKeyGeneratorTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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(
Expand Down
24 changes: 8 additions & 16 deletions ...rc/test/java/org/apache/polaris/service/auth/internal/broker/RSAKeyPairJWTBrokerTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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 =
Expand Down