Skip to content

Commit 1e05a56

Browse files
collado-mikecollado-mike
committed
Add additional test for loadResolvedEntities by id
1 parent 649e71a commit 1e05a56

File tree

4 files changed

+162
-31
lines changed

4 files changed

+162
-31
lines changed

‎polaris-core/src/main/java/org/apache/polaris/core/entity/PolarisBaseEntity.java‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -339,7 +339,10 @@ public boolean equals(Object o) {
339339
PolarisBaseEntity that = (PolarisBaseEntity) o;
340340
return subTypeCode == that.subTypeCode
341341
&& createTimestamp == that.createTimestamp
342+
&& dropTimestamp == that.dropTimestamp
343+
&& purgeTimestamp == that.purgeTimestamp
342344
&& toPurgeTimestamp == that.toPurgeTimestamp
345+
&& lastUpdateTimestamp == that.lastUpdateTimestamp
343346
&& grantRecordsVersion == that.grantRecordsVersion
344347
&& Objects.equals(properties, that.properties)
345348
&& Objects.equals(internalProperties, that.internalProperties);
@@ -356,7 +359,10 @@ public int hashCode() {
356359
entityVersion,
357360
subTypeCode,
358361
createTimestamp,
362+
dropTimestamp,
363+
purgeTimestamp,
359364
toPurgeTimestamp,
365+
lastUpdateTimestamp,
360366
properties,
361367
internalProperties,
362368
grantRecordsVersion);

‎polaris-core/src/main/java/org/apache/polaris/core/persistence/AtomicOperationMetaStoreManager.java‎

Lines changed: 32 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -1786,27 +1786,13 @@ public ResolvedEntitiesResult loadResolvedEntities(
17861786
@Nonnull PolarisCallContext callCtx,
17871787
@Nonnull List<EntityNameLookupRecord> entityLookupRecords) {
17881788
BasePersistence ms = callCtx.getMetaStore();
1789-
List<PolarisBaseEntity> entities =
1790-
ms.lookupEntities(
1791-
callCtx,
1792-
entityLookupRecords.stream()
1793-
.map(r -> new PolarisEntityId(r.getCatalogId(), r.getId()))
1794-
.collect(Collectors.toList()));
1795-
// mimic the behavior of loadEntity above, return null if not found or type mismatch
1796-
List<ResolvedPolarisEntity> ret =
1797-
IntStream.range(0, entityLookupRecords.size())
1798-
.mapToObj(
1799-
i -> {
1800-
if (entities.get(i) != null
1801-
&& !entities.get(i).getType().equals(entityLookupRecords.get(i).getType())) {
1802-
return null;
1803-
} else {
1804-
return entities.get(i);
1805-
}
1806-
})
1807-
.map(e -> toResolvedPolarisEntity(callCtx, e, ms))
1789+
List<PolarisEntityId> entityIds =
1790+
entityLookupRecords.stream()
1791+
.map(r -> new PolarisEntityId(r.getCatalogId(), r.getId()))
18081792
.collect(Collectors.toList());
1809-
return new ResolvedEntitiesResult(ret);
1793+
Function<Integer, PolarisEntityType> entityTypeForIndex =
1794+
i -> entityLookupRecords.get(i).getType();
1795+
return getResolvedEntitiesResult(callCtx, ms, entityIds, entityTypeForIndex);
18101796
}
18111797

18121798
@Nonnull
@@ -1816,20 +1802,43 @@ public ResolvedEntitiesResult loadResolvedEntities(
18161802
@Nonnull PolarisEntityType entityType,
18171803
@Nonnull List<PolarisEntityId> entityIds) {
18181804
BasePersistence ms = callCtx.getMetaStore();
1819-
List<PolarisBaseEntity> entities = ms.lookupEntities(callCtx, entityIds);
1805+
return getResolvedEntitiesResult(callCtx, ms, entityIds, i -> entityType);
1806+
}
18201807

1808+
private static ResolvedEntitiesResult getResolvedEntitiesResult(
1809+
PolarisCallContext callCtx,
1810+
BasePersistence ms,
1811+
List<PolarisEntityId> entityIds,
1812+
Function<Integer, PolarisEntityType> entityTypeForIndex) {
1813+
List<PolarisBaseEntity> entities = ms.lookupEntities(callCtx, entityIds);
18211814
// mimic the behavior of loadEntity above, return null if not found or type mismatch
18221815
List<ResolvedPolarisEntity> ret =
18231816
IntStream.range(0, entityIds.size())
18241817
.mapToObj(
18251818
i -> {
1826-
if (entities.get(i) != null && !entities.get(i).getType().equals(entityType)) {
1819+
if (entities.get(i) != null
1820+
&& !entities.get(i).getType().equals(entityTypeForIndex.apply(i))) {
18271821
return null;
18281822
} else {
18291823
return entities.get(i);
18301824
}
18311825
})
1832-
.map(e -> toResolvedPolarisEntity(callCtx, e, ms))
1826+
.map(
1827+
e -> {
1828+
if (e == null) {
1829+
return null;
1830+
} else {
1831+
// load the grant records
1832+
final List<PolarisGrantRecord> grantRecordsAsSecurable =
1833+
ms.loadAllGrantRecordsOnSecurable(callCtx, e.getCatalogId(), e.getId());
1834+
final List<PolarisGrantRecord> grantRecordsAsGrantee =
1835+
e.getType().isGrantee()
1836+
? ms.loadAllGrantRecordsOnGrantee(callCtx, e.getCatalogId(), e.getId())
1837+
: List.of();
1838+
return new ResolvedPolarisEntity(
1839+
PolarisEntity.of(e), grantRecordsAsGrantee, grantRecordsAsSecurable);
1840+
}
1841+
})
18331842
.collect(Collectors.toList());
18341843
return new ResolvedEntitiesResult(ret);
18351844
}

‎polaris-core/src/testFixtures/java/org/apache/polaris/core/persistence/BasePolarisMetaStoreManagerTest.java‎

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -240,8 +240,14 @@ protected void testLookup() {
240240

241241
/** test batch entity load */
242242
@Test
243-
protected void testBatchLoad() {
244-
polarisTestMetaStoreManager.testBatchLoad();
243+
protected void testLoadResolvedEntities() {
244+
polarisTestMetaStoreManager.testLoadResolvedEntities();
245+
}
246+
247+
/** test batch entity load */
248+
@Test
249+
protected void testLoadResolvedEntitiesById() {
250+
polarisTestMetaStoreManager.testLoadResolvedEntitiesById();
245251
}
246252

247253
/** Test the set of functions for the entity cache */

‎polaris-core/src/testFixtures/java/org/apache/polaris/core/persistence/PolarisTestMetaStoreManager.java‎

Lines changed: 116 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@
2626
import java.util.HashSet;
2727
import java.util.List;
2828
import java.util.Map;
29+
import java.util.Objects;
2930
import org.apache.commons.lang3.tuple.ImmutablePair;
3031
import org.apache.commons.lang3.tuple.Pair;
3132
import org.apache.polaris.core.PolarisCallContext;
@@ -2710,7 +2711,7 @@ public void testLookup() {
27102711
this.ensureNotExistsById(catalog.getId(), T1.getId(), PolarisEntityType.NAMESPACE);
27112712
}
27122713

2713-
public void testBatchLoad() {
2714+
public void testLoadResolvedEntities() {
27142715
// load all principals
27152716
List<EntityNameLookupRecord> principals =
27162717
polarisMetaStoreManager
@@ -2773,8 +2774,9 @@ public void testBatchLoad() {
27732774
InstanceOfAssertFactories.list(ResolvedPolarisEntity.class))
27742775
.hasSize(4)
27752776
.allSatisfy(entity -> Assertions.assertThat(entity).isNotNull())
2776-
.extracting(r -> (PolarisBaseEntity) r.getEntity())
2777-
.containsExactly(catalog, N1, N1_N2, T1);
2777+
.extracting(r -> getEntityCore(r.getEntity()))
2778+
.containsExactly(
2779+
getEntityCore(catalog), getEntityCore(N1), getEntityCore(N1_N2), getEntityCore(T1));
27782780

27792781
ResolvedPolarisEntity catalogEntity = entitiesResult.getResolvedEntities().get(0);
27802782
Assertions.assertThat(catalogEntity)
@@ -2866,10 +2868,118 @@ public void testBatchLoad() {
28662868
ResolvedEntitiesResult::getResolvedEntities,
28672869
InstanceOfAssertFactories.list(ResolvedPolarisEntity.class))
28682870
.hasSize(6)
2869-
.filteredOn(e -> e != null)
2871+
.filteredOn(Objects::nonNull)
28702872
.hasSize(2)
2871-
.extracting(r -> (PolarisBaseEntity) r.getEntity())
2872-
.containsExactly(catalog, T1);
2873+
.extracting(r -> getEntityCore(r.getEntity()))
2874+
.containsExactly(getEntityCore(catalog), getEntityCore(T1));
2875+
}
2876+
2877+
public void testLoadResolvedEntitiesById() {
2878+
// load all principals
2879+
List<EntityNameLookupRecord> principals =
2880+
polarisMetaStoreManager
2881+
.listEntities(
2882+
this.polarisCallContext,
2883+
null,
2884+
PolarisEntityType.PRINCIPAL,
2885+
PolarisEntitySubType.NULL_SUBTYPE,
2886+
PageToken.readEverything())
2887+
.getEntities();
2888+
2889+
// create new catalog
2890+
PolarisBaseEntity catalog =
2891+
new PolarisBaseEntity(
2892+
PolarisEntityConstants.getNullId(),
2893+
polarisMetaStoreManager.generateNewEntityId(this.polarisCallContext).getId(),
2894+
PolarisEntityType.CATALOG,
2895+
PolarisEntitySubType.NULL_SUBTYPE,
2896+
PolarisEntityConstants.getRootEntityId(),
2897+
"test");
2898+
CreateCatalogResult catalogCreated =
2899+
polarisMetaStoreManager.createCatalog(this.polarisCallContext, catalog, List.of());
2900+
Assertions.assertThat(catalogCreated).isNotNull();
2901+
2902+
// load the catalog again, since the grant versions are different
2903+
catalog =
2904+
polarisMetaStoreManager
2905+
.loadEntity(
2906+
polarisCallContext,
2907+
0L,
2908+
catalogCreated.getCatalog().getId(),
2909+
PolarisEntityType.CATALOG)
2910+
.getEntity();
2911+
2912+
// now create all objects
2913+
PolarisBaseEntity N1 = this.createEntity(List.of(catalog), PolarisEntityType.NAMESPACE, "N1");
2914+
PolarisBaseEntity N1_N2 =
2915+
this.createEntity(List.of(catalog, N1), PolarisEntityType.NAMESPACE, "N2");
2916+
PolarisBaseEntity T1 =
2917+
this.createEntity(
2918+
List.of(catalog, N1, N1_N2),
2919+
PolarisEntityType.TABLE_LIKE,
2920+
PolarisEntitySubType.ICEBERG_TABLE,
2921+
"T1");
2922+
2923+
// batch load all entities. They should all be present and non-null
2924+
ResolvedEntitiesResult entitiesResult =
2925+
polarisMetaStoreManager.loadResolvedEntities(
2926+
polarisCallContext,
2927+
PolarisEntityType.NAMESPACE,
2928+
List.of(
2929+
new PolarisEntityId(N1.getCatalogId(), N1.getId()),
2930+
new PolarisEntityId(N1_N2.getCatalogId(), N1_N2.getId())));
2931+
Assertions.assertThat(entitiesResult)
2932+
.isNotNull()
2933+
.returns(BaseResult.ReturnStatus.SUCCESS, ResolvedEntitiesResult::getReturnStatus)
2934+
.extracting(
2935+
ResolvedEntitiesResult::getResolvedEntities,
2936+
InstanceOfAssertFactories.list(ResolvedPolarisEntity.class))
2937+
.hasSize(2)
2938+
.allSatisfy(entity -> Assertions.assertThat(entity).isNotNull())
2939+
.extracting(r -> getEntityCore(r.getEntity()))
2940+
.containsExactly(getEntityCore(N1), getEntityCore(N1_N2));
2941+
2942+
// try entities which do not exist
2943+
entitiesResult =
2944+
polarisMetaStoreManager.loadResolvedEntities(
2945+
polarisCallContext,
2946+
PolarisEntityType.CATALOG,
2947+
List.of(
2948+
new PolarisEntityId(catalog.getId(), 27),
2949+
new PolarisEntityId(catalog.getId(), 35)));
2950+
Assertions.assertThat(entitiesResult)
2951+
.isNotNull()
2952+
.returns(BaseResult.ReturnStatus.SUCCESS, ResolvedEntitiesResult::getReturnStatus)
2953+
.extracting(
2954+
ResolvedEntitiesResult::getResolvedEntities,
2955+
InstanceOfAssertFactories.list(ResolvedPolarisEntity.class))
2956+
.hasSize(2)
2957+
.allSatisfy(entity -> Assertions.assertThat(entity).isNull());
2958+
2959+
// existing entities, some with wrong type
2960+
entitiesResult =
2961+
polarisMetaStoreManager.loadResolvedEntities(
2962+
polarisCallContext,
2963+
PolarisEntityType.NAMESPACE,
2964+
List.of(
2965+
new PolarisEntityId(catalog.getCatalogId(), catalog.getId()),
2966+
new PolarisEntityId(catalog.getId(), N1_N2.getId()),
2967+
new PolarisEntityId(catalog.getId(), T1.getId())));
2968+
Assertions.assertThat(entitiesResult)
2969+
.isNotNull()
2970+
.returns(BaseResult.ReturnStatus.SUCCESS, ResolvedEntitiesResult::getReturnStatus)
2971+
.extracting(
2972+
ResolvedEntitiesResult::getResolvedEntities,
2973+
InstanceOfAssertFactories.list(ResolvedPolarisEntity.class))
2974+
.hasSize(3)
2975+
.filteredOn(Objects::nonNull)
2976+
.hasSize(1)
2977+
.extracting(r -> getEntityCore(r.getEntity()))
2978+
.containsExactly(getEntityCore(N1_N2));
2979+
}
2980+
2981+
private static PolarisEntityCore getEntityCore(PolarisBaseEntity entity) {
2982+
return new PolarisEntityCore.Builder<>(entity).build();
28732983
}
28742984

28752985
/** Test the set of functions for the entity cache */

0 commit comments

Comments
 (0)