Add loadEntities batch call and rename listFullEntities (#2508)

* Add loadEntities batch call and rename listFullEntities

* Changed batch call to implement loadResolvedEntities instead

* Add loadResolvedEntities by id and entity cache support

* Add additional test for loadResolvedEntities by id

* Added additional test and updated comments in EntityCache interface

* Add additional constructor to ResolvedEntitiesResult

* Fixed unused method reference

* Removed loadResolvedEntities method with lookup record param

* Pulled out toResolvedPolarisEntity method per PR comment

Author: collado-mike

persistence/relational-jdbc/src/main/java/org/apache/polaris/persistence/relational/jdbc/JdbcBasePersistenceImpl.java

@@ -31,6 +31,7 @@
 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;
@@ -463,7 +464,12 @@ public List<PolarisBaseEntity> lookupEntities(
     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 +482,7 @@ public List<PolarisChangeTrackingVersions> lookupEntityVersions(
       @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 +577,7 @@ public Page<EntityNameLookupRecord> listEntities(
 
   @Nonnull
   @Override
-  public <T> Page<T> loadEntities(
+  public <T> Page<T> listFullEntities(
       @Nonnull PolarisCallContext callCtx,
       long catalogId,
       long parentId,

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;
@@ -67,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;
@@ -705,7 +707,7 @@ private void revokeGrantRecord(
 
   /** {@inheritDoc} */
   @Override
-  public @Nonnull Page<PolarisBaseEntity> loadEntities(
+  public @Nonnull Page<PolarisBaseEntity> listFullEntities(
       @Nonnull PolarisCallContext callCtx,
       @Nullable List<PolarisEntityCore> catalogPath,
       @Nonnull PolarisEntityType entityType,
@@ -728,7 +730,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 +1202,7 @@ private void revokeGrantRecord(
 
       // get the list of catalog roles, at most 2
       List<PolarisBaseEntity> catalogRoles =
-          ms.loadEntities(
+          ms.listFullEntities(
                   callCtx,
                   catalogId,
                   catalogId,
@@ -1520,7 +1522,7 @@ private void revokeGrantRecord(
 
     // find all available tasks
     Page<PolarisBaseEntity> availableTasks =
-        ms.loadEntities(
+        ms.listFullEntities(
             callCtx,
             PolarisEntityConstants.getRootEntityId(),
             PolarisEntityConstants.getRootEntityId(),
@@ -1760,6 +1762,56 @@ private void revokeGrantRecord(
     return result;
   }
 
+  @Nonnull
+  @Override
+  public ResolvedEntitiesResult loadResolvedEntities(
+      @Nonnull PolarisCallContext callCtx,
+      @Nonnull PolarisEntityType entityType,
+      @Nonnull List<PolarisEntityId> entityIds) {
+    BasePersistence ms = callCtx.getMetaStore();
+    return getResolvedEntitiesResult(callCtx, ms, entityIds, i -> entityType);
+  }
+
+  private static ResolvedEntitiesResult getResolvedEntitiesResult(
+      PolarisCallContext callCtx,
+      BasePersistence ms,
+      List<PolarisEntityId> entityIds,
+      Function<Integer, PolarisEntityType> entityTypeForIndex) {
+    List<PolarisBaseEntity> entities = ms.lookupEntities(callCtx, entityIds);
+    // mimic the behavior of loadEntity above, return null if not found or type mismatch
+    List<ResolvedPolarisEntity> ret =
+        IntStream.range(0, entityIds.size())
+            .mapToObj(
+                i -> {
+                  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))
+            .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<PolarisGrantRecord> grantRecordsAsSecurable =
+          ms.loadAllGrantRecordsOnSecurable(callCtx, e.getCatalogId(), e.getId());
+      final List<PolarisGrantRecord> 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(

polaris-core/src/main/java/org/apache/polaris/core/persistence/BasePersistence.java

@@ -279,7 +279,7 @@ List<PolarisChangeTrackingVersions> 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<EntityNameLookupRecord> listEntities(
    * @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,

polaris-core/src/main/java/org/apache/polaris/core/persistence/PolarisMetaStoreManager.java

@@ -47,6 +47,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;
@@ -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.
+   * #listFullEntitiesAll} 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<PolarisBaseEntity> loadEntities(
+  Page<PolarisBaseEntity> listFullEntities(
       @Nonnull PolarisCallContext callCtx,
       @Nullable List<PolarisEntityCore> catalogPath,
       @Nonnull PolarisEntityType entityType,
@@ -154,7 +155,7 @@ Page<PolarisBaseEntity> 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<PolarisBaseEntity> loadEntities(
    * @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> listFullEntitiesAll(
       @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();
   }
 
@@ -416,6 +418,23 @@ ResolvedEntityResult loadResolvedEntityByName(
       @Nonnull PolarisEntityType entityType,
       @Nonnull String entityName);
 
+  /**
+   * 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<PolarisEntityId> 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

polaris-core/src/main/java/org/apache/polaris/core/persistence/TransactionWorkspaceMetaStoreManager.java

@@ -53,6 +53,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;
@@ -133,7 +134,7 @@ public EntityResult readEntityByName(
   }
 
   @Override
-  public @Nonnull Page<PolarisBaseEntity> loadEntities(
+  public @Nonnull Page<PolarisBaseEntity> listFullEntities(
       @Nonnull PolarisCallContext callCtx,
       @Nullable List<PolarisEntityCore> catalogPath,
       @Nonnull PolarisEntityType entityType,
@@ -379,6 +380,16 @@ public ResolvedEntityResult loadResolvedEntityByName(
     return null;
   }
 
+  @Nonnull
+  @Override
+  public ResolvedEntitiesResult loadResolvedEntities(
+      @Nonnull PolarisCallContext callCtx,
+      @Nonnull PolarisEntityType entityType,
+      @Nonnull List<PolarisEntityId> entityIds) {
+    diagnostics.fail("illegal_method_in_transaction_workspace", "loadResolvedEntities");
+    return null;
+  }
+
   @Override
   public ResolvedEntityResult refreshResolvedEntity(
       @Nonnull PolarisCallContext callCtx,

polaris-core/src/main/java/org/apache/polaris/core/persistence/cache/EntityCache.java

@@ -20,8 +20,10 @@
 
 import jakarta.annotation.Nonnull;
 import jakarta.annotation.Nullable;
+import java.util.List;
 import org.apache.polaris.core.PolarisCallContext;
 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 +82,32 @@ 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.
+   *
+   * <p>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
+   * @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<EntityCacheLookupResult> getOrLoadResolvedEntities(
+      @Nonnull PolarisCallContext callCtx,
+      @Nonnull PolarisEntityType entityType,
+      @Nonnull List<PolarisEntityId> entityIds);
 }