From 7bf846287144b7090b4c98e7e6a7293a87b6f3b6 Mon Sep 17 00:00:00 2001 From: Dmitri Bourlatchkov Date: Tue, 4 Mar 2025 19:47:54 -0500 Subject: [PATCH 1/2] Added typeCode to BasePersistence.lookupEntity This is mainly needed for the upcoming NoSQL implementation. Also, this change harmonizes the API parameters across all lookup methods. This change induces changes in the MetaStore Manager API and related classes. In most top-level call sites the required extra information is already available. A minor inefficiency exists in PolarisAdminService, but it will only manifest in backends that do rely on the new parameter to perform lookups, plus the affected use case is not expected to be frequent (Admin API queries to obtain existing grants for a catalog role). --- ...olarisEclipseLinkMetaStoreSessionImpl.java | 2 +- .../core/persistence/BasePersistence.java | 3 +- .../persistence/PolarisMetaStoreManager.java | 11 ++- .../TransactionWorkspaceMetaStoreManager.java | 13 ++- .../core/persistence/cache/EntityCache.java | 10 +- .../core/persistence/resolver/Resolver.java | 4 +- .../AbstractTransactionalPersistence.java | 6 +- .../PolarisMetaStoreManagerImpl.java | 78 ++++++++++----- .../PolarisTreeMapMetaStoreSessionImpl.java | 2 +- .../core/storage/PolarisCredentialVendor.java | 3 + .../storage/cache/StorageCredentialCache.java | 1 + .../core/persistence/EntityCacheTest.java | 12 ++- .../core/persistence/ResolverTest.java | 2 +- .../cache/StorageCredentialCacheTest.java | 5 + .../PolarisTestMetaStoreManager.java | 95 +++++++++++++------ .../quarkus/admin/PolarisAuthzTestBase.java | 5 +- .../quarkus/auth/JWTRSAKeyPairTest.java | 3 +- .../auth/JWTSymmetricKeyGeneratorTest.java | 3 +- .../catalog/BasePolarisCatalogTest.java | 1 + .../service/admin/PolarisAdminService.java | 69 +++++++++++--- .../auth/BasePolarisAuthenticator.java | 5 +- .../auth/DefaultActiveRolesProvider.java | 6 +- .../polaris/service/auth/JWTBroker.java | 3 +- .../service/auth/TestOAuth2ApiService.java | 5 +- .../polaris/service/auth/TokenBroker.java | 5 +- .../service/task/TaskExecutorImpl.java | 4 +- .../auth/BasePolarisAuthenticatorTest.java | 7 +- 27 files changed, 267 insertions(+), 96 deletions(-) diff --git a/extension/persistence/eclipselink/src/main/java/org/apache/polaris/extension/persistence/impl/eclipselink/PolarisEclipseLinkMetaStoreSessionImpl.java b/extension/persistence/eclipselink/src/main/java/org/apache/polaris/extension/persistence/impl/eclipselink/PolarisEclipseLinkMetaStoreSessionImpl.java index 524d4701d3..5637a93279 100644 --- a/extension/persistence/eclipselink/src/main/java/org/apache/polaris/extension/persistence/impl/eclipselink/PolarisEclipseLinkMetaStoreSessionImpl.java +++ b/extension/persistence/eclipselink/src/main/java/org/apache/polaris/extension/persistence/impl/eclipselink/PolarisEclipseLinkMetaStoreSessionImpl.java @@ -351,7 +351,7 @@ public void deleteAll(@Nonnull PolarisCallContext callCtx) { /** {@inheritDoc} */ @Override public @Nullable PolarisBaseEntity lookupEntity( - @Nonnull PolarisCallContext callCtx, long catalogId, long entityId) { + @Nonnull PolarisCallContext callCtx, long catalogId, long entityId, int typeCode) { return ModelEntity.toEntity(this.store.lookupEntity(localSession.get(), catalogId, entityId)); } diff --git a/polaris-core/src/main/java/org/apache/polaris/core/persistence/BasePersistence.java b/polaris-core/src/main/java/org/apache/polaris/core/persistence/BasePersistence.java index b92a25af92..f5448cb65c 100644 --- a/polaris-core/src/main/java/org/apache/polaris/core/persistence/BasePersistence.java +++ b/polaris-core/src/main/java/org/apache/polaris/core/persistence/BasePersistence.java @@ -153,11 +153,12 @@ void deleteAllEntityGrantRecords( * @param callCtx call context * @param catalogId catalog id or NULL_ID * @param entityId entity id + * @param typeCode the PolarisEntityType code of the entity to lookup * @return null if the entity was not found, else the retrieved entity. */ @Nullable PolarisBaseEntity lookupEntity( - @Nonnull PolarisCallContext callCtx, long catalogId, long entityId); + @Nonnull PolarisCallContext callCtx, long catalogId, long entityId, int typeCode); /** * Lookup an entity given its catalogId, parentId, typeCode, and name. 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 0825c8f49d..ffaea3f127 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 @@ -285,7 +285,11 @@ DropEntityResult dropEntityIfExists( * @param entityId the id of the entity to load */ @Nonnull - EntityResult loadEntity(@Nonnull PolarisCallContext callCtx, long entityCatalogId, long entityId); + EntityResult loadEntity( + @Nonnull PolarisCallContext callCtx, + long entityCatalogId, + long entityId, + @Nonnull PolarisEntityType entityType); /** * Fetch a list of tasks to be completed. Tasks @@ -327,7 +331,10 @@ ChangeTrackingResult loadEntitiesChangeTracking( */ @Nonnull ResolvedEntityResult loadResolvedEntityById( - @Nonnull PolarisCallContext callCtx, long entityCatalogId, long entityId); + @Nonnull PolarisCallContext callCtx, + long entityCatalogId, + long entityId, + PolarisEntityType entityType); /** * Load a resolved entity, i.e. an entity definition and associated grant records, from the diff --git a/polaris-core/src/main/java/org/apache/polaris/core/persistence/TransactionWorkspaceMetaStoreManager.java b/polaris-core/src/main/java/org/apache/polaris/core/persistence/TransactionWorkspaceMetaStoreManager.java index 4a389f1c44..2691023a35 100644 --- a/polaris-core/src/main/java/org/apache/polaris/core/persistence/TransactionWorkspaceMetaStoreManager.java +++ b/polaris-core/src/main/java/org/apache/polaris/core/persistence/TransactionWorkspaceMetaStoreManager.java @@ -303,7 +303,10 @@ public ChangeTrackingResult loadEntitiesChangeTracking( @Override public EntityResult loadEntity( - @Nonnull PolarisCallContext callCtx, long entityCatalogId, long entityId) { + @Nonnull PolarisCallContext callCtx, + long entityCatalogId, + long entityId, + PolarisEntityType entityType) { callCtx.getDiagServices().fail("illegal_method_in_transaction_workspace", "loadEntity"); return null; } @@ -320,6 +323,7 @@ public ScopedCredentialsResult getSubscopedCredsForEntity( @Nonnull PolarisCallContext callCtx, long catalogId, long entityId, + PolarisEntityType entityType, boolean allowListOperation, @Nonnull Set allowedReadLocations, @Nonnull Set allowedWriteLocations) { @@ -327,6 +331,7 @@ public ScopedCredentialsResult getSubscopedCredsForEntity( callCtx, catalogId, entityId, + entityType, allowListOperation, allowedReadLocations, allowedWriteLocations); @@ -337,6 +342,7 @@ public ValidateAccessResult validateAccessToLocations( @Nonnull PolarisCallContext callCtx, long catalogId, long entityId, + PolarisEntityType entityType, @Nonnull Set actions, @Nonnull Set locations) { callCtx @@ -347,7 +353,10 @@ public ValidateAccessResult validateAccessToLocations( @Override public ResolvedEntityResult loadResolvedEntityById( - @Nonnull PolarisCallContext callCtx, long entityCatalogId, long entityId) { + @Nonnull PolarisCallContext callCtx, + long entityCatalogId, + long entityId, + PolarisEntityType entityType) { callCtx .getDiagServices() .fail("illegal_method_in_transaction_workspace", "loadResolvedEntityById"); diff --git a/polaris-core/src/main/java/org/apache/polaris/core/persistence/cache/EntityCache.java b/polaris-core/src/main/java/org/apache/polaris/core/persistence/cache/EntityCache.java index 7ca5de6213..2f2476a6bc 100644 --- a/polaris-core/src/main/java/org/apache/polaris/core/persistence/cache/EntityCache.java +++ b/polaris-core/src/main/java/org/apache/polaris/core/persistence/cache/EntityCache.java @@ -292,7 +292,7 @@ && isNewer(existingCacheEntry, existingCacheEntryByName)) { // try to load it refreshedCacheEntry = this.polarisMetaStoreManager.loadResolvedEntityById( - callContext, entityCatalogId, entityId); + callContext, entityCatalogId, entityId, entityType); if (refreshedCacheEntry.isSuccess()) { entity = refreshedCacheEntry.getEntity(); grantRecords = refreshedCacheEntry.getEntityGrantRecords(); @@ -361,7 +361,10 @@ && isNewer(existingCacheEntry, existingCacheEntryByName)) { * entity, either as found in the cache or loaded from the backend */ public @Nullable EntityCacheLookupResult getOrLoadEntityById( - @Nonnull PolarisCallContext callContext, long entityCatalogId, long entityId) { + @Nonnull PolarisCallContext callContext, + long entityCatalogId, + long entityId, + PolarisEntityType entityType) { // if it exists, we are set ResolvedPolarisEntity entry = this.getEntityById(entityId); @@ -374,7 +377,8 @@ && isNewer(existingCacheEntry, existingCacheEntryByName)) { // load it ResolvedEntityResult result = - polarisMetaStoreManager.loadResolvedEntityById(callContext, entityCatalogId, entityId); + polarisMetaStoreManager.loadResolvedEntityById( + callContext, entityCatalogId, entityId, entityType); // not found, exit if (!result.isSuccess()) { diff --git a/polaris-core/src/main/java/org/apache/polaris/core/persistence/resolver/Resolver.java b/polaris-core/src/main/java/org/apache/polaris/core/persistence/resolver/Resolver.java index b90ca7c3e2..bb5837297c 100644 --- a/polaris-core/src/main/java/org/apache/polaris/core/persistence/resolver/Resolver.java +++ b/polaris-core/src/main/java/org/apache/polaris/core/persistence/resolver/Resolver.java @@ -1027,7 +1027,7 @@ private ResolvedPolarisEntity resolveById( if (this.cache != null) { // get or load by name EntityCacheLookupResult lookupResult = - this.cache.getOrLoadEntityById(this.polarisCallContext, catalogId, entityId); + this.cache.getOrLoadEntityById(this.polarisCallContext, catalogId, entityId, entityType); // if not found, return null if (lookupResult == null) { @@ -1049,7 +1049,7 @@ private ResolvedPolarisEntity resolveById( // If no cache, load directly from metastore manager. ResolvedEntityResult result = polarisMetaStoreManager.loadResolvedEntityById( - this.polarisCallContext, catalogId, entityId); + this.polarisCallContext, catalogId, entityId, entityType); if (!result.isSuccess()) { // not found return null; diff --git a/polaris-core/src/main/java/org/apache/polaris/core/persistence/transactional/AbstractTransactionalPersistence.java b/polaris-core/src/main/java/org/apache/polaris/core/persistence/transactional/AbstractTransactionalPersistence.java index 4393e23455..6ed613ef3f 100644 --- a/polaris-core/src/main/java/org/apache/polaris/core/persistence/transactional/AbstractTransactionalPersistence.java +++ b/polaris-core/src/main/java/org/apache/polaris/core/persistence/transactional/AbstractTransactionalPersistence.java @@ -68,7 +68,11 @@ public PolarisBaseEntity lookupEntityByName( // lookup the entity, should be there PolarisBaseEntity entity = - lookupEntity(callCtx, entityActiveRecord.getCatalogId(), entityActiveRecord.getId()); + lookupEntity( + callCtx, + entityActiveRecord.getCatalogId(), + entityActiveRecord.getId(), + entityActiveRecord.getTypeCode()); callCtx .getDiagServices() .checkNotNull( diff --git a/polaris-core/src/main/java/org/apache/polaris/core/persistence/transactional/PolarisMetaStoreManagerImpl.java b/polaris-core/src/main/java/org/apache/polaris/core/persistence/transactional/PolarisMetaStoreManagerImpl.java index 30aa1d9392..5736af3c40 100644 --- a/polaris-core/src/main/java/org/apache/polaris/core/persistence/transactional/PolarisMetaStoreManagerImpl.java +++ b/polaris-core/src/main/java/org/apache/polaris/core/persistence/transactional/PolarisMetaStoreManagerImpl.java @@ -328,7 +328,7 @@ private void dropEntity( // load the grantee (either a catalog/principal role or a principal) and increment its grants // version PolarisBaseEntity granteeEntity = - ms.lookupEntity(callCtx, grantee.getCatalogId(), grantee.getId()); + ms.lookupEntity(callCtx, grantee.getCatalogId(), grantee.getId(), grantee.getTypeCode()); callCtx .getDiagServices() .checkNotNull(granteeEntity, "grantee_not_found", "grantee={}", grantee); @@ -341,7 +341,8 @@ private void dropEntity( // we also need to invalidate the grants on that securable so that we can reload them. // load the securable and increment its grants version PolarisBaseEntity securableEntity = - ms.lookupEntity(callCtx, securable.getCatalogId(), securable.getId()); + ms.lookupEntity( + callCtx, securable.getCatalogId(), securable.getId(), securable.getTypeCode()); callCtx .getDiagServices() .checkNotNull(securableEntity, "securable_not_found", "securable={}", securable); @@ -404,7 +405,7 @@ private void revokeGrantRecord( // load the grantee and increment its grants version PolarisBaseEntity refreshGrantee = - ms.lookupEntity(callCtx, grantee.getCatalogId(), grantee.getId()); + ms.lookupEntity(callCtx, grantee.getCatalogId(), grantee.getId(), grantee.getTypeCode()); callCtx .getDiagServices() .checkNotNull( @@ -418,7 +419,8 @@ private void revokeGrantRecord( // we also need to invalidate the grants on that securable so that we can reload them. // load the securable and increment its grants version PolarisBaseEntity refreshSecurable = - ms.lookupEntity(callCtx, securable.getCatalogId(), securable.getId()); + ms.lookupEntity( + callCtx, securable.getCatalogId(), securable.getId(), securable.getTypeCode()); callCtx .getDiagServices() .checkNotNull( @@ -460,7 +462,7 @@ private void revokeGrantRecord( // check if that catalog has already been created PolarisBaseEntity refreshCatalog = - ms.lookupEntity(callCtx, catalog.getCatalogId(), catalog.getId()); + ms.lookupEntity(callCtx, catalog.getCatalogId(), catalog.getId(), catalog.getTypeCode()); // if found, probably a retry, simply return the previously created catalog if (refreshCatalog != null) { @@ -822,7 +824,8 @@ public Map deserializeProperties(PolarisCallContext callCtx, Str // check if that catalog has already been created PolarisBaseEntity refreshPrincipal = - ms.lookupEntity(callCtx, principal.getCatalogId(), principal.getId()); + ms.lookupEntity( + callCtx, principal.getCatalogId(), principal.getId(), principal.getTypeCode()); // if found, probably a retry, simply return the previously created principal if (refreshPrincipal != null) { @@ -950,7 +953,12 @@ public Map deserializeProperties(PolarisCallContext callCtx, Str @Nonnull String oldSecretHash) { // if not found, the principal must have been dropped EntityResult loadEntityResult = - loadEntity(callCtx, ms, PolarisEntityConstants.getNullId(), principalId); + loadEntity( + callCtx, + ms, + PolarisEntityConstants.getNullId(), + principalId, + PolarisEntityType.PRINCIPAL.getCode()); if (loadEntityResult.getReturnStatus() != BaseResult.ReturnStatus.SUCCESS) { return null; } @@ -1061,7 +1069,8 @@ public Map deserializeProperties(PolarisCallContext callCtx, Str callCtx.getDiagServices().checkNotNull(entity.getName(), "unexpected_null_entity_name"); // first, check if the entity has already been created, in which case we will simply return it - PolarisBaseEntity entityFound = ms.lookupEntity(callCtx, entity.getCatalogId(), entity.getId()); + PolarisBaseEntity entityFound = + ms.lookupEntity(callCtx, entity.getCatalogId(), entity.getId(), entity.getTypeCode()); if (entityFound != null) { // probably the client retried, simply return it return new EntityResult(entityFound); @@ -1158,7 +1167,7 @@ public Map deserializeProperties(PolarisCallContext callCtx, Str // lookup the entity, cannot be null PolarisBaseEntity entityRefreshed = - ms.lookupEntity(callCtx, entity.getCatalogId(), entity.getId()); + ms.lookupEntity(callCtx, entity.getCatalogId(), entity.getId(), entity.getTypeCode()); callCtx .getDiagServices() .checkNotNull(entityRefreshed, "unexpected_entity_not_found", "entity={}", entity); @@ -1280,7 +1289,11 @@ public Map deserializeProperties(PolarisCallContext callCtx, Str // find the entity to rename PolarisBaseEntity refreshEntityToRename = - ms.lookupEntity(callCtx, entityToRename.getCatalogId(), entityToRename.getId()); + ms.lookupEntity( + callCtx, + entityToRename.getCatalogId(), + entityToRename.getId(), + entityToRename.getTypeCode()); // if this entity was not found, return failure. Not expected here because it was // resolved successfully (see above) @@ -1389,7 +1402,8 @@ public Map deserializeProperties(PolarisCallContext callCtx, Str // first find the entity to drop PolarisBaseEntity refreshEntityToDrop = - ms.lookupEntity(callCtx, entityToDrop.getCatalogId(), entityToDrop.getId()); + ms.lookupEntity( + callCtx, entityToDrop.getCatalogId(), entityToDrop.getId(), entityToDrop.getTypeCode()); // if this entity was not found, return failure if (refreshEntityToDrop == null) { @@ -1921,14 +1935,15 @@ private PolarisEntityResolver resolveSecurableToRoleGrant( callCtx, () -> this.loadEntitiesChangeTracking(callCtx, ms, entityIds)); } - /** Refer to {@link #loadEntity(PolarisCallContext, long, long)} */ + /** Refer to {@link #loadEntity(PolarisCallContext, long, long, PolarisEntityType)} */ private @Nonnull EntityResult loadEntity( @Nonnull PolarisCallContext callCtx, @Nonnull TransactionalPersistence ms, long entityCatalogId, - long entityId) { + long entityId, + int entityTypeCode) { // this is an easy one - PolarisBaseEntity entity = ms.lookupEntity(callCtx, entityCatalogId, entityId); + PolarisBaseEntity entity = ms.lookupEntity(callCtx, entityCatalogId, entityId, entityTypeCode); return (entity != null) ? new EntityResult(entity) : new EntityResult(BaseResult.ReturnStatus.ENTITY_NOT_FOUND, null); @@ -1937,13 +1952,17 @@ private PolarisEntityResolver resolveSecurableToRoleGrant( /** {@inheritDoc} */ @Override public @Nonnull EntityResult loadEntity( - @Nonnull PolarisCallContext callCtx, long entityCatalogId, long entityId) { + @Nonnull PolarisCallContext callCtx, + long entityCatalogId, + long entityId, + @Nonnull PolarisEntityType entityType) { // get metastore we should be using TransactionalPersistence ms = callCtx.getMetaStore(); // need to run inside a read transaction return ms.runInReadTransaction( - callCtx, () -> this.loadEntity(callCtx, ms, entityCatalogId, entityId)); + callCtx, + () -> this.loadEntity(callCtx, ms, entityCatalogId, entityId, entityType.getCode())); } /** Refer to {@link #loadTasks(PolarisCallContext, String, int)} */ @@ -2011,6 +2030,7 @@ private PolarisEntityResolver resolveSecurableToRoleGrant( @Nonnull PolarisCallContext callCtx, long catalogId, long entityId, + PolarisEntityType entityType, boolean allowListOperation, @Nonnull Set allowedReadLocations, @Nonnull Set allowedWriteLocations) { @@ -2024,7 +2044,7 @@ private PolarisEntityResolver resolveSecurableToRoleGrant( "allowed_locations_to_subscope_is_required"); // reload the entity, error out if not found - EntityResult reloadedEntity = loadEntity(callCtx, catalogId, entityId); + EntityResult reloadedEntity = loadEntity(callCtx, catalogId, entityId, entityType); if (reloadedEntity.getReturnStatus() != BaseResult.ReturnStatus.SUCCESS) { return new ScopedCredentialsResult( reloadedEntity.getReturnStatus(), reloadedEntity.getExtraInformation()); @@ -2067,6 +2087,7 @@ private PolarisEntityResolver resolveSecurableToRoleGrant( @Nonnull PolarisCallContext callCtx, long catalogId, long entityId, + PolarisEntityType entityType, @Nonnull Set actions, @Nonnull Set locations) { // get meta store we should be using @@ -2077,7 +2098,7 @@ private PolarisEntityResolver resolveSecurableToRoleGrant( !actions.isEmpty() && !locations.isEmpty(), "locations_and_operations_privileges_are_required"); // reload the entity, error out if not found - EntityResult reloadedEntity = loadEntity(callCtx, catalogId, entityId); + EntityResult reloadedEntity = loadEntity(callCtx, catalogId, entityId, entityType); if (reloadedEntity.getReturnStatus() != BaseResult.ReturnStatus.SUCCESS) { return new ValidateAccessResult( reloadedEntity.getReturnStatus(), reloadedEntity.getExtraInformation()); @@ -2129,15 +2150,16 @@ public Map getInternalPropertyMap( return deserializeProperties(callCtx, internalPropStr); } - /** {@link #loadResolvedEntityById(PolarisCallContext, long, long)} */ + /** {@link #loadResolvedEntityById(PolarisCallContext, long, long, PolarisEntityType)} */ private @Nonnull ResolvedEntityResult loadResolvedEntityById( @Nonnull PolarisCallContext callCtx, @Nonnull TransactionalPersistence ms, long entityCatalogId, - long entityId) { + long entityId, + int typeCode) { // load that entity - PolarisBaseEntity entity = ms.lookupEntity(callCtx, entityCatalogId, entityId); + PolarisBaseEntity entity = ms.lookupEntity(callCtx, entityCatalogId, entityId, typeCode); // if entity not found, return null if (entity == null) { @@ -2161,16 +2183,22 @@ public Map getInternalPropertyMap( /** {@inheritDoc} */ @Override public @Nonnull ResolvedEntityResult loadResolvedEntityById( - @Nonnull PolarisCallContext callCtx, long entityCatalogId, long entityId) { + @Nonnull PolarisCallContext callCtx, + long entityCatalogId, + long entityId, + PolarisEntityType entityType) { // get metastore we should be using TransactionalPersistence ms = callCtx.getMetaStore(); // need to run inside a read transaction return ms.runInReadTransaction( - callCtx, () -> this.loadResolvedEntityById(callCtx, ms, entityCatalogId, entityId)); + callCtx, + () -> + this.loadResolvedEntityById( + callCtx, ms, entityCatalogId, entityId, entityType.getCode())); } - /** {@link #loadResolvedEntityById(PolarisCallContext, long, long)} */ + /** {@link #loadResolvedEntityById(PolarisCallContext, long, long, PolarisEntityType)} */ private @Nonnull ResolvedEntityResult loadResolvedEntityByName( @Nonnull PolarisCallContext callCtx, @Nonnull TransactionalPersistence ms, @@ -2292,7 +2320,7 @@ public Map getInternalPropertyMap( // load the entity if something changed final PolarisBaseEntity entity; if (entityVersion != entityVersions.getEntityVersion()) { - entity = ms.lookupEntity(callCtx, entityCatalogId, entityId); + entity = ms.lookupEntity(callCtx, entityCatalogId, entityId, entityType.getCode()); // if not found, return null if (entity == null) { diff --git a/polaris-core/src/main/java/org/apache/polaris/core/persistence/transactional/PolarisTreeMapMetaStoreSessionImpl.java b/polaris-core/src/main/java/org/apache/polaris/core/persistence/transactional/PolarisTreeMapMetaStoreSessionImpl.java index dbd2ce8340..9894b82e69 100644 --- a/polaris-core/src/main/java/org/apache/polaris/core/persistence/transactional/PolarisTreeMapMetaStoreSessionImpl.java +++ b/polaris-core/src/main/java/org/apache/polaris/core/persistence/transactional/PolarisTreeMapMetaStoreSessionImpl.java @@ -214,7 +214,7 @@ public void deleteAll(@Nonnull PolarisCallContext callCtx) { /** {@inheritDoc} */ @Override public @Nullable PolarisBaseEntity lookupEntity( - @Nonnull PolarisCallContext callCtx, long catalogId, long entityId) { + @Nonnull PolarisCallContext callCtx, long catalogId, long entityId, int typeCode) { return this.store.getSliceEntities().read(this.store.buildKeyComposite(catalogId, entityId)); } diff --git a/polaris-core/src/main/java/org/apache/polaris/core/storage/PolarisCredentialVendor.java b/polaris-core/src/main/java/org/apache/polaris/core/storage/PolarisCredentialVendor.java index 9c3c1ce671..6a51fba8b2 100644 --- a/polaris-core/src/main/java/org/apache/polaris/core/storage/PolarisCredentialVendor.java +++ b/polaris-core/src/main/java/org/apache/polaris/core/storage/PolarisCredentialVendor.java @@ -26,6 +26,7 @@ import java.util.Map; import java.util.Set; import org.apache.polaris.core.PolarisCallContext; +import org.apache.polaris.core.entity.PolarisEntityType; import org.apache.polaris.core.persistence.dao.entity.BaseResult; /** Manage credentials for storage locations. */ @@ -48,6 +49,7 @@ ScopedCredentialsResult getSubscopedCredsForEntity( @Nonnull PolarisCallContext callCtx, long catalogId, long entityId, + PolarisEntityType entityType, boolean allowListOperation, @Nonnull Set allowedReadLocations, @Nonnull Set allowedWriteLocations); @@ -90,6 +92,7 @@ ValidateAccessResult validateAccessToLocations( @Nonnull PolarisCallContext callCtx, long catalogId, long entityId, + PolarisEntityType entityType, @Nonnull Set actions, @Nonnull Set locations); diff --git a/polaris-core/src/main/java/org/apache/polaris/core/storage/cache/StorageCredentialCache.java b/polaris-core/src/main/java/org/apache/polaris/core/storage/cache/StorageCredentialCache.java index 44a1b0fd46..e92053852b 100644 --- a/polaris-core/src/main/java/org/apache/polaris/core/storage/cache/StorageCredentialCache.java +++ b/polaris-core/src/main/java/org/apache/polaris/core/storage/cache/StorageCredentialCache.java @@ -148,6 +148,7 @@ public Map getOrGenerateSubScopeCreds( k.getCallContext(), k.getCatalogId(), k.getEntityId(), + polarisEntity.getType(), k.isAllowedListAction(), k.getAllowedReadLocations(), k.getAllowedWriteLocations()); diff --git a/polaris-core/src/test/java/org/apache/polaris/core/persistence/EntityCacheTest.java b/polaris-core/src/test/java/org/apache/polaris/core/persistence/EntityCacheTest.java index d6015160ee..0a3986357c 100644 --- a/polaris-core/src/test/java/org/apache/polaris/core/persistence/EntityCacheTest.java +++ b/polaris-core/src/test/java/org/apache/polaris/core/persistence/EntityCacheTest.java @@ -129,7 +129,9 @@ void testGetOrLoadEntityByName() { Assertions.assertThat(lookup.isCacheHit()).isTrue(); // do it again by id, should be found in the cache - lookup = cache.getOrLoadEntityById(this.callCtx, catalog.getCatalogId(), catalog.getId()); + lookup = + cache.getOrLoadEntityById( + this.callCtx, catalog.getCatalogId(), catalog.getId(), catalog.getType()); Assertions.assertThat(lookup).isNotNull(); Assertions.assertThat(lookup.isCacheHit()).isTrue(); Assertions.assertThat(lookup.getCacheEntry()).isNotNull(); @@ -148,7 +150,7 @@ void testGetOrLoadEntityByName() { Assertions.assertThat(cacheEntry).isNull(); // try to find it in the cache by id. Should not be there, i.e. no cache hit - lookup = cache.getOrLoadEntityById(this.callCtx, N1.getCatalogId(), N1.getId()); + lookup = cache.getOrLoadEntityById(this.callCtx, N1.getCatalogId(), N1.getId(), N1.getType()); Assertions.assertThat(lookup).isNotNull(); Assertions.assertThat(lookup.isCacheHit()).isFalse(); @@ -171,7 +173,7 @@ void testGetOrLoadEntityByName() { Assertions.assertThat(N1_entry.getGrantRecordsAsSecurable()).isNotNull(); // negative tests, load an entity which does not exist - lookup = cache.getOrLoadEntityById(this.callCtx, N1.getCatalogId(), 10000); + lookup = cache.getOrLoadEntityById(this.callCtx, N1.getCatalogId(), 10000, N1.getType()); Assertions.assertThat(lookup).isNull(); lookup = cache.getOrLoadEntityByName( @@ -429,7 +431,7 @@ void testRenameAndCacheDestinationBeforeLoadingSource() { PolarisBaseEntity N1 = this.tm.ensureExistsByName(List.of(catalog), PolarisEntityType.NAMESPACE, "N1"); - lookup = cache.getOrLoadEntityById(this.callCtx, N1.getCatalogId(), N1.getId()); + lookup = cache.getOrLoadEntityById(this.callCtx, N1.getCatalogId(), N1.getId(), N1.getType()); Assertions.assertThat(lookup).isNotNull(); EntityCacheByNameKey T4_name = @@ -457,7 +459,7 @@ void testRenameAndCacheDestinationBeforeLoadingSource() { // new entry if lookup by id EntityCacheLookupResult lookupResult = - cache.getOrLoadEntityById(callCtx, T4.getCatalogId(), T4.getId()); + cache.getOrLoadEntityById(callCtx, T4.getCatalogId(), T4.getId(), T4.getType()); Assertions.assertThat(lookupResult).isNotNull(); Assertions.assertThat(lookupResult.getCacheEntry()).isNotNull(); Assertions.assertThat(lookupResult.getCacheEntry().getEntity().getName()) diff --git a/polaris-core/src/test/java/org/apache/polaris/core/persistence/ResolverTest.java b/polaris-core/src/test/java/org/apache/polaris/core/persistence/ResolverTest.java index 7099fdf445..fbcb2f9c94 100644 --- a/polaris-core/src/test/java/org/apache/polaris/core/persistence/ResolverTest.java +++ b/polaris-core/src/test/java/org/apache/polaris/core/persistence/ResolverTest.java @@ -982,7 +982,7 @@ private void ensureResolved( // reload the cached entry from the backend ResolvedEntityResult refResolvedEntity = this.metaStoreManager.loadResolvedEntityById( - this.callCtx, refEntity.getCatalogId(), refEntity.getId()); + this.callCtx, refEntity.getCatalogId(), refEntity.getId(), refEntity.getType()); // should exist Assertions.assertThat(refResolvedEntity).isNotNull(); diff --git a/polaris-core/src/test/java/org/apache/polaris/core/storage/cache/StorageCredentialCacheTest.java b/polaris-core/src/test/java/org/apache/polaris/core/storage/cache/StorageCredentialCacheTest.java index 22e4c8b9d8..3672048752 100644 --- a/polaris-core/src/test/java/org/apache/polaris/core/storage/cache/StorageCredentialCacheTest.java +++ b/polaris-core/src/test/java/org/apache/polaris/core/storage/cache/StorageCredentialCacheTest.java @@ -84,6 +84,7 @@ public void testBadResult() { Mockito.any(), Mockito.anyLong(), Mockito.anyLong(), + Mockito.any(), Mockito.anyBoolean(), Mockito.anySet(), Mockito.anySet())) @@ -115,6 +116,7 @@ public void testCacheHit() { Mockito.any(), Mockito.anyLong(), Mockito.anyLong(), + Mockito.any(), Mockito.anyBoolean(), Mockito.anySet(), Mockito.anySet())) @@ -157,6 +159,7 @@ public void testCacheEvict() throws InterruptedException { Mockito.any(), Mockito.anyLong(), Mockito.anyLong(), + Mockito.any(), Mockito.anyBoolean(), Mockito.anySet(), Mockito.anySet())) @@ -214,6 +217,7 @@ public void testCacheGenerateNewEntries() { Mockito.any(), Mockito.anyLong(), Mockito.anyLong(), + Mockito.any(), Mockito.anyBoolean(), Mockito.anySet(), Mockito.anySet())) @@ -301,6 +305,7 @@ public void testCacheNotAffectedBy() { Mockito.any(), Mockito.anyLong(), Mockito.anyLong(), + Mockito.any(), Mockito.anyBoolean(), Mockito.anySet(), Mockito.anySet())) 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 0db9a6df44..81271305c1 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 @@ -118,7 +118,7 @@ private PolarisBaseEntity ensureExistsById( // make sure this entity was persisted PolarisBaseEntity entity = polarisMetaStoreManager - .loadEntity(this.polarisCallContext, catalogId, entityId) + .loadEntity(this.polarisCallContext, catalogId, entityId, expectedType) .getEntity(); // assert all expected values @@ -249,12 +249,17 @@ void ensureGrantRecordExists( // re-load both entities, ensure not null securable = polarisMetaStoreManager - .loadEntity(this.polarisCallContext, securable.getCatalogId(), securable.getId()) + .loadEntity( + this.polarisCallContext, + securable.getCatalogId(), + securable.getId(), + securable.getType()) .getEntity(); Assertions.assertThat(securable).isNotNull(); grantee = polarisMetaStoreManager - .loadEntity(this.polarisCallContext, grantee.getCatalogId(), grantee.getId()) + .loadEntity( + this.polarisCallContext, grantee.getCatalogId(), grantee.getId(), grantee.getType()) .getEntity(); Assertions.assertThat(grantee).isNotNull(); @@ -305,10 +310,15 @@ private void validateLoadedGrants(LoadGrantsResult loadGrantRecords, boolean isG long entityId = isGrantee ? grantRecord.getSecurableId() : grantRecord.getGranteeId(); // load that entity - PolarisBaseEntity entity = - polarisMetaStoreManager - .loadEntity(this.polarisCallContext, catalogId, entityId) - .getEntity(); + PolarisBaseEntity entity = null; + for (PolarisEntityType type : PolarisEntityType.values()) { + EntityResult entityResult = + polarisMetaStoreManager.loadEntity(this.polarisCallContext, catalogId, entityId, type); + if (entityResult.isSuccess()) { + entity = entityResult.getEntity(); + break; + } + } Assertions.assertThat(entity).isNotNull(); Assertions.assertThat(entities.get(entityId)).isEqualTo(entity); } @@ -326,12 +336,17 @@ void ensureGrantRecordRemoved( // re-load both entities, ensure not null securable = polarisMetaStoreManager - .loadEntity(this.polarisCallContext, securable.getCatalogId(), securable.getId()) + .loadEntity( + this.polarisCallContext, + securable.getCatalogId(), + securable.getId(), + securable.getType()) .getEntity(); Assertions.assertThat(securable).isNotNull(); grantee = polarisMetaStoreManager - .loadEntity(this.polarisCallContext, grantee.getCatalogId(), grantee.getId()) + .loadEntity( + this.polarisCallContext, grantee.getCatalogId(), grantee.getId(), grantee.getType()) .getEntity(); Assertions.assertThat(grantee).isNotNull(); @@ -451,7 +466,11 @@ PolarisBaseEntity createPrincipal(String name) { PolarisBaseEntity reloadPrincipal = polarisMetaStoreManager - .loadEntity(this.polarisCallContext, 0L, createPrincipalResult.getPrincipal().getId()) + .loadEntity( + this.polarisCallContext, + 0L, + createPrincipalResult.getPrincipal().getId(), + createPrincipalResult.getPrincipal().getType()) .getEntity(); internalProperties = PolarisObjectMapperUtil.deserializeProperties( @@ -508,7 +527,8 @@ PolarisBaseEntity createPrincipal(String name) { PolarisBaseEntity newPrincipal = polarisMetaStoreManager - .loadEntity(this.polarisCallContext, 0L, principalEntity.getId()) + .loadEntity( + this.polarisCallContext, 0L, principalEntity.getId(), principalEntity.getType()) .getEntity(); internalProperties = PolarisObjectMapperUtil.deserializeProperties( @@ -542,7 +562,8 @@ PolarisBaseEntity createPrincipal(String name) { PolarisBaseEntity finalPrincipal = polarisMetaStoreManager - .loadEntity(this.polarisCallContext, 0L, principalEntity.getId()) + .loadEntity( + this.polarisCallContext, 0L, principalEntity.getId(), principalEntity.getType()) .getEntity(); internalProperties = PolarisObjectMapperUtil.deserializeProperties( @@ -638,7 +659,11 @@ void dropEntity(List catalogPath, PolarisEntityCore entityToD // check if it exists PolarisBaseEntity entity = polarisMetaStoreManager - .loadEntity(this.polarisCallContext, entityToDrop.getCatalogId(), entityToDrop.getId()) + .loadEntity( + this.polarisCallContext, + entityToDrop.getCatalogId(), + entityToDrop.getId(), + entityToDrop.getType()) .getEntity(); if (entity != null) { EntityResult entityFound = @@ -746,7 +771,11 @@ void dropEntity(List catalogPath, PolarisEntityCore entityToD Assertions.assertThat(dropResult.getCleanupTaskId()).isNotNull(); PolarisBaseEntity cleanupTask = polarisMetaStoreManager - .loadEntity(this.polarisCallContext, 0L, dropResult.getCleanupTaskId()) + .loadEntity( + this.polarisCallContext, + 0L, + dropResult.getCleanupTaskId(), + PolarisEntityType.TASK) .getEntity(); Assertions.assertThat(cleanupTask).isNotNull(); Assertions.assertThat(cleanupTask.getType()).isEqualTo(PolarisEntityType.TASK); @@ -775,7 +804,10 @@ void dropEntity(List catalogPath, PolarisEntityCore entityToD PolarisBaseEntity entityAfterDrop = polarisMetaStoreManager .loadEntity( - this.polarisCallContext, entityToDrop.getCatalogId(), entityToDrop.getId()) + this.polarisCallContext, + entityToDrop.getCatalogId(), + entityToDrop.getId(), + entityToDrop.getType()) .getEntity(); // ensure dropped @@ -1093,7 +1125,8 @@ PolarisBaseEntity updateEntity( // lookup that entity, ensure it exists PolarisBaseEntity beforeUpdateEntity = polarisMetaStoreManager - .loadEntity(this.polarisCallContext, entity.getCatalogId(), entity.getId()) + .loadEntity( + this.polarisCallContext, entity.getCatalogId(), entity.getId(), entity.getType()) .getEntity(); // update that property @@ -1110,7 +1143,8 @@ PolarisBaseEntity updateEntity( // refresh catalog info entity = polarisMetaStoreManager - .loadEntity(this.polarisCallContext, entity.getCatalogId(), entity.getId()) + .loadEntity( + this.polarisCallContext, entity.getCatalogId(), entity.getId(), entity.getType()) .getEntity(); // ensure nothing has changed @@ -1248,7 +1282,7 @@ private void validateCacheEntryLoad(ResolvedEntityResult cacheEntry) { PolarisEntity refEntity = PolarisEntity.of( this.polarisMetaStoreManager.loadEntity( - this.polarisCallContext, entity.getCatalogId(), entity.getId())); + this.polarisCallContext, entity.getCatalogId(), entity.getId(), entity.getType())); Assertions.assertThat(refEntity).isNotNull(); // same entity @@ -1295,6 +1329,7 @@ private void validateCacheEntryRefresh( ResolvedEntityResult cacheEntry, long catalogId, long entityId, + PolarisEntityType entityType, int entityVersion, int entityGrantRecordsVersion) { // cannot be null @@ -1305,7 +1340,7 @@ private void validateCacheEntryRefresh( // reload the entity PolarisBaseEntity refEntity = this.polarisMetaStoreManager - .loadEntity(this.polarisCallContext, catalogId, entityId) + .loadEntity(this.polarisCallContext, catalogId, entityId, entityType) .getEntity(); Assertions.assertThat(refEntity).isNotNull(); @@ -1411,11 +1446,11 @@ private PolarisBaseEntity loadCacheEntryByName( * @return return just the entity */ private PolarisBaseEntity loadCacheEntryById( - long entityCatalogId, long entityId, boolean expectExists) { + long entityCatalogId, long entityId, PolarisEntityType entityType, boolean expectExists) { // load cached entry ResolvedEntityResult cacheEntry = this.polarisMetaStoreManager.loadResolvedEntityById( - this.polarisCallContext, entityCatalogId, entityId); + this.polarisCallContext, entityCatalogId, entityId, entityType); // if null, validate that indeed the entry does not exist Assertions.assertThat(cacheEntry.isSuccess()).isEqualTo(expectExists); @@ -1437,8 +1472,9 @@ private PolarisBaseEntity loadCacheEntryById( * @param entityId parent id of the entity * @return return just the entity */ - private PolarisBaseEntity loadCacheEntryById(long entityCatalogId, long entityId) { - return this.loadCacheEntryById(entityCatalogId, entityId, true); + private PolarisBaseEntity loadCacheEntryById( + long entityCatalogId, long entityId, PolarisEntityType entityType) { + return this.loadCacheEntryById(entityCatalogId, entityId, entityType, true); } /** @@ -1475,7 +1511,12 @@ private void refreshCacheEntry( // if not null, validate it if (cacheEntry.isSuccess()) { this.validateCacheEntryRefresh( - cacheEntry, entityCatalogId, entityId, entityVersion, entityGrantRecordsVersion); + cacheEntry, + entityCatalogId, + entityId, + entityType, + entityVersion, + entityGrantRecordsVersion); } } @@ -2226,7 +2267,7 @@ public void testEntityCache() { "test"); // and again by id - TEST = this.loadCacheEntryById(TEST.getCatalogId(), TEST.getId()); + TEST = this.loadCacheEntryById(TEST.getCatalogId(), TEST.getId(), TEST.getType()); // get namespace N1 PolarisBaseEntity N1 = @@ -2264,7 +2305,7 @@ public void testEntityCache() { // load role R1 PolarisBaseEntity R1 = this.loadCacheEntryByName(TEST.getId(), TEST.getId(), PolarisEntityType.CATALOG_ROLE, "R1"); - R1 = this.loadCacheEntryById(R1.getCatalogId(), R1.getId()); + R1 = this.loadCacheEntryById(R1.getCatalogId(), R1.getId(), R1.getType()); // add a grant record to N1 this.grantPrivilege(R1, List.of(TEST), N1, PolarisPrivilege.NAMESPACE_FULL_METADATA); @@ -2291,7 +2332,7 @@ public void testEntityCache() { // now validate that load something which does not exist, will also work this.loadCacheEntryByName( N1.getCatalogId(), N1.getId(), PolarisEntityType.TABLE_LIKE, "do_not_exists", false); - this.loadCacheEntryById(N1.getCatalogId() + 1000, N1.getId(), false); + this.loadCacheEntryById(N1.getCatalogId() + 1000, N1.getId(), N1.getType(), false); // refresh a purged entity this.refreshCacheEntry( diff --git a/quarkus/service/src/test/java/org/apache/polaris/service/quarkus/admin/PolarisAuthzTestBase.java b/quarkus/service/src/test/java/org/apache/polaris/service/quarkus/admin/PolarisAuthzTestBase.java index 65bd9edbf9..2b28e9f7f6 100644 --- a/quarkus/service/src/test/java/org/apache/polaris/service/quarkus/admin/PolarisAuthzTestBase.java +++ b/quarkus/service/src/test/java/org/apache/polaris/service/quarkus/admin/PolarisAuthzTestBase.java @@ -370,7 +370,10 @@ public void after() { .map( gr -> metaStoreManager.loadEntity( - callContext.getPolarisCallContext(), 0L, gr.getSecurableId())) + callContext.getPolarisCallContext(), + 0L, + gr.getSecurableId(), + PolarisEntityType.PRINCIPAL_ROLE)) .map(EntityResult::getEntity) .map(PolarisBaseEntity::getName) .collect(Collectors.toSet()); diff --git a/quarkus/service/src/test/java/org/apache/polaris/service/quarkus/auth/JWTRSAKeyPairTest.java b/quarkus/service/src/test/java/org/apache/polaris/service/quarkus/auth/JWTRSAKeyPairTest.java index 75fea55ad3..d48db61573 100644 --- a/quarkus/service/src/test/java/org/apache/polaris/service/quarkus/auth/JWTRSAKeyPairTest.java +++ b/quarkus/service/src/test/java/org/apache/polaris/service/quarkus/auth/JWTRSAKeyPairTest.java @@ -75,7 +75,8 @@ public void testSuccessfulTokenGeneration() throws Exception { PolarisEntitySubType.NULL_SUBTYPE, 0L, "principal"); - Mockito.when(metastoreManager.loadEntity(polarisCallContext, 0L, 1L)) + Mockito.when( + metastoreManager.loadEntity(polarisCallContext, 0L, 1L, PolarisEntityType.PRINCIPAL)) .thenReturn(new EntityResult(principal)); TokenBroker tokenBroker = new JWTRSAKeyPair(metastoreManager, 420, publicFileLocation, privateFileLocation); diff --git a/quarkus/service/src/test/java/org/apache/polaris/service/quarkus/auth/JWTSymmetricKeyGeneratorTest.java b/quarkus/service/src/test/java/org/apache/polaris/service/quarkus/auth/JWTSymmetricKeyGeneratorTest.java index 5151132e3e..b12fd9fc25 100644 --- a/quarkus/service/src/test/java/org/apache/polaris/service/quarkus/auth/JWTSymmetricKeyGeneratorTest.java +++ b/quarkus/service/src/test/java/org/apache/polaris/service/quarkus/auth/JWTSymmetricKeyGeneratorTest.java @@ -80,7 +80,8 @@ public Map contextVariables() { PolarisEntitySubType.NULL_SUBTYPE, 0L, "principal"); - Mockito.when(metastoreManager.loadEntity(polarisCallContext, 0L, 1L)) + Mockito.when( + metastoreManager.loadEntity(polarisCallContext, 0L, 1L, PolarisEntityType.PRINCIPAL)) .thenReturn(new EntityResult(principal)); TokenBroker generator = new JWTSymmetricKeyBroker(metastoreManager, 666, () -> "polaris"); TokenResponse token = diff --git a/quarkus/service/src/test/java/org/apache/polaris/service/quarkus/catalog/BasePolarisCatalogTest.java b/quarkus/service/src/test/java/org/apache/polaris/service/quarkus/catalog/BasePolarisCatalogTest.java index 84ac558b6e..5dcd140165 100644 --- a/quarkus/service/src/test/java/org/apache/polaris/service/quarkus/catalog/BasePolarisCatalogTest.java +++ b/quarkus/service/src/test/java/org/apache/polaris/service/quarkus/catalog/BasePolarisCatalogTest.java @@ -1429,6 +1429,7 @@ public void testDropTableWithPurge() { polarisContext, 0, taskEntity.getId(), + taskEntity.getType(), true, Set.of(tableMetadata.location()), Set.of(tableMetadata.location())) diff --git a/service/common/src/main/java/org/apache/polaris/service/admin/PolarisAdminService.java b/service/common/src/main/java/org/apache/polaris/service/admin/PolarisAdminService.java index 70b2e20efd..c05a6b2a3d 100644 --- a/service/common/src/main/java/org/apache/polaris/service/admin/PolarisAdminService.java +++ b/service/common/src/main/java/org/apache/polaris/service/admin/PolarisAdminService.java @@ -87,6 +87,7 @@ import org.apache.polaris.core.persistence.dao.entity.CreateCatalogResult; import org.apache.polaris.core.persistence.dao.entity.CreatePrincipalResult; import org.apache.polaris.core.persistence.dao.entity.DropEntityResult; +import org.apache.polaris.core.persistence.dao.entity.EntityResult; import org.apache.polaris.core.persistence.resolver.PolarisResolutionManifest; import org.apache.polaris.core.persistence.resolver.ResolverPath; import org.apache.polaris.core.persistence.resolver.ResolverStatus; @@ -761,7 +762,8 @@ private List listCatalogsUnsafe() { .map( nameAndId -> PolarisEntity.of( - metaStoreManager.loadEntity(getCurrentPolarisContext(), 0, nameAndId.getId()))) + metaStoreManager.loadEntity( + getCurrentPolarisContext(), 0, nameAndId.getId(), nameAndId.getType()))) .toList(); } @@ -887,7 +889,10 @@ public void deletePrincipal(String name) { PolarisEntity newPrincipal = PolarisEntity.of( metaStoreManager.loadEntity( - getCurrentPolarisContext(), 0L, currentPrincipalEntity.getId())); + getCurrentPolarisContext(), + 0L, + currentPrincipalEntity.getId(), + currentPrincipalEntity.getType())); return new PrincipalWithCredentials( PrincipalEntity.of(newPrincipal).asPrincipal(), new PrincipalWithCredentialsCredentials( @@ -923,7 +928,8 @@ public List listPrincipals() { .map( nameAndId -> PolarisEntity.of( - metaStoreManager.loadEntity(getCurrentPolarisContext(), 0, nameAndId.getId()))) + metaStoreManager.loadEntity( + getCurrentPolarisContext(), 0, nameAndId.getId(), nameAndId.getType()))) .toList(); } @@ -1031,7 +1037,8 @@ public List listPrincipalRoles() { .map( nameAndId -> PolarisEntity.of( - metaStoreManager.loadEntity(getCurrentPolarisContext(), 0, nameAndId.getId()))) + metaStoreManager.loadEntity( + getCurrentPolarisContext(), 0, nameAndId.getId(), nameAndId.getType()))) .toList(); } @@ -1159,7 +1166,10 @@ public List listCatalogRoles(String catalogName) { nameAndId -> PolarisEntity.of( metaStoreManager.loadEntity( - getCurrentPolarisContext(), catalogEntity.getId(), nameAndId.getId()))) + getCurrentPolarisContext(), + catalogEntity.getId(), + nameAndId.getId(), + nameAndId.getType()))) .toList(); } @@ -1209,7 +1219,7 @@ public List listPrincipalRolesAssigned(String principalName) { LoadGrantsResult grantList = metaStoreManager.loadGrantsToGrantee( getCurrentPolarisContext(), principalEntity.getCatalogId(), principalEntity.getId()); - return buildEntitiesFromGrantResults(grantList, false, null); + return buildEntitiesFromGrantResults(grantList, false, PolarisEntityType.PRINCIPAL_ROLE, null); } public boolean assignCatalogRoleToPrincipalRole( @@ -1275,7 +1285,7 @@ public List listAssigneePrincipalsForPrincipalRole(String princip getCurrentPolarisContext(), principalRoleEntity.getCatalogId(), principalRoleEntity.getId()); - return buildEntitiesFromGrantResults(grantList, true, null); + return buildEntitiesFromGrantResults(grantList, true, PolarisEntityType.PRINCIPAL, null); } /** @@ -1291,6 +1301,7 @@ public List listAssigneePrincipalsForPrincipalRole(String princip private List buildEntitiesFromGrantResults( @Nonnull LoadGrantsResult grantList, boolean grantees, + PolarisEntityType entityType, @Nullable Function grantFilter) { Map granteeMap = grantList.getEntitiesAsMap(); List toReturn = new ArrayList<>(grantList.getGrantRecords().size()); @@ -1300,7 +1311,8 @@ private List buildEntitiesFromGrantResults( grantees ? grantRecord.getGranteeCatalogId() : grantRecord.getSecurableCatalogId(); long entityId = grantees ? grantRecord.getGranteeId() : grantRecord.getSecurableId(); // get the entity associated with the grantee - PolarisBaseEntity entity = this.getOrLoadEntity(granteeMap, catalogId, entityId); + PolarisBaseEntity entity = + this.getOrLoadEntity(granteeMap, catalogId, entityId, entityType); if (entity != null) { toReturn.add(PolarisEntity.of(entity)); } @@ -1329,7 +1341,10 @@ public List listCatalogRolesForPrincipalRole( principalRoleEntity.getCatalogId(), principalRoleEntity.getId()); return buildEntitiesFromGrantResults( - grantList, false, grantRec -> grantRec.getSecurableCatalogId() == catalogEntity.getId()); + grantList, + false, + PolarisEntityType.CATALOG_ROLE, + grantRec -> grantRec.getSecurableCatalogId() == catalogEntity.getId()); } /** Adds a grant on the root container of this realm to {@code principalRoleName}. */ @@ -1554,7 +1569,7 @@ public List listAssigneePrincipalRolesForCatalogRole( getCurrentPolarisContext(), catalogRoleEntity.getCatalogId(), catalogRoleEntity.getId()); - return buildEntitiesFromGrantResults(grantList, true, null); + return buildEntitiesFromGrantResults(grantList, true, PolarisEntityType.PRINCIPAL_ROLE, null); } /** @@ -1580,8 +1595,7 @@ public List listGrantsForCatalogRole(String catalogName, String c Map entityMap = grantList.getEntitiesAsMap(); for (PolarisGrantRecord record : grantList.getGrantRecords()) { PolarisPrivilege privilege = PolarisPrivilege.fromCode(record.getPrivilegeCode()); - PolarisBaseEntity baseEntity = - this.getOrLoadEntity(entityMap, record.getSecurableCatalogId(), record.getSecurableId()); + PolarisBaseEntity baseEntity = this.getOrLoadEntityForGrant(entityMap, record); if (baseEntity != null) { switch (baseEntity.getType()) { case CATALOG: @@ -1650,15 +1664,42 @@ public List listGrantsForCatalogRole(String catalogName, String c * @param entitiesMap map of entities * @param catalogId the id of the catalog of the entity we are looking for * @param id id of the entity we are looking for + * @param entityType * @return null if the entity does not exist */ private @Nullable PolarisBaseEntity getOrLoadEntity( - @Nullable Map entitiesMap, long catalogId, long id) { + @Nullable Map entitiesMap, + long catalogId, + long id, + PolarisEntityType entityType) { return (entitiesMap == null) - ? metaStoreManager.loadEntity(getCurrentPolarisContext(), catalogId, id).getEntity() + ? metaStoreManager + .loadEntity(getCurrentPolarisContext(), catalogId, id, entityType) + .getEntity() : entitiesMap.get(id); } + private @Nullable PolarisBaseEntity getOrLoadEntityForGrant( + @Nullable Map entitiesMap, PolarisGrantRecord record) { + if (entitiesMap != null) { + return entitiesMap.get(record.getSecurableId()); + } + + for (PolarisEntityType type : PolarisEntityType.values()) { + EntityResult entityResult = + metaStoreManager.loadEntity( + getCurrentPolarisContext(), + record.getSecurableCatalogId(), + record.getSecurableId(), + type); + if (entityResult.isSuccess()) { + return entityResult.getEntity(); + } + } + + return null; + } + /** Adds a table-level or view-level grant on {@code identifier} to {@code catalogRoleName}. */ private boolean grantPrivilegeOnTableLikeToRole( String catalogName, diff --git a/service/common/src/main/java/org/apache/polaris/service/auth/BasePolarisAuthenticator.java b/service/common/src/main/java/org/apache/polaris/service/auth/BasePolarisAuthenticator.java index fb713659b2..5c284a9345 100644 --- a/service/common/src/main/java/org/apache/polaris/service/auth/BasePolarisAuthenticator.java +++ b/service/common/src/main/java/org/apache/polaris/service/auth/BasePolarisAuthenticator.java @@ -69,7 +69,10 @@ protected Optional getPrincipal(DecodedToken toke tokenInfo.getPrincipalId() > 0 ? PolarisEntity.of( metaStoreManager.loadEntity( - callContext.getPolarisCallContext(), 0L, tokenInfo.getPrincipalId())) + callContext.getPolarisCallContext(), + 0L, + tokenInfo.getPrincipalId(), + PolarisEntityType.PRINCIPAL)) : PolarisEntity.of( metaStoreManager.readEntityByName( callContext.getPolarisCallContext(), diff --git a/service/common/src/main/java/org/apache/polaris/service/auth/DefaultActiveRolesProvider.java b/service/common/src/main/java/org/apache/polaris/service/auth/DefaultActiveRolesProvider.java index c679cddc2a..95ed838ca5 100644 --- a/service/common/src/main/java/org/apache/polaris/service/auth/DefaultActiveRolesProvider.java +++ b/service/common/src/main/java/org/apache/polaris/service/auth/DefaultActiveRolesProvider.java @@ -32,6 +32,7 @@ import org.apache.polaris.core.context.CallContext; import org.apache.polaris.core.context.RealmContext; import org.apache.polaris.core.entity.PolarisEntity; +import org.apache.polaris.core.entity.PolarisEntityType; import org.apache.polaris.core.entity.PrincipalRoleEntity; import org.apache.polaris.core.persistence.MetaStoreManagerFactory; import org.apache.polaris.core.persistence.PolarisMetaStoreManager; @@ -90,7 +91,10 @@ protected List loadActivePrincipalRoles( .map( gr -> metaStoreManager.loadEntity( - polarisContext, gr.getSecurableCatalogId(), gr.getSecurableId())) + polarisContext, + gr.getSecurableCatalogId(), + gr.getSecurableId(), + PolarisEntityType.PRINCIPAL_ROLE)) .filter(EntityResult::isSuccess) .map(EntityResult::getEntity) .map(PrincipalRoleEntity::of) diff --git a/service/common/src/main/java/org/apache/polaris/service/auth/JWTBroker.java b/service/common/src/main/java/org/apache/polaris/service/auth/JWTBroker.java index 0e7ca3107a..43e4804aa4 100644 --- a/service/common/src/main/java/org/apache/polaris/service/auth/JWTBroker.java +++ b/service/common/src/main/java/org/apache/polaris/service/auth/JWTBroker.java @@ -114,7 +114,8 @@ public TokenResponse generateFromToken( metaStoreManager.loadEntity( CallContext.getCurrentContext().getPolarisCallContext(), 0L, - decodedToken.getPrincipalId()); + decodedToken.getPrincipalId(), + PolarisEntityType.PRINCIPAL); if (!principalLookup.isSuccess() || principalLookup.getEntity().getType() != PolarisEntityType.PRINCIPAL) { return new TokenResponse(OAuthTokenErrorResponse.Error.unauthorized_client); diff --git a/service/common/src/main/java/org/apache/polaris/service/auth/TestOAuth2ApiService.java b/service/common/src/main/java/org/apache/polaris/service/auth/TestOAuth2ApiService.java index 00330bc2c7..5b13617718 100644 --- a/service/common/src/main/java/org/apache/polaris/service/auth/TestOAuth2ApiService.java +++ b/service/common/src/main/java/org/apache/polaris/service/auth/TestOAuth2ApiService.java @@ -91,7 +91,10 @@ private String getPrincipalName(String clientId, RealmContext realmContext) { LOGGER.debug("Found principal secrets for client id {}", clientId); EntityResult principalResult = metaStoreManager.loadEntity( - polarisCallContext, 0L, secretsResult.getPrincipalSecrets().getPrincipalId()); + polarisCallContext, + 0L, + secretsResult.getPrincipalSecrets().getPrincipalId(), + PolarisEntityType.PRINCIPAL); if (!principalResult.isSuccess()) { throw new NotAuthorizedException("Failed to load principal entity"); } diff --git a/service/common/src/main/java/org/apache/polaris/service/auth/TokenBroker.java b/service/common/src/main/java/org/apache/polaris/service/auth/TokenBroker.java index aaf0949504..2896008028 100644 --- a/service/common/src/main/java/org/apache/polaris/service/auth/TokenBroker.java +++ b/service/common/src/main/java/org/apache/polaris/service/auth/TokenBroker.java @@ -127,7 +127,10 @@ TokenResponse generateFromToken( } EntityResult result = metaStoreManager.loadEntity( - polarisCallContext, 0L, principalSecrets.getPrincipalSecrets().getPrincipalId()); + polarisCallContext, + 0L, + principalSecrets.getPrincipalSecrets().getPrincipalId(), + PolarisEntityType.PRINCIPAL); if (!result.isSuccess() || result.getEntity().getType() != PolarisEntityType.PRINCIPAL) { return Optional.empty(); } diff --git a/service/common/src/main/java/org/apache/polaris/service/task/TaskExecutorImpl.java b/service/common/src/main/java/org/apache/polaris/service/task/TaskExecutorImpl.java index a404718cfd..a800682325 100644 --- a/service/common/src/main/java/org/apache/polaris/service/task/TaskExecutorImpl.java +++ b/service/common/src/main/java/org/apache/polaris/service/task/TaskExecutorImpl.java @@ -114,7 +114,9 @@ protected void handleTask(long taskEntityId, CallContext ctx, int attempt) { PolarisMetaStoreManager metaStoreManager = metaStoreManagerFactory.getOrCreateMetaStoreManager(ctx.getRealmContext()); PolarisBaseEntity taskEntity = - metaStoreManager.loadEntity(ctx.getPolarisCallContext(), 0L, taskEntityId).getEntity(); + metaStoreManager + .loadEntity(ctx.getPolarisCallContext(), 0L, taskEntityId, PolarisEntityType.TASK) + .getEntity(); if (!PolarisEntityType.TASK.equals(taskEntity.getType())) { throw new IllegalArgumentException("Provided taskId must be a task entity type"); } diff --git a/service/common/src/test/java/org/apache/polaris/service/auth/BasePolarisAuthenticatorTest.java b/service/common/src/test/java/org/apache/polaris/service/auth/BasePolarisAuthenticatorTest.java index 69e15bc4e0..ec9eda26a8 100644 --- a/service/common/src/test/java/org/apache/polaris/service/auth/BasePolarisAuthenticatorTest.java +++ b/service/common/src/test/java/org/apache/polaris/service/auth/BasePolarisAuthenticatorTest.java @@ -27,6 +27,7 @@ import org.apache.polaris.core.auth.AuthenticatedPolarisPrincipal; import org.apache.polaris.core.context.CallContext; import org.apache.polaris.core.context.RealmContext; +import org.apache.polaris.core.entity.PolarisEntityType; import org.apache.polaris.core.persistence.MetaStoreManagerFactory; import org.apache.polaris.core.persistence.PolarisMetaStoreManager; import org.apache.polaris.core.persistence.dao.entity.BaseResult; @@ -66,7 +67,8 @@ public void testFetchPrincipalThrowsServiceExceptionOnMetastoreException() { DecodedToken token = Mockito.mock(DecodedToken.class); long principalId = 100L; when(token.getPrincipalId()).thenReturn(principalId); - when(metaStoreManager.loadEntity(polarisCallContext, 0L, principalId)) + when(metaStoreManager.loadEntity( + polarisCallContext, 0L, principalId, PolarisEntityType.PRINCIPAL)) .thenThrow(new RuntimeException("Metastore exception")); Assertions.assertThatThrownBy(() -> authenticator.getPrincipal(token)) @@ -80,7 +82,8 @@ public void testFetchPrincipalThrowsNotAuthorizedWhenNotFound() { long principalId = 100L; when(token.getPrincipalId()).thenReturn(principalId); when(token.getClientId()).thenReturn("abc"); - when(metaStoreManager.loadEntity(polarisCallContext, 0L, principalId)) + when(metaStoreManager.loadEntity( + polarisCallContext, 0L, principalId, PolarisEntityType.PRINCIPAL)) .thenReturn(new EntityResult(BaseResult.ReturnStatus.ENTITY_NOT_FOUND, "")); Assertions.assertThatThrownBy(() -> authenticator.getPrincipal(token)) From 0a73252e35246aef2f6cebf905f42eb0873937a0 Mon Sep 17 00:00:00 2001 From: Dmitri Bourlatchkov Date: Thu, 6 Mar 2025 11:55:57 -0500 Subject: [PATCH 2/2] review: clarify javadoc --- .../polaris/core/persistence/BasePersistence.java | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/polaris-core/src/main/java/org/apache/polaris/core/persistence/BasePersistence.java b/polaris-core/src/main/java/org/apache/polaris/core/persistence/BasePersistence.java index f5448cb65c..75c30d430d 100644 --- a/polaris-core/src/main/java/org/apache/polaris/core/persistence/BasePersistence.java +++ b/polaris-core/src/main/java/org/apache/polaris/core/persistence/BasePersistence.java @@ -147,8 +147,14 @@ void deleteAllEntityGrantRecords( /** * Lookup an entity given its catalog id (which can be {@link - * org.apache.polaris.core.entity.PolarisEntityConstants#NULL_ID} for top-level entities) and its - * entityId. + * org.apache.polaris.core.entity.PolarisEntityConstants#NULL_ID} for top-level entities), its + * entityId and type code (from {@link PolarisEntityType#getCode()}. + * + *

The type code parameter is redundant but can be used to optimize implementations in some + * cases. All callers are required to provide a valid value for the type code parameter. If the + * given type code does not match the type code of the previously created entity with the + * specified {@code entityId}, implementations may still return the entity or may behave as if the + * entity were not found. * * @param callCtx call context * @param catalogId catalog id or NULL_ID