From f9637a6d9e364735977555a2919689dd362fd1cf Mon Sep 17 00:00:00 2001 From: Michael Collado Date: Thu, 4 Sep 2025 11:20:51 -0700 Subject: [PATCH 1/9] Add loadEntities batch call and rename listFullEntities --- .../jdbc/JdbcBasePersistenceImpl.java | 42 +++--- .../core/entity/PolarisBaseEntity.java | 6 - .../AtomicOperationMetaStoreManager.java | 37 ++++- .../core/persistence/BasePersistence.java | 4 +- .../persistence/PolarisMetaStoreManager.java | 29 +++- .../TransactionWorkspaceMetaStoreManager.java | 12 +- .../AbstractTransactionalPersistence.java | 2 +- .../TransactionalMetaStoreManagerImpl.java | 51 ++++++- .../TransactionalPersistence.java | 2 +- .../BasePolarisMetaStoreManagerTest.java | 8 +- .../PolarisTestMetaStoreManager.java | 140 ++++++++++++++++++ .../service/admin/PolarisAdminService.java | 8 +- .../service/catalog/policy/PolicyCatalog.java | 2 +- 13 files changed, 294 insertions(+), 49 deletions(-) diff --git a/persistence/relational-jdbc/src/main/java/org/apache/polaris/persistence/relational/jdbc/JdbcBasePersistenceImpl.java b/persistence/relational-jdbc/src/main/java/org/apache/polaris/persistence/relational/jdbc/JdbcBasePersistenceImpl.java index 8d6201453f..7a032cfac9 100644 --- a/persistence/relational-jdbc/src/main/java/org/apache/polaris/persistence/relational/jdbc/JdbcBasePersistenceImpl.java +++ b/persistence/relational-jdbc/src/main/java/org/apache/polaris/persistence/relational/jdbc/JdbcBasePersistenceImpl.java @@ -18,24 +18,9 @@ */ package org.apache.polaris.persistence.relational.jdbc; -import static org.apache.polaris.persistence.relational.jdbc.QueryGenerator.PreparedQuery; - import com.google.common.base.Preconditions; import jakarta.annotation.Nonnull; import jakarta.annotation.Nullable; -import java.sql.Connection; -import java.sql.SQLException; -import java.util.ArrayList; -import java.util.Collections; -import java.util.HashMap; -import java.util.LinkedHashMap; -import java.util.List; -import java.util.Map; -import java.util.Optional; -import java.util.concurrent.atomic.AtomicReference; -import java.util.function.Function; -import java.util.function.Predicate; -import java.util.stream.Collectors; import org.apache.polaris.core.PolarisCallContext; import org.apache.polaris.core.PolarisDiagnostics; import org.apache.polaris.core.entity.EntityNameLookupRecord; @@ -77,6 +62,23 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import java.sql.Connection; +import java.sql.SQLException; +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashMap; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.Objects; +import java.util.Optional; +import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Function; +import java.util.function.Predicate; +import java.util.stream.Collectors; + +import static org.apache.polaris.persistence.relational.jdbc.QueryGenerator.PreparedQuery; + public class JdbcBasePersistenceImpl implements BasePersistence, IntegrationPersistence { private static final Logger LOGGER = LoggerFactory.getLogger(JdbcBasePersistenceImpl.class); @@ -463,7 +465,12 @@ public List lookupEntities( PreparedQuery query = QueryGenerator.generateSelectQueryWithEntityIds(realmId, schemaVersion, entityIds); try { - return datasourceOperations.executeSelect(query, new ModelEntity(schemaVersion)); + Map idMap = + datasourceOperations.executeSelect(query, new ModelEntity(schemaVersion)).stream() + .collect( + Collectors.toMap( + e -> new PolarisEntityId(e.getCatalogId(), e.getId()), Function.identity())); + return entityIds.stream().map(idMap::get).collect(Collectors.toList()); } catch (SQLException e) { throw new RuntimeException( String.format("Failed to retrieve polaris entities due to %s", e.getMessage()), e); @@ -476,6 +483,7 @@ public List lookupEntityVersions( @Nonnull PolarisCallContext callCtx, List entityIds) { Map idToEntityMap = lookupEntities(callCtx, entityIds).stream() + .filter(Objects::nonNull) .collect( Collectors.toMap( entry -> new PolarisEntityId(entry.getCatalogId(), entry.getId()), @@ -570,7 +578,7 @@ public Page listEntities( @Nonnull @Override - public Page loadEntities( + public Page listFullEntities( @Nonnull PolarisCallContext callCtx, long catalogId, long parentId, diff --git a/polaris-core/src/main/java/org/apache/polaris/core/entity/PolarisBaseEntity.java b/polaris-core/src/main/java/org/apache/polaris/core/entity/PolarisBaseEntity.java index 80b0a96405..6c91517a1f 100644 --- a/polaris-core/src/main/java/org/apache/polaris/core/entity/PolarisBaseEntity.java +++ b/polaris-core/src/main/java/org/apache/polaris/core/entity/PolarisBaseEntity.java @@ -339,10 +339,7 @@ public boolean equals(Object o) { PolarisBaseEntity that = (PolarisBaseEntity) o; return subTypeCode == that.subTypeCode && createTimestamp == that.createTimestamp - && dropTimestamp == that.dropTimestamp - && purgeTimestamp == that.purgeTimestamp && toPurgeTimestamp == that.toPurgeTimestamp - && lastUpdateTimestamp == that.lastUpdateTimestamp && grantRecordsVersion == that.grantRecordsVersion && Objects.equals(properties, that.properties) && Objects.equals(internalProperties, that.internalProperties); @@ -359,10 +356,7 @@ public int hashCode() { entityVersion, subTypeCode, createTimestamp, - dropTimestamp, - purgeTimestamp, toPurgeTimestamp, - lastUpdateTimestamp, properties, internalProperties, grantRecordsVersion); diff --git a/polaris-core/src/main/java/org/apache/polaris/core/persistence/AtomicOperationMetaStoreManager.java b/polaris-core/src/main/java/org/apache/polaris/core/persistence/AtomicOperationMetaStoreManager.java index c3841486a7..684dd17c7f 100644 --- a/polaris-core/src/main/java/org/apache/polaris/core/persistence/AtomicOperationMetaStoreManager.java +++ b/polaris-core/src/main/java/org/apache/polaris/core/persistence/AtomicOperationMetaStoreManager.java @@ -32,6 +32,7 @@ import java.util.concurrent.atomic.AtomicInteger; import java.util.function.Function; import java.util.stream.Collectors; +import java.util.stream.IntStream; import org.apache.polaris.core.PolarisCallContext; import org.apache.polaris.core.PolarisDiagnostics; import org.apache.polaris.core.config.FeatureConfiguration; @@ -705,7 +706,7 @@ private void revokeGrantRecord( /** {@inheritDoc} */ @Override - public @Nonnull Page loadEntities( + public @Nonnull Page listFullEntities( @Nonnull PolarisCallContext callCtx, @Nullable List catalogPath, @Nonnull PolarisEntityType entityType, @@ -728,7 +729,7 @@ private void revokeGrantRecord( // with sensitive data; but the window of inconsistency is only the duration of a single // in-flight request (the cache-based resolution follows a different path entirely). - return ms.loadEntities( + return ms.listFullEntities( callCtx, catalogId, parentId, @@ -1200,7 +1201,7 @@ private void revokeGrantRecord( // get the list of catalog roles, at most 2 List catalogRoles = - ms.loadEntities( + ms.listFullEntities( callCtx, catalogId, catalogId, @@ -1513,6 +1514,34 @@ private void revokeGrantRecord( : new EntityResult(BaseResult.ReturnStatus.ENTITY_NOT_FOUND, null); } + @Nonnull + @Override + public EntitiesResult loadEntities( + @Nonnull PolarisCallContext callCtx, + @Nonnull List entityLookupRecords) { + BasePersistence ms = callCtx.getMetaStore(); + List entities = + ms.lookupEntities( + callCtx, + entityLookupRecords.stream() + .map(r -> new PolarisEntityId(r.getCatalogId(), r.getId())) + .collect(Collectors.toList())); + // mimic the behavior of loadEntity above, return null if not found or type mismatch + List ret = + IntStream.range(0, entityLookupRecords.size()) + .mapToObj( + i -> { + if (entities.get(i) != null + && !entities.get(i).getType().equals(entityLookupRecords.get(i).getType())) { + return null; + } else { + return entities.get(i); + } + }) + .collect(Collectors.toList()); + return new EntitiesResult(Page.fromItems(ret)); + } + @Override public @Nonnull EntitiesResult loadTasks( @Nonnull PolarisCallContext callCtx, String executorId, PageToken pageToken) { @@ -1520,7 +1549,7 @@ private void revokeGrantRecord( // find all available tasks Page availableTasks = - ms.loadEntities( + ms.listFullEntities( callCtx, PolarisEntityConstants.getRootEntityId(), PolarisEntityConstants.getRootEntityId(), 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 d238e490e0..afea79fda3 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 @@ -279,7 +279,7 @@ List lookupEntityVersions( /** * List lightweight information of entities matching the given criteria with pagination. If all - * properties of the entity are required,use {@link #loadEntities} instead. + * properties of the entity are required,use {@link #listFullEntities} instead. * * @param callCtx call context * @param catalogId catalog id for that entity, NULL_ID if the entity is top-level @@ -314,7 +314,7 @@ Page listEntities( * @return the paged list of matching entities after transformation */ @Nonnull - Page loadEntities( + Page listFullEntities( @Nonnull PolarisCallContext callCtx, long catalogId, long parentId, 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 cf3912fa95..e65d36d566 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 @@ -26,6 +26,7 @@ import org.apache.polaris.core.PolarisCallContext; import org.apache.polaris.core.auth.PolarisGrantManager; import org.apache.polaris.core.auth.PolarisSecretsManager; +import org.apache.polaris.core.entity.EntityNameLookupRecord; import org.apache.polaris.core.entity.LocationBasedEntity; import org.apache.polaris.core.entity.PolarisBaseEntity; import org.apache.polaris.core.entity.PolarisEntity; @@ -114,7 +115,7 @@ EntityResult readEntityByName( /** * List lightweight information about entities matching the given criteria. If all properties of - * the entity are required,use {@link #loadEntities} instead. + * the entity are required,use {@link #listFullEntities} instead. * * @param callCtx call context * @param catalogPath path inside a catalog. If null or empty, the entities to list are top-level, @@ -135,7 +136,7 @@ ListEntitiesResult listEntities( /** * Load full entities matching the given criteria with pagination. If only the entity name/id/type * is required, use {@link #listEntities} instead. If no pagination is required, use {@link - * #loadEntitiesAll} instead. + * #loadFullEntitiesAll} instead. * * @param callCtx call context * @param catalogPath path inside a catalog. If null or empty, the entities to list are top-level, @@ -145,7 +146,7 @@ ListEntitiesResult listEntities( * @return paged list of matching entities */ @Nonnull - Page loadEntities( + Page listFullEntities( @Nonnull PolarisCallContext callCtx, @Nullable List catalogPath, @Nonnull PolarisEntityType entityType, @@ -154,7 +155,7 @@ Page loadEntities( /** * Load full entities matching the given criteria into an unpaged list. If pagination is required - * use {@link #loadEntities} instead. If only the entity name/id/type is required, use {@link + * use {@link #listFullEntities} instead. If only the entity name/id/type is required, use {@link * #listEntities} instead. * * @param callCtx call context @@ -164,12 +165,13 @@ Page loadEntities( * @param entitySubType subType of entities to list (or ANY_SUBTYPE) * @return list of all matching entities */ - default @Nonnull List loadEntitiesAll( + default @Nonnull List loadFullEntitiesAll( @Nonnull PolarisCallContext callCtx, @Nullable List catalogPath, @Nonnull PolarisEntityType entityType, @Nonnull PolarisEntitySubType entitySubType) { - return loadEntities(callCtx, catalogPath, entityType, entitySubType, PageToken.readEverything()) + return listFullEntities( + callCtx, catalogPath, entityType, entitySubType, PageToken.readEverything()) .items(); } @@ -346,6 +348,21 @@ EntityResult loadEntity( long entityId, @Nonnull PolarisEntityType entityType); + /** + * Load a batch of entities given their {@link EntityNameLookupRecord}. Will return an empty list + * if the input list is empty. Order in that returned list is the same as the input list. Some + * elements might be NULL if the entity has been dropped. + * + * @param callCtx call context + * @param entityLookupRecords the list of entity lookup records to load + * @return a non-null list of entities corresponding to the lookup keys. Some elements might be + * NULL if the entity has been dropped. + */ + @Nonnull + EntitiesResult loadEntities( + @Nonnull PolarisCallContext callCtx, + @Nonnull List entityLookupRecords); + /** * Fetch a list of tasks to be completed. Tasks * 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 8729558935..bf16c7cb8a 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 @@ -28,6 +28,7 @@ import java.util.Set; import org.apache.polaris.core.PolarisCallContext; import org.apache.polaris.core.PolarisDiagnostics; +import org.apache.polaris.core.entity.EntityNameLookupRecord; import org.apache.polaris.core.entity.LocationBasedEntity; import org.apache.polaris.core.entity.PolarisBaseEntity; import org.apache.polaris.core.entity.PolarisEntity; @@ -133,7 +134,7 @@ public EntityResult readEntityByName( } @Override - public @Nonnull Page loadEntities( + public @Nonnull Page listFullEntities( @Nonnull PolarisCallContext callCtx, @Nullable List catalogPath, @Nonnull PolarisEntityType entityType, @@ -330,6 +331,15 @@ public EntityResult loadEntity( return null; } + @Nonnull + @Override + public EntitiesResult loadEntities( + @Nonnull PolarisCallContext callCtx, + @Nonnull List entityLookupRecords) { + diagnostics.fail("illegal_method_in_transaction_workspace", "loadEntities"); + return null; + } + @Override public EntitiesResult loadTasks( @Nonnull PolarisCallContext callCtx, String executorId, PageToken pageToken) { 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 902e2c8d6a..4eaf40fe57 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 @@ -389,7 +389,7 @@ public Page listEntities( /** {@inheritDoc} */ @Override @Nonnull - public Page loadEntities( + public Page listFullEntities( @Nonnull PolarisCallContext callCtx, long catalogId, long parentId, diff --git a/polaris-core/src/main/java/org/apache/polaris/core/persistence/transactional/TransactionalMetaStoreManagerImpl.java b/polaris-core/src/main/java/org/apache/polaris/core/persistence/transactional/TransactionalMetaStoreManagerImpl.java index 402cdc280e..2256a5614d 100644 --- a/polaris-core/src/main/java/org/apache/polaris/core/persistence/transactional/TransactionalMetaStoreManagerImpl.java +++ b/polaris-core/src/main/java/org/apache/polaris/core/persistence/transactional/TransactionalMetaStoreManagerImpl.java @@ -31,6 +31,7 @@ import java.util.Set; import java.util.function.Function; import java.util.stream.Collectors; +import java.util.stream.IntStream; import org.apache.polaris.core.PolarisCallContext; import org.apache.polaris.core.PolarisDiagnostics; import org.apache.polaris.core.config.FeatureConfiguration; @@ -723,10 +724,10 @@ private void bootstrapPolarisService( } /** - * See {@link PolarisMetaStoreManager#loadEntities(PolarisCallContext, List, PolarisEntityType, - * PolarisEntitySubType, PageToken)} + * See {@link PolarisMetaStoreManager#listFullEntities(PolarisCallContext, List, + * PolarisEntityType, PolarisEntitySubType, PageToken)} */ - private @Nonnull Page loadEntities( + private @Nonnull Page listFullEntities( @Nonnull PolarisCallContext callCtx, @Nonnull TransactionalPersistence ms, @Nullable List catalogPath, @@ -756,7 +757,7 @@ private void bootstrapPolarisService( /** {@inheritDoc} */ @Override - public @Nonnull Page loadEntities( + public @Nonnull Page listFullEntities( @Nonnull PolarisCallContext callCtx, @Nullable List catalogPath, @Nonnull PolarisEntityType entityType, @@ -768,7 +769,7 @@ private void bootstrapPolarisService( // run operation in a read transaction return ms.runInReadTransaction( callCtx, - () -> loadEntities(callCtx, ms, catalogPath, entityType, entitySubType, pageToken)); + () -> listFullEntities(callCtx, ms, catalogPath, entityType, entitySubType, pageToken)); } /** {@link #createPrincipal(PolarisCallContext, PrincipalEntity)} */ @@ -2011,6 +2012,46 @@ private PolarisEntityResolver resolveSecurableToRoleGrant( () -> this.loadEntity(callCtx, ms, entityCatalogId, entityId, entityType.getCode())); } + /** Refer to {@link #loadEntity(PolarisCallContext, long, long, PolarisEntityType)} */ + private @Nonnull EntitiesResult loadEntities( + @Nonnull PolarisCallContext callCtx, + @Nonnull TransactionalPersistence ms, + @Nonnull List entityLookupRecords) { + List entities = + ms.lookupEntitiesInCurrentTxn( + callCtx, + entityLookupRecords.stream() + .map(r -> new PolarisEntityId(r.getCatalogId(), r.getId())) + .collect(Collectors.toList())); + // mimic the behavior of loadEntity above, return null if not found or type mismatch + List ret = + IntStream.range(0, entityLookupRecords.size()) + .mapToObj( + i -> { + if (entities.get(i) != null + && !entities.get(i).getType().equals(entityLookupRecords.get(i).getType())) { + return null; + } else { + return entities.get(i); + } + }) + .collect(Collectors.toList()); + return new EntitiesResult(Page.fromItems(ret)); + } + + @Nonnull + @Override + public EntitiesResult loadEntities( + @Nonnull PolarisCallContext callCtx, + @Nonnull List entityLookupRecords) { + TransactionalPersistence ms = ((TransactionalPersistence) callCtx.getMetaStore()); + return ms.runInReadTransaction( + callCtx, + () -> + this.loadEntities( + callCtx, (TransactionalPersistence) callCtx.getMetaStore(), entityLookupRecords)); + } + /** Refer to {@link #loadTasks(PolarisCallContext, String, PageToken)} */ private @Nonnull EntitiesResult loadTasks( @Nonnull PolarisCallContext callCtx, diff --git a/polaris-core/src/main/java/org/apache/polaris/core/persistence/transactional/TransactionalPersistence.java b/polaris-core/src/main/java/org/apache/polaris/core/persistence/transactional/TransactionalPersistence.java index 3802908b82..dfcd5f9256 100644 --- a/polaris-core/src/main/java/org/apache/polaris/core/persistence/transactional/TransactionalPersistence.java +++ b/polaris-core/src/main/java/org/apache/polaris/core/persistence/transactional/TransactionalPersistence.java @@ -225,7 +225,7 @@ List lookupEntityVersionsInCurrentTxn( pageToken); } - /** See {@link org.apache.polaris.core.persistence.BasePersistence#loadEntities} */ + /** See {@link org.apache.polaris.core.persistence.BasePersistence#listFullEntities} */ @Nonnull Page loadEntitiesInCurrentTxn( @Nonnull PolarisCallContext callCtx, diff --git a/polaris-core/src/testFixtures/java/org/apache/polaris/core/persistence/BasePolarisMetaStoreManagerTest.java b/polaris-core/src/testFixtures/java/org/apache/polaris/core/persistence/BasePolarisMetaStoreManagerTest.java index 43a0c55815..21122bda61 100644 --- a/polaris-core/src/testFixtures/java/org/apache/polaris/core/persistence/BasePolarisMetaStoreManagerTest.java +++ b/polaris-core/src/testFixtures/java/org/apache/polaris/core/persistence/BasePolarisMetaStoreManagerTest.java @@ -118,7 +118,7 @@ protected void testCreateEntities() { .containsExactly(PolarisEntity.toCore(task1), PolarisEntity.toCore(task2)); List listedEntities = - metaStoreManager.loadEntitiesAll( + metaStoreManager.loadFullEntitiesAll( polarisTestMetaStoreManager.polarisCallContext, null, PolarisEntityType.TASK, @@ -238,6 +238,12 @@ protected void testLookup() { polarisTestMetaStoreManager.testLookup(); } + /** test batch entity load */ + @Test + protected void testBatchLoad() { + polarisTestMetaStoreManager.testBatchLoad(); + } + /** Test the set of functions for the entity cache */ @Test protected void testEntityCache() { 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 c8e486ee6e..d5660d8d00 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 @@ -44,6 +44,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.EntitiesResult; import org.apache.polaris.core.persistence.dao.entity.EntityResult; import org.apache.polaris.core.persistence.dao.entity.LoadGrantsResult; import org.apache.polaris.core.persistence.dao.entity.LoadPolicyMappingsResult; @@ -55,6 +56,7 @@ import org.apache.polaris.core.policy.PolicyType; import org.apache.polaris.core.policy.PredefinedPolicyTypes; import org.assertj.core.api.Assertions; +import org.assertj.core.api.InstanceOfAssertFactories; /** Test the Polaris persistence layer */ public class PolarisTestMetaStoreManager { @@ -650,6 +652,7 @@ PolarisBaseEntity createEntity( .parentId(parentId) .name(name) .propertiesAsMap(properties) + .internalPropertiesAsMap(Map.of()) .build(); PolarisBaseEntity entity = polarisMetaStoreManager @@ -2685,6 +2688,143 @@ public void testLookup() { this.ensureNotExistsById(catalog.getId(), T1.getId(), PolarisEntityType.NAMESPACE); } + public void testBatchLoad() { + // load all principals + List principals = + polarisMetaStoreManager + .listEntities( + this.polarisCallContext, + null, + PolarisEntityType.PRINCIPAL, + PolarisEntitySubType.NULL_SUBTYPE, + PageToken.readEverything()) + .getEntities(); + + // create new catalog + PolarisBaseEntity catalog = + new PolarisBaseEntity( + PolarisEntityConstants.getNullId(), + polarisMetaStoreManager.generateNewEntityId(this.polarisCallContext).getId(), + PolarisEntityType.CATALOG, + PolarisEntitySubType.NULL_SUBTYPE, + PolarisEntityConstants.getRootEntityId(), + "test"); + CreateCatalogResult catalogCreated = + polarisMetaStoreManager.createCatalog(this.polarisCallContext, catalog, List.of()); + Assertions.assertThat(catalogCreated).isNotNull(); + + // load the catalog again, since the grant versions are different + catalog = + polarisMetaStoreManager + .loadEntity( + polarisCallContext, + 0L, + catalogCreated.getCatalog().getId(), + PolarisEntityType.CATALOG) + .getEntity(); + + // now create all objects + PolarisBaseEntity N1 = this.createEntity(List.of(catalog), PolarisEntityType.NAMESPACE, "N1"); + PolarisBaseEntity N1_N2 = + this.createEntity(List.of(catalog, N1), PolarisEntityType.NAMESPACE, "N2"); + PolarisBaseEntity T1 = + this.createEntity( + List.of(catalog, N1, N1_N2), + PolarisEntityType.TABLE_LIKE, + PolarisEntitySubType.ICEBERG_TABLE, + "T1"); + + // batch load all entities. They should all be present and non-null + EntitiesResult entitiesResult = + polarisMetaStoreManager.loadEntities( + polarisCallContext, + List.of( + new EntityNameLookupRecord(catalog), + new EntityNameLookupRecord(N1), + new EntityNameLookupRecord(N1_N2), + new EntityNameLookupRecord(T1))); + Assertions.assertThat(entitiesResult) + .isNotNull() + .returns(BaseResult.ReturnStatus.SUCCESS, EntitiesResult::getReturnStatus) + .extracting( + EntitiesResult::getEntities, InstanceOfAssertFactories.list(PolarisBaseEntity.class)) + .hasSize(4) + .allSatisfy(entity -> Assertions.assertThat(entity).isNotNull()) + .containsExactly(catalog, N1, N1_N2, T1); + + // try entities which do not exist + entitiesResult = + polarisMetaStoreManager.loadEntities( + polarisCallContext, + List.of( + new EntityNameLookupRecord( + catalog.getId(), + 27, + 0L, + "CATALOG_DOES_NOT_EXIST", + PolarisEntityType.CATALOG.getCode(), + PolarisEntitySubType.NULL_SUBTYPE.getCode()), + new EntityNameLookupRecord( + catalog.getId(), + 35, + 0L, + "PRINCIPAL_DOES_NOT_EXIST", + PolarisEntityType.PRINCIPAL.getCode(), + PolarisEntitySubType.NULL_SUBTYPE.getCode()))); + Assertions.assertThat(entitiesResult) + .isNotNull() + .returns(BaseResult.ReturnStatus.SUCCESS, EntitiesResult::getReturnStatus) + .extracting( + EntitiesResult::getEntities, InstanceOfAssertFactories.list(PolarisBaseEntity.class)) + .hasSize(2) + .allSatisfy(entity -> Assertions.assertThat(entity).isNull()); + + // mix of existing entities and entities with wrong type + entitiesResult = + polarisMetaStoreManager.loadEntities( + polarisCallContext, + List.of( + new EntityNameLookupRecord( + catalog.getId(), + 27, + 0L, + "CATALOG_DOES_NOT_EXIST", + PolarisEntityType.CATALOG.getCode(), + PolarisEntitySubType.NULL_SUBTYPE.getCode()), + new EntityNameLookupRecord( + catalog.getId(), + 35, + 0L, + "PRINCIPAL_DOES_NOT_EXIST", + PolarisEntityType.PRINCIPAL.getCode(), + PolarisEntitySubType.NULL_SUBTYPE.getCode()), + new EntityNameLookupRecord(catalog), + new EntityNameLookupRecord( + N1.getCatalogId(), + N1.getId(), + N1.getParentId(), + N1.getName(), + PolarisEntityType.CATALOG_ROLE.getCode(), + PolarisEntitySubType.NULL_SUBTYPE.getCode()), + new EntityNameLookupRecord( + N1_N2.getCatalogId(), + N1_N2.getId(), + N1_N2.getParentId(), + N1_N2.getName(), + PolarisEntityType.TABLE_LIKE.getCode(), + PolarisEntitySubType.ANY_SUBTYPE.getCode()), + new EntityNameLookupRecord(T1))); + Assertions.assertThat(entitiesResult) + .isNotNull() + .returns(BaseResult.ReturnStatus.SUCCESS, EntitiesResult::getReturnStatus) + .extracting( + EntitiesResult::getEntities, InstanceOfAssertFactories.list(PolarisBaseEntity.class)) + .hasSize(6) + .filteredOn(e -> e != null) + .hasSize(2) + .containsExactly(catalog, T1); + } + /** Test the set of functions for the entity cache */ public void testEntityCache() { // create test catalog diff --git a/runtime/service/src/main/java/org/apache/polaris/service/admin/PolarisAdminService.java b/runtime/service/src/main/java/org/apache/polaris/service/admin/PolarisAdminService.java index 3aaeafbaf1..fcc8cd97d6 100644 --- a/runtime/service/src/main/java/org/apache/polaris/service/admin/PolarisAdminService.java +++ b/runtime/service/src/main/java/org/apache/polaris/service/admin/PolarisAdminService.java @@ -967,7 +967,7 @@ public List listCatalogs() { /** List all catalogs without checking for permission. */ private Stream listCatalogsUnsafe() { return metaStoreManager - .loadEntitiesAll( + .loadFullEntitiesAll( getCurrentPolarisContext(), null, PolarisEntityType.CATALOG, @@ -1214,7 +1214,7 @@ public List listPrincipals() { authorizeBasicRootOperationOrThrow(op); return metaStoreManager - .loadEntitiesAll( + .loadFullEntitiesAll( getCurrentPolarisContext(), null, PolarisEntityType.PRINCIPAL, @@ -1321,7 +1321,7 @@ public List listPrincipalRoles() { authorizeBasicRootOperationOrThrow(op); return metaStoreManager - .loadEntitiesAll( + .loadFullEntitiesAll( getCurrentPolarisContext(), null, PolarisEntityType.PRINCIPAL_ROLE, @@ -1445,7 +1445,7 @@ public List listCatalogRoles(String catalogName) { CatalogEntity catalogEntity = getCatalogByName(resolutionManifest, catalogName); List catalogPath = PolarisEntity.toCoreList(List.of(catalogEntity)); return metaStoreManager - .loadEntitiesAll( + .loadFullEntitiesAll( getCurrentPolarisContext(), catalogPath, PolarisEntityType.CATALOG_ROLE, diff --git a/runtime/service/src/main/java/org/apache/polaris/service/catalog/policy/PolicyCatalog.java b/runtime/service/src/main/java/org/apache/polaris/service/catalog/policy/PolicyCatalog.java index 8e4063b87b..451eba049d 100644 --- a/runtime/service/src/main/java/org/apache/polaris/service/catalog/policy/PolicyCatalog.java +++ b/runtime/service/src/main/java/org/apache/polaris/service/catalog/policy/PolicyCatalog.java @@ -188,7 +188,7 @@ public List listPolicies(Namespace namespace, @Nullable Policy } // with a policyType filter we need to load the full PolicyEntity to apply the filter return metaStoreManager - .loadEntitiesAll( + .loadFullEntitiesAll( callContext.getPolarisCallContext(), catalogPath, PolarisEntityType.POLICY, From 5981dbd0f4bda3d2ec0b08618d71c53b33f90843 Mon Sep 17 00:00:00 2001 From: Michael Collado Date: Fri, 5 Sep 2025 16:27:39 -0700 Subject: [PATCH 2/9] Changed batch call to implement loadResolvedEntities instead --- .../AtomicOperationMetaStoreManager.java | 73 ++++++++----- .../persistence/PolarisMetaStoreManager.java | 31 +++--- .../TransactionWorkspaceMetaStoreManager.java | 19 ++-- .../dao/entity/ResolvedEntitiesResult.java | 44 ++++++++ .../TransactionalMetaStoreManagerImpl.java | 100 +++++++++++------- .../PolarisTestMetaStoreManager.java | 47 ++++++-- 6 files changed, 211 insertions(+), 103 deletions(-) create mode 100644 polaris-core/src/main/java/org/apache/polaris/core/persistence/dao/entity/ResolvedEntitiesResult.java diff --git a/polaris-core/src/main/java/org/apache/polaris/core/persistence/AtomicOperationMetaStoreManager.java b/polaris-core/src/main/java/org/apache/polaris/core/persistence/AtomicOperationMetaStoreManager.java index 684dd17c7f..ee692657e7 100644 --- a/polaris-core/src/main/java/org/apache/polaris/core/persistence/AtomicOperationMetaStoreManager.java +++ b/polaris-core/src/main/java/org/apache/polaris/core/persistence/AtomicOperationMetaStoreManager.java @@ -68,6 +68,7 @@ import org.apache.polaris.core.persistence.dao.entity.PolicyAttachmentResult; import org.apache.polaris.core.persistence.dao.entity.PrincipalSecretsResult; import org.apache.polaris.core.persistence.dao.entity.PrivilegeResult; +import org.apache.polaris.core.persistence.dao.entity.ResolvedEntitiesResult; import org.apache.polaris.core.persistence.dao.entity.ResolvedEntityResult; import org.apache.polaris.core.persistence.dao.entity.ScopedCredentialsResult; import org.apache.polaris.core.persistence.pagination.Page; @@ -1514,34 +1515,6 @@ private void revokeGrantRecord( : new EntityResult(BaseResult.ReturnStatus.ENTITY_NOT_FOUND, null); } - @Nonnull - @Override - public EntitiesResult loadEntities( - @Nonnull PolarisCallContext callCtx, - @Nonnull List entityLookupRecords) { - BasePersistence ms = callCtx.getMetaStore(); - List entities = - ms.lookupEntities( - callCtx, - entityLookupRecords.stream() - .map(r -> new PolarisEntityId(r.getCatalogId(), r.getId())) - .collect(Collectors.toList())); - // mimic the behavior of loadEntity above, return null if not found or type mismatch - List ret = - IntStream.range(0, entityLookupRecords.size()) - .mapToObj( - i -> { - if (entities.get(i) != null - && !entities.get(i).getType().equals(entityLookupRecords.get(i).getType())) { - return null; - } else { - return entities.get(i); - } - }) - .collect(Collectors.toList()); - return new EntitiesResult(Page.fromItems(ret)); - } - @Override public @Nonnull EntitiesResult loadTasks( @Nonnull PolarisCallContext callCtx, String executorId, PageToken pageToken) { @@ -1789,6 +1762,50 @@ public EntitiesResult loadEntities( return result; } + @Nonnull + @Override + public ResolvedEntitiesResult loadResolvedEntities( + @Nonnull PolarisCallContext callCtx, + @Nonnull List entityLookupRecords) { + BasePersistence ms = callCtx.getMetaStore(); + List entities = + ms.lookupEntities( + callCtx, + entityLookupRecords.stream() + .map(r -> new PolarisEntityId(r.getCatalogId(), r.getId())) + .collect(Collectors.toList())); + // mimic the behavior of loadEntity above, return null if not found or type mismatch + List ret = + IntStream.range(0, entityLookupRecords.size()) + .mapToObj( + i -> { + if (entities.get(i) != null + && !entities.get(i).getType().equals(entityLookupRecords.get(i).getType())) { + return null; + } else { + return entities.get(i); + } + }) + .map( + e -> { + if (e == null) { + return null; + } else { + // load the grant records + final List grantRecordsAsSecurable = + ms.loadAllGrantRecordsOnSecurable(callCtx, e.getCatalogId(), e.getId()); + final List grantRecordsAsGrantee = + e.getType().isGrantee() + ? ms.loadAllGrantRecordsOnGrantee(callCtx, e.getCatalogId(), e.getId()) + : List.of(); + return new ResolvedPolarisEntity( + PolarisEntity.of(e), grantRecordsAsGrantee, grantRecordsAsSecurable); + } + }) + .collect(Collectors.toList()); + return new ResolvedEntitiesResult(ret); + } + /** {@inheritDoc} */ @Override public @Nonnull ResolvedEntityResult refreshResolvedEntity( 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 e65d36d566..d1154e1413 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 @@ -48,6 +48,7 @@ import org.apache.polaris.core.persistence.dao.entity.EntityWithPath; import org.apache.polaris.core.persistence.dao.entity.GenerateEntityIdResult; import org.apache.polaris.core.persistence.dao.entity.ListEntitiesResult; +import org.apache.polaris.core.persistence.dao.entity.ResolvedEntitiesResult; import org.apache.polaris.core.persistence.dao.entity.ResolvedEntityResult; import org.apache.polaris.core.persistence.pagination.Page; import org.apache.polaris.core.persistence.pagination.PageToken; @@ -348,21 +349,6 @@ EntityResult loadEntity( long entityId, @Nonnull PolarisEntityType entityType); - /** - * Load a batch of entities given their {@link EntityNameLookupRecord}. Will return an empty list - * if the input list is empty. Order in that returned list is the same as the input list. Some - * elements might be NULL if the entity has been dropped. - * - * @param callCtx call context - * @param entityLookupRecords the list of entity lookup records to load - * @return a non-null list of entities corresponding to the lookup keys. Some elements might be - * NULL if the entity has been dropped. - */ - @Nonnull - EntitiesResult loadEntities( - @Nonnull PolarisCallContext callCtx, - @Nonnull List entityLookupRecords); - /** * Fetch a list of tasks to be completed. Tasks * @@ -433,6 +419,21 @@ ResolvedEntityResult loadResolvedEntityByName( @Nonnull PolarisEntityType entityType, @Nonnull String entityName); + /** + * Load a batch of resolved entities given their {@link EntityNameLookupRecord}. Will return an + * empty list if the input list is empty. Order in that returned list is the same as the input + * list. Some elements might be NULL if the entity has been dropped. + * + * @param callCtx call context + * @param entityLookupRecords the list of entity lookup records to load + * @return a non-null list of entities corresponding to the lookup keys. Some elements might be + * NULL if the entity has been dropped. + */ + @Nonnull + ResolvedEntitiesResult loadResolvedEntities( + @Nonnull PolarisCallContext callCtx, + @Nonnull List entityLookupRecords); + /** * Refresh a resolved entity from the backend store. Will return NULL if the entity does not * exist, i.e. has been purged or dropped. Else, will determine what has changed based on 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 bf16c7cb8a..8914a77d96 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 @@ -54,6 +54,7 @@ import org.apache.polaris.core.persistence.dao.entity.PolicyAttachmentResult; import org.apache.polaris.core.persistence.dao.entity.PrincipalSecretsResult; import org.apache.polaris.core.persistence.dao.entity.PrivilegeResult; +import org.apache.polaris.core.persistence.dao.entity.ResolvedEntitiesResult; import org.apache.polaris.core.persistence.dao.entity.ResolvedEntityResult; import org.apache.polaris.core.persistence.dao.entity.ScopedCredentialsResult; import org.apache.polaris.core.persistence.pagination.Page; @@ -331,15 +332,6 @@ public EntityResult loadEntity( return null; } - @Nonnull - @Override - public EntitiesResult loadEntities( - @Nonnull PolarisCallContext callCtx, - @Nonnull List entityLookupRecords) { - diagnostics.fail("illegal_method_in_transaction_workspace", "loadEntities"); - return null; - } - @Override public EntitiesResult loadTasks( @Nonnull PolarisCallContext callCtx, String executorId, PageToken pageToken) { @@ -389,6 +381,15 @@ public ResolvedEntityResult loadResolvedEntityByName( return null; } + @Nonnull + @Override + public ResolvedEntitiesResult loadResolvedEntities( + @Nonnull PolarisCallContext callCtx, + @Nonnull List entityLookupRecords) { + diagnostics.fail("illegal_method_in_transaction_workspace", "loadEntities"); + return null; + } + @Override public ResolvedEntityResult refreshResolvedEntity( @Nonnull PolarisCallContext callCtx, diff --git a/polaris-core/src/main/java/org/apache/polaris/core/persistence/dao/entity/ResolvedEntitiesResult.java b/polaris-core/src/main/java/org/apache/polaris/core/persistence/dao/entity/ResolvedEntitiesResult.java new file mode 100644 index 0000000000..aec18414e8 --- /dev/null +++ b/polaris-core/src/main/java/org/apache/polaris/core/persistence/dao/entity/ResolvedEntitiesResult.java @@ -0,0 +1,44 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.polaris.core.persistence.dao.entity; + +import jakarta.annotation.Nonnull; +import jakarta.annotation.Nullable; +import java.util.List; +import org.apache.polaris.core.persistence.ResolvedPolarisEntity; + +public class ResolvedEntitiesResult extends BaseResult { + private final List resolvedEntities; + + public ResolvedEntitiesResult(List resolvedEntities) { + super(ReturnStatus.SUCCESS, null); + this.resolvedEntities = resolvedEntities; + } + + public ResolvedEntitiesResult( + @Nonnull ReturnStatus returnStatus, @Nullable String extraInformation) { + super(returnStatus, extraInformation); + this.resolvedEntities = null; + } + + public List getResolvedEntities() { + return resolvedEntities; + } +} diff --git a/polaris-core/src/main/java/org/apache/polaris/core/persistence/transactional/TransactionalMetaStoreManagerImpl.java b/polaris-core/src/main/java/org/apache/polaris/core/persistence/transactional/TransactionalMetaStoreManagerImpl.java index 2256a5614d..11b5b5dd71 100644 --- a/polaris-core/src/main/java/org/apache/polaris/core/persistence/transactional/TransactionalMetaStoreManagerImpl.java +++ b/polaris-core/src/main/java/org/apache/polaris/core/persistence/transactional/TransactionalMetaStoreManagerImpl.java @@ -57,6 +57,7 @@ import org.apache.polaris.core.persistence.PolarisMetaStoreManager; import org.apache.polaris.core.persistence.PolarisObjectMapperUtil; import org.apache.polaris.core.persistence.PolicyMappingAlreadyExistsException; +import org.apache.polaris.core.persistence.ResolvedPolarisEntity; import org.apache.polaris.core.persistence.RetryOnConcurrencyException; import org.apache.polaris.core.persistence.dao.entity.BaseResult; import org.apache.polaris.core.persistence.dao.entity.ChangeTrackingResult; @@ -72,6 +73,7 @@ import org.apache.polaris.core.persistence.dao.entity.PolicyAttachmentResult; import org.apache.polaris.core.persistence.dao.entity.PrincipalSecretsResult; import org.apache.polaris.core.persistence.dao.entity.PrivilegeResult; +import org.apache.polaris.core.persistence.dao.entity.ResolvedEntitiesResult; import org.apache.polaris.core.persistence.dao.entity.ResolvedEntityResult; import org.apache.polaris.core.persistence.dao.entity.ScopedCredentialsResult; import org.apache.polaris.core.persistence.pagination.Page; @@ -2012,46 +2014,6 @@ private PolarisEntityResolver resolveSecurableToRoleGrant( () -> this.loadEntity(callCtx, ms, entityCatalogId, entityId, entityType.getCode())); } - /** Refer to {@link #loadEntity(PolarisCallContext, long, long, PolarisEntityType)} */ - private @Nonnull EntitiesResult loadEntities( - @Nonnull PolarisCallContext callCtx, - @Nonnull TransactionalPersistence ms, - @Nonnull List entityLookupRecords) { - List entities = - ms.lookupEntitiesInCurrentTxn( - callCtx, - entityLookupRecords.stream() - .map(r -> new PolarisEntityId(r.getCatalogId(), r.getId())) - .collect(Collectors.toList())); - // mimic the behavior of loadEntity above, return null if not found or type mismatch - List ret = - IntStream.range(0, entityLookupRecords.size()) - .mapToObj( - i -> { - if (entities.get(i) != null - && !entities.get(i).getType().equals(entityLookupRecords.get(i).getType())) { - return null; - } else { - return entities.get(i); - } - }) - .collect(Collectors.toList()); - return new EntitiesResult(Page.fromItems(ret)); - } - - @Nonnull - @Override - public EntitiesResult loadEntities( - @Nonnull PolarisCallContext callCtx, - @Nonnull List entityLookupRecords) { - TransactionalPersistence ms = ((TransactionalPersistence) callCtx.getMetaStore()); - return ms.runInReadTransaction( - callCtx, - () -> - this.loadEntities( - callCtx, (TransactionalPersistence) callCtx.getMetaStore(), entityLookupRecords)); - } - /** Refer to {@link #loadTasks(PolarisCallContext, String, PageToken)} */ private @Nonnull EntitiesResult loadTasks( @Nonnull PolarisCallContext callCtx, @@ -2334,6 +2296,64 @@ public EntitiesResult loadEntities( return result; } + /** Refer to {@link #loadEntity(PolarisCallContext, long, long, PolarisEntityType)} */ + private @Nonnull ResolvedEntitiesResult loadResolvedEntities( + @Nonnull PolarisCallContext callCtx, + @Nonnull TransactionalPersistence ms, + @Nonnull List entityLookupRecords) { + List entities = + ms.lookupEntitiesInCurrentTxn( + callCtx, + entityLookupRecords.stream() + .map(r -> new PolarisEntityId(r.getCatalogId(), r.getId())) + .collect(Collectors.toList())); + // mimic the behavior of loadEntity above, return null if not found or type mismatch + List ret = + IntStream.range(0, entityLookupRecords.size()) + .mapToObj( + i -> { + if (entities.get(i) != null + && !entities.get(i).getType().equals(entityLookupRecords.get(i).getType())) { + return null; + } else { + return entities.get(i); + } + }) + .map( + e -> { + if (e == null) { + return null; + } else { + // load the grant records + final List grantRecordsAsSecurable = + ms.loadAllGrantRecordsOnSecurableInCurrentTxn( + callCtx, e.getCatalogId(), e.getId()); + final List grantRecordsAsGrantee = + e.getType().isGrantee() + ? ms.loadAllGrantRecordsOnGranteeInCurrentTxn( + callCtx, e.getCatalogId(), e.getId()) + : List.of(); + return new ResolvedPolarisEntity( + PolarisEntity.of(e), grantRecordsAsGrantee, grantRecordsAsSecurable); + } + }) + .collect(Collectors.toList()); + return new ResolvedEntitiesResult(ret); + } + + @Nonnull + @Override + public ResolvedEntitiesResult loadResolvedEntities( + @Nonnull PolarisCallContext callCtx, + @Nonnull List entityLookupRecords) { + TransactionalPersistence ms = ((TransactionalPersistence) callCtx.getMetaStore()); + return ms.runInReadTransaction( + callCtx, + () -> + this.loadResolvedEntities( + callCtx, (TransactionalPersistence) callCtx.getMetaStore(), entityLookupRecords)); + } + /** {@inheritDoc} */ private @Nonnull ResolvedEntityResult refreshResolvedEntity( @Nonnull PolarisCallContext callCtx, 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 d5660d8d00..f6062d06fb 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 @@ -44,11 +44,11 @@ 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.EntitiesResult; import org.apache.polaris.core.persistence.dao.entity.EntityResult; import org.apache.polaris.core.persistence.dao.entity.LoadGrantsResult; import org.apache.polaris.core.persistence.dao.entity.LoadPolicyMappingsResult; import org.apache.polaris.core.persistence.dao.entity.PolicyAttachmentResult; +import org.apache.polaris.core.persistence.dao.entity.ResolvedEntitiesResult; import org.apache.polaris.core.persistence.dao.entity.ResolvedEntityResult; import org.apache.polaris.core.persistence.pagination.PageToken; import org.apache.polaris.core.policy.PolarisPolicyMappingRecord; @@ -2735,8 +2735,8 @@ public void testBatchLoad() { "T1"); // batch load all entities. They should all be present and non-null - EntitiesResult entitiesResult = - polarisMetaStoreManager.loadEntities( + ResolvedEntitiesResult entitiesResult = + polarisMetaStoreManager.loadResolvedEntities( polarisCallContext, List.of( new EntityNameLookupRecord(catalog), @@ -2745,16 +2745,38 @@ public void testBatchLoad() { new EntityNameLookupRecord(T1))); Assertions.assertThat(entitiesResult) .isNotNull() - .returns(BaseResult.ReturnStatus.SUCCESS, EntitiesResult::getReturnStatus) + .returns(BaseResult.ReturnStatus.SUCCESS, ResolvedEntitiesResult::getReturnStatus) .extracting( - EntitiesResult::getEntities, InstanceOfAssertFactories.list(PolarisBaseEntity.class)) + ResolvedEntitiesResult::getResolvedEntities, + InstanceOfAssertFactories.list(ResolvedPolarisEntity.class)) .hasSize(4) .allSatisfy(entity -> Assertions.assertThat(entity).isNotNull()) + .extracting(r -> (PolarisBaseEntity) r.getEntity()) .containsExactly(catalog, N1, N1_N2, T1); + ResolvedPolarisEntity catalogEntity = entitiesResult.getResolvedEntities().get(0); + Assertions.assertThat(catalogEntity) + .extracting(ResolvedPolarisEntity::getAllGrantRecords) + .isNotNull() + .asInstanceOf(InstanceOfAssertFactories.list(PolarisGrantRecord.class)) + .hasSize(2) + .containsExactlyInAnyOrder( + new PolarisGrantRecord( + 0L, + catalog.getId(), + catalogCreated.getCatalogAdminRole().getCatalogId(), + catalogCreated.getCatalogAdminRole().getId(), + PolarisPrivilege.CATALOG_MANAGE_ACCESS.getCode()), + new PolarisGrantRecord( + 0L, + catalog.getId(), + catalogCreated.getCatalogAdminRole().getCatalogId(), + catalogCreated.getCatalogAdminRole().getId(), + PolarisPrivilege.CATALOG_MANAGE_METADATA.getCode())); + // try entities which do not exist entitiesResult = - polarisMetaStoreManager.loadEntities( + polarisMetaStoreManager.loadResolvedEntities( polarisCallContext, List.of( new EntityNameLookupRecord( @@ -2773,15 +2795,16 @@ public void testBatchLoad() { PolarisEntitySubType.NULL_SUBTYPE.getCode()))); Assertions.assertThat(entitiesResult) .isNotNull() - .returns(BaseResult.ReturnStatus.SUCCESS, EntitiesResult::getReturnStatus) + .returns(BaseResult.ReturnStatus.SUCCESS, ResolvedEntitiesResult::getReturnStatus) .extracting( - EntitiesResult::getEntities, InstanceOfAssertFactories.list(PolarisBaseEntity.class)) + ResolvedEntitiesResult::getResolvedEntities, + InstanceOfAssertFactories.list(ResolvedPolarisEntity.class)) .hasSize(2) .allSatisfy(entity -> Assertions.assertThat(entity).isNull()); // mix of existing entities and entities with wrong type entitiesResult = - polarisMetaStoreManager.loadEntities( + polarisMetaStoreManager.loadResolvedEntities( polarisCallContext, List.of( new EntityNameLookupRecord( @@ -2816,12 +2839,14 @@ public void testBatchLoad() { new EntityNameLookupRecord(T1))); Assertions.assertThat(entitiesResult) .isNotNull() - .returns(BaseResult.ReturnStatus.SUCCESS, EntitiesResult::getReturnStatus) + .returns(BaseResult.ReturnStatus.SUCCESS, ResolvedEntitiesResult::getReturnStatus) .extracting( - EntitiesResult::getEntities, InstanceOfAssertFactories.list(PolarisBaseEntity.class)) + ResolvedEntitiesResult::getResolvedEntities, + InstanceOfAssertFactories.list(ResolvedPolarisEntity.class)) .hasSize(6) .filteredOn(e -> e != null) .hasSize(2) + .extracting(r -> (PolarisBaseEntity) r.getEntity()) .containsExactly(catalog, T1); } From e5bed8274b9fbc1452bc173650ad0d4e23d198d2 Mon Sep 17 00:00:00 2001 From: Michael Collado Date: Tue, 23 Sep 2025 13:25:42 -0700 Subject: [PATCH 3/9] Add loadResolvedEntities by id and entity cache support --- .../AtomicOperationMetaStoreManager.java | 51 +- .../persistence/PolarisMetaStoreManager.java | 21 +- .../TransactionWorkspaceMetaStoreManager.java | 10 + .../core/persistence/cache/EntityCache.java | 35 + .../cache/InMemoryEntityCache.java | 174 +++- .../TransactionalMetaStoreManagerImpl.java | 43 +- .../cache/InMemoryEntityCacheTest.java | 845 +++++++++++++++--- .../BasePolarisMetaStoreManagerTest.java | 2 +- .../service/admin/PolarisAdminService.java | 8 +- .../service/catalog/policy/PolicyCatalog.java | 2 +- 10 files changed, 1016 insertions(+), 175 deletions(-) diff --git a/polaris-core/src/main/java/org/apache/polaris/core/persistence/AtomicOperationMetaStoreManager.java b/polaris-core/src/main/java/org/apache/polaris/core/persistence/AtomicOperationMetaStoreManager.java index ee692657e7..6e52902d49 100644 --- a/polaris-core/src/main/java/org/apache/polaris/core/persistence/AtomicOperationMetaStoreManager.java +++ b/polaris-core/src/main/java/org/apache/polaris/core/persistence/AtomicOperationMetaStoreManager.java @@ -1786,26 +1786,53 @@ public ResolvedEntitiesResult loadResolvedEntities( return entities.get(i); } }) - .map( - e -> { - if (e == null) { + .map(e -> toResolvedPolarisEntity(callCtx, e, ms)) + .collect(Collectors.toList()); + return new ResolvedEntitiesResult(ret); + } + + @Nonnull + @Override + public ResolvedEntitiesResult loadResolvedEntities( + @Nonnull PolarisCallContext callCtx, + @Nonnull PolarisEntityType entityType, + @Nonnull List entityIds) { + BasePersistence ms = callCtx.getMetaStore(); + List entities = ms.lookupEntities(callCtx, entityIds); + + // mimic the behavior of loadEntity above, return null if not found or type mismatch + List ret = + IntStream.range(0, entityIds.size()) + .mapToObj( + i -> { + if (entities.get(i) != null && !entities.get(i).getType().equals(entityType)) { return null; } else { - // load the grant records - final List grantRecordsAsSecurable = - ms.loadAllGrantRecordsOnSecurable(callCtx, e.getCatalogId(), e.getId()); - final List grantRecordsAsGrantee = - e.getType().isGrantee() - ? ms.loadAllGrantRecordsOnGrantee(callCtx, e.getCatalogId(), e.getId()) - : List.of(); - return new ResolvedPolarisEntity( - PolarisEntity.of(e), grantRecordsAsGrantee, grantRecordsAsSecurable); + return entities.get(i); } }) + .map(e -> toResolvedPolarisEntity(callCtx, e, ms)) .collect(Collectors.toList()); return new ResolvedEntitiesResult(ret); } + private static ResolvedPolarisEntity toResolvedPolarisEntity( + PolarisCallContext callCtx, PolarisBaseEntity e, BasePersistence ms) { + if (e == null) { + return null; + } else { + // load the grant records + final List grantRecordsAsSecurable = + ms.loadAllGrantRecordsOnSecurable(callCtx, e.getCatalogId(), e.getId()); + final List grantRecordsAsGrantee = + e.getType().isGrantee() + ? ms.loadAllGrantRecordsOnGrantee(callCtx, e.getCatalogId(), e.getId()) + : List.of(); + return new ResolvedPolarisEntity( + PolarisEntity.of(e), grantRecordsAsGrantee, grantRecordsAsSecurable); + } + } + /** {@inheritDoc} */ @Override public @Nonnull ResolvedEntityResult refreshResolvedEntity( 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 d1154e1413..bab8530e0e 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 @@ -137,7 +137,7 @@ ListEntitiesResult listEntities( /** * Load full entities matching the given criteria with pagination. If only the entity name/id/type * is required, use {@link #listEntities} instead. If no pagination is required, use {@link - * #loadFullEntitiesAll} instead. + * #listFullEntitiesAll} instead. * * @param callCtx call context * @param catalogPath path inside a catalog. If null or empty, the entities to list are top-level, @@ -166,7 +166,7 @@ Page listFullEntities( * @param entitySubType subType of entities to list (or ANY_SUBTYPE) * @return list of all matching entities */ - default @Nonnull List loadFullEntitiesAll( + default @Nonnull List listFullEntitiesAll( @Nonnull PolarisCallContext callCtx, @Nullable List catalogPath, @Nonnull PolarisEntityType entityType, @@ -434,6 +434,23 @@ ResolvedEntitiesResult loadResolvedEntities( @Nonnull PolarisCallContext callCtx, @Nonnull List entityLookupRecords); + /** + * Load a batch of resolved entities of a specified entity type given their {@link + * PolarisEntityId}. Will return an empty list if the input list is empty. Order in that returned + * list is the same as the input list. Some elements might be NULL if the entity has been dropped. + * + * @param callCtx call context + * @param entityType the type of entities to load + * @param entityIds the list of entity ids to load + * @return a non-null list of entities corresponding to the lookup keys. Some elements might be + * NULL if the entity has been dropped. + */ + @Nonnull + ResolvedEntitiesResult loadResolvedEntities( + @Nonnull PolarisCallContext callCtx, + @Nonnull PolarisEntityType entityType, + @Nonnull List entityIds); + /** * Refresh a resolved entity from the backend store. Will return NULL if the entity does not * exist, i.e. has been purged or dropped. Else, will determine what has changed based on 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 8914a77d96..39c1f908f7 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 @@ -390,6 +390,16 @@ public ResolvedEntitiesResult loadResolvedEntities( return null; } + @Nonnull + @Override + public ResolvedEntitiesResult loadResolvedEntities( + @Nonnull PolarisCallContext callCtx, + @Nonnull PolarisEntityType entityType, + @Nonnull List entityIds) { + diagnostics.fail("illegal_method_in_transaction_workspace", "loadResolvedEntities"); + return null; + } + @Override public ResolvedEntityResult refreshResolvedEntity( @Nonnull PolarisCallContext callCtx, 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 cd438c9950..93dc87ac97 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 @@ -20,8 +20,11 @@ import jakarta.annotation.Nonnull; import jakarta.annotation.Nullable; +import java.util.List; import org.apache.polaris.core.PolarisCallContext; +import org.apache.polaris.core.entity.EntityNameLookupRecord; import org.apache.polaris.core.entity.PolarisBaseEntity; +import org.apache.polaris.core.entity.PolarisEntityId; import org.apache.polaris.core.entity.PolarisEntityType; import org.apache.polaris.core.persistence.ResolvedPolarisEntity; @@ -80,4 +83,36 @@ EntityCacheLookupResult getOrLoadEntityById( @Nullable EntityCacheLookupResult getOrLoadEntityByName( @Nonnull PolarisCallContext callContext, @Nonnull EntityCacheByNameKey entityNameKey); + + /** + * Load multiple entities by id, returning those found in the cache and loading those not found. + * + * @param callCtx the Polaris call context + * @param entityType the entity type + * @param entityIds the list of entity ids to load + * @return the list of resolved entities, in the same order as the requested entity ids. As in + * {@link + * org.apache.polaris.core.persistence.PolarisMetaStoreManager#loadResolvedEntities(PolarisCallContext, + * PolarisEntityType, List)}, elements in the returned list may be null if the corresponding + * entity id does not exist. + */ + List getOrLoadResolvedEntities( + @Nonnull PolarisCallContext callCtx, + @Nonnull PolarisEntityType entityType, + @Nonnull List entityIds); + + /** + * Load multiple entities by {@link EntityNameLookupRecord}, returning those found in the cache + * and loading those not found. + * + * @param callCtx the Polaris call context + * @param lookupRecords the list of entity name to load + * @return the list of resolved entities, in the same order as the requested entity records. As in + * {@link + * org.apache.polaris.core.persistence.PolarisMetaStoreManager#loadResolvedEntities(PolarisCallContext, + * PolarisEntityType, List)}, elements in the returned list may be null if the corresponding + * entity id does not exist. + */ + List getOrLoadResolvedEntities( + @Nonnull PolarisCallContext callCtx, @Nonnull List lookupRecords); } diff --git a/polaris-core/src/main/java/org/apache/polaris/core/persistence/cache/InMemoryEntityCache.java b/polaris-core/src/main/java/org/apache/polaris/core/persistence/cache/InMemoryEntityCache.java index c30b996f14..a10a1f5ead 100644 --- a/polaris-core/src/main/java/org/apache/polaris/core/persistence/cache/InMemoryEntityCache.java +++ b/polaris-core/src/main/java/org/apache/polaris/core/persistence/cache/InMemoryEntityCache.java @@ -23,25 +23,41 @@ import com.github.benmanes.caffeine.cache.RemovalListener; import jakarta.annotation.Nonnull; import jakarta.annotation.Nullable; -import java.util.AbstractMap; -import java.util.List; -import java.util.concurrent.ConcurrentHashMap; -import java.util.concurrent.TimeUnit; import org.apache.polaris.core.PolarisCallContext; import org.apache.polaris.core.PolarisDiagnostics; import org.apache.polaris.core.config.BehaviorChangeConfiguration; import org.apache.polaris.core.config.FeatureConfiguration; import org.apache.polaris.core.config.RealmConfig; +import org.apache.polaris.core.entity.EntityNameLookupRecord; import org.apache.polaris.core.entity.PolarisBaseEntity; +import org.apache.polaris.core.entity.PolarisChangeTrackingVersions; +import org.apache.polaris.core.entity.PolarisEntityId; import org.apache.polaris.core.entity.PolarisEntityType; import org.apache.polaris.core.entity.PolarisGrantRecord; import org.apache.polaris.core.persistence.PolarisMetaStoreManager; import org.apache.polaris.core.persistence.ResolvedPolarisEntity; +import org.apache.polaris.core.persistence.dao.entity.ChangeTrackingResult; +import org.apache.polaris.core.persistence.dao.entity.ResolvedEntitiesResult; import org.apache.polaris.core.persistence.dao.entity.ResolvedEntityResult; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.AbstractMap; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.Iterator; +import java.util.List; +import java.util.Map; +import java.util.Objects; +import java.util.Optional; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.TimeUnit; +import java.util.function.Function; +import java.util.stream.Collectors; /** An in-memory entity cache with a limit of 100k entities and a 1h TTL. */ public class InMemoryEntityCache implements EntityCache { - + private static final Logger LOGGER = LoggerFactory.getLogger(InMemoryEntityCache.class); private final PolarisDiagnostics diagnostics; private final PolarisMetaStoreManager polarisMetaStoreManager; private final Cache byId; @@ -451,4 +467,152 @@ && isNewer(existingCacheEntry, existingCacheEntryByName)) { // return what we found return new EntityCacheLookupResult(entry, cacheHit); } + + @Override + public List getOrLoadResolvedEntities( + @Nonnull PolarisCallContext callCtx, + @Nonnull PolarisEntityType entityType, + @Nonnull List entityIds) { + // use a map to collect cached entries to avoid concurrency problems in case a second thread is + // trying to populate + // the cache from a different snapshot + Map resolvedEntities = new HashMap<>(); + for (int i = 0; i < 100; i++) { + Function, ResolvedEntitiesResult> loaderFunc = + idsToLoad -> polarisMetaStoreManager.loadResolvedEntities(callCtx, entityType, idsToLoad); + if (isCacheStateValid(callCtx, resolvedEntities, entityIds, loaderFunc)) { + break; + } + } + + return entityIds.stream() + .map( + id -> { + ResolvedPolarisEntity entity = resolvedEntities.get(id); + return entity == null ? null : new EntityCacheLookupResult(entity, true); + }) + .collect(Collectors.toList()); + } + + @Override + public List getOrLoadResolvedEntities( + @Nonnull PolarisCallContext callCtx, @Nonnull List lookupRecords) { + Map entityIdMap = + lookupRecords.stream() + .collect( + Collectors.toMap( + e -> new PolarisEntityId(e.getCatalogId(), e.getId()), + Function.identity(), + (a, b) -> a)); + Function, ResolvedEntitiesResult> loaderFunc = + idsToLoad -> + polarisMetaStoreManager.loadResolvedEntities( + callCtx, idsToLoad.stream().map(entityIdMap::get).collect(Collectors.toList())); + + // use a map to collect cached entries to avoid concurrency problems in case a second thread is + // trying to populate + // the cache from a different snapshot + Map resolvedEntities = new HashMap<>(); + List entityIds = + lookupRecords.stream() + .map(e -> new PolarisEntityId(e.getCatalogId(), e.getId())) + .collect(Collectors.toList()); + for (int i = 0; i < 100; i++) { + if (isCacheStateValid(callCtx, resolvedEntities, entityIds, loaderFunc)) { + break; + } + } + + return lookupRecords.stream() + .map( + lookupRecord -> { + ResolvedPolarisEntity entity = + resolvedEntities.get( + new PolarisEntityId(lookupRecord.getCatalogId(), lookupRecord.getId())); + return entity == null ? null : new EntityCacheLookupResult(entity, true); + }) + .collect(Collectors.toList()); + } + + private boolean isCacheStateValid( + @Nonnull PolarisCallContext callCtx, + @Nonnull Map resolvedEntities, + @Nonnull List entityIds, + @Nonnull Function, ResolvedEntitiesResult> loaderFunc) { + ChangeTrackingResult changeTrackingResult = + polarisMetaStoreManager.loadEntitiesChangeTracking(callCtx, entityIds); + List idsToLoad = new ArrayList<>(); + if (changeTrackingResult.isSuccess()) { + idsToLoad.addAll(validateCacheEntries(entityIds, resolvedEntities, changeTrackingResult)); + } else { + idsToLoad.addAll(entityIds); + } + if (!idsToLoad.isEmpty()) { + ResolvedEntitiesResult resolvedEntitiesResult = loaderFunc.apply(idsToLoad); + if (resolvedEntitiesResult.isSuccess()) { + LOGGER.debug("Resolved entities - validating cache"); + resolvedEntitiesResult.getResolvedEntities().stream() + .filter(Objects::nonNull) + .forEach( + e -> { + this.cacheNewEntry(e); + resolvedEntities.put( + new PolarisEntityId(e.getEntity().getCatalogId(), e.getEntity().getId()), e); + }); + } + } + + // the loader function should always return a batch of results from the same "snapshot" of the + // persistence, so + // if the changeTracking call above failed, we should have loaded the entire batch in one shot. + // There should be no + // need to revalidate the entities. + List idsToReload = + changeTrackingResult.isSuccess() + ? validateCacheEntries(entityIds, resolvedEntities, changeTrackingResult) + : List.of(); + return idsToReload.isEmpty(); + } + + private List validateCacheEntries( + List entityIds, + Map resolvedEntities, + ChangeTrackingResult changeTrackingResult) { + List idsToReload = new ArrayList<>(); + Iterator idIterator = entityIds.iterator(); + Iterator changeTrackingIterator = + changeTrackingResult.getChangeTrackingVersions().iterator(); + while (idIterator.hasNext() && changeTrackingIterator.hasNext()) { + PolarisEntityId entityId = idIterator.next(); + PolarisChangeTrackingVersions changeTrackingVersions = changeTrackingIterator.next(); + if (changeTrackingVersions == null) { + // entity has been purged + ResolvedPolarisEntity cachedEntity = getEntityById(entityId.getId()); + if (cachedEntity != null || resolvedEntities.containsKey(entityId)) { + LOGGER.debug("Entity {} has been purged, removing from cache", entityId); + Optional.ofNullable(cachedEntity).ifPresent(this::removeCacheEntry); + resolvedEntities.remove(entityId); + } + continue; + } + // compare versions using equals rather than less than so we can use the same function to + // validate that the cache + // entries are consistent with a single call to the change tracking table, rather than some + // grants ahead and some + // grants behind + ResolvedPolarisEntity cachedEntity = + resolvedEntities.computeIfAbsent(entityId, id -> this.getEntityById(id.getId())); + if (cachedEntity == null + || cachedEntity.getEntity().getEntityVersion() + != changeTrackingVersions.getEntityVersion() + || cachedEntity.getEntity().getGrantRecordsVersion() + != changeTrackingVersions.getGrantRecordsVersion()) { + idsToReload.add(entityId); + } else { + resolvedEntities.put(entityId, cachedEntity); + } + } + LOGGER.debug("Cache entries {} need to be reloaded", idsToReload); + return idsToReload; + } } diff --git a/polaris-core/src/main/java/org/apache/polaris/core/persistence/transactional/TransactionalMetaStoreManagerImpl.java b/polaris-core/src/main/java/org/apache/polaris/core/persistence/transactional/TransactionalMetaStoreManagerImpl.java index 11b5b5dd71..87fa9ac044 100644 --- a/polaris-core/src/main/java/org/apache/polaris/core/persistence/transactional/TransactionalMetaStoreManagerImpl.java +++ b/polaris-core/src/main/java/org/apache/polaris/core/persistence/transactional/TransactionalMetaStoreManagerImpl.java @@ -2296,24 +2296,19 @@ private PolarisEntityResolver resolveSecurableToRoleGrant( return result; } - /** Refer to {@link #loadEntity(PolarisCallContext, long, long, PolarisEntityType)} */ - private @Nonnull ResolvedEntitiesResult loadResolvedEntities( - @Nonnull PolarisCallContext callCtx, - @Nonnull TransactionalPersistence ms, - @Nonnull List entityLookupRecords) { - List entities = - ms.lookupEntitiesInCurrentTxn( - callCtx, - entityLookupRecords.stream() - .map(r -> new PolarisEntityId(r.getCatalogId(), r.getId())) - .collect(Collectors.toList())); + private static ResolvedEntitiesResult getResolvedEntitiesResult( + PolarisCallContext callCtx, + TransactionalPersistence ms, + List entityIds, + Function entityTypeForIndex) { + List entities = ms.lookupEntitiesInCurrentTxn(callCtx, entityIds); // mimic the behavior of loadEntity above, return null if not found or type mismatch List ret = - IntStream.range(0, entityLookupRecords.size()) + IntStream.range(0, entityIds.size()) .mapToObj( i -> { if (entities.get(i) != null - && !entities.get(i).getType().equals(entityLookupRecords.get(i).getType())) { + && !entities.get(i).getType().equals(entityTypeForIndex.apply(i))) { return null; } else { return entities.get(i); @@ -2347,11 +2342,25 @@ public ResolvedEntitiesResult loadResolvedEntities( @Nonnull PolarisCallContext callCtx, @Nonnull List entityLookupRecords) { TransactionalPersistence ms = ((TransactionalPersistence) callCtx.getMetaStore()); + List entityIds = + entityLookupRecords.stream() + .map(r -> new PolarisEntityId(r.getCatalogId(), r.getId())) + .collect(Collectors.toList()); + Function entityTypeForIndex = + i -> entityLookupRecords.get(i).getType(); return ms.runInReadTransaction( - callCtx, - () -> - this.loadResolvedEntities( - callCtx, (TransactionalPersistence) callCtx.getMetaStore(), entityLookupRecords)); + callCtx, () -> getResolvedEntitiesResult(callCtx, ms, entityIds, entityTypeForIndex)); + } + + @Nonnull + @Override + public ResolvedEntitiesResult loadResolvedEntities( + @Nonnull PolarisCallContext callCtx, + @Nonnull PolarisEntityType entityType, + @Nonnull List entityIds) { + TransactionalPersistence ms = ((TransactionalPersistence) callCtx.getMetaStore()); + return ms.runInReadTransaction( + callCtx, () -> getResolvedEntitiesResult(callCtx, ms, entityIds, i -> entityType)); } /** {@inheritDoc} */ diff --git a/polaris-core/src/test/java/org/apache/polaris/core/persistence/cache/InMemoryEntityCacheTest.java b/polaris-core/src/test/java/org/apache/polaris/core/persistence/cache/InMemoryEntityCacheTest.java index ce34320504..dd90f5027a 100644 --- a/polaris-core/src/test/java/org/apache/polaris/core/persistence/cache/InMemoryEntityCacheTest.java +++ b/polaris-core/src/test/java/org/apache/polaris/core/persistence/cache/InMemoryEntityCacheTest.java @@ -19,16 +19,26 @@ package org.apache.polaris.core.persistence.cache; import static org.apache.polaris.core.persistence.PrincipalSecretsGenerator.RANDOM_SECRETS; +import static org.assertj.core.api.Assertions.assertThat; import java.time.Clock; import java.util.List; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.Semaphore; +import java.util.concurrent.atomic.AtomicReference; import java.util.stream.Collectors; import org.apache.iceberg.catalog.TableIdentifier; import org.apache.polaris.core.PolarisCallContext; import org.apache.polaris.core.PolarisDefaultDiagServiceImpl; import org.apache.polaris.core.PolarisDiagnostics; +import org.apache.polaris.core.entity.EntityNameLookupRecord; import org.apache.polaris.core.entity.PolarisBaseEntity; +import org.apache.polaris.core.entity.PolarisChangeTrackingVersions; import org.apache.polaris.core.entity.PolarisEntity; +import org.apache.polaris.core.entity.PolarisEntityId; import org.apache.polaris.core.entity.PolarisEntitySubType; import org.apache.polaris.core.entity.PolarisEntityType; import org.apache.polaris.core.entity.PolarisGrantRecord; @@ -37,17 +47,21 @@ import org.apache.polaris.core.persistence.PolarisMetaStoreManager; import org.apache.polaris.core.persistence.PolarisTestMetaStoreManager; import org.apache.polaris.core.persistence.ResolvedPolarisEntity; +import org.apache.polaris.core.persistence.dao.entity.ChangeTrackingResult; +import org.apache.polaris.core.persistence.dao.entity.ResolvedEntitiesResult; import org.apache.polaris.core.persistence.transactional.TransactionalMetaStoreManagerImpl; import org.apache.polaris.core.persistence.transactional.TransactionalPersistence; import org.apache.polaris.core.persistence.transactional.TreeMapMetaStore; import org.apache.polaris.core.persistence.transactional.TreeMapTransactionalPersistenceImpl; -import org.assertj.core.api.Assertions; import org.junit.jupiter.api.Test; import org.mockito.Mockito; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; /** Unit testing of the entity cache */ public class InMemoryEntityCacheTest { + public static final Logger LOGGER = LoggerFactory.getLogger(InMemoryEntityCache.class); private final PolarisDiagnostics diagServices; private final PolarisCallContext callCtx; private final PolarisTestMetaStoreManager tm; @@ -105,31 +119,31 @@ void testGetOrLoadEntityByName() { EntityCacheLookupResult lookup = cache.getOrLoadEntityByName( this.callCtx, new EntityCacheByNameKey(PolarisEntityType.CATALOG, "test")); - Assertions.assertThat(lookup).isNotNull(); - Assertions.assertThat(lookup.isCacheHit()).isFalse(); - Assertions.assertThat(lookup.getCacheEntry()).isNotNull(); + assertThat(lookup).isNotNull(); + assertThat(lookup.isCacheHit()).isFalse(); + assertThat(lookup.getCacheEntry()).isNotNull(); // validate the cache entry PolarisBaseEntity catalog = lookup.getCacheEntry().getEntity(); - Assertions.assertThat(catalog).isNotNull(); - Assertions.assertThat(catalog.getType()).isEqualTo(PolarisEntityType.CATALOG); + assertThat(catalog).isNotNull(); + assertThat(catalog.getType()).isEqualTo(PolarisEntityType.CATALOG); // do it again, should be found in the cache lookup = cache.getOrLoadEntityByName( this.callCtx, new EntityCacheByNameKey(PolarisEntityType.CATALOG, "test")); - Assertions.assertThat(lookup).isNotNull(); - Assertions.assertThat(lookup.isCacheHit()).isTrue(); + assertThat(lookup).isNotNull(); + assertThat(lookup.isCacheHit()).isTrue(); // do it again by id, should be found in the cache 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(); - Assertions.assertThat(lookup.getCacheEntry().getEntity()).isNotNull(); - Assertions.assertThat(lookup.getCacheEntry().getGrantRecordsAsSecurable()).isNotNull(); + assertThat(lookup).isNotNull(); + assertThat(lookup.isCacheHit()).isTrue(); + assertThat(lookup.getCacheEntry()).isNotNull(); + assertThat(lookup.getCacheEntry().getEntity()).isNotNull(); + assertThat(lookup.getCacheEntry().getGrantRecordsAsSecurable()).isNotNull(); // get N1 PolarisBaseEntity N1 = @@ -140,128 +154,119 @@ void testGetOrLoadEntityByName() { new EntityCacheByNameKey( catalog.getId(), catalog.getId(), PolarisEntityType.NAMESPACE, "N1"); ResolvedPolarisEntity cacheEntry = cache.getEntityByName(N1_name); - Assertions.assertThat(cacheEntry).isNull(); + 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(), N1.getType()); - Assertions.assertThat(lookup).isNotNull(); - Assertions.assertThat(lookup.isCacheHit()).isFalse(); + assertThat(lookup).isNotNull(); + assertThat(lookup.isCacheHit()).isFalse(); // should be there now, by name cacheEntry = cache.getEntityByName(N1_name); - Assertions.assertThat(cacheEntry).isNotNull(); - Assertions.assertThat(cacheEntry.getEntity()).isNotNull(); - Assertions.assertThat(cacheEntry.getGrantRecordsAsSecurable()).isNotNull(); + assertThat(cacheEntry).isNotNull(); + assertThat(cacheEntry.getEntity()).isNotNull(); + assertThat(cacheEntry.getGrantRecordsAsSecurable()).isNotNull(); // should be there now, by id cacheEntry = cache.getEntityById(N1.getId()); - Assertions.assertThat(cacheEntry).isNotNull(); - Assertions.assertThat(cacheEntry.getEntity()).isNotNull(); - Assertions.assertThat(cacheEntry.getGrantRecordsAsSecurable()).isNotNull(); + assertThat(cacheEntry).isNotNull(); + assertThat(cacheEntry.getEntity()).isNotNull(); + assertThat(cacheEntry.getGrantRecordsAsSecurable()).isNotNull(); // lookup N1 ResolvedPolarisEntity N1_entry = cache.getEntityById(N1.getId()); - Assertions.assertThat(N1_entry).isNotNull(); - Assertions.assertThat(N1_entry.getEntity()).isNotNull(); - Assertions.assertThat(N1_entry.getGrantRecordsAsSecurable()).isNotNull(); + assertThat(N1_entry).isNotNull(); + assertThat(N1_entry.getEntity()).isNotNull(); + assertThat(N1_entry.getGrantRecordsAsSecurable()).isNotNull(); // negative tests, load an entity which does not exist lookup = cache.getOrLoadEntityById(this.callCtx, N1.getCatalogId(), 10000, N1.getType()); - Assertions.assertThat(lookup).isNull(); + assertThat(lookup).isNull(); lookup = cache.getOrLoadEntityByName( this.callCtx, new EntityCacheByNameKey(PolarisEntityType.CATALOG, "non_existant_catalog")); - Assertions.assertThat(lookup).isNull(); + assertThat(lookup).isNull(); // lookup N2 to validate grants EntityCacheByNameKey N2_name = new EntityCacheByNameKey(catalog.getId(), N1.getId(), PolarisEntityType.NAMESPACE, "N2"); lookup = cache.getOrLoadEntityByName(callCtx, N2_name); - Assertions.assertThat(lookup).isNotNull(); + assertThat(lookup).isNotNull(); ResolvedPolarisEntity cacheEntry_N1 = lookup.getCacheEntry(); - Assertions.assertThat(cacheEntry_N1).isNotNull(); - Assertions.assertThat(cacheEntry_N1.getEntity()).isNotNull(); - Assertions.assertThat(cacheEntry_N1.getGrantRecordsAsSecurable()).isNotNull(); + assertThat(cacheEntry_N1).isNotNull(); + assertThat(cacheEntry_N1.getEntity()).isNotNull(); + assertThat(cacheEntry_N1.getGrantRecordsAsSecurable()).isNotNull(); // lookup catalog role R1 EntityCacheByNameKey R1_name = new EntityCacheByNameKey( catalog.getId(), catalog.getId(), PolarisEntityType.CATALOG_ROLE, "R1"); lookup = cache.getOrLoadEntityByName(callCtx, R1_name); - Assertions.assertThat(lookup).isNotNull(); + assertThat(lookup).isNotNull(); ResolvedPolarisEntity cacheEntry_R1 = lookup.getCacheEntry(); - Assertions.assertThat(cacheEntry_R1).isNotNull(); - Assertions.assertThat(cacheEntry_R1.getEntity()).isNotNull(); - Assertions.assertThat(cacheEntry_R1.getGrantRecordsAsSecurable()).isNotNull(); - Assertions.assertThat(cacheEntry_R1.getGrantRecordsAsGrantee()).isNotNull(); + assertThat(cacheEntry_R1).isNotNull(); + assertThat(cacheEntry_R1.getEntity()).isNotNull(); + assertThat(cacheEntry_R1.getGrantRecordsAsSecurable()).isNotNull(); + assertThat(cacheEntry_R1.getGrantRecordsAsGrantee()).isNotNull(); // we expect one TABLE_READ grant on that securable granted to the catalog role R1 - Assertions.assertThat(cacheEntry_N1.getGrantRecordsAsSecurable()).hasSize(1); + assertThat(cacheEntry_N1.getGrantRecordsAsSecurable()).hasSize(1); PolarisGrantRecord gr = cacheEntry_N1.getGrantRecordsAsSecurable().get(0); // securable is N1, grantee is R1 - Assertions.assertThat(gr.getGranteeId()).isEqualTo(cacheEntry_R1.getEntity().getId()); - Assertions.assertThat(gr.getGranteeCatalogId()) - .isEqualTo(cacheEntry_R1.getEntity().getCatalogId()); - Assertions.assertThat(gr.getSecurableId()).isEqualTo(cacheEntry_N1.getEntity().getId()); - Assertions.assertThat(gr.getSecurableCatalogId()) - .isEqualTo(cacheEntry_N1.getEntity().getCatalogId()); - Assertions.assertThat(gr.getPrivilegeCode()) - .isEqualTo(PolarisPrivilege.TABLE_READ_DATA.getCode()); + assertThat(gr.getGranteeId()).isEqualTo(cacheEntry_R1.getEntity().getId()); + assertThat(gr.getGranteeCatalogId()).isEqualTo(cacheEntry_R1.getEntity().getCatalogId()); + assertThat(gr.getSecurableId()).isEqualTo(cacheEntry_N1.getEntity().getId()); + assertThat(gr.getSecurableCatalogId()).isEqualTo(cacheEntry_N1.getEntity().getCatalogId()); + assertThat(gr.getPrivilegeCode()).isEqualTo(PolarisPrivilege.TABLE_READ_DATA.getCode()); // R1 should have 4 privileges granted to it - Assertions.assertThat(cacheEntry_R1.getGrantRecordsAsGrantee()).hasSize(4); + assertThat(cacheEntry_R1.getGrantRecordsAsGrantee()).hasSize(4); List matchPriv = cacheEntry_R1.getGrantRecordsAsGrantee().stream() .filter( grantRecord -> grantRecord.getPrivilegeCode() == PolarisPrivilege.TABLE_READ_DATA.getCode()) .collect(Collectors.toList()); - Assertions.assertThat(matchPriv).hasSize(1); + assertThat(matchPriv).hasSize(1); gr = matchPriv.get(0); - Assertions.assertThat(gr.getGranteeId()).isEqualTo(cacheEntry_R1.getEntity().getId()); - Assertions.assertThat(gr.getGranteeCatalogId()) - .isEqualTo(cacheEntry_R1.getEntity().getCatalogId()); - Assertions.assertThat(gr.getSecurableId()).isEqualTo(cacheEntry_N1.getEntity().getId()); - Assertions.assertThat(gr.getSecurableCatalogId()) - .isEqualTo(cacheEntry_N1.getEntity().getCatalogId()); - Assertions.assertThat(gr.getPrivilegeCode()) - .isEqualTo(PolarisPrivilege.TABLE_READ_DATA.getCode()); + assertThat(gr.getGranteeId()).isEqualTo(cacheEntry_R1.getEntity().getId()); + assertThat(gr.getGranteeCatalogId()).isEqualTo(cacheEntry_R1.getEntity().getCatalogId()); + assertThat(gr.getSecurableId()).isEqualTo(cacheEntry_N1.getEntity().getId()); + assertThat(gr.getSecurableCatalogId()).isEqualTo(cacheEntry_N1.getEntity().getCatalogId()); + assertThat(gr.getPrivilegeCode()).isEqualTo(PolarisPrivilege.TABLE_READ_DATA.getCode()); // lookup principal role PR1 EntityCacheByNameKey PR1_name = new EntityCacheByNameKey(PolarisEntityType.PRINCIPAL_ROLE, "PR1"); lookup = cache.getOrLoadEntityByName(callCtx, PR1_name); - Assertions.assertThat(lookup).isNotNull(); + assertThat(lookup).isNotNull(); ResolvedPolarisEntity cacheEntry_PR1 = lookup.getCacheEntry(); - Assertions.assertThat(cacheEntry_PR1).isNotNull(); - Assertions.assertThat(cacheEntry_PR1.getEntity()).isNotNull(); - Assertions.assertThat(cacheEntry_PR1.getGrantRecordsAsSecurable()).isNotNull(); - Assertions.assertThat(cacheEntry_PR1.getGrantRecordsAsGrantee()).isNotNull(); + assertThat(cacheEntry_PR1).isNotNull(); + assertThat(cacheEntry_PR1.getEntity()).isNotNull(); + assertThat(cacheEntry_PR1.getGrantRecordsAsSecurable()).isNotNull(); + assertThat(cacheEntry_PR1.getGrantRecordsAsGrantee()).isNotNull(); // R1 should have 1 CATALOG_ROLE_USAGE privilege granted *on* it to PR1 - Assertions.assertThat(cacheEntry_R1.getGrantRecordsAsSecurable()).hasSize(1); + assertThat(cacheEntry_R1.getGrantRecordsAsSecurable()).hasSize(1); gr = cacheEntry_R1.getGrantRecordsAsSecurable().get(0); - Assertions.assertThat(gr.getSecurableId()).isEqualTo(cacheEntry_R1.getEntity().getId()); - Assertions.assertThat(gr.getSecurableCatalogId()) - .isEqualTo(cacheEntry_R1.getEntity().getCatalogId()); - Assertions.assertThat(gr.getGranteeId()).isEqualTo(cacheEntry_PR1.getEntity().getId()); - Assertions.assertThat(gr.getGranteeCatalogId()) - .isEqualTo(cacheEntry_PR1.getEntity().getCatalogId()); - Assertions.assertThat(gr.getPrivilegeCode()) - .isEqualTo(PolarisPrivilege.CATALOG_ROLE_USAGE.getCode()); + assertThat(gr.getSecurableId()).isEqualTo(cacheEntry_R1.getEntity().getId()); + assertThat(gr.getSecurableCatalogId()).isEqualTo(cacheEntry_R1.getEntity().getCatalogId()); + assertThat(gr.getGranteeId()).isEqualTo(cacheEntry_PR1.getEntity().getId()); + assertThat(gr.getGranteeCatalogId()).isEqualTo(cacheEntry_PR1.getEntity().getCatalogId()); + assertThat(gr.getPrivilegeCode()).isEqualTo(PolarisPrivilege.CATALOG_ROLE_USAGE.getCode()); // PR1 should have 1 grant on it to P1. - Assertions.assertThat(cacheEntry_PR1.getGrantRecordsAsSecurable()).hasSize(1); - Assertions.assertThat(cacheEntry_PR1.getGrantRecordsAsSecurable().get(0).getPrivilegeCode()) + assertThat(cacheEntry_PR1.getGrantRecordsAsSecurable()).hasSize(1); + assertThat(cacheEntry_PR1.getGrantRecordsAsSecurable().get(0).getPrivilegeCode()) .isEqualTo(PolarisPrivilege.PRINCIPAL_ROLE_USAGE.getCode()); // PR1 should have 2 grants to it, on R1 and R2 - Assertions.assertThat(cacheEntry_PR1.getGrantRecordsAsGrantee()).hasSize(2); - Assertions.assertThat(cacheEntry_PR1.getGrantRecordsAsGrantee().get(0).getPrivilegeCode()) + assertThat(cacheEntry_PR1.getGrantRecordsAsGrantee()).hasSize(2); + assertThat(cacheEntry_PR1.getGrantRecordsAsGrantee().get(0).getPrivilegeCode()) .isEqualTo(PolarisPrivilege.CATALOG_ROLE_USAGE.getCode()); - Assertions.assertThat(cacheEntry_PR1.getGrantRecordsAsGrantee().get(1).getPrivilegeCode()) + assertThat(cacheEntry_PR1.getGrantRecordsAsGrantee().get(1).getPrivilegeCode()) .isEqualTo(PolarisPrivilege.CATALOG_ROLE_USAGE.getCode()); } @@ -274,14 +279,14 @@ void testRefresh() { EntityCacheLookupResult lookup = cache.getOrLoadEntityByName( this.callCtx, new EntityCacheByNameKey(PolarisEntityType.CATALOG, "test")); - Assertions.assertThat(lookup).isNotNull(); - Assertions.assertThat(lookup.isCacheHit()).isFalse(); + assertThat(lookup).isNotNull(); + assertThat(lookup.isCacheHit()).isFalse(); // the catalog - Assertions.assertThat(lookup.getCacheEntry()).isNotNull(); + assertThat(lookup.getCacheEntry()).isNotNull(); PolarisBaseEntity catalog = lookup.getCacheEntry().getEntity(); - Assertions.assertThat(catalog).isNotNull(); - Assertions.assertThat(catalog.getType()).isEqualTo(PolarisEntityType.CATALOG); + assertThat(catalog).isNotNull(); + assertThat(catalog.getType()).isEqualTo(PolarisEntityType.CATALOG); // find table N5/N6/T6 PolarisBaseEntity N5 = @@ -294,23 +299,23 @@ void testRefresh() { PolarisEntityType.TABLE_LIKE, PolarisEntitySubType.ICEBERG_TABLE, "T6"); - Assertions.assertThat(T6v1).isNotNull(); + assertThat(T6v1).isNotNull(); // that table is not in the cache ResolvedPolarisEntity cacheEntry = cache.getEntityById(T6v1.getId()); - Assertions.assertThat(cacheEntry).isNull(); + assertThat(cacheEntry).isNull(); // now load that table in the cache cacheEntry = cache.getAndRefreshIfNeeded( this.callCtx, T6v1, T6v1.getEntityVersion(), T6v1.getGrantRecordsVersion()); - Assertions.assertThat(cacheEntry).isNotNull(); - Assertions.assertThat(cacheEntry.getEntity()).isNotNull(); - Assertions.assertThat(cacheEntry.getGrantRecordsAsSecurable()).isNotNull(); + assertThat(cacheEntry).isNotNull(); + assertThat(cacheEntry.getEntity()).isNotNull(); + assertThat(cacheEntry.getGrantRecordsAsSecurable()).isNotNull(); PolarisBaseEntity table = cacheEntry.getEntity(); - Assertions.assertThat(table.getId()).isEqualTo(T6v1.getId()); - Assertions.assertThat(table.getEntityVersion()).isEqualTo(T6v1.getEntityVersion()); - Assertions.assertThat(table.getGrantRecordsVersion()).isEqualTo(T6v1.getGrantRecordsVersion()); + assertThat(table.getId()).isEqualTo(T6v1.getId()); + assertThat(table.getEntityVersion()).isEqualTo(T6v1.getEntityVersion()); + assertThat(table.getGrantRecordsVersion()).isEqualTo(T6v1.getGrantRecordsVersion()); // update the entity PolarisBaseEntity T6v2 = @@ -319,31 +324,31 @@ void testRefresh() { T6v1, "{\"v2_properties\": \"some value\"}", "{\"v2_internal_properties\": \"internal value\"}"); - Assertions.assertThat(T6v2).isNotNull(); + assertThat(T6v2).isNotNull(); // now refresh that entity. But because we don't change the versions, nothing should be reloaded cacheEntry = cache.getAndRefreshIfNeeded( this.callCtx, T6v1, T6v1.getEntityVersion(), T6v1.getGrantRecordsVersion()); - Assertions.assertThat(cacheEntry).isNotNull(); - Assertions.assertThat(cacheEntry.getEntity()).isNotNull(); - Assertions.assertThat(cacheEntry.getGrantRecordsAsSecurable()).isNotNull(); + assertThat(cacheEntry).isNotNull(); + assertThat(cacheEntry.getEntity()).isNotNull(); + assertThat(cacheEntry.getGrantRecordsAsSecurable()).isNotNull(); table = cacheEntry.getEntity(); - Assertions.assertThat(table.getId()).isEqualTo(T6v1.getId()); - Assertions.assertThat(table.getEntityVersion()).isEqualTo(T6v1.getEntityVersion()); - Assertions.assertThat(table.getGrantRecordsVersion()).isEqualTo(T6v1.getGrantRecordsVersion()); + assertThat(table.getId()).isEqualTo(T6v1.getId()); + assertThat(table.getEntityVersion()).isEqualTo(T6v1.getEntityVersion()); + assertThat(table.getGrantRecordsVersion()).isEqualTo(T6v1.getGrantRecordsVersion()); // now refresh again, this time with the new versions. Should be reloaded cacheEntry = cache.getAndRefreshIfNeeded( this.callCtx, T6v2, T6v2.getEntityVersion(), T6v2.getGrantRecordsVersion()); - Assertions.assertThat(cacheEntry).isNotNull(); - Assertions.assertThat(cacheEntry.getEntity()).isNotNull(); - Assertions.assertThat(cacheEntry.getGrantRecordsAsSecurable()).isNotNull(); + assertThat(cacheEntry).isNotNull(); + assertThat(cacheEntry.getEntity()).isNotNull(); + assertThat(cacheEntry.getGrantRecordsAsSecurable()).isNotNull(); table = cacheEntry.getEntity(); - Assertions.assertThat(table.getId()).isEqualTo(T6v2.getId()); - Assertions.assertThat(table.getEntityVersion()).isEqualTo(T6v2.getEntityVersion()); - Assertions.assertThat(table.getGrantRecordsVersion()).isEqualTo(T6v2.getGrantRecordsVersion()); + assertThat(table.getId()).isEqualTo(T6v2.getId()); + assertThat(table.getEntityVersion()).isEqualTo(T6v2.getEntityVersion()); + assertThat(table.getGrantRecordsVersion()).isEqualTo(T6v2.getGrantRecordsVersion()); // update it again PolarisBaseEntity T6v3 = @@ -352,7 +357,7 @@ void testRefresh() { T6v2, "{\"v3_properties\": \"some value\"}", "{\"v3_internal_properties\": \"internal value\"}"); - Assertions.assertThat(T6v3).isNotNull(); + assertThat(T6v3).isNotNull(); // the two catalog roles PolarisBaseEntity R1 = @@ -368,9 +373,9 @@ void testRefresh() { this.callCtx, N2, N2.getEntityVersion(), N2.getGrantRecordsVersion()); // should have one single grant - Assertions.assertThat(cacheEntry).isNotNull(); - Assertions.assertThat(cacheEntry.getGrantRecordsAsSecurable()).isNotNull(); - Assertions.assertThat(cacheEntry.getGrantRecordsAsSecurable()).hasSize(1); + assertThat(cacheEntry).isNotNull(); + assertThat(cacheEntry.getGrantRecordsAsSecurable()).isNotNull(); + assertThat(cacheEntry.getGrantRecordsAsSecurable()).hasSize(1); // perform an additional grant to R1 this.tm.grantPrivilege(R1, List.of(catalog, N1), N2, PolarisPrivilege.NAMESPACE_FULL_METADATA); @@ -380,8 +385,8 @@ void testRefresh() { this.tm.ensureExistsByName(List.of(catalog, N1), PolarisEntityType.NAMESPACE, "N2"); // same entity version but different grant records - Assertions.assertThat(N2v2).isNotNull(); - Assertions.assertThat(N2v2.getGrantRecordsVersion()).isEqualTo(N2.getGrantRecordsVersion() + 1); + assertThat(N2v2).isNotNull(); + assertThat(N2v2.getGrantRecordsVersion()).isEqualTo(N2.getGrantRecordsVersion() + 1); // the cache is outdated now lookup = @@ -389,24 +394,24 @@ void testRefresh() { this.callCtx, new EntityCacheByNameKey( catalog.getId(), N1.getId(), PolarisEntityType.NAMESPACE, "N2")); - Assertions.assertThat(lookup).isNotNull(); + assertThat(lookup).isNotNull(); cacheEntry = lookup.getCacheEntry(); - Assertions.assertThat(cacheEntry).isNotNull(); - Assertions.assertThat(cacheEntry.getEntity()).isNotNull(); - Assertions.assertThat(cacheEntry.getGrantRecordsAsSecurable()).isNotNull(); - Assertions.assertThat(cacheEntry.getGrantRecordsAsSecurable()).hasSize(1); - Assertions.assertThat(cacheEntry.getEntity().getGrantRecordsVersion()) + assertThat(cacheEntry).isNotNull(); + assertThat(cacheEntry.getEntity()).isNotNull(); + assertThat(cacheEntry.getGrantRecordsAsSecurable()).isNotNull(); + assertThat(cacheEntry.getGrantRecordsAsSecurable()).hasSize(1); + assertThat(cacheEntry.getEntity().getGrantRecordsVersion()) .isEqualTo(N2.getGrantRecordsVersion()); // now refresh cacheEntry = cache.getAndRefreshIfNeeded( this.callCtx, N2, N2v2.getEntityVersion(), N2v2.getGrantRecordsVersion()); - Assertions.assertThat(cacheEntry).isNotNull(); - Assertions.assertThat(cacheEntry.getEntity()).isNotNull(); - Assertions.assertThat(cacheEntry.getGrantRecordsAsSecurable()).isNotNull(); - Assertions.assertThat(cacheEntry.getGrantRecordsAsSecurable()).hasSize(2); - Assertions.assertThat(cacheEntry.getEntity().getGrantRecordsVersion()) + assertThat(cacheEntry).isNotNull(); + assertThat(cacheEntry.getEntity()).isNotNull(); + assertThat(cacheEntry.getGrantRecordsAsSecurable()).isNotNull(); + assertThat(cacheEntry.getGrantRecordsAsSecurable()).hasSize(2); + assertThat(cacheEntry.getEntity().getGrantRecordsVersion()) .isEqualTo(N2v2.getGrantRecordsVersion()); } @@ -418,23 +423,23 @@ void testRenameAndCacheDestinationBeforeLoadingSource() { EntityCacheLookupResult lookup = cache.getOrLoadEntityByName( this.callCtx, new EntityCacheByNameKey(PolarisEntityType.CATALOG, "test")); - Assertions.assertThat(lookup).isNotNull(); - Assertions.assertThat(lookup.getCacheEntry()).isNotNull(); + assertThat(lookup).isNotNull(); + assertThat(lookup.getCacheEntry()).isNotNull(); PolarisBaseEntity catalog = lookup.getCacheEntry().getEntity(); PolarisBaseEntity N1 = this.tm.ensureExistsByName(List.of(catalog), PolarisEntityType.NAMESPACE, "N1"); lookup = cache.getOrLoadEntityById(this.callCtx, N1.getCatalogId(), N1.getId(), N1.getType()); - Assertions.assertThat(lookup).isNotNull(); + assertThat(lookup).isNotNull(); EntityCacheByNameKey T4_name = new EntityCacheByNameKey(N1.getCatalogId(), N1.getId(), PolarisEntityType.TABLE_LIKE, "T4"); lookup = cache.getOrLoadEntityByName(callCtx, T4_name); - Assertions.assertThat(lookup).isNotNull(); + assertThat(lookup).isNotNull(); ResolvedPolarisEntity cacheEntry_T4 = lookup.getCacheEntry(); - Assertions.assertThat(cacheEntry_T4).isNotNull(); - Assertions.assertThat(cacheEntry_T4.getEntity()).isNotNull(); - Assertions.assertThat(cacheEntry_T4.getGrantRecordsAsSecurable()).isNotNull(); + assertThat(cacheEntry_T4).isNotNull(); + assertThat(cacheEntry_T4.getEntity()).isNotNull(); + assertThat(cacheEntry_T4.getGrantRecordsAsSecurable()).isNotNull(); PolarisBaseEntity T4 = cacheEntry_T4.getEntity(); @@ -445,18 +450,17 @@ void testRenameAndCacheDestinationBeforeLoadingSource() { new EntityCacheByNameKey( N1.getCatalogId(), N1.getId(), PolarisEntityType.TABLE_LIKE, "T4_renamed"); lookup = cache.getOrLoadEntityByName(callCtx, T4_renamed); - Assertions.assertThat(lookup).isNotNull(); + assertThat(lookup).isNotNull(); ResolvedPolarisEntity cacheEntry_T4_renamed = lookup.getCacheEntry(); - Assertions.assertThat(cacheEntry_T4_renamed).isNotNull(); + assertThat(cacheEntry_T4_renamed).isNotNull(); PolarisBaseEntity T4_renamed_entity = cacheEntry_T4_renamed.getEntity(); // new entry if lookup by id EntityCacheLookupResult lookupResult = cache.getOrLoadEntityById(callCtx, T4.getCatalogId(), T4.getId(), T4.getType()); - Assertions.assertThat(lookupResult).isNotNull(); - Assertions.assertThat(lookupResult.getCacheEntry()).isNotNull(); - Assertions.assertThat(lookupResult.getCacheEntry().getEntity().getName()) - .isEqualTo("T4_renamed"); + assertThat(lookupResult).isNotNull(); + assertThat(lookupResult.getCacheEntry()).isNotNull(); + assertThat(lookupResult.getCacheEntry().getEntity().getName()).isEqualTo("T4_renamed"); // old name is gone, replaced by new name // Assertions.assertNull(cache.getOrLoadEntityByName(callCtx, T4_name)); @@ -469,7 +473,7 @@ void testRenameAndCacheDestinationBeforeLoadingSource() { T4_renamed_entity.getGrantRecordsVersion()); // now the loading by the old name should return null - Assertions.assertThat(cache.getOrLoadEntityByName(callCtx, T4_name)).isNull(); + assertThat(cache.getOrLoadEntityByName(callCtx, T4_name)).isNull(); } /* Helper for `testEntityWeigher` */ @@ -495,7 +499,582 @@ void testEntityWeigher() { .setMetadataLocation("a".repeat(1000 * 1000)) .build(); - Assertions.assertThat(getEntityWeight(smallEntity)).isLessThan(getEntityWeight(mediumEntity)); - Assertions.assertThat(getEntityWeight(mediumEntity)).isLessThan(getEntityWeight(largeEntity)); + assertThat(getEntityWeight(smallEntity)).isLessThan(getEntityWeight(mediumEntity)); + assertThat(getEntityWeight(mediumEntity)).isLessThan(getEntityWeight(largeEntity)); + } + + @Test + public void testBatchLoadByEntityIds() { + // get a new cache + InMemoryEntityCache cache = this.allocateNewCache(); + + // Load catalog into cache + EntityCacheLookupResult lookup = + cache.getOrLoadEntityByName( + this.callCtx, new EntityCacheByNameKey(PolarisEntityType.CATALOG, "test")); + assertThat(lookup).isNotNull(); + PolarisBaseEntity catalog = lookup.getCacheEntry().getEntity(); + + // Get some entities that exist in the test setup + PolarisBaseEntity N1 = + this.tm.ensureExistsByName(List.of(catalog), PolarisEntityType.NAMESPACE, "N1"); + PolarisBaseEntity N5 = + this.tm.ensureExistsByName(List.of(catalog), PolarisEntityType.NAMESPACE, "N5"); + PolarisBaseEntity N5_N6 = + this.tm.ensureExistsByName(List.of(catalog, N5), PolarisEntityType.NAMESPACE, "N6"); + + // Pre-load N1 into cache + cache.getOrLoadEntityById(this.callCtx, N1.getCatalogId(), N1.getId(), N1.getType()); + + // Create list of entity IDs - N1 is already cached, N5, N5_N6 are not + List entityIds = + List.of(getPolarisEntityId(N1), getPolarisEntityId(N5), getPolarisEntityId(N5_N6)); + + // Test batch loading by entity IDs (all are namespaces) + List results = + cache.getOrLoadResolvedEntities(this.callCtx, PolarisEntityType.NAMESPACE, entityIds); + + // Verify all entities were found + assertThat(results).hasSize(3); + assertThat(results.get(0)).isNotNull(); // N1 - was cached + assertThat(results.get(1)).isNotNull(); // N5 - was loaded + assertThat(results.get(2)).isNotNull(); // N5_N6 - was loaded + + // Verify the entities are correct + assertThat(results.get(0).getCacheEntry().getEntity().getId()).isEqualTo(N1.getId()); + assertThat(results.get(1).getCacheEntry().getEntity().getId()).isEqualTo(N5.getId()); + assertThat(results.get(2).getCacheEntry().getEntity().getId()).isEqualTo(N5_N6.getId()); + + // All should be cache hits now since they were loaded in the previous call + assertThat(results.get(0).isCacheHit()).isTrue(); + assertThat(results.get(1).isCacheHit()).isTrue(); + assertThat(results.get(2).isCacheHit()).isTrue(); + + // Test with a non-existent entity ID + List nonExistentIds = + List.of(new PolarisEntityId(catalog.getCatalogId(), 99999L)); + List nonExistentResults = + cache.getOrLoadResolvedEntities(this.callCtx, PolarisEntityType.NAMESPACE, nonExistentIds); + + assertThat(nonExistentResults).hasSize(1); + assertThat(nonExistentResults.get(0)).isNull(); + + // Test with table entities separately + PolarisBaseEntity T6 = + this.tm.ensureExistsByName( + List.of(catalog, N5, N5_N6), + PolarisEntityType.TABLE_LIKE, + PolarisEntitySubType.ICEBERG_TABLE, + "T6"); + + List tableIds = List.of(getPolarisEntityId(T6)); + + List tableResults = + cache.getOrLoadResolvedEntities(this.callCtx, PolarisEntityType.TABLE_LIKE, tableIds); + + assertThat(tableResults).hasSize(1); + assertThat(tableResults.get(0)).isNotNull(); + assertThat(tableResults.get(0).getCacheEntry().getEntity().getId()).isEqualTo(T6.getId()); + } + + @Test + public void testBatchLoadByNameLookupRecords() { + // get a new cache + InMemoryEntityCache cache = this.allocateNewCache(); + + // Load catalog into cache + EntityCacheLookupResult lookup = + cache.getOrLoadEntityByName( + this.callCtx, new EntityCacheByNameKey(PolarisEntityType.CATALOG, "test")); + assertThat(lookup).isNotNull(); + PolarisBaseEntity catalog = lookup.getCacheEntry().getEntity(); + + // Get some entities that exist in the test setup + PolarisBaseEntity N1 = + this.tm.ensureExistsByName(List.of(catalog), PolarisEntityType.NAMESPACE, "N1"); + PolarisBaseEntity N2 = + this.tm.ensureExistsByName(List.of(catalog, N1), PolarisEntityType.NAMESPACE, "N2"); + PolarisBaseEntity R1 = + this.tm.ensureExistsByName(List.of(catalog), PolarisEntityType.CATALOG_ROLE, "R1"); + + // Pre-load N1 into cache by name + cache.getOrLoadEntityByName( + this.callCtx, + new EntityCacheByNameKey( + catalog.getId(), catalog.getId(), PolarisEntityType.NAMESPACE, "N1")); + + // Create list of EntityNameLookupRecords + List lookupRecords = + List.of( + new EntityNameLookupRecord(N1), // already cached + new EntityNameLookupRecord(N2), // not cached + new EntityNameLookupRecord(R1) // not cached + ); + + // Test batch loading by name lookup records + List results = + cache.getOrLoadResolvedEntities(this.callCtx, lookupRecords); + + // Verify all entities were found + assertThat(results).hasSize(3); + assertThat(results.get(0)).isNotNull(); // N1 - was cached + assertThat(results.get(1)).isNotNull(); // N2 - was loaded + assertThat(results.get(2)).isNotNull(); // R1 - was loaded + + // Verify the entities are correct + assertThat(results.get(0).getCacheEntry().getEntity().getId()).isEqualTo(N1.getId()); + assertThat(results.get(1).getCacheEntry().getEntity().getId()).isEqualTo(N2.getId()); + assertThat(results.get(2).getCacheEntry().getEntity().getId()).isEqualTo(R1.getId()); + + // All should be cache hits now + assertThat(results.get(0).isCacheHit()).isTrue(); + assertThat(results.get(1).isCacheHit()).isTrue(); + assertThat(results.get(2).isCacheHit()).isTrue(); + } + + @Test + public void testBatchLoadWithStaleVersions() { + // get a new cache + InMemoryEntityCache cache = this.allocateNewCache(); + + // Load catalog into cache + EntityCacheLookupResult lookup = + cache.getOrLoadEntityByName( + this.callCtx, new EntityCacheByNameKey(PolarisEntityType.CATALOG, "test")); + assertThat(lookup).isNotNull(); + PolarisBaseEntity catalog = lookup.getCacheEntry().getEntity(); + + // Get table T6 that we can update + PolarisBaseEntity N5 = + this.tm.ensureExistsByName(List.of(catalog), PolarisEntityType.NAMESPACE, "N5"); + PolarisBaseEntity N5_N6 = + this.tm.ensureExistsByName(List.of(catalog, N5), PolarisEntityType.NAMESPACE, "N6"); + PolarisBaseEntity T6v1 = + this.tm.ensureExistsByName( + List.of(catalog, N5, N5_N6), + PolarisEntityType.TABLE_LIKE, + PolarisEntitySubType.ICEBERG_TABLE, + "T6"); + + // Load T6 into cache initially + cache.getOrLoadEntityById(this.callCtx, T6v1.getCatalogId(), T6v1.getId(), T6v1.getType()); + + // Verify it's in cache with original version + ResolvedPolarisEntity cachedT6 = cache.getEntityById(T6v1.getId()); + assertThat(cachedT6).isNotNull(); + assertThat(cachedT6.getEntity().getEntityVersion()).isEqualTo(T6v1.getEntityVersion()); + + // Update the entity to create a new version + PolarisBaseEntity T6v2 = + this.tm.updateEntity( + List.of(catalog, N5, N5_N6), + T6v1, + "{\"v2_properties\": \"some value\"}", + "{\"v2_internal_properties\": \"internal value\"}"); + assertThat(T6v2).isNotNull(); + assertThat(T6v2.getEntityVersion()).isGreaterThan(T6v1.getEntityVersion()); + + // Create entity ID list with the updated entity + List entityIds = List.of(getPolarisEntityId(T6v2)); + + // Call batch load - this should detect the stale version and reload + List results = + cache.getOrLoadResolvedEntities(this.callCtx, PolarisEntityType.TABLE_LIKE, entityIds); + + // Verify the entity was reloaded with the new version + assertThat(results).hasSize(1); + assertThat(results.get(0)).isNotNull(); + + ResolvedPolarisEntity reloadedT6 = results.get(0).getCacheEntry(); + assertThat(reloadedT6.getEntity().getId()).isEqualTo(T6v2.getId()); + assertThat(reloadedT6.getEntity().getEntityVersion()).isEqualTo(T6v2.getEntityVersion()); + + // Verify the cache now contains the updated version + cachedT6 = cache.getEntityById(T6v2.getId()); + assertThat(cachedT6).isNotNull(); + assertThat(cachedT6.getEntity().getEntityVersion()).isEqualTo(T6v2.getEntityVersion()); + } + + @Test + public void testBatchLoadWithStaleGrantVersions() { + // get a new cache + InMemoryEntityCache cache = this.allocateNewCache(); + + // Load catalog into cache + EntityCacheLookupResult lookup = + cache.getOrLoadEntityByName( + this.callCtx, new EntityCacheByNameKey(PolarisEntityType.CATALOG, "test")); + assertThat(lookup).isNotNull(); + PolarisBaseEntity catalog = lookup.getCacheEntry().getEntity(); + + // Get entities we'll work with + PolarisBaseEntity N1 = + this.tm.ensureExistsByName(List.of(catalog), PolarisEntityType.NAMESPACE, "N1"); + PolarisBaseEntity N2 = + this.tm.ensureExistsByName(List.of(catalog, N1), PolarisEntityType.NAMESPACE, "N2"); + PolarisBaseEntity R1 = + this.tm.ensureExistsByName(List.of(catalog), PolarisEntityType.CATALOG_ROLE, "R1"); + + // Load N2 into cache initially + cache.getOrLoadEntityByName( + this.callCtx, + new EntityCacheByNameKey(catalog.getId(), N1.getId(), PolarisEntityType.NAMESPACE, "N2")); + + // Verify it's in cache with original grant version + ResolvedPolarisEntity cachedN2 = cache.getEntityById(N2.getId()); + assertThat(cachedN2).isNotNull(); + int originalGrantVersion = cachedN2.getEntity().getGrantRecordsVersion(); + + // Grant additional privilege to change grant version + this.tm.grantPrivilege(R1, List.of(catalog, N1), N2, PolarisPrivilege.NAMESPACE_FULL_METADATA); + + // Get the updated entity with new grant version + PolarisBaseEntity N2Updated = + this.tm.ensureExistsByName(List.of(catalog, N1), PolarisEntityType.NAMESPACE, "N2"); + assertThat(N2Updated.getGrantRecordsVersion()).isGreaterThan(originalGrantVersion); + + // Create entity ID list + List entityIds = List.of(getPolarisEntityId(N2Updated)); + + // Call batch load - this should detect the stale grant version and reload + List results = + cache.getOrLoadResolvedEntities(this.callCtx, PolarisEntityType.NAMESPACE, entityIds); + + // Verify the entity was reloaded with the new grant version + assertThat(results).hasSize(1); + assertThat(results.get(0)).isNotNull(); + + ResolvedPolarisEntity reloadedN2 = results.get(0).getCacheEntry(); + assertThat(reloadedN2.getEntity().getId()).isEqualTo(N2Updated.getId()); + assertThat(reloadedN2.getEntity().getGrantRecordsVersion()) + .isEqualTo(N2Updated.getGrantRecordsVersion()); + + // Should now have more grant records + assertThat(reloadedN2.getGrantRecordsAsSecurable().size()).isGreaterThan(1); + + // Verify the cache now contains the updated grant version + cachedN2 = cache.getEntityById(N2Updated.getId()); + assertThat(cachedN2).isNotNull(); + assertThat(cachedN2.getEntity().getGrantRecordsVersion()) + .isEqualTo(N2Updated.getGrantRecordsVersion()); + } + + @Test + public void testBatchLoadVersionRetryLogic() { + // get a new cache + PolarisMetaStoreManager metaStoreManager = Mockito.spy(this.metaStoreManager); + InMemoryEntityCache cache = + new InMemoryEntityCache(diagServices, callCtx.getRealmConfig(), metaStoreManager); + + // Load catalog into cache + EntityCacheLookupResult lookup = + cache.getOrLoadEntityByName( + this.callCtx, new EntityCacheByNameKey(PolarisEntityType.CATALOG, "test")); + assertThat(lookup).isNotNull(); + PolarisBaseEntity catalog = lookup.getCacheEntry().getEntity(); + + // Get entities that we can work with + PolarisBaseEntity N1 = + this.tm.ensureExistsByName(List.of(catalog), PolarisEntityType.NAMESPACE, "N1"); + PolarisBaseEntity N2 = + this.tm.ensureExistsByName(List.of(catalog, N1), PolarisEntityType.NAMESPACE, "N2"); + + // Load N2 into cache initially + cache.getOrLoadEntityByName( + this.callCtx, + new EntityCacheByNameKey(catalog.getId(), N1.getId(), PolarisEntityType.NAMESPACE, "N2")); + + // Verify it's in cache with original version + ResolvedPolarisEntity cachedN2 = cache.getEntityById(N2.getId()); + assertThat(cachedN2).isNotNull(); + int originalEntityVersion = cachedN2.getEntity().getEntityVersion(); + + // Update the entity multiple times to create version skew + PolarisBaseEntity N2v2 = + this.tm.updateEntity(List.of(catalog, N1), N2, "{\"v2\": \"value\"}", null); + + // the first call should return the first version, then we call the real method to get the + // latest + Mockito.doReturn( + new ChangeTrackingResult( + List.of( + changeTrackingFor(catalog), changeTrackingFor(N1), changeTrackingFor(N2v2)))) + .when(metaStoreManager) + .loadEntitiesChangeTracking(Mockito.any(), Mockito.any()); + Mockito.doCallRealMethod() + .when(metaStoreManager) + .loadEntitiesChangeTracking(Mockito.any(), Mockito.any()); + + // update again to create v3, which isn't returned in the change tracking result + PolarisBaseEntity N2v3 = + this.tm.updateEntity(List.of(catalog, N1), N2v2, "{\"v3\": \"value\"}", null); + + // Verify versions increased + assertThat(N2v2.getEntityVersion()).isGreaterThan(originalEntityVersion); + assertThat(N2v3.getEntityVersion()).isGreaterThan(N2v2.getEntityVersion()); + + // Create entity ID list + List entityIds = + List.of(getPolarisEntityId(catalog), getPolarisEntityId(N1), getPolarisEntityId(N2)); + + // Call batch load - this should detect the stale versions and reload until consistent + List results = + cache.getOrLoadResolvedEntities(this.callCtx, PolarisEntityType.NAMESPACE, entityIds); + + // Verify the entity was reloaded with the latest version + assertThat(results).hasSize(3); + assertThat(results) + .doesNotContainNull() + .extracting(EntityCacheLookupResult::getCacheEntry) + .doesNotContainNull() + .extracting(e -> e.getEntity().getId()) + .containsExactly(catalog.getId(), N1.getId(), N2.getId()); + + ResolvedPolarisEntity reloadedN2 = results.get(2).getCacheEntry(); + assertThat(reloadedN2.getEntity().getId()).isEqualTo(N2v3.getId()); + assertThat(reloadedN2.getEntity().getEntityVersion()).isEqualTo(N2v3.getEntityVersion()); + + // Verify the cache now contains the latest version + cachedN2 = cache.getEntityById(N2v3.getId()); + assertThat(cachedN2).isNotNull(); + assertThat(cachedN2.getEntity().getEntityVersion()).isEqualTo(N2v3.getEntityVersion()); + } + + private static PolarisEntityId getPolarisEntityId(PolarisBaseEntity catalog) { + return new PolarisEntityId(catalog.getCatalogId(), catalog.getId()); + } + + @Test + public void testConcurrentClientLoadingBehavior() throws Exception { + // Load catalog into cache + EntityCacheLookupResult lookup = + allocateNewCache() + .getOrLoadEntityByName( + this.callCtx, new EntityCacheByNameKey(PolarisEntityType.CATALOG, "test")); + assertThat(lookup).isNotNull(); + PolarisBaseEntity catalog = lookup.getCacheEntry().getEntity(); + + // Get multiple entities to create a larger list for concurrent processing + PolarisBaseEntity N1 = + this.tm.ensureExistsByName(List.of(catalog), PolarisEntityType.NAMESPACE, "N1"); + PolarisBaseEntity N2 = + this.tm.ensureExistsByName(List.of(catalog, N1), PolarisEntityType.NAMESPACE, "N2"); + PolarisBaseEntity N3 = + this.tm.ensureExistsByName(List.of(catalog, N1), PolarisEntityType.NAMESPACE, "N3"); + PolarisBaseEntity N5 = + this.tm.ensureExistsByName(List.of(catalog), PolarisEntityType.NAMESPACE, "N5"); + PolarisBaseEntity N5_N6 = + this.tm.ensureExistsByName(List.of(catalog, N5), PolarisEntityType.NAMESPACE, "N6"); + + // Create entity IDs list for both clients + List entityIds = + List.of( + getPolarisEntityId(N1), + getPolarisEntityId(N2), + getPolarisEntityId(N3), + getPolarisEntityId(N5), + getPolarisEntityId(N5_N6)); + + // Update one of the entities to create version differences + PolarisBaseEntity N2v2 = + this.tm.updateEntity( + List.of(catalog, N1), N2, "{\"concurrent_test\": \"client1_version\"}", null); + PolarisBaseEntity N2v3 = + this.tm.updateEntity( + List.of(catalog, N1), N2v2, "{\"concurrent_test\": \"client2_version\"}", null); + + // Mock the metastore manager to control the timing of method calls + PolarisMetaStoreManager mockedMetaStoreManager = Mockito.spy(this.metaStoreManager); + + // Create caches with the mocked metastore manager + InMemoryEntityCache cache = + new InMemoryEntityCache(diagServices, callCtx.getRealmConfig(), mockedMetaStoreManager); + + // Synchronization primitives for controlling execution order + CountDownLatch client1ChangeTrackingResult = new CountDownLatch(1); + Semaphore client1ResolvedEntitiesBlocker = new Semaphore(0); + + // Atomic references to capture results from both threads + AtomicReference client1Exception = new AtomicReference<>(); + AtomicReference client2Exception = new AtomicReference<>(); + + // Configure mock behavior: + // 1. Allow both threads to call loadEntitiesChangeTracking() with different versions + // 2. Block client1's loadResolvedEntities() until client2's loadResolvedEntities() completes + + // Mock loadEntitiesChangeTracking to return different versions for each client + Mockito.doAnswer( + invocation -> { + // First call (client1) - returns older version for N2 + LOGGER.debug("Returning change tracking for client1"); + return new ChangeTrackingResult( + List.of( + changeTrackingFor(N1), + changeTrackingFor(N2v2), // older version + changeTrackingFor(N3), + changeTrackingFor(N5), + changeTrackingFor(N5_N6))); + }) + .doAnswer( + invocation -> { + // Second call (client2) - returns newer version for N2 + LOGGER.debug("Returning change tracking for client2"); + return new ChangeTrackingResult( + List.of( + changeTrackingFor(N1), + changeTrackingFor(N2v3), // newer version + changeTrackingFor(N3), + changeTrackingFor(N5), + changeTrackingFor(N5_N6))); + }) + .when(mockedMetaStoreManager) + .loadEntitiesChangeTracking(Mockito.any(), Mockito.any()); + + // Mock loadResolvedEntities to control timing - client1 blocks until client2 completes + // client1 receives the older version of all entities, while client2 receives the newer version + // of N2 + Mockito.doAnswer( + invocation -> { + // This is client1's loadResolvedEntities call - block until client2 completes + try { + client1ChangeTrackingResult.countDown(); + LOGGER.debug("Awaiting client2 to complete resolved entities load"); + client1ResolvedEntitiesBlocker.acquire(); // Block until client2 signals completion + List resolvedEntities = + List.of( + getResolvedPolarisEntity(N1), + getResolvedPolarisEntity(N2v2), + getResolvedPolarisEntity(N3), + getResolvedPolarisEntity(N5), + getResolvedPolarisEntity(N5_N6)); + LOGGER.debug("Client1 returning results {}", resolvedEntities); + return new ResolvedEntitiesResult(resolvedEntities); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new RuntimeException(e); + } + }) + .doAnswer( + invocation -> { + // This is client2's loadResolvedEntities call - execute normally and signal client1 + try { + LOGGER.debug("Client2 loading resolved entities"); + var result = + new ResolvedEntitiesResult( + List.of( + getResolvedPolarisEntity(N1), + getResolvedPolarisEntity(N2v3), + getResolvedPolarisEntity(N3), + getResolvedPolarisEntity(N5), + getResolvedPolarisEntity(N5_N6))); + client1ResolvedEntitiesBlocker.release(); // Allow client1 to proceed + LOGGER.debug("Client2 returning results {}", result.getResolvedEntities()); + return result; + } catch (Exception e) { + client1ResolvedEntitiesBlocker.release(); // Release in case of error + throw e; + } + }) + .when(mockedMetaStoreManager) + .loadResolvedEntities(Mockito.any(), Mockito.any(), Mockito.any()); + + // ExecutorService isn't AutoCloseable in JDK 11 :( + ExecutorService executorService = Executors.newFixedThreadPool(2); + try { + // Client 1 task - should get older version and be blocked during loadResolvedEntities + Future> client1Task = + executorService.submit( + () -> { + try { + // Client1 calls getOrLoadResolvedEntities + // - loadEntitiesChangeTracking returns older version for N2 + // - loadResolvedEntities will block until client2 completes + List results = + cache.getOrLoadResolvedEntities( + this.callCtx, PolarisEntityType.NAMESPACE, entityIds); + + return results; + } catch (Exception e) { + client1Exception.set(e); + return null; + } + }); + + // Client 2 task - should get newer version and complete first + Future> client2Task = + executorService.submit( + () -> { + try { + // Client2 calls getOrLoadResolvedEntities + // - loadEntitiesChangeTracking returns newer version for N2 + // - loadResolvedEntities executes normally and signals client1 when done + client1ChangeTrackingResult.await(); + List results = + cache.getOrLoadResolvedEntities( + this.callCtx, PolarisEntityType.NAMESPACE, entityIds); + + return results; + } catch (Exception e) { + client2Exception.set(e); + client1ResolvedEntitiesBlocker.release(); // Release in case of error + return null; + } + }); + + // Wait for both tasks to complete + List client1Results = client1Task.get(); + List client2Results = client2Task.get(); + + // Verify no exceptions occurred + assertThat(client1Exception.get()).isNull(); + assertThat(client2Exception.get()).isNull(); + + // Verify both clients got results + assertThat(client1Results).isNotNull(); + assertThat(client2Results).isNotNull(); + assertThat(client1Results).hasSize(5); + assertThat(client2Results).hasSize(5); + + // All entities should be found + assertThat(client1Results).doesNotContainNull(); + assertThat(client2Results).doesNotContainNull(); + + // Verify that client1 got the older version of N2 (index 1 in the list) + ResolvedPolarisEntity client1N2 = client1Results.get(1).getCacheEntry(); + assertThat(client1N2.getEntity().getId()).isEqualTo(N2.getId()); + assertThat(client1N2.getEntity().getEntityVersion()).isEqualTo(N2v2.getEntityVersion()); + + // Verify that client2 got the newer version of N2 + ResolvedPolarisEntity client2N2 = client2Results.get(1).getCacheEntry(); + assertThat(client2N2.getEntity().getId()).isEqualTo(N2.getId()); + assertThat(client2N2.getEntity().getEntityVersion()).isEqualTo(N2v3.getEntityVersion()); + + // Verify that both clients got consistent versions for other entities + for (int i = 0; i < 5; i++) { + if (i != 1) { // Skip N2 which we expect to be different + ResolvedPolarisEntity client1Entity = client1Results.get(i).getCacheEntry(); + ResolvedPolarisEntity client2Entity = client2Results.get(i).getCacheEntry(); + + assertThat(client1Entity.getEntity().getId()) + .isEqualTo(client2Entity.getEntity().getId()); + assertThat(client1Entity.getEntity().getEntityVersion()) + .isEqualTo(client2Entity.getEntity().getEntityVersion()); + assertThat(client1Entity.getEntity().getGrantRecordsVersion()) + .isEqualTo(client2Entity.getEntity().getGrantRecordsVersion()); + } + } + assertThat(entityIds).extracting(id -> cache.getEntityById(id.getId())).doesNotContainNull(); + } finally { + executorService.shutdown(); + } + } + + private static ResolvedPolarisEntity getResolvedPolarisEntity(PolarisBaseEntity catalog) { + return new ResolvedPolarisEntity(PolarisEntity.of(catalog), List.of(), List.of()); + } + + private static PolarisChangeTrackingVersions changeTrackingFor(PolarisBaseEntity entity) { + return new PolarisChangeTrackingVersions( + entity.getEntityVersion(), entity.getGrantRecordsVersion()); } } diff --git a/polaris-core/src/testFixtures/java/org/apache/polaris/core/persistence/BasePolarisMetaStoreManagerTest.java b/polaris-core/src/testFixtures/java/org/apache/polaris/core/persistence/BasePolarisMetaStoreManagerTest.java index 21122bda61..c65f81a6e2 100644 --- a/polaris-core/src/testFixtures/java/org/apache/polaris/core/persistence/BasePolarisMetaStoreManagerTest.java +++ b/polaris-core/src/testFixtures/java/org/apache/polaris/core/persistence/BasePolarisMetaStoreManagerTest.java @@ -118,7 +118,7 @@ protected void testCreateEntities() { .containsExactly(PolarisEntity.toCore(task1), PolarisEntity.toCore(task2)); List listedEntities = - metaStoreManager.loadFullEntitiesAll( + metaStoreManager.listFullEntitiesAll( polarisTestMetaStoreManager.polarisCallContext, null, PolarisEntityType.TASK, diff --git a/runtime/service/src/main/java/org/apache/polaris/service/admin/PolarisAdminService.java b/runtime/service/src/main/java/org/apache/polaris/service/admin/PolarisAdminService.java index fcc8cd97d6..dd6839fd83 100644 --- a/runtime/service/src/main/java/org/apache/polaris/service/admin/PolarisAdminService.java +++ b/runtime/service/src/main/java/org/apache/polaris/service/admin/PolarisAdminService.java @@ -967,7 +967,7 @@ public List listCatalogs() { /** List all catalogs without checking for permission. */ private Stream listCatalogsUnsafe() { return metaStoreManager - .loadFullEntitiesAll( + .listFullEntitiesAll( getCurrentPolarisContext(), null, PolarisEntityType.CATALOG, @@ -1214,7 +1214,7 @@ public List listPrincipals() { authorizeBasicRootOperationOrThrow(op); return metaStoreManager - .loadFullEntitiesAll( + .listFullEntitiesAll( getCurrentPolarisContext(), null, PolarisEntityType.PRINCIPAL, @@ -1321,7 +1321,7 @@ public List listPrincipalRoles() { authorizeBasicRootOperationOrThrow(op); return metaStoreManager - .loadFullEntitiesAll( + .listFullEntitiesAll( getCurrentPolarisContext(), null, PolarisEntityType.PRINCIPAL_ROLE, @@ -1445,7 +1445,7 @@ public List listCatalogRoles(String catalogName) { CatalogEntity catalogEntity = getCatalogByName(resolutionManifest, catalogName); List catalogPath = PolarisEntity.toCoreList(List.of(catalogEntity)); return metaStoreManager - .loadFullEntitiesAll( + .listFullEntitiesAll( getCurrentPolarisContext(), catalogPath, PolarisEntityType.CATALOG_ROLE, diff --git a/runtime/service/src/main/java/org/apache/polaris/service/catalog/policy/PolicyCatalog.java b/runtime/service/src/main/java/org/apache/polaris/service/catalog/policy/PolicyCatalog.java index 451eba049d..9b2f7c8a62 100644 --- a/runtime/service/src/main/java/org/apache/polaris/service/catalog/policy/PolicyCatalog.java +++ b/runtime/service/src/main/java/org/apache/polaris/service/catalog/policy/PolicyCatalog.java @@ -188,7 +188,7 @@ public List listPolicies(Namespace namespace, @Nullable Policy } // with a policyType filter we need to load the full PolicyEntity to apply the filter return metaStoreManager - .loadFullEntitiesAll( + .listFullEntitiesAll( callContext.getPolarisCallContext(), catalogPath, PolarisEntityType.POLICY, From 0b0523e0497432d1360e1d7e437916acfa716648 Mon Sep 17 00:00:00 2001 From: Michael Collado Date: Tue, 23 Sep 2025 15:27:11 -0700 Subject: [PATCH 4/9] Add additional test for loadResolvedEntities by id --- .../core/entity/PolarisBaseEntity.java | 6 + .../AtomicOperationMetaStoreManager.java | 55 ++++---- .../BasePolarisMetaStoreManagerTest.java | 10 +- .../PolarisTestMetaStoreManager.java | 122 +++++++++++++++++- 4 files changed, 162 insertions(+), 31 deletions(-) diff --git a/polaris-core/src/main/java/org/apache/polaris/core/entity/PolarisBaseEntity.java b/polaris-core/src/main/java/org/apache/polaris/core/entity/PolarisBaseEntity.java index 6c91517a1f..80b0a96405 100644 --- a/polaris-core/src/main/java/org/apache/polaris/core/entity/PolarisBaseEntity.java +++ b/polaris-core/src/main/java/org/apache/polaris/core/entity/PolarisBaseEntity.java @@ -339,7 +339,10 @@ public boolean equals(Object o) { PolarisBaseEntity that = (PolarisBaseEntity) o; return subTypeCode == that.subTypeCode && createTimestamp == that.createTimestamp + && dropTimestamp == that.dropTimestamp + && purgeTimestamp == that.purgeTimestamp && toPurgeTimestamp == that.toPurgeTimestamp + && lastUpdateTimestamp == that.lastUpdateTimestamp && grantRecordsVersion == that.grantRecordsVersion && Objects.equals(properties, that.properties) && Objects.equals(internalProperties, that.internalProperties); @@ -356,7 +359,10 @@ public int hashCode() { entityVersion, subTypeCode, createTimestamp, + dropTimestamp, + purgeTimestamp, toPurgeTimestamp, + lastUpdateTimestamp, properties, internalProperties, grantRecordsVersion); diff --git a/polaris-core/src/main/java/org/apache/polaris/core/persistence/AtomicOperationMetaStoreManager.java b/polaris-core/src/main/java/org/apache/polaris/core/persistence/AtomicOperationMetaStoreManager.java index 6e52902d49..88d1c73bbf 100644 --- a/polaris-core/src/main/java/org/apache/polaris/core/persistence/AtomicOperationMetaStoreManager.java +++ b/polaris-core/src/main/java/org/apache/polaris/core/persistence/AtomicOperationMetaStoreManager.java @@ -1768,27 +1768,13 @@ public ResolvedEntitiesResult loadResolvedEntities( @Nonnull PolarisCallContext callCtx, @Nonnull List entityLookupRecords) { BasePersistence ms = callCtx.getMetaStore(); - List entities = - ms.lookupEntities( - callCtx, - entityLookupRecords.stream() - .map(r -> new PolarisEntityId(r.getCatalogId(), r.getId())) - .collect(Collectors.toList())); - // mimic the behavior of loadEntity above, return null if not found or type mismatch - List ret = - IntStream.range(0, entityLookupRecords.size()) - .mapToObj( - i -> { - if (entities.get(i) != null - && !entities.get(i).getType().equals(entityLookupRecords.get(i).getType())) { - return null; - } else { - return entities.get(i); - } - }) - .map(e -> toResolvedPolarisEntity(callCtx, e, ms)) + List entityIds = + entityLookupRecords.stream() + .map(r -> new PolarisEntityId(r.getCatalogId(), r.getId())) .collect(Collectors.toList()); - return new ResolvedEntitiesResult(ret); + Function entityTypeForIndex = + i -> entityLookupRecords.get(i).getType(); + return getResolvedEntitiesResult(callCtx, ms, entityIds, entityTypeForIndex); } @Nonnull @@ -1798,20 +1784,43 @@ public ResolvedEntitiesResult loadResolvedEntities( @Nonnull PolarisEntityType entityType, @Nonnull List entityIds) { BasePersistence ms = callCtx.getMetaStore(); - List entities = ms.lookupEntities(callCtx, entityIds); + return getResolvedEntitiesResult(callCtx, ms, entityIds, i -> entityType); + } + private static ResolvedEntitiesResult getResolvedEntitiesResult( + PolarisCallContext callCtx, + BasePersistence ms, + List entityIds, + Function entityTypeForIndex) { + List entities = ms.lookupEntities(callCtx, entityIds); // mimic the behavior of loadEntity above, return null if not found or type mismatch List ret = IntStream.range(0, entityIds.size()) .mapToObj( i -> { - if (entities.get(i) != null && !entities.get(i).getType().equals(entityType)) { + if (entities.get(i) != null + && !entities.get(i).getType().equals(entityTypeForIndex.apply(i))) { return null; } else { return entities.get(i); } }) - .map(e -> toResolvedPolarisEntity(callCtx, e, ms)) + .map( + e -> { + if (e == null) { + return null; + } else { + // load the grant records + final List grantRecordsAsSecurable = + ms.loadAllGrantRecordsOnSecurable(callCtx, e.getCatalogId(), e.getId()); + final List grantRecordsAsGrantee = + e.getType().isGrantee() + ? ms.loadAllGrantRecordsOnGrantee(callCtx, e.getCatalogId(), e.getId()) + : List.of(); + return new ResolvedPolarisEntity( + PolarisEntity.of(e), grantRecordsAsGrantee, grantRecordsAsSecurable); + } + }) .collect(Collectors.toList()); return new ResolvedEntitiesResult(ret); } diff --git a/polaris-core/src/testFixtures/java/org/apache/polaris/core/persistence/BasePolarisMetaStoreManagerTest.java b/polaris-core/src/testFixtures/java/org/apache/polaris/core/persistence/BasePolarisMetaStoreManagerTest.java index c65f81a6e2..0675e7bf1e 100644 --- a/polaris-core/src/testFixtures/java/org/apache/polaris/core/persistence/BasePolarisMetaStoreManagerTest.java +++ b/polaris-core/src/testFixtures/java/org/apache/polaris/core/persistence/BasePolarisMetaStoreManagerTest.java @@ -240,8 +240,14 @@ protected void testLookup() { /** test batch entity load */ @Test - protected void testBatchLoad() { - polarisTestMetaStoreManager.testBatchLoad(); + protected void testLoadResolvedEntities() { + polarisTestMetaStoreManager.testLoadResolvedEntities(); + } + + /** test batch entity load */ + @Test + protected void testLoadResolvedEntitiesById() { + polarisTestMetaStoreManager.testLoadResolvedEntitiesById(); } /** Test the set of functions for the entity cache */ 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 f6062d06fb..6c7ff7d6fb 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 @@ -23,6 +23,7 @@ import java.util.HashSet; import java.util.List; import java.util.Map; +import java.util.Objects; import org.apache.commons.lang3.tuple.ImmutablePair; import org.apache.commons.lang3.tuple.Pair; import org.apache.polaris.core.PolarisCallContext; @@ -2688,7 +2689,7 @@ public void testLookup() { this.ensureNotExistsById(catalog.getId(), T1.getId(), PolarisEntityType.NAMESPACE); } - public void testBatchLoad() { + public void testLoadResolvedEntities() { // load all principals List principals = polarisMetaStoreManager @@ -2751,8 +2752,9 @@ public void testBatchLoad() { InstanceOfAssertFactories.list(ResolvedPolarisEntity.class)) .hasSize(4) .allSatisfy(entity -> Assertions.assertThat(entity).isNotNull()) - .extracting(r -> (PolarisBaseEntity) r.getEntity()) - .containsExactly(catalog, N1, N1_N2, T1); + .extracting(r -> getEntityCore(r.getEntity())) + .containsExactly( + getEntityCore(catalog), getEntityCore(N1), getEntityCore(N1_N2), getEntityCore(T1)); ResolvedPolarisEntity catalogEntity = entitiesResult.getResolvedEntities().get(0); Assertions.assertThat(catalogEntity) @@ -2844,10 +2846,118 @@ public void testBatchLoad() { ResolvedEntitiesResult::getResolvedEntities, InstanceOfAssertFactories.list(ResolvedPolarisEntity.class)) .hasSize(6) - .filteredOn(e -> e != null) + .filteredOn(Objects::nonNull) .hasSize(2) - .extracting(r -> (PolarisBaseEntity) r.getEntity()) - .containsExactly(catalog, T1); + .extracting(r -> getEntityCore(r.getEntity())) + .containsExactly(getEntityCore(catalog), getEntityCore(T1)); + } + + public void testLoadResolvedEntitiesById() { + // load all principals + List principals = + polarisMetaStoreManager + .listEntities( + this.polarisCallContext, + null, + PolarisEntityType.PRINCIPAL, + PolarisEntitySubType.NULL_SUBTYPE, + PageToken.readEverything()) + .getEntities(); + + // create new catalog + PolarisBaseEntity catalog = + new PolarisBaseEntity( + PolarisEntityConstants.getNullId(), + polarisMetaStoreManager.generateNewEntityId(this.polarisCallContext).getId(), + PolarisEntityType.CATALOG, + PolarisEntitySubType.NULL_SUBTYPE, + PolarisEntityConstants.getRootEntityId(), + "test"); + CreateCatalogResult catalogCreated = + polarisMetaStoreManager.createCatalog(this.polarisCallContext, catalog, List.of()); + Assertions.assertThat(catalogCreated).isNotNull(); + + // load the catalog again, since the grant versions are different + catalog = + polarisMetaStoreManager + .loadEntity( + polarisCallContext, + 0L, + catalogCreated.getCatalog().getId(), + PolarisEntityType.CATALOG) + .getEntity(); + + // now create all objects + PolarisBaseEntity N1 = this.createEntity(List.of(catalog), PolarisEntityType.NAMESPACE, "N1"); + PolarisBaseEntity N1_N2 = + this.createEntity(List.of(catalog, N1), PolarisEntityType.NAMESPACE, "N2"); + PolarisBaseEntity T1 = + this.createEntity( + List.of(catalog, N1, N1_N2), + PolarisEntityType.TABLE_LIKE, + PolarisEntitySubType.ICEBERG_TABLE, + "T1"); + + // batch load all entities. They should all be present and non-null + ResolvedEntitiesResult entitiesResult = + polarisMetaStoreManager.loadResolvedEntities( + polarisCallContext, + PolarisEntityType.NAMESPACE, + List.of( + new PolarisEntityId(N1.getCatalogId(), N1.getId()), + new PolarisEntityId(N1_N2.getCatalogId(), N1_N2.getId()))); + Assertions.assertThat(entitiesResult) + .isNotNull() + .returns(BaseResult.ReturnStatus.SUCCESS, ResolvedEntitiesResult::getReturnStatus) + .extracting( + ResolvedEntitiesResult::getResolvedEntities, + InstanceOfAssertFactories.list(ResolvedPolarisEntity.class)) + .hasSize(2) + .allSatisfy(entity -> Assertions.assertThat(entity).isNotNull()) + .extracting(r -> getEntityCore(r.getEntity())) + .containsExactly(getEntityCore(N1), getEntityCore(N1_N2)); + + // try entities which do not exist + entitiesResult = + polarisMetaStoreManager.loadResolvedEntities( + polarisCallContext, + PolarisEntityType.CATALOG, + List.of( + new PolarisEntityId(catalog.getId(), 27), + new PolarisEntityId(catalog.getId(), 35))); + Assertions.assertThat(entitiesResult) + .isNotNull() + .returns(BaseResult.ReturnStatus.SUCCESS, ResolvedEntitiesResult::getReturnStatus) + .extracting( + ResolvedEntitiesResult::getResolvedEntities, + InstanceOfAssertFactories.list(ResolvedPolarisEntity.class)) + .hasSize(2) + .allSatisfy(entity -> Assertions.assertThat(entity).isNull()); + + // existing entities, some with wrong type + entitiesResult = + polarisMetaStoreManager.loadResolvedEntities( + polarisCallContext, + PolarisEntityType.NAMESPACE, + List.of( + new PolarisEntityId(catalog.getCatalogId(), catalog.getId()), + new PolarisEntityId(catalog.getId(), N1_N2.getId()), + new PolarisEntityId(catalog.getId(), T1.getId()))); + Assertions.assertThat(entitiesResult) + .isNotNull() + .returns(BaseResult.ReturnStatus.SUCCESS, ResolvedEntitiesResult::getReturnStatus) + .extracting( + ResolvedEntitiesResult::getResolvedEntities, + InstanceOfAssertFactories.list(ResolvedPolarisEntity.class)) + .hasSize(3) + .filteredOn(Objects::nonNull) + .hasSize(1) + .extracting(r -> getEntityCore(r.getEntity())) + .containsExactly(getEntityCore(N1_N2)); + } + + private static PolarisEntityCore getEntityCore(PolarisBaseEntity entity) { + return new PolarisEntityCore.Builder<>(entity).build(); } /** Test the set of functions for the entity cache */ From cda9847fb6e599516c3c1d2f3525b4804884250e Mon Sep 17 00:00:00 2001 From: Michael Collado Date: Tue, 23 Sep 2025 15:55:48 -0700 Subject: [PATCH 5/9] Added additional test and updated comments in EntityCache interface --- .../core/persistence/cache/EntityCache.java | 16 + .../cache/InMemoryEntityCache.java | 23 +- .../cache/InMemoryEntityCacheTest.java | 284 +++++++++++++----- 3 files changed, 244 insertions(+), 79 deletions(-) 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 93dc87ac97..7cfea2a84b 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 @@ -87,6 +87,17 @@ EntityCacheLookupResult getOrLoadEntityByName( /** * Load multiple entities by id, returning those found in the cache and loading those not found. * + *

Cached entity versions and grant versions must be verified against the versions returned by + * the {@link + * org.apache.polaris.core.persistence.PolarisMetaStoreManager#loadEntitiesChangeTracking(PolarisCallContext, + * List)} API to ensure the returned entities are consistent with the current state of the + * metastore. Cache implementations must never return a mix of stale entities and fresh entities, + * as authorization or table conflict decisions could be made based on inconsistent data. For + * example, a Principal may have a grant to a Principal Role in a cached entry, but that grant may + * be revoked prior to the Principal Role being granted a privilege on a Catalog. If the Principal + * record is stale, but the Principal Role is refreshed, the Principal may be incorrectly + * authorized to access the Catalog. + * * @param callCtx the Polaris call context * @param entityType the entity type * @param entityIds the list of entity ids to load @@ -105,6 +116,11 @@ List getOrLoadResolvedEntities( * Load multiple entities by {@link EntityNameLookupRecord}, returning those found in the cache * and loading those not found. * + *

see note in {@link #getOrLoadResolvedEntities(PolarisCallContext, PolarisEntityType, List)} + * re: the need to validate cache contents against the {@link + * org.apache.polaris.core.persistence.PolarisMetaStoreManager#loadEntitiesChangeTracking(PolarisCallContext, + * List)} to avoid invalid authorization or conflict detection decisions based on stale entries. + * * @param callCtx the Polaris call context * @param lookupRecords the list of entity name to load * @return the list of resolved entities, in the same order as the requested entity records. As in diff --git a/polaris-core/src/main/java/org/apache/polaris/core/persistence/cache/InMemoryEntityCache.java b/polaris-core/src/main/java/org/apache/polaris/core/persistence/cache/InMemoryEntityCache.java index a10a1f5ead..e8b4163e79 100644 --- a/polaris-core/src/main/java/org/apache/polaris/core/persistence/cache/InMemoryEntityCache.java +++ b/polaris-core/src/main/java/org/apache/polaris/core/persistence/cache/InMemoryEntityCache.java @@ -58,6 +58,7 @@ /** An in-memory entity cache with a limit of 100k entities and a 1h TTL. */ public class InMemoryEntityCache implements EntityCache { private static final Logger LOGGER = LoggerFactory.getLogger(InMemoryEntityCache.class); + public static final int MAX_CACHE_REFRESH_ATTEMPTS = 100; private final PolarisDiagnostics diagnostics; private final PolarisMetaStoreManager polarisMetaStoreManager; private final Cache byId; @@ -477,13 +478,22 @@ public List getOrLoadResolvedEntities( // trying to populate // the cache from a different snapshot Map resolvedEntities = new HashMap<>(); - for (int i = 0; i < 100; i++) { + boolean stateResolved = false; + for (int i = 0; i < MAX_CACHE_REFRESH_ATTEMPTS; i++) { Function, ResolvedEntitiesResult> loaderFunc = idsToLoad -> polarisMetaStoreManager.loadResolvedEntities(callCtx, entityType, idsToLoad); if (isCacheStateValid(callCtx, resolvedEntities, entityIds, loaderFunc)) { + stateResolved = true; break; } } + if (!stateResolved) { + LOGGER.warn( + "Unable to resolve entities in cache after multiple attempts {} - resolved: {}", + entityIds, + resolvedEntities); + diagnostics.fail("cannot_resolve_all_entities", "Unable to resolve entities in cache"); + } return entityIds.stream() .map( @@ -517,11 +527,20 @@ public List getOrLoadResolvedEntities( lookupRecords.stream() .map(e -> new PolarisEntityId(e.getCatalogId(), e.getId())) .collect(Collectors.toList()); - for (int i = 0; i < 100; i++) { + boolean stateResolved = false; + for (int i = 0; i < MAX_CACHE_REFRESH_ATTEMPTS; i++) { if (isCacheStateValid(callCtx, resolvedEntities, entityIds, loaderFunc)) { + stateResolved = true; break; } } + if (!stateResolved) { + LOGGER.warn( + "Unable to resolve entities in cache after multiple attempts {} - resolved: {}", + entityIds, + resolvedEntities); + diagnostics.fail("cannot_resolve_all_entities", "Unable to resolve entities in cache"); + } return lookupRecords.stream() .map( diff --git a/polaris-core/src/test/java/org/apache/polaris/core/persistence/cache/InMemoryEntityCacheTest.java b/polaris-core/src/test/java/org/apache/polaris/core/persistence/cache/InMemoryEntityCacheTest.java index dd90f5027a..43d9c88206 100644 --- a/polaris-core/src/test/java/org/apache/polaris/core/persistence/cache/InMemoryEntityCacheTest.java +++ b/polaris-core/src/test/java/org/apache/polaris/core/persistence/cache/InMemoryEntityCacheTest.java @@ -53,8 +53,12 @@ import org.apache.polaris.core.persistence.transactional.TransactionalPersistence; import org.apache.polaris.core.persistence.transactional.TreeMapMetaStore; import org.apache.polaris.core.persistence.transactional.TreeMapTransactionalPersistenceImpl; +import org.assertj.core.api.Assertions; import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.ValueSource; import org.mockito.Mockito; +import org.mockito.stubbing.Answer; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -632,8 +636,9 @@ public void testBatchLoadByNameLookupRecords() { assertThat(results.get(2).isCacheHit()).isTrue(); } - @Test - public void testBatchLoadWithStaleVersions() { + @ParameterizedTest + @ValueSource(strings = {"id", "name"}) + public void testBatchLoadWithStaleVersions(String loadType) { // get a new cache InMemoryEntityCache cache = this.allocateNewCache(); @@ -674,12 +679,18 @@ public void testBatchLoadWithStaleVersions() { assertThat(T6v2).isNotNull(); assertThat(T6v2.getEntityVersion()).isGreaterThan(T6v1.getEntityVersion()); - // Create entity ID list with the updated entity - List entityIds = List.of(getPolarisEntityId(T6v2)); - - // Call batch load - this should detect the stale version and reload - List results = - cache.getOrLoadResolvedEntities(this.callCtx, PolarisEntityType.TABLE_LIKE, entityIds); + List results; + if (loadType.equals("id")) { + // Create entity ID list with the updated entity + List entityIds = List.of(getPolarisEntityId(T6v2)); + + // Call batch load - this should detect the stale version and reload + results = + cache.getOrLoadResolvedEntities(this.callCtx, PolarisEntityType.TABLE_LIKE, entityIds); + } else { + results = + cache.getOrLoadResolvedEntities(this.callCtx, List.of(new EntityNameLookupRecord(T6v2))); + } // Verify the entity was reloaded with the new version assertThat(results).hasSize(1); @@ -695,8 +706,9 @@ public void testBatchLoadWithStaleVersions() { assertThat(cachedT6.getEntity().getEntityVersion()).isEqualTo(T6v2.getEntityVersion()); } - @Test - public void testBatchLoadWithStaleGrantVersions() { + @ParameterizedTest + @ValueSource(strings = {"id", "name"}) + public void testBatchLoadWithStaleGrantVersions(String loadType) { // get a new cache InMemoryEntityCache cache = this.allocateNewCache(); @@ -733,12 +745,19 @@ public void testBatchLoadWithStaleGrantVersions() { this.tm.ensureExistsByName(List.of(catalog, N1), PolarisEntityType.NAMESPACE, "N2"); assertThat(N2Updated.getGrantRecordsVersion()).isGreaterThan(originalGrantVersion); - // Create entity ID list - List entityIds = List.of(getPolarisEntityId(N2Updated)); - - // Call batch load - this should detect the stale grant version and reload - List results = - cache.getOrLoadResolvedEntities(this.callCtx, PolarisEntityType.NAMESPACE, entityIds); + List results; + if (loadType.equals("id")) { + // Create entity ID list + List entityIds = List.of(getPolarisEntityId(N2Updated)); + + // Call batch load - this should detect the stale grant version and reload + results = + cache.getOrLoadResolvedEntities(this.callCtx, PolarisEntityType.NAMESPACE, entityIds); + } else { + results = + cache.getOrLoadResolvedEntities( + this.callCtx, List.of(new EntityNameLookupRecord(N2Updated))); + } // Verify the entity was reloaded with the new grant version assertThat(results).hasSize(1); @@ -759,8 +778,9 @@ public void testBatchLoadWithStaleGrantVersions() { .isEqualTo(N2Updated.getGrantRecordsVersion()); } - @Test - public void testBatchLoadVersionRetryLogic() { + @ParameterizedTest + @ValueSource(strings = {"id", "name"}) + public void testBatchLoadVersionRetryLogic(String loadType) { // get a new cache PolarisMetaStoreManager metaStoreManager = Mockito.spy(this.metaStoreManager); InMemoryEntityCache cache = @@ -813,13 +833,24 @@ public void testBatchLoadVersionRetryLogic() { assertThat(N2v2.getEntityVersion()).isGreaterThan(originalEntityVersion); assertThat(N2v3.getEntityVersion()).isGreaterThan(N2v2.getEntityVersion()); - // Create entity ID list - List entityIds = - List.of(getPolarisEntityId(catalog), getPolarisEntityId(N1), getPolarisEntityId(N2)); - - // Call batch load - this should detect the stale versions and reload until consistent - List results = - cache.getOrLoadResolvedEntities(this.callCtx, PolarisEntityType.NAMESPACE, entityIds); + List results; + if (loadType.equals("id")) { + // Create entity ID list + List entityIds = + List.of(getPolarisEntityId(catalog), getPolarisEntityId(N1), getPolarisEntityId(N2)); + + // Call batch load - this should detect the stale versions and reload until consistent + results = + cache.getOrLoadResolvedEntities(this.callCtx, PolarisEntityType.NAMESPACE, entityIds); + } else { + results = + cache.getOrLoadResolvedEntities( + this.callCtx, + List.of( + new EntityNameLookupRecord(catalog), + new EntityNameLookupRecord(N1), + new EntityNameLookupRecord(N2))); + } // Verify the entity was reloaded with the latest version assertThat(results).hasSize(3); @@ -840,12 +871,86 @@ public void testBatchLoadVersionRetryLogic() { assertThat(cachedN2.getEntity().getEntityVersion()).isEqualTo(N2v3.getEntityVersion()); } + @ParameterizedTest + @ValueSource(strings = {"id", "name"}) + public void testBatchLoadVersionRetryFailsAfterMaxAttempts(String loadType) { + // get a new cache + PolarisMetaStoreManager metaStoreManager = Mockito.spy(this.metaStoreManager); + InMemoryEntityCache cache = + new InMemoryEntityCache(diagServices, callCtx.getRealmConfig(), metaStoreManager); + + // Load catalog into cache + EntityCacheLookupResult lookup = + cache.getOrLoadEntityByName( + this.callCtx, new EntityCacheByNameKey(PolarisEntityType.CATALOG, "test")); + assertThat(lookup).isNotNull(); + PolarisBaseEntity catalog = lookup.getCacheEntry().getEntity(); + + // Get entities that we can work with + PolarisBaseEntity N1 = + this.tm.ensureExistsByName(List.of(catalog), PolarisEntityType.NAMESPACE, "N1"); + PolarisBaseEntity N2 = + this.tm.ensureExistsByName(List.of(catalog, N1), PolarisEntityType.NAMESPACE, "N2"); + + // Load N2 into cache initially + cache.getOrLoadEntityByName( + this.callCtx, + new EntityCacheByNameKey(catalog.getId(), N1.getId(), PolarisEntityType.NAMESPACE, "N2")); + + // Verify it's in cache with original version + ResolvedPolarisEntity cachedN2 = cache.getEntityById(N2.getId()); + assertThat(cachedN2).isNotNull(); + int originalEntityVersion = cachedN2.getEntity().getEntityVersion(); + + // Update the entity multiple times to create version skew + PolarisBaseEntity N2v2 = + this.tm.updateEntity(List.of(catalog, N1), N2, "{\"v2\": \"value\"}", null); + + // return v2 for the change tracking + Mockito.doReturn( + new ChangeTrackingResult( + List.of( + changeTrackingFor(catalog), changeTrackingFor(N1), changeTrackingFor(N2v2)))) + .when(metaStoreManager) + .loadEntitiesChangeTracking(Mockito.any(), Mockito.any()); + + // update again to create v3, which isn't returned in the change tracking result + PolarisBaseEntity N2v3 = + this.tm.updateEntity(List.of(catalog, N1), N2v2, "{\"v3\": \"value\"}", null); + + // Verify versions increased + assertThat(N2v2.getEntityVersion()).isGreaterThan(originalEntityVersion); + assertThat(N2v3.getEntityVersion()).isGreaterThan(N2v2.getEntityVersion()); + + if (loadType.equals("id")) { + // Create entity ID list + List entityIds = + List.of(getPolarisEntityId(catalog), getPolarisEntityId(N1), getPolarisEntityId(N2)); + Assertions.assertThatThrownBy( + () -> + cache.getOrLoadResolvedEntities( + this.callCtx, PolarisEntityType.NAMESPACE, entityIds)) + .isInstanceOf(RuntimeException.class); + } else { + Assertions.assertThatThrownBy( + () -> + cache.getOrLoadResolvedEntities( + this.callCtx, + List.of( + new EntityNameLookupRecord(catalog), + new EntityNameLookupRecord(N1), + new EntityNameLookupRecord(N2)))) + .isInstanceOf(RuntimeException.class); + } + } + private static PolarisEntityId getPolarisEntityId(PolarisBaseEntity catalog) { return new PolarisEntityId(catalog.getCatalogId(), catalog.getId()); } - @Test - public void testConcurrentClientLoadingBehavior() throws Exception { + @ParameterizedTest + @ValueSource(strings = {"id", "name"}) + public void testConcurrentClientLoadingBehavior(String loadType) throws Exception { // Load catalog into cache EntityCacheLookupResult lookup = allocateNewCache() @@ -933,50 +1038,59 @@ public void testConcurrentClientLoadingBehavior() throws Exception { // Mock loadResolvedEntities to control timing - client1 blocks until client2 completes // client1 receives the older version of all entities, while client2 receives the newer version // of N2 - Mockito.doAnswer( - invocation -> { - // This is client1's loadResolvedEntities call - block until client2 completes - try { - client1ChangeTrackingResult.countDown(); - LOGGER.debug("Awaiting client2 to complete resolved entities load"); - client1ResolvedEntitiesBlocker.acquire(); // Block until client2 signals completion - List resolvedEntities = + Answer client1Answer = + invocation -> { + // This is client1's loadResolvedEntities call - block until client2 completes + try { + client1ChangeTrackingResult.countDown(); + LOGGER.debug("Awaiting client2 to complete resolved entities load"); + client1ResolvedEntitiesBlocker.acquire(); // Block until client2 signals completion + List resolvedEntities = + List.of( + getResolvedPolarisEntity(N1), + getResolvedPolarisEntity(N2v2), + getResolvedPolarisEntity(N3), + getResolvedPolarisEntity(N5), + getResolvedPolarisEntity(N5_N6)); + LOGGER.debug("Client1 returning results {}", resolvedEntities); + return new ResolvedEntitiesResult(resolvedEntities); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new RuntimeException(e); + } + }; + Answer client2Answer = + invocation -> { + // This is client2's loadResolvedEntities call - execute normally and signal client1 + try { + LOGGER.debug("Client2 loading resolved entities"); + var result = + new ResolvedEntitiesResult( List.of( getResolvedPolarisEntity(N1), - getResolvedPolarisEntity(N2v2), + getResolvedPolarisEntity(N2v3), getResolvedPolarisEntity(N3), getResolvedPolarisEntity(N5), - getResolvedPolarisEntity(N5_N6)); - LOGGER.debug("Client1 returning results {}", resolvedEntities); - return new ResolvedEntitiesResult(resolvedEntities); - } catch (InterruptedException e) { - Thread.currentThread().interrupt(); - throw new RuntimeException(e); - } - }) - .doAnswer( - invocation -> { - // This is client2's loadResolvedEntities call - execute normally and signal client1 - try { - LOGGER.debug("Client2 loading resolved entities"); - var result = - new ResolvedEntitiesResult( - List.of( - getResolvedPolarisEntity(N1), - getResolvedPolarisEntity(N2v3), - getResolvedPolarisEntity(N3), - getResolvedPolarisEntity(N5), - getResolvedPolarisEntity(N5_N6))); - client1ResolvedEntitiesBlocker.release(); // Allow client1 to proceed - LOGGER.debug("Client2 returning results {}", result.getResolvedEntities()); - return result; - } catch (Exception e) { - client1ResolvedEntitiesBlocker.release(); // Release in case of error - throw e; - } - }) - .when(mockedMetaStoreManager) - .loadResolvedEntities(Mockito.any(), Mockito.any(), Mockito.any()); + getResolvedPolarisEntity(N5_N6))); + client1ResolvedEntitiesBlocker.release(); // Allow client1 to proceed + LOGGER.debug("Client2 returning results {}", result.getResolvedEntities()); + return result; + } catch (Exception e) { + client1ResolvedEntitiesBlocker.release(); // Release in case of error + throw e; + } + }; + if (loadType.equals("id")) { + Mockito.doAnswer(client1Answer) + .doAnswer(client2Answer) + .when(mockedMetaStoreManager) + .loadResolvedEntities(Mockito.any(), Mockito.any(), Mockito.any()); + } else { + Mockito.doAnswer(client1Answer) + .doAnswer(client2Answer) + .when(mockedMetaStoreManager) + .loadResolvedEntities(Mockito.any(), Mockito.any()); + } // ExecutorService isn't AutoCloseable in JDK 11 :( ExecutorService executorService = Executors.newFixedThreadPool(2); @@ -989,11 +1103,19 @@ public void testConcurrentClientLoadingBehavior() throws Exception { // Client1 calls getOrLoadResolvedEntities // - loadEntitiesChangeTracking returns older version for N2 // - loadResolvedEntities will block until client2 completes - List results = - cache.getOrLoadResolvedEntities( - this.callCtx, PolarisEntityType.NAMESPACE, entityIds); - - return results; + if (loadType.equals("id")) { + return cache.getOrLoadResolvedEntities( + this.callCtx, PolarisEntityType.NAMESPACE, entityIds); + } else { + return cache.getOrLoadResolvedEntities( + this.callCtx, + List.of( + new EntityNameLookupRecord(N1), + new EntityNameLookupRecord(N2), + new EntityNameLookupRecord(N3), + new EntityNameLookupRecord(N5), + new EntityNameLookupRecord(N5_N6))); + } } catch (Exception e) { client1Exception.set(e); return null; @@ -1009,11 +1131,19 @@ public void testConcurrentClientLoadingBehavior() throws Exception { // - loadEntitiesChangeTracking returns newer version for N2 // - loadResolvedEntities executes normally and signals client1 when done client1ChangeTrackingResult.await(); - List results = - cache.getOrLoadResolvedEntities( - this.callCtx, PolarisEntityType.NAMESPACE, entityIds); - - return results; + if (loadType.equals("id")) { + return cache.getOrLoadResolvedEntities( + this.callCtx, PolarisEntityType.NAMESPACE, entityIds); + } else { + return cache.getOrLoadResolvedEntities( + this.callCtx, + List.of( + new EntityNameLookupRecord(N1), + new EntityNameLookupRecord(N2), + new EntityNameLookupRecord(N3), + new EntityNameLookupRecord(N5), + new EntityNameLookupRecord(N5_N6))); + } } catch (Exception e) { client2Exception.set(e); client1ResolvedEntitiesBlocker.release(); // Release in case of error From e830c883cdf573f82a047c16a975f863562e717f Mon Sep 17 00:00:00 2001 From: Michael Collado Date: Tue, 23 Sep 2025 16:28:18 -0700 Subject: [PATCH 6/9] Add additional constructor to ResolvedEntitiesResult --- .../dao/entity/ResolvedEntitiesResult.java | 16 +++++++++++++++- 1 file changed, 15 insertions(+), 1 deletion(-) diff --git a/polaris-core/src/main/java/org/apache/polaris/core/persistence/dao/entity/ResolvedEntitiesResult.java b/polaris-core/src/main/java/org/apache/polaris/core/persistence/dao/entity/ResolvedEntitiesResult.java index aec18414e8..91fc4673a5 100644 --- a/polaris-core/src/main/java/org/apache/polaris/core/persistence/dao/entity/ResolvedEntitiesResult.java +++ b/polaris-core/src/main/java/org/apache/polaris/core/persistence/dao/entity/ResolvedEntitiesResult.java @@ -19,11 +19,17 @@ package org.apache.polaris.core.persistence.dao.entity; +import com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonProperty; import jakarta.annotation.Nonnull; import jakarta.annotation.Nullable; -import java.util.List; import org.apache.polaris.core.persistence.ResolvedPolarisEntity; +import java.util.List; + +/** + * Response object for the loadResolvedEntities call. + */ public class ResolvedEntitiesResult extends BaseResult { private final List resolvedEntities; @@ -38,6 +44,14 @@ public ResolvedEntitiesResult( this.resolvedEntities = null; } + @JsonCreator + private ResolvedEntitiesResult(@JsonProperty("returnStatus") ReturnStatus returnStatus, + @JsonProperty("extraInformation") String extraInformation, + @JsonProperty("resolvedEntities") List resolvedEntities) { + super(returnStatus, extraInformation); + this.resolvedEntities = resolvedEntities; + } + public List getResolvedEntities() { return resolvedEntities; } From 20ef26810e74eea3f5c09fc779ee5a672300f2e8 Mon Sep 17 00:00:00 2001 From: Michael Collado Date: Tue, 23 Sep 2025 16:35:46 -0700 Subject: [PATCH 7/9] Fixed unused method reference --- .../AtomicOperationMetaStoreManager.java | 17 +---------------- .../dao/entity/ResolvedEntitiesResult.java | 14 ++++++-------- 2 files changed, 7 insertions(+), 24 deletions(-) diff --git a/polaris-core/src/main/java/org/apache/polaris/core/persistence/AtomicOperationMetaStoreManager.java b/polaris-core/src/main/java/org/apache/polaris/core/persistence/AtomicOperationMetaStoreManager.java index 88d1c73bbf..9dfd3b4f00 100644 --- a/polaris-core/src/main/java/org/apache/polaris/core/persistence/AtomicOperationMetaStoreManager.java +++ b/polaris-core/src/main/java/org/apache/polaris/core/persistence/AtomicOperationMetaStoreManager.java @@ -1805,22 +1805,7 @@ private static ResolvedEntitiesResult getResolvedEntitiesResult( return entities.get(i); } }) - .map( - e -> { - if (e == null) { - return null; - } else { - // load the grant records - final List grantRecordsAsSecurable = - ms.loadAllGrantRecordsOnSecurable(callCtx, e.getCatalogId(), e.getId()); - final List grantRecordsAsGrantee = - e.getType().isGrantee() - ? ms.loadAllGrantRecordsOnGrantee(callCtx, e.getCatalogId(), e.getId()) - : List.of(); - return new ResolvedPolarisEntity( - PolarisEntity.of(e), grantRecordsAsGrantee, grantRecordsAsSecurable); - } - }) + .map(e -> toResolvedPolarisEntity(callCtx, e, ms)) .collect(Collectors.toList()); return new ResolvedEntitiesResult(ret); } diff --git a/polaris-core/src/main/java/org/apache/polaris/core/persistence/dao/entity/ResolvedEntitiesResult.java b/polaris-core/src/main/java/org/apache/polaris/core/persistence/dao/entity/ResolvedEntitiesResult.java index 91fc4673a5..61eb27da1f 100644 --- a/polaris-core/src/main/java/org/apache/polaris/core/persistence/dao/entity/ResolvedEntitiesResult.java +++ b/polaris-core/src/main/java/org/apache/polaris/core/persistence/dao/entity/ResolvedEntitiesResult.java @@ -23,13 +23,10 @@ import com.fasterxml.jackson.annotation.JsonProperty; import jakarta.annotation.Nonnull; import jakarta.annotation.Nullable; -import org.apache.polaris.core.persistence.ResolvedPolarisEntity; - import java.util.List; +import org.apache.polaris.core.persistence.ResolvedPolarisEntity; -/** - * Response object for the loadResolvedEntities call. - */ +/** Response object for the loadResolvedEntities call. */ public class ResolvedEntitiesResult extends BaseResult { private final List resolvedEntities; @@ -45,9 +42,10 @@ public ResolvedEntitiesResult( } @JsonCreator - private ResolvedEntitiesResult(@JsonProperty("returnStatus") ReturnStatus returnStatus, - @JsonProperty("extraInformation") String extraInformation, - @JsonProperty("resolvedEntities") List resolvedEntities) { + private ResolvedEntitiesResult( + @JsonProperty("returnStatus") ReturnStatus returnStatus, + @JsonProperty("extraInformation") String extraInformation, + @JsonProperty("resolvedEntities") List resolvedEntities) { super(returnStatus, extraInformation); this.resolvedEntities = resolvedEntities; } From d1d7b2c4a8adcf08502c1425cc4e2079c692d38c Mon Sep 17 00:00:00 2001 From: Michael Collado Date: Wed, 8 Oct 2025 14:44:20 -0700 Subject: [PATCH 8/9] Removed loadResolvedEntities method with lookup record param --- .../AtomicOperationMetaStoreManager.java | 15 -- .../persistence/PolarisMetaStoreManager.java | 16 -- .../TransactionWorkspaceMetaStoreManager.java | 10 - .../core/persistence/cache/EntityCache.java | 21 -- .../cache/InMemoryEntityCache.java | 50 ---- .../TransactionalMetaStoreManagerImpl.java | 16 -- .../cache/InMemoryEntityCacheTest.java | 215 ++++-------------- .../BasePolarisMetaStoreManagerTest.java | 6 - .../PolarisTestMetaStoreManager.java | 163 ------------- 9 files changed, 44 insertions(+), 468 deletions(-) diff --git a/polaris-core/src/main/java/org/apache/polaris/core/persistence/AtomicOperationMetaStoreManager.java b/polaris-core/src/main/java/org/apache/polaris/core/persistence/AtomicOperationMetaStoreManager.java index 9dfd3b4f00..1ec8c89d41 100644 --- a/polaris-core/src/main/java/org/apache/polaris/core/persistence/AtomicOperationMetaStoreManager.java +++ b/polaris-core/src/main/java/org/apache/polaris/core/persistence/AtomicOperationMetaStoreManager.java @@ -1762,21 +1762,6 @@ private void revokeGrantRecord( return result; } - @Nonnull - @Override - public ResolvedEntitiesResult loadResolvedEntities( - @Nonnull PolarisCallContext callCtx, - @Nonnull List entityLookupRecords) { - BasePersistence ms = callCtx.getMetaStore(); - List entityIds = - entityLookupRecords.stream() - .map(r -> new PolarisEntityId(r.getCatalogId(), r.getId())) - .collect(Collectors.toList()); - Function entityTypeForIndex = - i -> entityLookupRecords.get(i).getType(); - return getResolvedEntitiesResult(callCtx, ms, entityIds, entityTypeForIndex); - } - @Nonnull @Override public ResolvedEntitiesResult loadResolvedEntities( 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 bab8530e0e..efa73a60b7 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 @@ -26,7 +26,6 @@ import org.apache.polaris.core.PolarisCallContext; import org.apache.polaris.core.auth.PolarisGrantManager; import org.apache.polaris.core.auth.PolarisSecretsManager; -import org.apache.polaris.core.entity.EntityNameLookupRecord; import org.apache.polaris.core.entity.LocationBasedEntity; import org.apache.polaris.core.entity.PolarisBaseEntity; import org.apache.polaris.core.entity.PolarisEntity; @@ -419,21 +418,6 @@ ResolvedEntityResult loadResolvedEntityByName( @Nonnull PolarisEntityType entityType, @Nonnull String entityName); - /** - * Load a batch of resolved entities given their {@link EntityNameLookupRecord}. Will return an - * empty list if the input list is empty. Order in that returned list is the same as the input - * list. Some elements might be NULL if the entity has been dropped. - * - * @param callCtx call context - * @param entityLookupRecords the list of entity lookup records to load - * @return a non-null list of entities corresponding to the lookup keys. Some elements might be - * NULL if the entity has been dropped. - */ - @Nonnull - ResolvedEntitiesResult loadResolvedEntities( - @Nonnull PolarisCallContext callCtx, - @Nonnull List entityLookupRecords); - /** * Load a batch of resolved entities of a specified entity type given their {@link * PolarisEntityId}. Will return an empty list if the input list is empty. Order in that returned 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 39c1f908f7..99c1f81624 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 @@ -28,7 +28,6 @@ import java.util.Set; import org.apache.polaris.core.PolarisCallContext; import org.apache.polaris.core.PolarisDiagnostics; -import org.apache.polaris.core.entity.EntityNameLookupRecord; import org.apache.polaris.core.entity.LocationBasedEntity; import org.apache.polaris.core.entity.PolarisBaseEntity; import org.apache.polaris.core.entity.PolarisEntity; @@ -381,15 +380,6 @@ public ResolvedEntityResult loadResolvedEntityByName( return null; } - @Nonnull - @Override - public ResolvedEntitiesResult loadResolvedEntities( - @Nonnull PolarisCallContext callCtx, - @Nonnull List entityLookupRecords) { - diagnostics.fail("illegal_method_in_transaction_workspace", "loadEntities"); - return null; - } - @Nonnull @Override public ResolvedEntitiesResult loadResolvedEntities( 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 7cfea2a84b..481fa7a9e3 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 @@ -22,7 +22,6 @@ import jakarta.annotation.Nullable; import java.util.List; import org.apache.polaris.core.PolarisCallContext; -import org.apache.polaris.core.entity.EntityNameLookupRecord; import org.apache.polaris.core.entity.PolarisBaseEntity; import org.apache.polaris.core.entity.PolarisEntityId; import org.apache.polaris.core.entity.PolarisEntityType; @@ -111,24 +110,4 @@ List getOrLoadResolvedEntities( @Nonnull PolarisCallContext callCtx, @Nonnull PolarisEntityType entityType, @Nonnull List entityIds); - - /** - * Load multiple entities by {@link EntityNameLookupRecord}, returning those found in the cache - * and loading those not found. - * - *

see note in {@link #getOrLoadResolvedEntities(PolarisCallContext, PolarisEntityType, List)} - * re: the need to validate cache contents against the {@link - * org.apache.polaris.core.persistence.PolarisMetaStoreManager#loadEntitiesChangeTracking(PolarisCallContext, - * List)} to avoid invalid authorization or conflict detection decisions based on stale entries. - * - * @param callCtx the Polaris call context - * @param lookupRecords the list of entity name to load - * @return the list of resolved entities, in the same order as the requested entity records. As in - * {@link - * org.apache.polaris.core.persistence.PolarisMetaStoreManager#loadResolvedEntities(PolarisCallContext, - * PolarisEntityType, List)}, elements in the returned list may be null if the corresponding - * entity id does not exist. - */ - List getOrLoadResolvedEntities( - @Nonnull PolarisCallContext callCtx, @Nonnull List lookupRecords); } diff --git a/polaris-core/src/main/java/org/apache/polaris/core/persistence/cache/InMemoryEntityCache.java b/polaris-core/src/main/java/org/apache/polaris/core/persistence/cache/InMemoryEntityCache.java index e8b4163e79..784bcf2a41 100644 --- a/polaris-core/src/main/java/org/apache/polaris/core/persistence/cache/InMemoryEntityCache.java +++ b/polaris-core/src/main/java/org/apache/polaris/core/persistence/cache/InMemoryEntityCache.java @@ -28,7 +28,6 @@ import org.apache.polaris.core.config.BehaviorChangeConfiguration; import org.apache.polaris.core.config.FeatureConfiguration; import org.apache.polaris.core.config.RealmConfig; -import org.apache.polaris.core.entity.EntityNameLookupRecord; import org.apache.polaris.core.entity.PolarisBaseEntity; import org.apache.polaris.core.entity.PolarisChangeTrackingVersions; import org.apache.polaris.core.entity.PolarisEntityId; @@ -504,55 +503,6 @@ public List getOrLoadResolvedEntities( .collect(Collectors.toList()); } - @Override - public List getOrLoadResolvedEntities( - @Nonnull PolarisCallContext callCtx, @Nonnull List lookupRecords) { - Map entityIdMap = - lookupRecords.stream() - .collect( - Collectors.toMap( - e -> new PolarisEntityId(e.getCatalogId(), e.getId()), - Function.identity(), - (a, b) -> a)); - Function, ResolvedEntitiesResult> loaderFunc = - idsToLoad -> - polarisMetaStoreManager.loadResolvedEntities( - callCtx, idsToLoad.stream().map(entityIdMap::get).collect(Collectors.toList())); - - // use a map to collect cached entries to avoid concurrency problems in case a second thread is - // trying to populate - // the cache from a different snapshot - Map resolvedEntities = new HashMap<>(); - List entityIds = - lookupRecords.stream() - .map(e -> new PolarisEntityId(e.getCatalogId(), e.getId())) - .collect(Collectors.toList()); - boolean stateResolved = false; - for (int i = 0; i < MAX_CACHE_REFRESH_ATTEMPTS; i++) { - if (isCacheStateValid(callCtx, resolvedEntities, entityIds, loaderFunc)) { - stateResolved = true; - break; - } - } - if (!stateResolved) { - LOGGER.warn( - "Unable to resolve entities in cache after multiple attempts {} - resolved: {}", - entityIds, - resolvedEntities); - diagnostics.fail("cannot_resolve_all_entities", "Unable to resolve entities in cache"); - } - - return lookupRecords.stream() - .map( - lookupRecord -> { - ResolvedPolarisEntity entity = - resolvedEntities.get( - new PolarisEntityId(lookupRecord.getCatalogId(), lookupRecord.getId())); - return entity == null ? null : new EntityCacheLookupResult(entity, true); - }) - .collect(Collectors.toList()); - } - private boolean isCacheStateValid( @Nonnull PolarisCallContext callCtx, @Nonnull Map resolvedEntities, diff --git a/polaris-core/src/main/java/org/apache/polaris/core/persistence/transactional/TransactionalMetaStoreManagerImpl.java b/polaris-core/src/main/java/org/apache/polaris/core/persistence/transactional/TransactionalMetaStoreManagerImpl.java index 87fa9ac044..0bbaf1e6a9 100644 --- a/polaris-core/src/main/java/org/apache/polaris/core/persistence/transactional/TransactionalMetaStoreManagerImpl.java +++ b/polaris-core/src/main/java/org/apache/polaris/core/persistence/transactional/TransactionalMetaStoreManagerImpl.java @@ -2336,22 +2336,6 @@ private static ResolvedEntitiesResult getResolvedEntitiesResult( return new ResolvedEntitiesResult(ret); } - @Nonnull - @Override - public ResolvedEntitiesResult loadResolvedEntities( - @Nonnull PolarisCallContext callCtx, - @Nonnull List entityLookupRecords) { - TransactionalPersistence ms = ((TransactionalPersistence) callCtx.getMetaStore()); - List entityIds = - entityLookupRecords.stream() - .map(r -> new PolarisEntityId(r.getCatalogId(), r.getId())) - .collect(Collectors.toList()); - Function entityTypeForIndex = - i -> entityLookupRecords.get(i).getType(); - return ms.runInReadTransaction( - callCtx, () -> getResolvedEntitiesResult(callCtx, ms, entityIds, entityTypeForIndex)); - } - @Nonnull @Override public ResolvedEntitiesResult loadResolvedEntities( diff --git a/polaris-core/src/test/java/org/apache/polaris/core/persistence/cache/InMemoryEntityCacheTest.java b/polaris-core/src/test/java/org/apache/polaris/core/persistence/cache/InMemoryEntityCacheTest.java index 43d9c88206..95774364d9 100644 --- a/polaris-core/src/test/java/org/apache/polaris/core/persistence/cache/InMemoryEntityCacheTest.java +++ b/polaris-core/src/test/java/org/apache/polaris/core/persistence/cache/InMemoryEntityCacheTest.java @@ -34,7 +34,6 @@ import org.apache.polaris.core.PolarisCallContext; import org.apache.polaris.core.PolarisDefaultDiagServiceImpl; import org.apache.polaris.core.PolarisDiagnostics; -import org.apache.polaris.core.entity.EntityNameLookupRecord; import org.apache.polaris.core.entity.PolarisBaseEntity; import org.apache.polaris.core.entity.PolarisChangeTrackingVersions; import org.apache.polaris.core.entity.PolarisEntity; @@ -55,8 +54,6 @@ import org.apache.polaris.core.persistence.transactional.TreeMapTransactionalPersistenceImpl; import org.assertj.core.api.Assertions; import org.junit.jupiter.api.Test; -import org.junit.jupiter.params.ParameterizedTest; -import org.junit.jupiter.params.provider.ValueSource; import org.mockito.Mockito; import org.mockito.stubbing.Answer; import org.slf4j.Logger; @@ -582,63 +579,7 @@ public void testBatchLoadByEntityIds() { } @Test - public void testBatchLoadByNameLookupRecords() { - // get a new cache - InMemoryEntityCache cache = this.allocateNewCache(); - - // Load catalog into cache - EntityCacheLookupResult lookup = - cache.getOrLoadEntityByName( - this.callCtx, new EntityCacheByNameKey(PolarisEntityType.CATALOG, "test")); - assertThat(lookup).isNotNull(); - PolarisBaseEntity catalog = lookup.getCacheEntry().getEntity(); - - // Get some entities that exist in the test setup - PolarisBaseEntity N1 = - this.tm.ensureExistsByName(List.of(catalog), PolarisEntityType.NAMESPACE, "N1"); - PolarisBaseEntity N2 = - this.tm.ensureExistsByName(List.of(catalog, N1), PolarisEntityType.NAMESPACE, "N2"); - PolarisBaseEntity R1 = - this.tm.ensureExistsByName(List.of(catalog), PolarisEntityType.CATALOG_ROLE, "R1"); - - // Pre-load N1 into cache by name - cache.getOrLoadEntityByName( - this.callCtx, - new EntityCacheByNameKey( - catalog.getId(), catalog.getId(), PolarisEntityType.NAMESPACE, "N1")); - - // Create list of EntityNameLookupRecords - List lookupRecords = - List.of( - new EntityNameLookupRecord(N1), // already cached - new EntityNameLookupRecord(N2), // not cached - new EntityNameLookupRecord(R1) // not cached - ); - - // Test batch loading by name lookup records - List results = - cache.getOrLoadResolvedEntities(this.callCtx, lookupRecords); - - // Verify all entities were found - assertThat(results).hasSize(3); - assertThat(results.get(0)).isNotNull(); // N1 - was cached - assertThat(results.get(1)).isNotNull(); // N2 - was loaded - assertThat(results.get(2)).isNotNull(); // R1 - was loaded - - // Verify the entities are correct - assertThat(results.get(0).getCacheEntry().getEntity().getId()).isEqualTo(N1.getId()); - assertThat(results.get(1).getCacheEntry().getEntity().getId()).isEqualTo(N2.getId()); - assertThat(results.get(2).getCacheEntry().getEntity().getId()).isEqualTo(R1.getId()); - - // All should be cache hits now - assertThat(results.get(0).isCacheHit()).isTrue(); - assertThat(results.get(1).isCacheHit()).isTrue(); - assertThat(results.get(2).isCacheHit()).isTrue(); - } - - @ParameterizedTest - @ValueSource(strings = {"id", "name"}) - public void testBatchLoadWithStaleVersions(String loadType) { + public void testBatchLoadWithStaleVersions() { // get a new cache InMemoryEntityCache cache = this.allocateNewCache(); @@ -680,17 +621,12 @@ public void testBatchLoadWithStaleVersions(String loadType) { assertThat(T6v2.getEntityVersion()).isGreaterThan(T6v1.getEntityVersion()); List results; - if (loadType.equals("id")) { - // Create entity ID list with the updated entity - List entityIds = List.of(getPolarisEntityId(T6v2)); - - // Call batch load - this should detect the stale version and reload - results = - cache.getOrLoadResolvedEntities(this.callCtx, PolarisEntityType.TABLE_LIKE, entityIds); - } else { - results = - cache.getOrLoadResolvedEntities(this.callCtx, List.of(new EntityNameLookupRecord(T6v2))); - } + // Create entity ID list with the updated entity + List entityIds = List.of(getPolarisEntityId(T6v2)); + + // Call batch load - this should detect the stale version and reload + results = + cache.getOrLoadResolvedEntities(this.callCtx, PolarisEntityType.TABLE_LIKE, entityIds); // Verify the entity was reloaded with the new version assertThat(results).hasSize(1); @@ -706,9 +642,8 @@ public void testBatchLoadWithStaleVersions(String loadType) { assertThat(cachedT6.getEntity().getEntityVersion()).isEqualTo(T6v2.getEntityVersion()); } - @ParameterizedTest - @ValueSource(strings = {"id", "name"}) - public void testBatchLoadWithStaleGrantVersions(String loadType) { + @Test + public void testBatchLoadWithStaleGrantVersions() { // get a new cache InMemoryEntityCache cache = this.allocateNewCache(); @@ -745,19 +680,12 @@ public void testBatchLoadWithStaleGrantVersions(String loadType) { this.tm.ensureExistsByName(List.of(catalog, N1), PolarisEntityType.NAMESPACE, "N2"); assertThat(N2Updated.getGrantRecordsVersion()).isGreaterThan(originalGrantVersion); - List results; - if (loadType.equals("id")) { - // Create entity ID list - List entityIds = List.of(getPolarisEntityId(N2Updated)); - - // Call batch load - this should detect the stale grant version and reload - results = - cache.getOrLoadResolvedEntities(this.callCtx, PolarisEntityType.NAMESPACE, entityIds); - } else { - results = - cache.getOrLoadResolvedEntities( - this.callCtx, List.of(new EntityNameLookupRecord(N2Updated))); - } + // Create entity ID list + List entityIds = List.of(getPolarisEntityId(N2Updated)); + + // Call batch load - this should detect the stale grant version and reload + List results = + cache.getOrLoadResolvedEntities(this.callCtx, PolarisEntityType.NAMESPACE, entityIds); // Verify the entity was reloaded with the new grant version assertThat(results).hasSize(1); @@ -778,9 +706,8 @@ public void testBatchLoadWithStaleGrantVersions(String loadType) { .isEqualTo(N2Updated.getGrantRecordsVersion()); } - @ParameterizedTest - @ValueSource(strings = {"id", "name"}) - public void testBatchLoadVersionRetryLogic(String loadType) { + @Test + public void testBatchLoadVersionRetryLogic() { // get a new cache PolarisMetaStoreManager metaStoreManager = Mockito.spy(this.metaStoreManager); InMemoryEntityCache cache = @@ -833,24 +760,13 @@ public void testBatchLoadVersionRetryLogic(String loadType) { assertThat(N2v2.getEntityVersion()).isGreaterThan(originalEntityVersion); assertThat(N2v3.getEntityVersion()).isGreaterThan(N2v2.getEntityVersion()); - List results; - if (loadType.equals("id")) { - // Create entity ID list - List entityIds = - List.of(getPolarisEntityId(catalog), getPolarisEntityId(N1), getPolarisEntityId(N2)); - - // Call batch load - this should detect the stale versions and reload until consistent - results = - cache.getOrLoadResolvedEntities(this.callCtx, PolarisEntityType.NAMESPACE, entityIds); - } else { - results = - cache.getOrLoadResolvedEntities( - this.callCtx, - List.of( - new EntityNameLookupRecord(catalog), - new EntityNameLookupRecord(N1), - new EntityNameLookupRecord(N2))); - } + // Create entity ID list + List entityIds = + List.of(getPolarisEntityId(catalog), getPolarisEntityId(N1), getPolarisEntityId(N2)); + + // Call batch load - this should detect the stale versions and reload until consistent + List results = + cache.getOrLoadResolvedEntities(this.callCtx, PolarisEntityType.NAMESPACE, entityIds); // Verify the entity was reloaded with the latest version assertThat(results).hasSize(3); @@ -871,9 +787,8 @@ public void testBatchLoadVersionRetryLogic(String loadType) { assertThat(cachedN2.getEntity().getEntityVersion()).isEqualTo(N2v3.getEntityVersion()); } - @ParameterizedTest - @ValueSource(strings = {"id", "name"}) - public void testBatchLoadVersionRetryFailsAfterMaxAttempts(String loadType) { + @Test + public void testBatchLoadVersionRetryFailsAfterMaxAttempts() { // get a new cache PolarisMetaStoreManager metaStoreManager = Mockito.spy(this.metaStoreManager); InMemoryEntityCache cache = @@ -922,35 +837,22 @@ public void testBatchLoadVersionRetryFailsAfterMaxAttempts(String loadType) { assertThat(N2v2.getEntityVersion()).isGreaterThan(originalEntityVersion); assertThat(N2v3.getEntityVersion()).isGreaterThan(N2v2.getEntityVersion()); - if (loadType.equals("id")) { - // Create entity ID list - List entityIds = - List.of(getPolarisEntityId(catalog), getPolarisEntityId(N1), getPolarisEntityId(N2)); - Assertions.assertThatThrownBy( - () -> - cache.getOrLoadResolvedEntities( - this.callCtx, PolarisEntityType.NAMESPACE, entityIds)) - .isInstanceOf(RuntimeException.class); - } else { - Assertions.assertThatThrownBy( - () -> - cache.getOrLoadResolvedEntities( - this.callCtx, - List.of( - new EntityNameLookupRecord(catalog), - new EntityNameLookupRecord(N1), - new EntityNameLookupRecord(N2)))) - .isInstanceOf(RuntimeException.class); - } + // Create entity ID list + List entityIds = + List.of(getPolarisEntityId(catalog), getPolarisEntityId(N1), getPolarisEntityId(N2)); + Assertions.assertThatThrownBy( + () -> + cache.getOrLoadResolvedEntities( + this.callCtx, PolarisEntityType.NAMESPACE, entityIds)) + .isInstanceOf(RuntimeException.class); } private static PolarisEntityId getPolarisEntityId(PolarisBaseEntity catalog) { return new PolarisEntityId(catalog.getCatalogId(), catalog.getId()); } - @ParameterizedTest - @ValueSource(strings = {"id", "name"}) - public void testConcurrentClientLoadingBehavior(String loadType) throws Exception { + @Test + public void testConcurrentClientLoadingBehavior() throws Exception { // Load catalog into cache EntityCacheLookupResult lookup = allocateNewCache() @@ -1080,17 +982,10 @@ public void testConcurrentClientLoadingBehavior(String loadType) throws Exceptio throw e; } }; - if (loadType.equals("id")) { - Mockito.doAnswer(client1Answer) - .doAnswer(client2Answer) - .when(mockedMetaStoreManager) - .loadResolvedEntities(Mockito.any(), Mockito.any(), Mockito.any()); - } else { - Mockito.doAnswer(client1Answer) - .doAnswer(client2Answer) - .when(mockedMetaStoreManager) - .loadResolvedEntities(Mockito.any(), Mockito.any()); - } + Mockito.doAnswer(client1Answer) + .doAnswer(client2Answer) + .when(mockedMetaStoreManager) + .loadResolvedEntities(Mockito.any(), Mockito.any(), Mockito.any()); // ExecutorService isn't AutoCloseable in JDK 11 :( ExecutorService executorService = Executors.newFixedThreadPool(2); @@ -1103,19 +998,8 @@ public void testConcurrentClientLoadingBehavior(String loadType) throws Exceptio // Client1 calls getOrLoadResolvedEntities // - loadEntitiesChangeTracking returns older version for N2 // - loadResolvedEntities will block until client2 completes - if (loadType.equals("id")) { - return cache.getOrLoadResolvedEntities( - this.callCtx, PolarisEntityType.NAMESPACE, entityIds); - } else { - return cache.getOrLoadResolvedEntities( - this.callCtx, - List.of( - new EntityNameLookupRecord(N1), - new EntityNameLookupRecord(N2), - new EntityNameLookupRecord(N3), - new EntityNameLookupRecord(N5), - new EntityNameLookupRecord(N5_N6))); - } + return cache.getOrLoadResolvedEntities( + this.callCtx, PolarisEntityType.NAMESPACE, entityIds); } catch (Exception e) { client1Exception.set(e); return null; @@ -1131,19 +1015,8 @@ public void testConcurrentClientLoadingBehavior(String loadType) throws Exceptio // - loadEntitiesChangeTracking returns newer version for N2 // - loadResolvedEntities executes normally and signals client1 when done client1ChangeTrackingResult.await(); - if (loadType.equals("id")) { - return cache.getOrLoadResolvedEntities( - this.callCtx, PolarisEntityType.NAMESPACE, entityIds); - } else { - return cache.getOrLoadResolvedEntities( - this.callCtx, - List.of( - new EntityNameLookupRecord(N1), - new EntityNameLookupRecord(N2), - new EntityNameLookupRecord(N3), - new EntityNameLookupRecord(N5), - new EntityNameLookupRecord(N5_N6))); - } + return cache.getOrLoadResolvedEntities( + this.callCtx, PolarisEntityType.NAMESPACE, entityIds); } catch (Exception e) { client2Exception.set(e); client1ResolvedEntitiesBlocker.release(); // Release in case of error diff --git a/polaris-core/src/testFixtures/java/org/apache/polaris/core/persistence/BasePolarisMetaStoreManagerTest.java b/polaris-core/src/testFixtures/java/org/apache/polaris/core/persistence/BasePolarisMetaStoreManagerTest.java index 0675e7bf1e..f8816260a2 100644 --- a/polaris-core/src/testFixtures/java/org/apache/polaris/core/persistence/BasePolarisMetaStoreManagerTest.java +++ b/polaris-core/src/testFixtures/java/org/apache/polaris/core/persistence/BasePolarisMetaStoreManagerTest.java @@ -238,12 +238,6 @@ protected void testLookup() { polarisTestMetaStoreManager.testLookup(); } - /** test batch entity load */ - @Test - protected void testLoadResolvedEntities() { - polarisTestMetaStoreManager.testLoadResolvedEntities(); - } - /** test batch entity load */ @Test protected void testLoadResolvedEntitiesById() { 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 6c7ff7d6fb..b32a4059da 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 @@ -2689,169 +2689,6 @@ public void testLookup() { this.ensureNotExistsById(catalog.getId(), T1.getId(), PolarisEntityType.NAMESPACE); } - public void testLoadResolvedEntities() { - // load all principals - List principals = - polarisMetaStoreManager - .listEntities( - this.polarisCallContext, - null, - PolarisEntityType.PRINCIPAL, - PolarisEntitySubType.NULL_SUBTYPE, - PageToken.readEverything()) - .getEntities(); - - // create new catalog - PolarisBaseEntity catalog = - new PolarisBaseEntity( - PolarisEntityConstants.getNullId(), - polarisMetaStoreManager.generateNewEntityId(this.polarisCallContext).getId(), - PolarisEntityType.CATALOG, - PolarisEntitySubType.NULL_SUBTYPE, - PolarisEntityConstants.getRootEntityId(), - "test"); - CreateCatalogResult catalogCreated = - polarisMetaStoreManager.createCatalog(this.polarisCallContext, catalog, List.of()); - Assertions.assertThat(catalogCreated).isNotNull(); - - // load the catalog again, since the grant versions are different - catalog = - polarisMetaStoreManager - .loadEntity( - polarisCallContext, - 0L, - catalogCreated.getCatalog().getId(), - PolarisEntityType.CATALOG) - .getEntity(); - - // now create all objects - PolarisBaseEntity N1 = this.createEntity(List.of(catalog), PolarisEntityType.NAMESPACE, "N1"); - PolarisBaseEntity N1_N2 = - this.createEntity(List.of(catalog, N1), PolarisEntityType.NAMESPACE, "N2"); - PolarisBaseEntity T1 = - this.createEntity( - List.of(catalog, N1, N1_N2), - PolarisEntityType.TABLE_LIKE, - PolarisEntitySubType.ICEBERG_TABLE, - "T1"); - - // batch load all entities. They should all be present and non-null - ResolvedEntitiesResult entitiesResult = - polarisMetaStoreManager.loadResolvedEntities( - polarisCallContext, - List.of( - new EntityNameLookupRecord(catalog), - new EntityNameLookupRecord(N1), - new EntityNameLookupRecord(N1_N2), - new EntityNameLookupRecord(T1))); - Assertions.assertThat(entitiesResult) - .isNotNull() - .returns(BaseResult.ReturnStatus.SUCCESS, ResolvedEntitiesResult::getReturnStatus) - .extracting( - ResolvedEntitiesResult::getResolvedEntities, - InstanceOfAssertFactories.list(ResolvedPolarisEntity.class)) - .hasSize(4) - .allSatisfy(entity -> Assertions.assertThat(entity).isNotNull()) - .extracting(r -> getEntityCore(r.getEntity())) - .containsExactly( - getEntityCore(catalog), getEntityCore(N1), getEntityCore(N1_N2), getEntityCore(T1)); - - ResolvedPolarisEntity catalogEntity = entitiesResult.getResolvedEntities().get(0); - Assertions.assertThat(catalogEntity) - .extracting(ResolvedPolarisEntity::getAllGrantRecords) - .isNotNull() - .asInstanceOf(InstanceOfAssertFactories.list(PolarisGrantRecord.class)) - .hasSize(2) - .containsExactlyInAnyOrder( - new PolarisGrantRecord( - 0L, - catalog.getId(), - catalogCreated.getCatalogAdminRole().getCatalogId(), - catalogCreated.getCatalogAdminRole().getId(), - PolarisPrivilege.CATALOG_MANAGE_ACCESS.getCode()), - new PolarisGrantRecord( - 0L, - catalog.getId(), - catalogCreated.getCatalogAdminRole().getCatalogId(), - catalogCreated.getCatalogAdminRole().getId(), - PolarisPrivilege.CATALOG_MANAGE_METADATA.getCode())); - - // try entities which do not exist - entitiesResult = - polarisMetaStoreManager.loadResolvedEntities( - polarisCallContext, - List.of( - new EntityNameLookupRecord( - catalog.getId(), - 27, - 0L, - "CATALOG_DOES_NOT_EXIST", - PolarisEntityType.CATALOG.getCode(), - PolarisEntitySubType.NULL_SUBTYPE.getCode()), - new EntityNameLookupRecord( - catalog.getId(), - 35, - 0L, - "PRINCIPAL_DOES_NOT_EXIST", - PolarisEntityType.PRINCIPAL.getCode(), - PolarisEntitySubType.NULL_SUBTYPE.getCode()))); - Assertions.assertThat(entitiesResult) - .isNotNull() - .returns(BaseResult.ReturnStatus.SUCCESS, ResolvedEntitiesResult::getReturnStatus) - .extracting( - ResolvedEntitiesResult::getResolvedEntities, - InstanceOfAssertFactories.list(ResolvedPolarisEntity.class)) - .hasSize(2) - .allSatisfy(entity -> Assertions.assertThat(entity).isNull()); - - // mix of existing entities and entities with wrong type - entitiesResult = - polarisMetaStoreManager.loadResolvedEntities( - polarisCallContext, - List.of( - new EntityNameLookupRecord( - catalog.getId(), - 27, - 0L, - "CATALOG_DOES_NOT_EXIST", - PolarisEntityType.CATALOG.getCode(), - PolarisEntitySubType.NULL_SUBTYPE.getCode()), - new EntityNameLookupRecord( - catalog.getId(), - 35, - 0L, - "PRINCIPAL_DOES_NOT_EXIST", - PolarisEntityType.PRINCIPAL.getCode(), - PolarisEntitySubType.NULL_SUBTYPE.getCode()), - new EntityNameLookupRecord(catalog), - new EntityNameLookupRecord( - N1.getCatalogId(), - N1.getId(), - N1.getParentId(), - N1.getName(), - PolarisEntityType.CATALOG_ROLE.getCode(), - PolarisEntitySubType.NULL_SUBTYPE.getCode()), - new EntityNameLookupRecord( - N1_N2.getCatalogId(), - N1_N2.getId(), - N1_N2.getParentId(), - N1_N2.getName(), - PolarisEntityType.TABLE_LIKE.getCode(), - PolarisEntitySubType.ANY_SUBTYPE.getCode()), - new EntityNameLookupRecord(T1))); - Assertions.assertThat(entitiesResult) - .isNotNull() - .returns(BaseResult.ReturnStatus.SUCCESS, ResolvedEntitiesResult::getReturnStatus) - .extracting( - ResolvedEntitiesResult::getResolvedEntities, - InstanceOfAssertFactories.list(ResolvedPolarisEntity.class)) - .hasSize(6) - .filteredOn(Objects::nonNull) - .hasSize(2) - .extracting(r -> getEntityCore(r.getEntity())) - .containsExactly(getEntityCore(catalog), getEntityCore(T1)); - } - public void testLoadResolvedEntitiesById() { // load all principals List principals = From beabf6bfebfc23f47bd829bd64b1af456b799cb4 Mon Sep 17 00:00:00 2001 From: Michael Collado Date: Wed, 8 Oct 2025 15:36:44 -0700 Subject: [PATCH 9/9] Pulled out toResolvedPolarisEntity method per PR comment --- .../jdbc/JdbcBasePersistenceImpl.java | 33 +++++++++-------- .../cache/InMemoryEntityCache.java | 25 +++++++------ .../TransactionalMetaStoreManagerImpl.java | 36 +++++++++---------- 3 files changed, 46 insertions(+), 48 deletions(-) diff --git a/persistence/relational-jdbc/src/main/java/org/apache/polaris/persistence/relational/jdbc/JdbcBasePersistenceImpl.java b/persistence/relational-jdbc/src/main/java/org/apache/polaris/persistence/relational/jdbc/JdbcBasePersistenceImpl.java index 7a032cfac9..9401df2dd0 100644 --- a/persistence/relational-jdbc/src/main/java/org/apache/polaris/persistence/relational/jdbc/JdbcBasePersistenceImpl.java +++ b/persistence/relational-jdbc/src/main/java/org/apache/polaris/persistence/relational/jdbc/JdbcBasePersistenceImpl.java @@ -18,9 +18,25 @@ */ package org.apache.polaris.persistence.relational.jdbc; +import static org.apache.polaris.persistence.relational.jdbc.QueryGenerator.PreparedQuery; + import com.google.common.base.Preconditions; import jakarta.annotation.Nonnull; import jakarta.annotation.Nullable; +import java.sql.Connection; +import java.sql.SQLException; +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashMap; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.Objects; +import java.util.Optional; +import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Function; +import java.util.function.Predicate; +import java.util.stream.Collectors; import org.apache.polaris.core.PolarisCallContext; import org.apache.polaris.core.PolarisDiagnostics; import org.apache.polaris.core.entity.EntityNameLookupRecord; @@ -62,23 +78,6 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; -import java.sql.Connection; -import java.sql.SQLException; -import java.util.ArrayList; -import java.util.Collections; -import java.util.HashMap; -import java.util.LinkedHashMap; -import java.util.List; -import java.util.Map; -import java.util.Objects; -import java.util.Optional; -import java.util.concurrent.atomic.AtomicReference; -import java.util.function.Function; -import java.util.function.Predicate; -import java.util.stream.Collectors; - -import static org.apache.polaris.persistence.relational.jdbc.QueryGenerator.PreparedQuery; - public class JdbcBasePersistenceImpl implements BasePersistence, IntegrationPersistence { private static final Logger LOGGER = LoggerFactory.getLogger(JdbcBasePersistenceImpl.class); diff --git a/polaris-core/src/main/java/org/apache/polaris/core/persistence/cache/InMemoryEntityCache.java b/polaris-core/src/main/java/org/apache/polaris/core/persistence/cache/InMemoryEntityCache.java index 784bcf2a41..cf5ad0c213 100644 --- a/polaris-core/src/main/java/org/apache/polaris/core/persistence/cache/InMemoryEntityCache.java +++ b/polaris-core/src/main/java/org/apache/polaris/core/persistence/cache/InMemoryEntityCache.java @@ -23,6 +23,18 @@ import com.github.benmanes.caffeine.cache.RemovalListener; import jakarta.annotation.Nonnull; import jakarta.annotation.Nullable; +import java.util.AbstractMap; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.Iterator; +import java.util.List; +import java.util.Map; +import java.util.Objects; +import java.util.Optional; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.TimeUnit; +import java.util.function.Function; +import java.util.stream.Collectors; import org.apache.polaris.core.PolarisCallContext; import org.apache.polaris.core.PolarisDiagnostics; import org.apache.polaris.core.config.BehaviorChangeConfiguration; @@ -41,19 +53,6 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; -import java.util.AbstractMap; -import java.util.ArrayList; -import java.util.HashMap; -import java.util.Iterator; -import java.util.List; -import java.util.Map; -import java.util.Objects; -import java.util.Optional; -import java.util.concurrent.ConcurrentHashMap; -import java.util.concurrent.TimeUnit; -import java.util.function.Function; -import java.util.stream.Collectors; - /** An in-memory entity cache with a limit of 100k entities and a 1h TTL. */ public class InMemoryEntityCache implements EntityCache { private static final Logger LOGGER = LoggerFactory.getLogger(InMemoryEntityCache.class); diff --git a/polaris-core/src/main/java/org/apache/polaris/core/persistence/transactional/TransactionalMetaStoreManagerImpl.java b/polaris-core/src/main/java/org/apache/polaris/core/persistence/transactional/TransactionalMetaStoreManagerImpl.java index 0bbaf1e6a9..db3ccd0f35 100644 --- a/polaris-core/src/main/java/org/apache/polaris/core/persistence/transactional/TransactionalMetaStoreManagerImpl.java +++ b/polaris-core/src/main/java/org/apache/polaris/core/persistence/transactional/TransactionalMetaStoreManagerImpl.java @@ -2314,28 +2314,28 @@ private static ResolvedEntitiesResult getResolvedEntitiesResult( return entities.get(i); } }) - .map( - e -> { - if (e == null) { - return null; - } else { - // load the grant records - final List grantRecordsAsSecurable = - ms.loadAllGrantRecordsOnSecurableInCurrentTxn( - callCtx, e.getCatalogId(), e.getId()); - final List grantRecordsAsGrantee = - e.getType().isGrantee() - ? ms.loadAllGrantRecordsOnGranteeInCurrentTxn( - callCtx, e.getCatalogId(), e.getId()) - : List.of(); - return new ResolvedPolarisEntity( - PolarisEntity.of(e), grantRecordsAsGrantee, grantRecordsAsSecurable); - } - }) + .map(e -> toResolvedPolarisEntity(callCtx, e, ms)) .collect(Collectors.toList()); return new ResolvedEntitiesResult(ret); } + private static ResolvedPolarisEntity toResolvedPolarisEntity( + PolarisCallContext callCtx, PolarisBaseEntity e, TransactionalPersistence ms) { + if (e == null) { + return null; + } else { + // load the grant records + final List grantRecordsAsSecurable = + ms.loadAllGrantRecordsOnSecurableInCurrentTxn(callCtx, e.getCatalogId(), e.getId()); + final List grantRecordsAsGrantee = + e.getType().isGrantee() + ? ms.loadAllGrantRecordsOnSecurableInCurrentTxn(callCtx, e.getCatalogId(), e.getId()) + : List.of(); + return new ResolvedPolarisEntity( + PolarisEntity.of(e), grantRecordsAsGrantee, grantRecordsAsSecurable); + } + } + @Nonnull @Override public ResolvedEntitiesResult loadResolvedEntities(