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..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,17 +147,24 @@ 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 * @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))