This is an automated email from the ASF dual-hosted git repository. dimas pushed a commit to branch dmitri-tmp1 in repository https://gitbox.apache.org/repos/asf/polaris.git
commit e8fee9be7ac2c8f18b5f9d62968b20f39318ab58 Author: Michael Collado <[email protected]> AuthorDate: Thu Sep 4 11:20:51 2025 -0700 Add loadEntities batch call and rename listFullEntities --- .../PolarisEclipseLinkMetaStoreSessionImpl.java | 10 +- .../relational/jdbc/JdbcBasePersistenceImpl.java | 42 ++++--- .../polaris/core/entity/PolarisBaseEntity.java | 6 - .../AtomicOperationMetaStoreManager.java | 37 +++++- .../polaris/core/persistence/BasePersistence.java | 4 +- .../core/persistence/PolarisMetaStoreManager.java | 29 ++++- .../TransactionWorkspaceMetaStoreManager.java | 12 +- .../AbstractTransactionalPersistence.java | 2 +- .../TransactionalMetaStoreManagerImpl.java | 51 +++++++- .../transactional/TransactionalPersistence.java | 2 +- .../BasePolarisMetaStoreManagerTest.java | 8 +- .../persistence/PolarisTestMetaStoreManager.java | 140 +++++++++++++++++++++ .../polaris/service/admin/PolarisAdminService.java | 8 +- .../service/catalog/policy/PolicyCatalog.java | 2 +- 14 files changed, 301 insertions(+), 52 deletions(-) diff --git a/persistence/eclipselink/src/main/java/org/apache/polaris/extension/persistence/impl/eclipselink/PolarisEclipseLinkMetaStoreSessionImpl.java b/persistence/eclipselink/src/main/java/org/apache/polaris/extension/persistence/impl/eclipselink/PolarisEclipseLinkMetaStoreSessionImpl.java index ccaf16cf0..1fa3e8667 100644 --- a/persistence/eclipselink/src/main/java/org/apache/polaris/extension/persistence/impl/eclipselink/PolarisEclipseLinkMetaStoreSessionImpl.java +++ b/persistence/eclipselink/src/main/java/org/apache/polaris/extension/persistence/impl/eclipselink/PolarisEclipseLinkMetaStoreSessionImpl.java @@ -378,9 +378,13 @@ public class PolarisEclipseLinkMetaStoreSessionImpl extends AbstractTransactiona @Override public @Nonnull List<PolarisBaseEntity> lookupEntitiesInCurrentTxn( @Nonnull PolarisCallContext callCtx, List<PolarisEntityId> entityIds) { - return this.store.lookupEntities(localSession.get(), entityIds).stream() - .map(ModelEntity::toEntity) - .toList(); + Map<PolarisEntityId, PolarisBaseEntity> idMap = + this.store.lookupEntities(localSession.get(), entityIds).stream() + .map(ModelEntity::toEntity) + .collect( + Collectors.toMap( + e -> new PolarisEntityId(e.getCatalogId(), e.getId()), Function.identity())); + return entityIds.stream().map(idMap::get).collect(Collectors.toList()); } /** {@inheritDoc} */ 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 8d6201453..7a032cfac 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.apache.polaris.persistence.relational.jdbc.models.SchemaVersion; 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 class JdbcBasePersistenceImpl implements BasePersistence, IntegrationPers PreparedQuery query = QueryGenerator.generateSelectQueryWithEntityIds(realmId, schemaVersion, entityIds); try { - return datasourceOperations.executeSelect(query, new ModelEntity(schemaVersion)); + Map<PolarisEntityId, PolarisBaseEntity> 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 class JdbcBasePersistenceImpl implements BasePersistence, IntegrationPers @Nonnull PolarisCallContext callCtx, List<PolarisEntityId> entityIds) { Map<PolarisEntityId, ModelEntity> idToEntityMap = lookupEntities(callCtx, entityIds).stream() + .filter(Objects::nonNull) .collect( Collectors.toMap( entry -> new PolarisEntityId(entry.getCatalogId(), entry.getId()), @@ -570,7 +578,7 @@ public class JdbcBasePersistenceImpl implements BasePersistence, IntegrationPers @Nonnull @Override - public <T> Page<T> loadEntities( + public <T> Page<T> 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 80b0a9640..6c91517a1 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 class PolarisBaseEntity extends PolarisEntityCore { 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 class PolarisBaseEntity extends PolarisEntityCore { 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 e2c46c151..e1628ded8 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.Set; 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 @@ public class AtomicOperationMetaStoreManager extends BaseMetaStoreManager { /** {@inheritDoc} */ @Override - public @Nonnull Page<PolarisBaseEntity> loadEntities( + public @Nonnull Page<PolarisBaseEntity> listFullEntities( @Nonnull PolarisCallContext callCtx, @Nullable List<PolarisEntityCore> catalogPath, @Nonnull PolarisEntityType entityType, @@ -728,7 +729,7 @@ public class AtomicOperationMetaStoreManager extends BaseMetaStoreManager { // 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, @@ -1203,7 +1204,7 @@ public class AtomicOperationMetaStoreManager extends BaseMetaStoreManager { // get the list of catalog roles, at most 2 List<PolarisBaseEntity> catalogRoles = - ms.loadEntities( + ms.listFullEntities( callCtx, catalogId, catalogId, @@ -1516,6 +1517,34 @@ public class AtomicOperationMetaStoreManager extends BaseMetaStoreManager { : new EntityResult(BaseResult.ReturnStatus.ENTITY_NOT_FOUND, null); } + @Nonnull + @Override + public EntitiesResult loadEntities( + @Nonnull PolarisCallContext callCtx, + @Nonnull List<EntityNameLookupRecord> entityLookupRecords) { + BasePersistence ms = callCtx.getMetaStore(); + List<PolarisBaseEntity> 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<PolarisBaseEntity> 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) { @@ -1523,7 +1552,7 @@ public class AtomicOperationMetaStoreManager extends BaseMetaStoreManager { // find all available tasks Page<PolarisBaseEntity> 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 4fa60e8c5..05aefa392 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 @@ public interface BasePersistence extends PolicyMappingPersistence { /** * 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 @@ public interface BasePersistence extends PolicyMappingPersistence { * @return the paged list of matching entities after transformation */ @Nonnull - <T> Page<T> loadEntities( + <T> Page<T> 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 566b10e64..f1b4109cf 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 java.util.Optional; 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 @@ public interface PolarisMetaStoreManager /** * 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 @@ public interface PolarisMetaStoreManager /** * 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 @@ public interface PolarisMetaStoreManager * @return paged list of matching entities */ @Nonnull - Page<PolarisBaseEntity> loadEntities( + Page<PolarisBaseEntity> listFullEntities( @Nonnull PolarisCallContext callCtx, @Nullable List<PolarisEntityCore> catalogPath, @Nonnull PolarisEntityType entityType, @@ -154,7 +155,7 @@ public interface PolarisMetaStoreManager /** * 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 @@ public interface PolarisMetaStoreManager * @param entitySubType subType of entities to list (or ANY_SUBTYPE) * @return list of all matching entities */ - default @Nonnull List<PolarisBaseEntity> loadEntitiesAll( + default @Nonnull List<PolarisBaseEntity> loadFullEntitiesAll( @Nonnull PolarisCallContext callCtx, @Nullable List<PolarisEntityCore> 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 @@ public interface PolarisMetaStoreManager 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<EntityNameLookupRecord> 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 872955893..bf16c7cb8 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.Optional; 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 class TransactionWorkspaceMetaStoreManager implements PolarisMetaStoreMan } @Override - public @Nonnull Page<PolarisBaseEntity> loadEntities( + public @Nonnull Page<PolarisBaseEntity> listFullEntities( @Nonnull PolarisCallContext callCtx, @Nullable List<PolarisEntityCore> catalogPath, @Nonnull PolarisEntityType entityType, @@ -330,6 +331,15 @@ public class TransactionWorkspaceMetaStoreManager implements PolarisMetaStoreMan return null; } + @Nonnull + @Override + public EntitiesResult loadEntities( + @Nonnull PolarisCallContext callCtx, + @Nonnull List<EntityNameLookupRecord> 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 848a8421e..a4d5024b6 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 abstract class AbstractTransactionalPersistence implements TransactionalP /** {@inheritDoc} */ @Override @Nonnull - public <T> Page<T> loadEntities( + public <T> Page<T> 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 402cdc280..2256a5614 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.Optional; 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 @@ public class TransactionalMetaStoreManagerImpl extends BaseMetaStoreManager { } /** - * See {@link PolarisMetaStoreManager#loadEntities(PolarisCallContext, List, PolarisEntityType, - * PolarisEntitySubType, PageToken)} + * See {@link PolarisMetaStoreManager#listFullEntities(PolarisCallContext, List, + * PolarisEntityType, PolarisEntitySubType, PageToken)} */ - private @Nonnull Page<PolarisBaseEntity> loadEntities( + private @Nonnull Page<PolarisBaseEntity> listFullEntities( @Nonnull PolarisCallContext callCtx, @Nonnull TransactionalPersistence ms, @Nullable List<PolarisEntityCore> catalogPath, @@ -756,7 +757,7 @@ public class TransactionalMetaStoreManagerImpl extends BaseMetaStoreManager { /** {@inheritDoc} */ @Override - public @Nonnull Page<PolarisBaseEntity> loadEntities( + public @Nonnull Page<PolarisBaseEntity> listFullEntities( @Nonnull PolarisCallContext callCtx, @Nullable List<PolarisEntityCore> catalogPath, @Nonnull PolarisEntityType entityType, @@ -768,7 +769,7 @@ public class TransactionalMetaStoreManagerImpl extends BaseMetaStoreManager { // 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 @@ public class TransactionalMetaStoreManagerImpl extends BaseMetaStoreManager { () -> 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<EntityNameLookupRecord> entityLookupRecords) { + List<PolarisBaseEntity> 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<PolarisBaseEntity> 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<EntityNameLookupRecord> 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 3802908b8..dfcd5f925 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 @@ public interface TransactionalPersistence pageToken); } - /** See {@link org.apache.polaris.core.persistence.BasePersistence#loadEntities} */ + /** See {@link org.apache.polaris.core.persistence.BasePersistence#listFullEntities} */ @Nonnull <T> Page<T> 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 43a0c5581..21122bda6 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 @@ public abstract class BasePolarisMetaStoreManagerTest { .containsExactly(PolarisEntity.toCore(task1), PolarisEntity.toCore(task2)); List<PolarisBaseEntity> listedEntities = - metaStoreManager.loadEntitiesAll( + metaStoreManager.loadFullEntitiesAll( polarisTestMetaStoreManager.polarisCallContext, null, PolarisEntityType.TASK, @@ -238,6 +238,12 @@ public abstract class BasePolarisMetaStoreManagerTest { 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 514a14c30..66bbbe82e 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.BaseResult; 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.PolicyEntity; 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 { @@ -655,6 +657,7 @@ public class PolarisTestMetaStoreManager { .parentId(parentId) .name(name) .propertiesAsMap(properties) + .internalPropertiesAsMap(Map.of()) .build(); PolarisBaseEntity entity = polarisMetaStoreManager @@ -2690,6 +2693,143 @@ public class PolarisTestMetaStoreManager { this.ensureNotExistsById(catalog.getId(), T1.getId(), PolarisEntityType.NAMESPACE); } + public void testBatchLoad() { + // load all principals + List<EntityNameLookupRecord> 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 9c9074dba..620009c78 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 @@ -979,7 +979,7 @@ public class PolarisAdminService { /** List all catalogs without checking for permission. */ private Stream<CatalogEntity> listCatalogsUnsafe() { return metaStoreManager - .loadEntitiesAll( + .loadFullEntitiesAll( getCurrentPolarisContext(), null, PolarisEntityType.CATALOG, @@ -1224,7 +1224,7 @@ public class PolarisAdminService { authorizeBasicRootOperationOrThrow(op); return metaStoreManager - .loadEntitiesAll( + .loadFullEntitiesAll( getCurrentPolarisContext(), null, PolarisEntityType.PRINCIPAL, @@ -1331,7 +1331,7 @@ public class PolarisAdminService { authorizeBasicRootOperationOrThrow(op); return metaStoreManager - .loadEntitiesAll( + .loadFullEntitiesAll( getCurrentPolarisContext(), null, PolarisEntityType.PRINCIPAL_ROLE, @@ -1455,7 +1455,7 @@ public class PolarisAdminService { CatalogEntity catalogEntity = getCatalogByName(resolutionManifest, catalogName); List<PolarisEntityCore> 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 8e4063b87..451eba049 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 class PolicyCatalog { } // with a policyType filter we need to load the full PolicyEntity to apply the filter return metaStoreManager - .loadEntitiesAll( + .loadFullEntitiesAll( callContext.getPolarisCallContext(), catalogPath, PolarisEntityType.POLICY,
