From 62dda815e0f07d17b55d8b4c606958ed65b2d161 Mon Sep 17 00:00:00 2001 From: Chris Lee Date: Wed, 28 Aug 2024 21:37:09 -0700 Subject: [PATCH] feat: adopt the internal snapshot/aspect union and asset in resources. (#396) * feat: adopt the internal snapshot/aspect union and asset in resources. * create internal model mode. * add experimental functions. * create lix interface for experimental controllers. * simplify the interface params and put link in java doc. * include ingestAsset in experiment. * integrate with ingestionParams * remove the ingestAsset resource lix * remove unused lix controllers. --- .../linkedin/metadata/dao/BaseLocalDAO.java | 4 +- .../metadata/dao/utils/ModelUtils.java | 4 + .../metadata/dao/utils/ModelUtilsTest.java | 10 + .../restli/BaseAspectRoutingResource.java | 235 ++++++++- .../restli/BaseBrowsableEntityResource.java | 36 +- .../metadata/restli/BaseEntityResource.java | 454 +++++++++++++---- .../restli/BaseSearchableEntityResource.java | 84 ++- .../BaseSingleAspectEntityResource.java | 27 +- ...eSingleAspectSearchableEntityResource.java | 37 +- .../metadata/restli/RestliConstants.java | 2 +- .../metadata/restli/lix/DummyResourceLix.java | 10 - .../metadata/restli/lix/ResourceLix.java | 18 - .../restli/BaseAspectRoutingResourceTest.java | 11 +- .../BaseBrowsableEntityResourceTest.java | 11 +- .../restli/BaseEntityResourceTest.java | 478 +++++++++++++++++- .../BaseEntitySimpleKeyResourceTest.java | 261 +++++++++- .../BaseSearchableEntityResourceTest.java | 13 +- ...SearchableEntitySimpleKeyResourceTest.java | 15 +- ...ngleAspectEntitySimpleKeyResourceTest.java | 13 +- ...titySimpleKeyVersionedSubResourceTest.java | 19 +- ...SearchableEntitySimpleKeyResourceTest.java | 15 +- .../InternalEntitySnapshot.pdl | 2 +- version.properties | 2 +- 23 files changed, 1522 insertions(+), 239 deletions(-) diff --git a/dao-api/src/main/java/com/linkedin/metadata/dao/BaseLocalDAO.java b/dao-api/src/main/java/com/linkedin/metadata/dao/BaseLocalDAO.java index 3fbbaf646..683a56c7d 100644 --- a/dao-api/src/main/java/com/linkedin/metadata/dao/BaseLocalDAO.java +++ b/dao-api/src/main/java/com/linkedin/metadata/dao/BaseLocalDAO.java @@ -822,7 +822,9 @@ public ASPECT add(@Nonnull URN urn, @Nonnull ASP @Nonnull public ASPECT add(@Nonnull URN urn, @Nonnull ASPECT newValue, @Nonnull AuditStamp auditStamp, @Nullable IngestionTrackingContext trackingContext, @Nullable IngestionParams ingestionParams) { - final IngestionParams nonNullIngestionParams = ingestionParams == null ? new IngestionParams().setIngestionMode(IngestionMode.LIVE) : ingestionParams; + final IngestionParams nonNullIngestionParams = + ingestionParams == null ? new IngestionParams().setIngestionMode(IngestionMode.LIVE).setTestMode(false) + : ingestionParams; return add(urn, (Class) newValue.getClass(), ignored -> newValue, auditStamp, trackingContext, nonNullIngestionParams); } diff --git a/dao-api/src/main/java/com/linkedin/metadata/dao/utils/ModelUtils.java b/dao-api/src/main/java/com/linkedin/metadata/dao/utils/ModelUtils.java index e4d6185f4..2175c9401 100644 --- a/dao-api/src/main/java/com/linkedin/metadata/dao/utils/ModelUtils.java +++ b/dao-api/src/main/java/com/linkedin/metadata/dao/utils/ModelUtils.java @@ -945,4 +945,8 @@ public static String getEntityType(@Nullable URN urn) { + return urn == null ? null : urn.getEntityType(); + } } diff --git a/dao-api/src/test/java/com/linkedin/metadata/dao/utils/ModelUtilsTest.java b/dao-api/src/test/java/com/linkedin/metadata/dao/utils/ModelUtilsTest.java index 2ea40a66a..fb59e5d56 100644 --- a/dao-api/src/test/java/com/linkedin/metadata/dao/utils/ModelUtilsTest.java +++ b/dao-api/src/test/java/com/linkedin/metadata/dao/utils/ModelUtilsTest.java @@ -751,4 +751,14 @@ public void testDecorateValue() { assertEquals(decoratedValue.getFoo(), entityFoo); assertEquals(decoratedValue.getBar(), expectedBar); } + + @Test + public void testGetEntityType() { + Urn expectedUrn = makeUrn(1); + + String entityType = ModelUtils.getEntityType(expectedUrn); + + assertEquals(entityType, expectedUrn.getEntityType()); + assertNull(ModelUtils.getEntityType(null)); + } } diff --git a/restli-resources/src/main/java/com/linkedin/metadata/restli/BaseAspectRoutingResource.java b/restli-resources/src/main/java/com/linkedin/metadata/restli/BaseAspectRoutingResource.java index 64ed75e56..1af7ee660 100644 --- a/restli-resources/src/main/java/com/linkedin/metadata/restli/BaseAspectRoutingResource.java +++ b/restli-resources/src/main/java/com/linkedin/metadata/restli/BaseAspectRoutingResource.java @@ -10,6 +10,7 @@ import com.linkedin.metadata.dao.utils.ModelUtils; import com.linkedin.metadata.events.IngestionTrackingContext; import com.linkedin.metadata.internal.IngestionParams; +import com.linkedin.metadata.restli.lix.ResourceLix; import com.linkedin.parseq.Task; import com.linkedin.restli.common.HttpStatus; import com.linkedin.restli.server.RestLiServiceException; @@ -55,21 +56,47 @@ public abstract class BaseAspectRoutingResource< URN extends Urn, SNAPSHOT extends RecordTemplate, ASPECT_UNION extends UnionTemplate, - DOCUMENT extends RecordTemplate> + DOCUMENT extends RecordTemplate, + INTERNAL_SNAPSHOT extends RecordTemplate, + INTERNAL_ASPECT_UNION extends UnionTemplate, + ASSET extends RecordTemplate> // @formatter:on - extends BaseBrowsableEntityResource { + extends + BaseBrowsableEntityResource { private final Class _valueClass; private final Class _aspectUnionClass; private final Class _snapshotClass; + private final Class _internalSnapshotClass; + private final Class _internalAspectUnionClass; + private final Class _assetClass; + + public BaseAspectRoutingResource(@Nullable Class snapshotClass, + @Nullable Class aspectUnionClass, @Nonnull Class urnClass, @Nonnull Class valueClass, + @Nonnull Class internalSnapshotClass, + @Nonnull Class internalAspectUnionClass, @Nonnull Class assetClass) { + super(snapshotClass, aspectUnionClass, urnClass, internalSnapshotClass, internalAspectUnionClass, assetClass); + _valueClass = valueClass; + _aspectUnionClass = aspectUnionClass; + _snapshotClass = snapshotClass; + _internalSnapshotClass = internalSnapshotClass; + _internalAspectUnionClass = internalAspectUnionClass; + _assetClass = assetClass; + } - public BaseAspectRoutingResource(@Nonnull Class snapshotClass, - @Nonnull Class aspectUnionClass, @Nonnull Class urnClass, - @Nonnull Class valueClass) { - super(snapshotClass, aspectUnionClass, urnClass); + public BaseAspectRoutingResource(@Nullable Class snapshotClass, + @Nullable Class aspectUnionClass, @Nonnull Class urnClass, @Nonnull Class valueClass, + @Nonnull Class internalSnapshotClass, + @Nonnull Class internalAspectUnionClass, @Nonnull Class assetClass, + @Nonnull ResourceLix resourceLix) { + super(snapshotClass, aspectUnionClass, urnClass, internalSnapshotClass, internalAspectUnionClass, assetClass, + resourceLix); _valueClass = valueClass; _aspectUnionClass = aspectUnionClass; _snapshotClass = snapshotClass; + _internalSnapshotClass = internalSnapshotClass; + _internalAspectUnionClass = internalAspectUnionClass; + _assetClass = assetClass; } /** @@ -85,9 +112,17 @@ public BaseAspectRoutingResource(@Nonnull Class snapshotClass, @Nonnull @Override public Task get(@Nonnull KEY id, @QueryParam(PARAM_ASPECTS) @Optional @Nullable String[] aspectNames) { + final URN urn = toUrn(id); + return get(id, aspectNames, _resourceLix.testGet(String.valueOf(urn), urn.getEntityType())); + } + + @Nonnull + @Override + protected Task get(@Nonnull KEY id, @QueryParam(PARAM_ASPECTS) @Optional @Nullable String[] aspectNames, + boolean isInternalModelsEnabled) { return RestliUtils.toTask(() -> { - final Set> aspectClasses = parseAspectsParam(aspectNames); + final Set> aspectClasses = parseAspectsParam(aspectNames, isInternalModelsEnabled); // The assumption is main GMS must have this entity. // If entity only has routing aspect, resourceNotFoundException will be thrown. @@ -98,40 +133,103 @@ public Task get(@Nonnull KEY id, @QueryParam(PARAM_ASPECTS) @Optional @Nu final Set> nonRoutingAspects = getNonRoutingAspects(aspectClasses); final VALUE valueFromLocalDao; if (nonRoutingAspects.isEmpty()) { - valueFromLocalDao = toValue(newSnapshot(urn)); + valueFromLocalDao = + isInternalModelsEnabled ? toInternalValue(newInternalSnapshot(urn)) : toValue(newSnapshot(urn)); } else { - valueFromLocalDao = getValueFromLocalDao(id, nonRoutingAspects); + valueFromLocalDao = getValueFromLocalDao(id, nonRoutingAspects, isInternalModelsEnabled); } return merge(valueFromLocalDao, getValueFromRoutingGms(toUrn(id), getRoutingAspects(aspectClasses))); }); } /** + * Deprecated to use {@link #getAsset(String, String[])} . * An action method for getting a snapshot of aspects for an entity. */ + @Deprecated @Action(name = ACTION_GET_SNAPSHOT) @Nonnull @Override public Task getSnapshot(@ActionParam(PARAM_URN) @Nonnull String urnString, @ActionParam(PARAM_ASPECTS) @Optional @Nullable String[] aspectNames) { + final URN urn = parseUrnParam(urnString); + return getSnapshot(urnString, aspectNames, + _resourceLix.testGetSnapshot(String.valueOf(urn), ModelUtils.getEntityType(urn))); + } + + @Nonnull + @Override + protected Task getSnapshot(@ActionParam(PARAM_URN) @Nonnull String urnString, + @ActionParam(PARAM_ASPECTS) @Optional @Nullable String[] aspectNames, boolean isInternalModelsEnabled) { + final URN urn = parseUrnParam(urnString); + final Set> aspectClasses = parseAspectsParam(aspectNames, isInternalModelsEnabled); + if (isInternalModelsEnabled) { + return RestliUtils.toTask(() -> { + if (!containsRoutingAspect(aspectClasses)) { + // Get snapshot from Local DAO. + final List aspectUnions = getInternalAspectsFromLocalDao(urn, aspectClasses); + return ModelUtils.newSnapshot(_snapshotClass, urn, + convertInternalAspectUnionToAspectUnion(_aspectUnionClass, aspectUnions)); + } else { + final Set> nonRoutingAspects = getNonRoutingAspects(aspectClasses); + final List aspectsFromLocalDao = + getInternalAspectsFromLocalDao(urn, nonRoutingAspects); + final Set> routingAspects = getRoutingAspects(aspectClasses); + final List aspectsFromGms = routingAspects.stream() + .map(routingAspect -> getInternalAspectsFromGms(urn, routingAspect)) + .flatMap(List::stream) + .collect(Collectors.toList()); + return ModelUtils.newSnapshot(_snapshotClass, urn, convertInternalAspectUnionToAspectUnion(_aspectUnionClass, + Stream.concat(aspectsFromGms.stream(), aspectsFromLocalDao.stream()).collect(Collectors.toList()))); + } + }); + } else { + return RestliUtils.toTask(() -> { + if (!containsRoutingAspect(aspectClasses)) { + // Get snapshot from Local DAO. + final List aspectUnions = getAspectsFromLocalDao(urn, aspectClasses); + return ModelUtils.newSnapshot(_snapshotClass, urn, aspectUnions); + } else { + final Set> nonRoutingAspects = getNonRoutingAspects(aspectClasses); + final List aspectsFromLocalDao = getAspectsFromLocalDao(urn, nonRoutingAspects); + final Set> routingAspects = getRoutingAspects(aspectClasses); + final List aspectsFromGms = routingAspects.stream() + .map(routingAspect -> getAspectsFromGms(urn, routingAspect)) + .flatMap(List::stream) + .collect(Collectors.toList()); + return ModelUtils.newSnapshot(_snapshotClass, urn, + Stream.concat(aspectsFromGms.stream(), aspectsFromLocalDao.stream()).collect(Collectors.toList())); + } + }); + } + } + + /** + * An action method for getting an asset of aspects for an entity. + */ + @Action(name = ACTION_GET_ASSET) + @Nonnull + @Override + public Task getAsset(@ActionParam(PARAM_URN) @Nonnull String urnString, + @ActionParam(PARAM_ASPECTS) @Optional @Nullable String[] aspectNames) { return RestliUtils.toTask(() -> { final URN urn = parseUrnParam(urnString); - final Set> aspectClasses = parseAspectsParam(aspectNames); + final Set> aspectClasses = parseAspectsParam(aspectNames, true); if (!containsRoutingAspect(aspectClasses)) { // Get snapshot from Local DAO. - final List aspectUnions = getAspectsFromLocalDao(urn, aspectClasses); - return ModelUtils.newSnapshot(_snapshotClass, urn, aspectUnions); + final List aspectUnions = getInternalAspectsFromLocalDao(urn, aspectClasses); + return ModelUtils.newAsset(_assetClass, urn, aspectUnions); } else { final Set> nonRoutingAspects = getNonRoutingAspects(aspectClasses); - final List aspectsFromLocalDao = getAspectsFromLocalDao(urn, nonRoutingAspects); + final List aspectsFromLocalDao = getInternalAspectsFromLocalDao(urn, nonRoutingAspects); final Set> routingAspects = getRoutingAspects(aspectClasses); - final List aspectsFromGms = routingAspects.stream() - .map(routingAspect -> getAspectsFromGms(urn, routingAspect)) + final List aspectsFromGms = routingAspects.stream() + .map(routingAspect -> getInternalAspectsFromGms(urn, routingAspect)) .flatMap(List::stream) .collect(Collectors.toList()); - return ModelUtils.newSnapshot(_snapshotClass, urn, + return ModelUtils.newAsset(_assetClass, urn, Stream.concat(aspectsFromGms.stream(), aspectsFromLocalDao.stream()).collect(Collectors.toList())); } }); @@ -144,6 +242,14 @@ public Task getSnapshot(@ActionParam(PARAM_URN) @Nonnull String urnStr @Nonnull public Task backfill(@ActionParam(PARAM_URNS) @Nonnull String[] urns, @ActionParam(PARAM_ASPECTS) @Optional @Nullable String[] aspectNames) { + final String urnString = urns[0]; + final URN urn = parseUrnParam(urnString); + return backfill(urns, aspectNames, _resourceLix.testBackfillWithUrns(urnString, ModelUtils.getEntityType(urn))); + } + + @Nonnull + private Task backfill(@ActionParam(PARAM_URNS) @Nonnull String[] urns, + @ActionParam(PARAM_ASPECTS) @Optional @Nullable String[] aspectNames, boolean isInternalModelsEnabled) { if (this._urnClass == null) { throw new IllegalStateException("URN class is not set for this resource"); } @@ -151,7 +257,8 @@ public Task backfill(@ActionParam(PARAM_URNS) @Nonnull String[] return RestliUtils.toTask(() -> { final Set urnSet = Arrays.stream(urns).map(this::parseUrnParam).collect(Collectors.toSet()); - final Set> aspectClasses = parseAspectsParam(aspectNames); + final Set> aspectClasses = + parseAspectsParam(aspectNames, isInternalModelsEnabled); Map, java.util.Optional>> urnToAspect = new HashMap<>(); @@ -185,19 +292,29 @@ public Task backfill(@ActionParam(PARAM_URNS) @Nonnull String[] @Override public Task backfillWithNewValue(@ActionParam(PARAM_URNS) @Nonnull String[] urns, @ActionParam(PARAM_ASPECTS) @Optional @Nullable String[] aspectNames) { + final String urnString = urns[0]; + final URN urn = parseUrnParam(urnString); + return backfillWithNewValue(urns, aspectNames, + _resourceLix.testBackfillWithNewValue(urnString, ModelUtils.getEntityType(urn))); + } + + private Task backfillWithNewValue(@ActionParam(PARAM_URNS) @Nonnull String[] urns, + @ActionParam(PARAM_ASPECTS) @Optional @Nullable String[] aspectNames, boolean isInternalModelsEnabled) { return RestliUtils.toTask(() -> { final Set urnSet = Arrays.stream(urns).map(this::parseUrnParam).collect(Collectors.toSet()); - final Set> aspectClasses = parseAspectsParam(aspectNames); - return RestliUtils.buildBackfillResult(getLocalDAO().backfillWithNewValue(getNonRoutingAspects(aspectClasses), urnSet)); + final Set> aspectClasses = + parseAspectsParam(aspectNames, isInternalModelsEnabled); + return RestliUtils.buildBackfillResult( + getLocalDAO().backfillWithNewValue(getNonRoutingAspects(aspectClasses), urnSet)); }); } @Nonnull @Override protected Task ingestInternal(@Nonnull SNAPSHOT snapshot, - @Nonnull Set> aspectsToIgnore, - @Nullable IngestionTrackingContext trackingContext, @Nullable IngestionParams ingestionParams) { + @Nonnull Set> aspectsToIgnore, @Nullable IngestionTrackingContext trackingContext, + @Nullable IngestionParams ingestionParams) { // TODO: META-18950: add trackingContext to BaseAspectRoutingResource. currently the param is unused. return RestliUtils.toTask(() -> { final URN urn = (URN) ModelUtils.getUrnFromSnapshot(snapshot); @@ -225,6 +342,37 @@ protected Task ingestInternal(@Nonnull SNAPSHOT snapshot, }); } + @Nonnull + @Override + protected Task ingestInternalAsset(@Nonnull ASSET asset, + @Nonnull Set> aspectsToIgnore, @Nullable IngestionTrackingContext trackingContext, + @Nullable IngestionParams ingestionParams) { + // TODO: META-18950: add trackingContext to BaseAspectRoutingResource. currently the param is unused. + return RestliUtils.toTask(() -> { + final URN urn = (URN) ModelUtils.getUrnFromAsset(asset); + final AuditStamp auditStamp = getAuditor().requestAuditStamp(getContext().getRawRequestContext()); + ModelUtils.getAspectsFromAsset(asset).forEach(aspect -> { + if (!aspectsToIgnore.contains(aspect.getClass())) { + if (getAspectRoutingGmsClientManager().hasRegistered(aspect.getClass())) { + try { + if (trackingContext != null) { + getAspectRoutingGmsClientManager().getRoutingGmsClient(aspect.getClass()) + .ingestWithTracking(urn, aspect, trackingContext, ingestionParams); + } else { + getAspectRoutingGmsClientManager().getRoutingGmsClient(aspect.getClass()).ingest(urn, aspect); + } + } catch (Exception exception) { + log.error("Couldn't ingest routing aspect {} for {}", aspect.getClass().getSimpleName(), urn, exception); + } + } else { + getLocalDAO().add(urn, aspect, auditStamp, trackingContext, ingestionParams); + } + } + }); + return null; + }); + } + /** * Whether given set of aspect classes contains routing aspect class. * @param aspectClasses A set of aspect classes @@ -265,9 +413,10 @@ private Set> getRoutingAspects(Set> aspectClasses) { + private VALUE getValueFromLocalDao(KEY id, Set> aspectClasses, + boolean isInternalModelsEnabled) { final URN urn = toUrn(id); - final VALUE value = getInternal(Collections.singleton(urn), aspectClasses).get(urn); + final VALUE value = getInternal(Collections.singleton(urn), aspectClasses, isInternalModelsEnabled).get(urn); if (value == null) { throw RestliUtils.resourceNotFoundException(); } @@ -284,6 +433,30 @@ private VALUE getValueFromLocalDao(KEY id, Set> @ParametersAreNonnullByDefault private List getAspectsFromLocalDao(URN urn, Set> aspectClasses) { + if (aspectClasses == null || aspectClasses.isEmpty()) { + return Collections.emptyList(); + } + final Set> keys = aspectClasses.stream() + .map(aspectClass -> new AspectKey<>(aspectClass, urn, LATEST_VERSION)) + .collect(Collectors.toSet()); + return getLocalDAO().get(keys) + .values() + .stream() + .filter(java.util.Optional::isPresent) + .map(aspect -> ModelUtils.newAspectUnion(_aspectUnionClass, aspect.get())) + .collect(Collectors.toList()); + } + + /** + * Get internal aspect values from local DAO for specified aspect classes. + * @param urn identifier of the entity. + * @param aspectClasses Aspects to be decorated on the entity + * @return A list of aspects. + */ + @Nonnull + @ParametersAreNonnullByDefault + private List getInternalAspectsFromLocalDao(URN urn, Set> aspectClasses) { + if (aspectClasses == null || aspectClasses.isEmpty()) { return Collections.emptyList(); } @@ -296,7 +469,7 @@ private List getAspectsFromLocalDao(URN urn, Set ModelUtils.newAspectUnion(_aspectUnionClass, aspect.get())) + .map(aspect -> ModelUtils.newAspectUnion(_internalAspectUnionClass, aspect.get())) .collect(Collectors.toList()); } @@ -314,6 +487,20 @@ private List getAspectsFromGms(URN urn, Class aspectClass) { return Collections.singletonList(ModelUtils.newAspectUnion(_aspectUnionClass, routingAspects.get(0))); } + /** + * Get internal aspect value from routing aspect GMS. + */ + @Nonnull + @ParametersAreNonnullByDefault + private List getInternalAspectsFromGms(URN urn, Class aspectClass) { + final List routingAspects = + getValueFromRoutingGms(urn, Collections.singletonList(aspectClass)); + if (routingAspects.isEmpty()) { + return new ArrayList<>(); + } + return Collections.singletonList(ModelUtils.newAspectUnion(_internalAspectUnionClass, routingAspects.get(0))); + } + /** * Merge routing aspect value from GMS into entity value retrieved from Local DAO. * @param valueFromLocalDao Entity value retrieved from Local DAO diff --git a/restli-resources/src/main/java/com/linkedin/metadata/restli/BaseBrowsableEntityResource.java b/restli-resources/src/main/java/com/linkedin/metadata/restli/BaseBrowsableEntityResource.java index 3ea2ebad8..2ef288cb8 100644 --- a/restli-resources/src/main/java/com/linkedin/metadata/restli/BaseBrowsableEntityResource.java +++ b/restli-resources/src/main/java/com/linkedin/metadata/restli/BaseBrowsableEntityResource.java @@ -8,6 +8,7 @@ import com.linkedin.metadata.dao.utils.QueryUtils; import com.linkedin.metadata.query.BrowseResult; import com.linkedin.metadata.query.Filter; +import com.linkedin.metadata.restli.lix.ResourceLix; import com.linkedin.parseq.Task; import com.linkedin.restli.server.annotations.Action; import com.linkedin.restli.server.annotations.ActionParam; @@ -29,6 +30,9 @@ * @param must be a valid snapshot type defined in com.linkedin.metadata.snapshot * @param must be a valid aspect union type supported by the snapshot * @param must be a valid search document type defined in com.linkedin.metadata.search + * @param must be a valid internal snapshot type defined in com.linkedin.metadata.snapshot + * @param must be a valid internal aspect union type supported by the internal snapshot + * @param must be a valid asset type defined in com.linkedin.metadata.asset */ public abstract class BaseBrowsableEntityResource< // @formatter:off @@ -37,18 +41,34 @@ public abstract class BaseBrowsableEntityResource< URN extends Urn, SNAPSHOT extends RecordTemplate, ASPECT_UNION extends UnionTemplate, - DOCUMENT extends RecordTemplate> + DOCUMENT extends RecordTemplate, + INTERNAL_SNAPSHOT extends RecordTemplate, + INTERNAL_ASPECT_UNION extends UnionTemplate, + ASSET extends RecordTemplate> // @formatter:on - extends BaseSearchableEntityResource { + extends + BaseSearchableEntityResource { - public BaseBrowsableEntityResource(@Nonnull Class snapshotClass, - @Nonnull Class aspectUnionClass) { - super(snapshotClass, aspectUnionClass); + public BaseBrowsableEntityResource(@Nullable Class snapshotClass, + @Nullable Class aspectUnionClass, @Nonnull Class internalSnapshotClass, + @Nonnull Class internalAspectUnionClass, @Nonnull Class assetClass) { + super(snapshotClass, aspectUnionClass, internalSnapshotClass, internalAspectUnionClass, assetClass); } - public BaseBrowsableEntityResource(@Nonnull Class snapshotClass, - @Nonnull Class aspectUnionClass, @Nonnull Class urnClass) { - super(snapshotClass, aspectUnionClass, urnClass); + public BaseBrowsableEntityResource(@Nullable Class snapshotClass, + @Nullable Class aspectUnionClass, @Nonnull Class urnClass, + @Nonnull Class internalSnapshotClass, + @Nonnull Class internalAspectUnionClass, @Nonnull Class assetClass) { + super(snapshotClass, aspectUnionClass, urnClass, internalSnapshotClass, internalAspectUnionClass, assetClass); + } + + public BaseBrowsableEntityResource(@Nullable Class snapshotClass, + @Nullable Class aspectUnionClass, @Nonnull Class urnClass, + @Nonnull Class internalSnapshotClass, + @Nonnull Class internalAspectUnionClass, @Nonnull Class assetClass, + @Nonnull ResourceLix resourceLix) { + super(snapshotClass, aspectUnionClass, urnClass, internalSnapshotClass, internalAspectUnionClass, assetClass, + resourceLix); } /** diff --git a/restli-resources/src/main/java/com/linkedin/metadata/restli/BaseEntityResource.java b/restli-resources/src/main/java/com/linkedin/metadata/restli/BaseEntityResource.java index 39cdf39b0..05f2588cd 100644 --- a/restli-resources/src/main/java/com/linkedin/metadata/restli/BaseEntityResource.java +++ b/restli-resources/src/main/java/com/linkedin/metadata/restli/BaseEntityResource.java @@ -21,6 +21,8 @@ import com.linkedin.metadata.query.IndexSortCriterion; import com.linkedin.metadata.query.ListResultMetadata; import com.linkedin.metadata.query.MapMetadata; +import com.linkedin.metadata.restli.lix.DummyResourceLix; +import com.linkedin.metadata.restli.lix.ResourceLix; import com.linkedin.parseq.Task; import com.linkedin.restli.common.EmptyRecord; import com.linkedin.restli.common.HttpStatus; @@ -66,6 +68,9 @@ * @param must be a valid {@link Urn} type for the snapshot * @param must be a valid snapshot type defined in com.linkedin.metadata.snapshot * @param must be a valid aspect union type supported by the snapshot + * @param must be a valid internal snapshot type defined in com.linkedin.metadata.snapshot + * @param must be a valid internal aspect union type supported by the internal snapshot + * @param must be a valid asset type defined in com.linkedin.metadata.asset */ public abstract class BaseEntityResource< // @formatter:off @@ -73,7 +78,10 @@ public abstract class BaseEntityResource< VALUE extends RecordTemplate, URN extends Urn, SNAPSHOT extends RecordTemplate, - ASPECT_UNION extends UnionTemplate> + ASPECT_UNION extends UnionTemplate, + INTERNAL_SNAPSHOT extends RecordTemplate, + INTERNAL_ASPECT_UNION extends UnionTemplate, + ASSET extends RecordTemplate> // @formatter:on extends CollectionResourceTaskTemplate { @@ -82,27 +90,70 @@ public abstract class BaseEntityResource< private final Class _snapshotClass; private final Class _aspectUnionClass; private final Set> _supportedAspectClasses; + private final Set> _supportedInternalAspectClasses; + private final Class _internalSnapshotClass; + private final Class _internalAspectUnionClass; + private final Class _assetClass; protected final Class _urnClass; + protected final ResourceLix _resourceLix; protected BaseTrackingManager _trackingManager = null; - public BaseEntityResource(@Nonnull Class snapshotClass, @Nonnull Class aspectUnionClass) { - this(snapshotClass, aspectUnionClass, null); + public BaseEntityResource(@Nullable Class snapshotClass, @Nullable Class aspectUnionClass, + @Nonnull Class internalSnapshotClass, + @Nonnull Class internalAspectUnionClass, @Nonnull Class assetClass) { + this(snapshotClass, aspectUnionClass, null, internalSnapshotClass, internalAspectUnionClass, assetClass, + new DummyResourceLix()); } - public BaseEntityResource(@Nonnull Class snapshotClass, @Nonnull Class aspectUnionClass, - @Nullable Class urnClass) { + public BaseEntityResource(@Nullable Class snapshotClass, @Nullable Class aspectUnionClass, + @Nonnull Class internalSnapshotClass, + @Nonnull Class internalAspectUnionClass, @Nonnull Class assetClass, + @Nonnull ResourceLix resourceLix) { + this(snapshotClass, aspectUnionClass, null, internalSnapshotClass, internalAspectUnionClass, assetClass, + resourceLix); + } + + public BaseEntityResource(@Nullable Class snapshotClass, @Nullable Class aspectUnionClass, + @Nullable Class urnClass, @Nonnull Class internalSnapshotClass, + @Nonnull Class internalAspectUnionClass, @Nonnull Class assetClass) { + this(snapshotClass, aspectUnionClass, urnClass, internalSnapshotClass, internalAspectUnionClass, assetClass, + new DummyResourceLix()); + } + + public BaseEntityResource(@Nullable Class snapshotClass, @Nullable Class aspectUnionClass, + @Nullable Class urnClass, @Nonnull Class internalSnapshotClass, + @Nonnull Class internalAspectUnionClass, @Nonnull Class assetClass, + @Nonnull ResourceLix resourceLix) { super(); - ModelUtils.validateSnapshotAspect(snapshotClass, aspectUnionClass); + ModelUtils.validateSnapshotAspect(internalSnapshotClass, internalAspectUnionClass); _snapshotClass = snapshotClass; _aspectUnionClass = aspectUnionClass; - _supportedAspectClasses = ModelUtils.getValidAspectTypes(_aspectUnionClass); _urnClass = urnClass; + _internalSnapshotClass = internalSnapshotClass; + _internalAspectUnionClass = internalAspectUnionClass; + _supportedAspectClasses = ModelUtils.getValidAspectTypes(_aspectUnionClass); + _supportedInternalAspectClasses = ModelUtils.getValidAspectTypes(_internalAspectUnionClass); + _assetClass = assetClass; + _resourceLix = resourceLix; } - public BaseEntityResource(@Nonnull Class snapshotClass, @Nonnull Class aspectUnionClass, - @Nullable Class urnClass, @Nullable BaseTrackingManager trackingManager) { - this(snapshotClass, aspectUnionClass, urnClass); + public BaseEntityResource(@Nullable Class snapshotClass, @Nullable Class aspectUnionClass, + @Nullable Class urnClass, @Nullable BaseTrackingManager trackingManager, + @Nonnull Class internalSnapshotClass, + @Nonnull Class internalAspectUnionClass, @Nonnull Class assetClass, + @Nonnull ResourceLix resourceLix) { + this(snapshotClass, aspectUnionClass, urnClass, internalSnapshotClass, internalAspectUnionClass, assetClass, + resourceLix); + _trackingManager = trackingManager; + } + + public BaseEntityResource(@Nullable Class snapshotClass, @Nullable Class aspectUnionClass, + @Nullable Class urnClass, @Nullable BaseTrackingManager trackingManager, + @Nonnull Class internalSnapshotClass, + @Nonnull Class internalAspectUnionClass, @Nonnull Class assetClass) { + this(snapshotClass, aspectUnionClass, urnClass, internalSnapshotClass, internalAspectUnionClass, assetClass, + new DummyResourceLix()); _trackingManager = trackingManager; } @@ -118,7 +169,7 @@ protected BaseRestliAuditor getAuditor() { * Returns an aspect-specific {@link BaseLocalDAO}. */ @Nonnull - protected abstract BaseLocalDAO getLocalDAO(); + protected abstract BaseLocalDAO getLocalDAO(); /** * Creates an URN from its string representation. @@ -144,6 +195,15 @@ protected BaseRestliAuditor getAuditor() { @Nonnull protected abstract VALUE toValue(@Nonnull SNAPSHOT snapshot); + /** + * Converts an internal snapshot to resource's value. + */ + @Nonnull + protected VALUE toInternalValue(@Nonnull INTERNAL_SNAPSHOT internalSnapshot) { + final SNAPSHOT snapshot = ModelUtils.convertSnapshots(_snapshotClass, internalSnapshot); + return ModelUtils.decorateValue(internalSnapshot, toValue(snapshot)); + } + /** * Converts a resource's value to a snapshot. */ @@ -155,15 +215,22 @@ protected BaseRestliAuditor getAuditor() { */ @RestMethod.Get @Nonnull - public Task get(@Nonnull KEY id, - @QueryParam(PARAM_ASPECTS) @Optional @Nullable String[] aspectNames) { + public Task get(@Nonnull KEY id, @QueryParam(PARAM_ASPECTS) @Optional @Nullable String[] aspectNames) { + final URN urn = toUrn(id); + return get(id, aspectNames, _resourceLix.testGet(String.valueOf(urn), urn.getEntityType())); + } + + protected Task get(@Nonnull KEY id, @QueryParam(PARAM_ASPECTS) @Optional @Nullable String[] aspectNames, + boolean isInternalModelsEnabled) { return RestliUtils.toTask(() -> { final URN urn = toUrn(id); if (!getLocalDAO().exists(urn)) { throw RestliUtils.resourceNotFoundException(); } - final VALUE value = getInternal(Collections.singleton(urn), parseAspectsParam(aspectNames)).get(urn); + final VALUE value = + getInternal(Collections.singleton(urn), parseAspectsParam(aspectNames, isInternalModelsEnabled), + isInternalModelsEnabled).get(urn); if (value == null) { throw RestliUtils.resourceNotFoundException(); } @@ -180,16 +247,22 @@ public Task get(@Nonnull KEY id, @RestMethod.BatchGet @Deprecated @Nonnull - public Task> batchGet( - @Nonnull Set ids, + public Task> batchGet(@Nonnull Set ids, @QueryParam(PARAM_ASPECTS) @Optional @Nullable String[] aspectNames) { + final URN urn = ids.stream().findFirst().isPresent() ? toUrn(ids.stream().findFirst().get()) : null; + return batchGet(ids, aspectNames, _resourceLix.testBatchGet(String.valueOf(urn), ModelUtils.getEntityType(urn))); + } + + @Deprecated + @Nonnull + private Task> batchGet(@Nonnull Set ids, + @QueryParam(PARAM_ASPECTS) @Optional @Nullable String[] aspectNames, boolean isInternalModelsEnabled) { return RestliUtils.toTask(() -> { - final Map urnMap = - ids.stream().collect(Collectors.toMap(this::toUrn, Function.identity())); - return getInternal(urnMap.keySet(), parseAspectsParam(aspectNames)).entrySet() + final Map urnMap = ids.stream().collect(Collectors.toMap(this::toUrn, Function.identity())); + return getInternal(urnMap.keySet(), parseAspectsParam(aspectNames, isInternalModelsEnabled), + isInternalModelsEnabled).entrySet() .stream() - .collect( - Collectors.toMap(e -> urnMap.get(e.getKey()), Map.Entry::getValue)); + .collect(Collectors.toMap(e -> urnMap.get(e.getKey()), Map.Entry::getValue)); }); } @@ -200,15 +273,23 @@ public Task> batchGet( */ @RestMethod.BatchGet @Nonnull - public Task> batchGetWithErrors( - @Nonnull Set ids, + public Task> batchGetWithErrors(@Nonnull Set ids, @QueryParam(PARAM_ASPECTS) @Optional @Nullable String[] aspectNames) { + final URN urn = ids.stream().findFirst().isPresent() ? toUrn(ids.stream().findFirst().get()) : null; + return batchGetWithErrors(ids, aspectNames, + _resourceLix.testBatchGetWithErrors(String.valueOf(urn), ModelUtils.getEntityType(urn))); + } + + @Nonnull + private Task> batchGetWithErrors(@Nonnull Set ids, + @QueryParam(PARAM_ASPECTS) @Optional @Nullable String[] aspectNames, boolean isInternalModelsEnabled) { return RestliUtils.toTask(() -> { final Map errors = new HashMap<>(); final Map statuses = new HashMap<>(); - final Map urnMap = - ids.stream().collect(Collectors.toMap(this::toUrn, Function.identity())); - final Map batchResult = getInternal(urnMap.keySet(), parseAspectsParam(aspectNames)); + final Map urnMap = ids.stream().collect(Collectors.toMap(this::toUrn, Function.identity())); + final Map batchResult = + getInternal(urnMap.keySet(), parseAspectsParam(aspectNames, isInternalModelsEnabled), + isInternalModelsEnabled); batchResult.entrySet().removeIf(entry -> { if (!entry.getValue().data().isEmpty()) { // don't remove if there is a non-empty value associated with the key @@ -228,8 +309,10 @@ public Task> batchGetWithErrors( } /** + * Deprecated to use {@link #ingestAsset(RecordTemplate, IngestionTrackingContext, IngestionParams)} instead. * An action method for automated ingestion pipeline. */ + @Deprecated @Action(name = ACTION_INGEST) @Nonnull public Task ingest(@ActionParam(PARAM_SNAPSHOT) @Nonnull SNAPSHOT snapshot) { @@ -237,11 +320,13 @@ public Task ingest(@ActionParam(PARAM_SNAPSHOT) @Nonnull SNAPSHOT snapshot } /** + * Deprecated to use {@link #ingestAsset(RecordTemplate, IngestionTrackingContext, IngestionParams)} instead. * Same as {@link #ingest(RecordTemplate)} but with tracking context attached. * @param snapshot Snapshot of the metadata change to be ingested * @param trackingContext {@link IngestionTrackingContext} to 1) track DAO-level metrics and 2) to pass on to MAE emission * @return ingest task */ + @Deprecated @Action(name = ACTION_INGEST_WITH_TRACKING) @Nonnull public Task ingestWithTracking(@ActionParam(PARAM_SNAPSHOT) @Nonnull SNAPSHOT snapshot, @@ -250,10 +335,24 @@ public Task ingestWithTracking(@ActionParam(PARAM_SNAPSHOT) @Nonnull SNAPS return ingestInternal(snapshot, Collections.emptySet(), trackingContext, ingestionParams); } + /** + * An action method for automated ingestion pipeline. + * @param asset Asset of the metadata change to be ingested + * @param trackingContext {@link IngestionTrackingContext} to 1) track DAO-level metrics and 2) to pass on to MAE emission + * @return ingest task + */ + @Action(name = ACTION_INGEST_ASSET) + @Nonnull + public Task ingestAsset(@ActionParam(PARAM_ASSET) @Nonnull ASSET asset, + @ActionParam(PARAM_TRACKING_CONTEXT) @Nonnull IngestionTrackingContext trackingContext, + @Optional @ActionParam(PARAM_INGESTION_PARAMS) IngestionParams ingestionParams) { + return ingestInternalAsset(asset, Collections.emptySet(), trackingContext, ingestionParams); + } + @Nonnull protected Task ingestInternal(@Nonnull SNAPSHOT snapshot, - @Nonnull Set> aspectsToIgnore, - @Nullable IngestionTrackingContext trackingContext, @Nullable IngestionParams ingestionParams) { + @Nonnull Set> aspectsToIgnore, @Nullable IngestionTrackingContext trackingContext, + @Nullable IngestionParams ingestionParams) { return RestliUtils.toTask(() -> { final URN urn = (URN) ModelUtils.getUrnFromSnapshot(snapshot); final AuditStamp auditStamp = getAuditor().requestAuditStamp(getContext().getRawRequestContext()); @@ -266,17 +365,83 @@ protected Task ingestInternal(@Nonnull SNAPSHOT snapshot, }); } + @Nonnull + protected Task ingestInternalAsset(@Nonnull ASSET asset, + @Nonnull Set> aspectsToIgnore, @Nullable IngestionTrackingContext trackingContext, + @Nullable IngestionParams ingestionParams) { + return RestliUtils.toTask(() -> { + final URN urn = (URN) ModelUtils.getUrnFromAsset(asset); + final AuditStamp auditStamp = getAuditor().requestAuditStamp(getContext().getRawRequestContext()); + ModelUtils.getAspectsFromAsset(asset).stream().forEach(aspect -> { + if (!aspectsToIgnore.contains(aspect.getClass())) { + getLocalDAO().add(urn, aspect, auditStamp, trackingContext, ingestionParams); + } + }); + return null; + }); + } + /** + * Deprecated to use {@link #getAsset(String, String[])} instead. * An action method for getting a snapshot of aspects for an entity. */ + @Deprecated @Action(name = ACTION_GET_SNAPSHOT) @Nonnull public Task getSnapshot(@ActionParam(PARAM_URN) @Nonnull String urnString, @ActionParam(PARAM_ASPECTS) @Optional @Nullable String[] aspectNames) { + final URN urn = parseUrnParam(urnString); + return getSnapshot(urnString, aspectNames, + _resourceLix.testGetSnapshot(String.valueOf(urn), ModelUtils.getEntityType(urn))); + } + + @Deprecated + @Nonnull + protected Task getSnapshot(@ActionParam(PARAM_URN) @Nonnull String urnString, + @ActionParam(PARAM_ASPECTS) @Optional @Nullable String[] aspectNames, boolean isInternalModelsEnabled) { + final URN urn = parseUrnParam(urnString); + final Set> keys = + parseAspectsParam(aspectNames, isInternalModelsEnabled).stream() + .map(aspectClass -> new AspectKey<>(aspectClass, urn, LATEST_VERSION)) + .collect(Collectors.toSet()); + if (isInternalModelsEnabled) { + return RestliUtils.toTask(() -> { + final List aspects = getLocalDAO().get(keys) + .values() + .stream() + .filter(java.util.Optional::isPresent) + .map(aspect -> ModelUtils.newAspectUnion(_internalAspectUnionClass, aspect.get())) + .collect(Collectors.toList()); + + return ModelUtils.newSnapshot(_snapshotClass, urn, + ModelUtils.convertInternalAspectUnionToAspectUnion(_aspectUnionClass, aspects)); + }); + } else { + return RestliUtils.toTask(() -> { + final List aspects = getLocalDAO().get(keys) + .values() + .stream() + .filter(java.util.Optional::isPresent) + .filter(aspect -> _supportedAspectClasses.contains(aspect.get().getClass())) + .map(aspect -> ModelUtils.newAspectUnion(_aspectUnionClass, aspect.get())) + .collect(Collectors.toList()); + + return ModelUtils.newSnapshot(_snapshotClass, urn, aspects); + }); + } + } + + /** + * An action method for getting an asset of aspects for an entity. + */ + @Action(name = ACTION_GET_ASSET) + @Nonnull + public Task getAsset(@ActionParam(PARAM_URN) @Nonnull String urnString, + @ActionParam(PARAM_ASPECTS) @Optional @Nullable String[] aspectNames) { return RestliUtils.toTask(() -> { final URN urn = parseUrnParam(urnString); - final Set> keys = parseAspectsParam(aspectNames).stream() + final Set> keys = parseAspectsParam(aspectNames, true).stream() .map(aspectClass -> new AspectKey<>(aspectClass, urn, LATEST_VERSION)) .collect(Collectors.toSet()); @@ -284,10 +449,10 @@ public Task getSnapshot(@ActionParam(PARAM_URN) @Nonnull String urnStr .values() .stream() .filter(java.util.Optional::isPresent) - .map(aspect -> ModelUtils.newAspectUnion(_aspectUnionClass, aspect.get())) + .map(aspect -> ModelUtils.newAspectUnion(_internalAspectUnionClass, aspect.get())) .collect(Collectors.toList()); - return ModelUtils.newSnapshot(_snapshotClass, urn, aspects); + return ModelUtils.newAsset(_assetClass, urn, aspects); }); } @@ -300,10 +465,18 @@ public Task getSnapshot(@ActionParam(PARAM_URN) @Nonnull String urnStr @Nonnull public Task backfill(@ActionParam(PARAM_URN) @Nonnull String urnString, @ActionParam(PARAM_ASPECTS) @Optional @Nullable String[] aspectNames) { + final URN urn = parseUrnParam(urnString); + return backfill(urnString, aspectNames, + _resourceLix.testBackfillLegacy(String.valueOf(urn), ModelUtils.getEntityType(urn))); + } + + @Nonnull + private Task backfill(@ActionParam(PARAM_URN) @Nonnull String urnString, + @ActionParam(PARAM_ASPECTS) @Optional @Nullable String[] aspectNames, boolean isInternalModelsEnabled) { return RestliUtils.toTask(() -> { final URN urn = parseUrnParam(urnString); - final List backfilledAspects = parseAspectsParam(aspectNames).stream() + final List backfilledAspects = parseAspectsParam(aspectNames, isInternalModelsEnabled).stream() .map(aspectClass -> getLocalDAO().backfill(aspectClass, urn)) .filter(optionalAspect -> optionalAspect.isPresent()) .map(optionalAspect -> ModelUtils.getAspectName(optionalAspect.get().getClass())) @@ -320,11 +493,20 @@ public Task backfill(@ActionParam(PARAM_URN) @Nonnull String urn @Action(name = ACTION_BACKFILL_WITH_URNS) @Nonnull public Task backfill(@ActionParam(PARAM_URNS) @Nonnull String[] urns, - @ActionParam(PARAM_ASPECTS) @Optional @Nullable String[] aspectNames) { + @ActionParam(PARAM_ASPECTS) @Optional @Nullable String[] aspectNames) { + final String urnString = urns[0]; + final URN urn = parseUrnParam(urnString); + return backfill(urns, aspectNames, _resourceLix.testBackfillWithUrns(urnString, ModelUtils.getEntityType(urn))); + } + + private Task backfill(@ActionParam(PARAM_URNS) @Nonnull String[] urns, + @ActionParam(PARAM_ASPECTS) @Optional @Nullable String[] aspectNames, boolean isInternalModelsEnabled) { return RestliUtils.toTask(() -> { - final Set urnSet = Arrays.stream(urns).map(urnString -> parseUrnParam(urnString)).collect(Collectors.toSet()); - return RestliUtils.buildBackfillResult(getLocalDAO().backfill(parseAspectsParam(aspectNames), urnSet)); + final Set urnSet = + Arrays.stream(urns).map(urnString -> parseUrnParam(urnString)).collect(Collectors.toSet()); + return RestliUtils.buildBackfillResult( + getLocalDAO().backfill(parseAspectsParam(aspectNames, isInternalModelsEnabled), urnSet)); }); } @@ -339,13 +521,24 @@ public Task backfill(@ActionParam(PARAM_URNS) @Nonnull String[] public Task emitNoChangeMetadataAuditEvent(@ActionParam(PARAM_URNS) @Nonnull String[] urns, @ActionParam(PARAM_ASPECTS) @Optional @Nullable String[] aspectNames, @ActionParam(PARAM_INGESTION_MODE) @Nonnull IngestionMode ingestionMode) { + final String urnString = urns[0]; + final URN urn = parseUrnParam(urnString); + return emitNoChangeMetadataAuditEvent(urns, aspectNames, ingestionMode, + _resourceLix.testEmitNoChangeMetadataAuditEvent(urnString, ModelUtils.getEntityType(urn))); + } + + @Nonnull + private Task emitNoChangeMetadataAuditEvent(@ActionParam(PARAM_URNS) @Nonnull String[] urns, + @ActionParam(PARAM_ASPECTS) @Optional @Nullable String[] aspectNames, + @ActionParam(PARAM_INGESTION_MODE) @Nonnull IngestionMode ingestionMode, boolean isInternalModelsEnabled) { BackfillMode backfillMode = ALLOWED_INGESTION_BACKFILL_BIMAP.get(ingestionMode); if (backfillMode == null) { return RestliUtils.toTask(BackfillResult::new); } return RestliUtils.toTask(() -> { final Set urnSet = Arrays.stream(urns).map(urnString -> parseUrnParam(urnString)).collect(Collectors.toSet()); - return RestliUtils.buildBackfillResult(getLocalDAO().backfill(backfillMode, parseAspectsParam(aspectNames), urnSet)); + return RestliUtils.buildBackfillResult( + getLocalDAO().backfill(backfillMode, parseAspectsParam(aspectNames, isInternalModelsEnabled), urnSet)); }); } @@ -359,10 +552,19 @@ public Task emitNoChangeMetadataAuditEvent(@ActionParam(PARAM_UR @Nonnull public Task backfillWithNewValue(@ActionParam(PARAM_URNS) @Nonnull String[] urns, @ActionParam(PARAM_ASPECTS) @Optional @Nullable String[] aspectNames) { + final String urnString = urns[0]; + final URN urn = parseUrnParam(urnString); + return backfillWithNewValue(urns, aspectNames, + _resourceLix.testBackfillWithNewValue(urnString, ModelUtils.getEntityType(urn))); + } - return RestliUtils.toTask(() -> { + private Task backfillWithNewValue(@ActionParam(PARAM_URNS) @Nonnull String[] urns, + @ActionParam(PARAM_ASPECTS) @Optional @Nullable String[] aspectNames, boolean isInternalModelsEnabled) { + + return RestliUtils.toTask(() -> { final Set urnSet = Arrays.stream(urns).map(urnString -> parseUrnParam(urnString)).collect(Collectors.toSet()); - return RestliUtils.buildBackfillResult(getLocalDAO().backfillWithNewValue(parseAspectsParam(aspectNames), urnSet)); + return RestliUtils.buildBackfillResult( + getLocalDAO().backfillWithNewValue(parseAspectsParam(aspectNames, isInternalModelsEnabled), urnSet)); }); } @@ -373,10 +575,19 @@ public Task backfillWithNewValue(@ActionParam(PARAM_URNS) @Nonnu @Nonnull public Task backfillEntityTables(@ActionParam(PARAM_URNS) @Nonnull String[] urns, @ActionParam(PARAM_ASPECTS) @Optional @Nullable String[] aspectNames) { + final String urnString = urns[0]; + final URN urn = parseUrnParam(urnString); + return backfillEntityTables(urns, aspectNames, + _resourceLix.testBackfillEntityTables(urnString, ModelUtils.getEntityType(urn))); + } + + private Task backfillEntityTables(@ActionParam(PARAM_URNS) @Nonnull String[] urns, + @ActionParam(PARAM_ASPECTS) @Optional @Nullable String[] aspectNames, boolean isInternalModelsEnabled) { return RestliUtils.toTask(() -> { final Set urnSet = Arrays.stream(urns).map(urnString -> parseUrnParam(urnString)).collect(Collectors.toSet()); - return RestliUtils.buildBackfillResult(getLocalDAO().backfillEntityTables(parseAspectsParam(aspectNames), urnSet)); + return RestliUtils.buildBackfillResult( + getLocalDAO().backfillEntityTables(parseAspectsParam(aspectNames, isInternalModelsEnabled), urnSet)); }); } @@ -387,12 +598,20 @@ public Task backfillEntityTables(@ActionParam(PARAM_URNS) @Nonnu @Nonnull public Task backfillRelationshipTables(@ActionParam(PARAM_URNS) @Nonnull String[] urns, @ActionParam(PARAM_ASPECTS) @Nonnull String[] aspectNames) { + final String urnString = urns[0]; + final URN urn = parseUrnParam(urnString); + return backfillRelationshipTables(urns, aspectNames, + _resourceLix.testBackfillRelationshipTables(urnString, ModelUtils.getEntityType(urn))); + } + + private Task backfillRelationshipTables(@ActionParam(PARAM_URNS) @Nonnull String[] urns, + @ActionParam(PARAM_ASPECTS) @Nonnull String[] aspectNames, boolean isInternalModelsEnabled) { final BackfillResult backfillResult = new BackfillResult() .setEntities(new BackfillResultEntityArray()) .setRelationships(new BackfillResultRelationshipArray()); for (String urn : urns) { - for (Class aspect : parseAspectsParam(aspectNames)) { + for (Class aspect : parseAspectsParam(aspectNames, isInternalModelsEnabled)) { getLocalDAO().backfillLocalRelationships(parseUrnParam(urn), aspect).forEach(relationshipUpdates -> { relationshipUpdates.getRelationships().forEach(relationship -> { try { @@ -425,12 +644,21 @@ public Task backfill(@ActionParam(PARAM_MODE) @Nonnull BackfillM @ActionParam(PARAM_ASPECTS) @Optional @Nullable String[] aspectNames, @ActionParam(PARAM_URN) @Optional @Nullable String lastUrn, @ActionParam(PARAM_LIMIT) int limit) { + return backfill(mode, aspectNames, lastUrn, limit, + _resourceLix.testBackfill(_assetClass.getSimpleName(), mode.name())); + } - return RestliUtils.toTask(() -> - RestliUtils.buildBackfillResult(getLocalDAO().backfill(mode, parseAspectsParam(aspectNames), - _urnClass, - parseUrnParam(lastUrn), - limit))); + @Nonnull + private Task backfill(@ActionParam(PARAM_MODE) @Nonnull BackfillMode mode, + @ActionParam(PARAM_ASPECTS) @Optional @Nullable String[] aspectNames, + @ActionParam(PARAM_URN) @Optional @Nullable String lastUrn, + @ActionParam(PARAM_LIMIT) int limit, boolean isInternalModelsEnabled) { + + return RestliUtils.toTask(() -> RestliUtils.buildBackfillResult( + getLocalDAO().backfill(mode, parseAspectsParam(aspectNames, isInternalModelsEnabled), + _urnClass, + parseUrnParam(lastUrn), + limit))); } /** @@ -465,7 +693,7 @@ public Task listUrnsFromIndex(@ActionParam(PARAM_FILTER) @Optional @Nu * @return ordered list of values of multiple entities */ @Nonnull - private List getUrnAspectValues(List> urnAspectEntries) { + private List getUrnAspectValues(List> urnAspectEntries, boolean isInternalModelsEnabled) { final Map> urnAspectsMap = new LinkedHashMap<>(); for (UrnAspectEntry entry : urnAspectEntries) { urnAspectsMap.compute(entry.getUrn(), (k, v) -> { @@ -474,7 +702,8 @@ private List getUrnAspectValues(List> urnAspectEntrie } v.addAll(entry.getAspects() .stream() - .map(recordTemplate -> ModelUtils.newAspectUnion(_aspectUnionClass, recordTemplate)) + .map(recordTemplate -> isInternalModelsEnabled ? ModelUtils.newAspectUnion(_internalAspectUnionClass, + recordTemplate) : ModelUtils.newAspectUnion(_aspectUnionClass, recordTemplate)) .collect(Collectors.toList())); return v; }); @@ -482,7 +711,8 @@ private List getUrnAspectValues(List> urnAspectEntrie return urnAspectsMap.entrySet() .stream() - .map(e -> toValue(newSnapshot(e.getKey(), e.getValue()))) + .map(e -> isInternalModelsEnabled ? toInternalValue(newInternalSnapshot(e.getKey(), e.getValue())) + : toValue(newSnapshot(e.getKey(), e.getValue()))) .collect(Collectors.toList()); } @@ -500,18 +730,18 @@ private List getUrnAspectValues(List> urnAspectEntrie * @return ordered list of values of multiple entities */ @Nonnull - private List filterAspects( - @Nonnull Set> aspectClasses, @Nullable IndexFilter filter, - @Nullable IndexSortCriterion indexSortCriterion, @Nullable String lastUrn, int count) { + private List filterAspects(@Nonnull Set> aspectClasses, + @Nullable IndexFilter filter, @Nullable IndexSortCriterion indexSortCriterion, @Nullable String lastUrn, + int count, boolean isInternalModelsEnabled) { final List> urnAspectEntries = getLocalDAO().getAspects(aspectClasses, filter, indexSortCriterion, parseUrnParam(lastUrn), count); - return getUrnAspectValues(urnAspectEntries); + return getUrnAspectValues(urnAspectEntries, isInternalModelsEnabled); } /** - * Similar to {@link #filterAspects(Set, IndexFilter, IndexSortCriterion, String, int)} but + * Similar to {@link #filterAspects(Set, IndexFilter, IndexSortCriterion, String, int, boolean)} but * takes in a start offset and returns a list result with pagination information. * * @param start defining the paging start @@ -519,14 +749,14 @@ private List filterAspects( * @return a {@link ListResult} containing a list of version numbers and other pagination information */ @Nonnull - private ListResult filterAspects( - @Nonnull Set> aspectClasses, @Nullable IndexFilter filter, - @Nullable IndexSortCriterion indexSortCriterion, int start, int count) { + private ListResult filterAspects(@Nonnull Set> aspectClasses, + @Nullable IndexFilter filter, @Nullable IndexSortCriterion indexSortCriterion, int start, int count, + boolean isInternalModelsEnabled) { final ListResult> listResult = getLocalDAO().getAspects(aspectClasses, filter, indexSortCriterion, start, count); final List> urnAspectEntries = listResult.getValues(); - final List values = getUrnAspectValues(urnAspectEntries); + final List values = getUrnAspectValues(urnAspectEntries, isInternalModelsEnabled); return ListResult.builder() .values(values) @@ -554,14 +784,16 @@ private ListResult filterAspects( */ @Nonnull private List filterUrns(@Nullable IndexFilter filter, @Nullable IndexSortCriterion indexSortCriterion, - @Nullable String lastUrn, int count) { + @Nullable String lastUrn, int count, boolean isInternalModelsEnabled) { final List urns = getLocalDAO().listUrns(filter, indexSortCriterion, parseUrnParam(lastUrn), count); - return urns.stream().map(urn -> toValue(newSnapshot(urn))).collect(Collectors.toList()); + return urns.stream() + .map(urn -> isInternalModelsEnabled ? toInternalValue(newInternalSnapshot(urn)) : toValue(newSnapshot(urn))) + .collect(Collectors.toList()); } /** - * Similar to {@link #filterUrns(IndexFilter, IndexSortCriterion, String, int)} but + * Similar to {@link #filterUrns(IndexFilter, IndexSortCriterion, String, int, boolean)} but * takes in a start offset and returns a list result with pagination information. * * @param start defining the paging start @@ -570,11 +802,13 @@ private List filterUrns(@Nullable IndexFilter filter, @Nullable IndexSort */ @Nonnull private ListResult filterUrns(@Nullable IndexFilter filter, @Nullable IndexSortCriterion indexSortCriterion, - int start, int count) { + int start, int count, boolean isInternalModelsEnabled) { final ListResult listResult = getLocalDAO().listUrns(filter, indexSortCriterion, start, count); final List urns = listResult.getValues(); - final List urnValues = urns.stream().map(urn -> toValue(newSnapshot(urn))).collect(Collectors.toList()); + final List urnValues = urns.stream() + .map(urn -> isInternalModelsEnabled ? toInternalValue(newInternalSnapshot(urn)) : toValue(newSnapshot(urn))) + .collect(Collectors.toList()); return ListResult.builder() .values(urnValues) @@ -605,19 +839,28 @@ private ListResult filterUrns(@Nullable IndexFilter filter, @Nullable Ind */ @Finder(FINDER_FILTER) @Nonnull - public Task> filter( - @QueryParam(PARAM_FILTER) @Optional @Nullable IndexFilter indexFilter, + public Task> filter(@QueryParam(PARAM_FILTER) @Optional @Nullable IndexFilter indexFilter, @QueryParam(PARAM_SORT) @Optional @Nullable IndexSortCriterion indexSortCriterion, @QueryParam(PARAM_ASPECTS) @Optional @Nullable String[] aspectNames, - @QueryParam(PARAM_URN) @Optional @Nullable String lastUrn, - @QueryParam(PARAM_COUNT) @Optional("10") int count) { + @QueryParam(PARAM_URN) @Optional @Nullable String lastUrn, @QueryParam(PARAM_COUNT) @Optional("10") int count) { + + return filter(indexFilter, indexSortCriterion, aspectNames, lastUrn, count, + _resourceLix.testFilter(_assetClass.getSimpleName())); + } + + @Nonnull + private Task> filter(@QueryParam(PARAM_FILTER) @Optional @Nullable IndexFilter indexFilter, + @QueryParam(PARAM_SORT) @Optional @Nullable IndexSortCriterion indexSortCriterion, + @QueryParam(PARAM_ASPECTS) @Optional @Nullable String[] aspectNames, + @QueryParam(PARAM_URN) @Optional @Nullable String lastUrn, @QueryParam(PARAM_COUNT) @Optional("10") int count, + boolean isInternalModelsEnabled) { return RestliUtils.toTask(() -> { - final Set> aspectClasses = parseAspectsParam(aspectNames); + final Set> aspectClasses = parseAspectsParam(aspectNames, isInternalModelsEnabled); if (aspectClasses.isEmpty()) { - return filterUrns(indexFilter, indexSortCriterion, lastUrn, count); + return filterUrns(indexFilter, indexSortCriterion, lastUrn, count, isInternalModelsEnabled); } else { - return filterAspects(aspectClasses, indexFilter, indexSortCriterion, lastUrn, count); + return filterAspects(aspectClasses, indexFilter, indexSortCriterion, lastUrn, count, isInternalModelsEnabled); } }); } @@ -652,18 +895,29 @@ public Task> filter( */ @Finder(FINDER_FILTER) @Nonnull - public Task> filter( - @QueryParam(PARAM_FILTER) @Optional @Nullable IndexFilter indexFilter, + public Task> filter(@QueryParam(PARAM_FILTER) @Optional @Nullable IndexFilter indexFilter, @QueryParam(PARAM_SORT) @Optional @Nullable IndexSortCriterion indexSortCriterion, @QueryParam(PARAM_ASPECTS) @Optional @Nullable String[] aspectNames, @PagingContextParam @Nonnull PagingContext pagingContext) { + return filter(indexFilter, indexSortCriterion, aspectNames, pagingContext, + _resourceLix.testFilter(_assetClass.getSimpleName())); + } + + @Nonnull + private Task> filter(@QueryParam(PARAM_FILTER) @Optional @Nullable IndexFilter indexFilter, + @QueryParam(PARAM_SORT) @Optional @Nullable IndexSortCriterion indexSortCriterion, + @QueryParam(PARAM_ASPECTS) @Optional @Nullable String[] aspectNames, + @PagingContextParam @Nonnull PagingContext pagingContext, boolean isInternalModelsEnabled) { return RestliUtils.toTask(() -> { - final Set> aspectClasses = parseAspectsParam(aspectNames); + final Set> aspectClasses = + parseAspectsParam(aspectNames, isInternalModelsEnabled); if (aspectClasses.isEmpty()) { - return filterUrns(indexFilter, indexSortCriterion, pagingContext.getStart(), pagingContext.getCount()); + return filterUrns(indexFilter, indexSortCriterion, pagingContext.getStart(), pagingContext.getCount(), + isInternalModelsEnabled); } else { - return filterAspects(aspectClasses, indexFilter, indexSortCriterion, pagingContext.getStart(), pagingContext.getCount()); + return filterAspects(aspectClasses, indexFilter, indexSortCriterion, pagingContext.getStart(), + pagingContext.getCount(), isInternalModelsEnabled); } }); } @@ -708,9 +962,10 @@ public Task> countAggregate( } @Nonnull - protected Set> parseAspectsParam(@Nullable String[] aspectNames) { + protected Set> parseAspectsParam(@Nullable String[] aspectNames, + boolean isInternalModelsEnabled) { if (aspectNames == null) { - return _supportedAspectClasses; + return isInternalModelsEnabled ? _supportedInternalAspectClasses : _supportedAspectClasses; } return Arrays.asList(aspectNames).stream().map(ModelUtils::getAspectClass).collect(Collectors.toSet()); } @@ -720,31 +975,36 @@ protected Set> parseAspectsParam(@Nullable Strin * * @param urns collection of urns * @param aspectClasses set of aspect classes + * @param isInternalModelsEnabled flag to switch the internal models * @return All {@link VALUE} objects keyed by {@link URN} obtained from DB */ @Nonnull protected Map getInternal(@Nonnull Collection urns, - @Nonnull Set> aspectClasses) { - return getUrnAspectMap(urns, aspectClasses).entrySet() + @Nonnull Set> aspectClasses, boolean isInternalModelsEnabled) { + return getUrnAspectMap(urns, aspectClasses, isInternalModelsEnabled).entrySet() .stream() - .collect(Collectors.toMap(Map.Entry::getKey, e -> toValue(newSnapshot(e.getKey(), e.getValue())))); + .collect(Collectors.toMap(Map.Entry::getKey, + e -> isInternalModelsEnabled ? toInternalValue(newInternalSnapshot(e.getKey(), e.getValue())) + : toValue(newSnapshot(e.getKey(), e.getValue())))); } /** - * Similar to {@link #getInternal(Collection, Set)} but filter out {@link URN}s which are not in the DB. + * Similar to {@link #getInternal(Collection, Set, boolean)} but filter out {@link URN}s which are not in the DB. */ @Nonnull protected Map getInternalNonEmpty(@Nonnull Collection urns, - @Nonnull Set> aspectClasses) { - return getUrnAspectMap(urns, aspectClasses).entrySet() + @Nonnull Set> aspectClasses, boolean isInternalModelsEnabled) { + return getUrnAspectMap(urns, aspectClasses, isInternalModelsEnabled).entrySet() .stream() .filter(e -> !e.getValue().isEmpty()) - .collect(Collectors.toMap(Map.Entry::getKey, e -> toValue(newSnapshot(e.getKey(), e.getValue())))); + .collect(Collectors.toMap(Map.Entry::getKey, + e -> isInternalModelsEnabled ? toInternalValue(newInternalSnapshot(e.getKey(), e.getValue())) + : toValue(newSnapshot(e.getKey(), e.getValue())))); } @Nonnull private Map> getUrnAspectMap(@Nonnull Collection urns, - @Nonnull Set> aspectClasses) { + @Nonnull Set> aspectClasses, boolean isInternalModelsEnabled) { // Construct the keys to retrieve latest version of all supported aspects for all URNs. final Set> keys = urns.stream() .map(urn -> aspectClasses.stream() @@ -756,10 +1016,15 @@ private Map> getUrnAspectMap(@Nonnull Collection u final Map> urnAspectsMap = urns.stream().collect(Collectors.toMap(Function.identity(), urn -> new ArrayList<>())); - getLocalDAO().get(keys) - .forEach((key, aspect) -> aspect.ifPresent( - metadata -> urnAspectsMap.get(key.getUrn()).add(ModelUtils.newAspectUnion(_aspectUnionClass, metadata)))); - + if (isInternalModelsEnabled) { + getLocalDAO().get(keys) + .forEach((key, aspect) -> aspect.ifPresent(metadata -> urnAspectsMap.get(key.getUrn()) + .add(ModelUtils.newAspectUnion(_internalAspectUnionClass, metadata)))); + } else { + getLocalDAO().get(keys) + .forEach((key, aspect) -> aspect.ifPresent( + metadata -> urnAspectsMap.get(key.getUrn()).add(ModelUtils.newAspectUnion(_aspectUnionClass, metadata)))); + } return urnAspectsMap; } @@ -776,6 +1041,19 @@ protected SNAPSHOT newSnapshot(@Nonnull URN urn) { return ModelUtils.newSnapshot(_snapshotClass, urn, Collections.emptyList()); } + @Nonnull + private INTERNAL_SNAPSHOT newInternalSnapshot(@Nonnull URN urn, @Nonnull List aspects) { + return ModelUtils.newSnapshot(_internalSnapshotClass, urn, aspects); + } + + /** + * Creates an Internal snapshot of the entity with no aspects set, just the URN. + */ + @Nonnull + protected INTERNAL_SNAPSHOT newInternalSnapshot(@Nonnull URN urn) { + return ModelUtils.newSnapshot(_internalSnapshotClass, urn, Collections.emptyList()); + } + @Nullable protected URN parseUrnParam(@Nullable String urnString) { if (urnString == null) { diff --git a/restli-resources/src/main/java/com/linkedin/metadata/restli/BaseSearchableEntityResource.java b/restli-resources/src/main/java/com/linkedin/metadata/restli/BaseSearchableEntityResource.java index b52a891ed..a86dd61fc 100644 --- a/restli-resources/src/main/java/com/linkedin/metadata/restli/BaseSearchableEntityResource.java +++ b/restli-resources/src/main/java/com/linkedin/metadata/restli/BaseSearchableEntityResource.java @@ -13,6 +13,8 @@ import com.linkedin.metadata.query.SearchResultMetadata; import com.linkedin.metadata.query.SortCriterion; import com.linkedin.metadata.query.SortOrder; +import com.linkedin.metadata.restli.lix.DummyResourceLix; +import com.linkedin.metadata.restli.lix.ResourceLix; import com.linkedin.parseq.Task; import com.linkedin.restli.server.CollectionResult; import com.linkedin.restli.server.PagingContext; @@ -43,6 +45,9 @@ * @param must be a valid snapshot type defined in com.linkedin.metadata.snapshot * @param must be a valid aspect union type supported by the snapshot * @param must be a valid search document type defined in com.linkedin.metadata.search + * @param must be a valid internal snapshot type defined in com.linkedin.metadata.snapshot + * @param must be a valid internal aspect union type supported by the internal snapshot + * @param must be a valid asset type defined in com.linkedin.metadata.asset */ public abstract class BaseSearchableEntityResource< // @formatter:off @@ -51,22 +56,40 @@ public abstract class BaseSearchableEntityResource< URN extends Urn, SNAPSHOT extends RecordTemplate, ASPECT_UNION extends UnionTemplate, - DOCUMENT extends RecordTemplate> + DOCUMENT extends RecordTemplate, + INTERNAL_SNAPSHOT extends RecordTemplate, + INTERNAL_ASPECT_UNION extends UnionTemplate, + ASSET extends RecordTemplate> // @formatter:on - extends BaseEntityResource { + extends + BaseEntityResource { private static final String DEFAULT_SORT_CRITERION_FIELD = "urn"; - public BaseSearchableEntityResource(@Nonnull Class snapshotClass, - @Nonnull Class aspectUnionClass) { - super(snapshotClass, aspectUnionClass); + public BaseSearchableEntityResource(@Nullable Class snapshotClass, + @Nullable Class aspectUnionClass, @Nonnull Class internalSnapshotClass, + @Nonnull Class internalAspectUnionClass, @Nonnull Class assetClass) { + super(snapshotClass, aspectUnionClass, internalSnapshotClass, internalAspectUnionClass, assetClass); } - public BaseSearchableEntityResource(@Nonnull Class snapshotClass, - @Nonnull Class aspectUnionClass, @Nonnull Class urnClass) { - super(snapshotClass, aspectUnionClass, urnClass); + public BaseSearchableEntityResource(@Nullable Class snapshotClass, + @Nullable Class aspectUnionClass, @Nonnull Class urnClass, + @Nonnull Class internalSnapshotClass, + @Nonnull Class internalAspectUnionClass, @Nonnull Class assetClass) { + super(snapshotClass, aspectUnionClass, urnClass, internalSnapshotClass, internalAspectUnionClass, assetClass, + new DummyResourceLix()); } + public BaseSearchableEntityResource(@Nullable Class snapshotClass, + @Nullable Class aspectUnionClass, @Nonnull Class urnClass, + @Nonnull Class internalSnapshotClass, + @Nonnull Class internalAspectUnionClass, @Nonnull Class assetClass, + @Nonnull ResourceLix resourceLix) { + super(snapshotClass, aspectUnionClass, urnClass, internalSnapshotClass, internalAspectUnionClass, assetClass, + resourceLix); + } + + /** * Returns a document-specific {@link BaseSearchDAO}. */ @@ -89,6 +112,15 @@ public Task> getAll(@Nonnull PagingContext pagingContext, @QueryParam(PARAM_ASPECTS) @Optional @Nullable String[] aspectNames, @QueryParam(PARAM_FILTER) @Optional @Nullable Filter filter, @QueryParam(PARAM_SORT) @Optional @Nullable SortCriterion sortCriterion) { + final String urnType = _urnClass == null ? null : _urnClass.getSimpleName(); + return getAll(pagingContext, aspectNames, filter, sortCriterion, _resourceLix.testGetAll(urnType)); + } + + @Nonnull + protected Task> getAll(@Nonnull PagingContext pagingContext, + @QueryParam(PARAM_ASPECTS) @Optional @Nullable String[] aspectNames, + @QueryParam(PARAM_FILTER) @Optional @Nullable Filter filter, + @QueryParam(PARAM_SORT) @Optional @Nullable SortCriterion sortCriterion, boolean isInternalModelsEnabled) { final Filter searchFilter = filter != null ? filter : QueryUtils.EMPTY_FILTER; final SortCriterion searchSortCriterion = sortCriterion != null ? sortCriterion @@ -96,7 +128,7 @@ public Task> getAll(@Nonnull PagingContext pagingContext, final SearchResult filterResult = getSearchDAO().filter(searchFilter, searchSortCriterion, pagingContext.getStart(), pagingContext.getCount()); return RestliUtils.toTask( - () -> getSearchQueryCollectionResult(filterResult, aspectNames).getElements()); + () -> getSearchQueryCollectionResult(filterResult, aspectNames, isInternalModelsEnabled).getElements()); } @Finder(FINDER_SEARCH) @@ -106,12 +138,22 @@ public Task> search(@QueryParam(PA @QueryParam(PARAM_FILTER) @Optional @Nullable Filter filter, @QueryParam(PARAM_SORT) @Optional @Nullable SortCriterion sortCriterion, @PagingContextParam @Nonnull PagingContext pagingContext) { + final String urnType = _urnClass == null ? null : _urnClass.getSimpleName(); + return search(input, aspectNames, filter, sortCriterion, pagingContext, _resourceLix.testSearch(urnType)); + } + + @Nonnull + private Task> search(@QueryParam(PARAM_INPUT) @Nonnull String input, + @QueryParam(PARAM_ASPECTS) @Optional @Nullable String[] aspectNames, + @QueryParam(PARAM_FILTER) @Optional @Nullable Filter filter, + @QueryParam(PARAM_SORT) @Optional @Nullable SortCriterion sortCriterion, + @PagingContextParam @Nonnull PagingContext pagingContext, boolean isInternalModelsEnabled) { final Filter searchFilter = filter != null ? filter : QueryUtils.EMPTY_FILTER; final SearchResult searchResult = getSearchDAO().search(input, searchFilter, sortCriterion, pagingContext.getStart(), pagingContext.getCount()); return RestliUtils.toTask( - () -> getSearchQueryCollectionResult(searchResult, aspectNames)); + () -> getSearchQueryCollectionResult(searchResult, aspectNames, isInternalModelsEnabled)); } /** @@ -133,12 +175,24 @@ public Task> searchV2(@QueryParam( @QueryParam(PARAM_SORT) @Optional @Nullable SortCriterion sortCriterion, @QueryParam(PARAM_PREFERENCE) @Optional @Nullable String preference, @PagingContextParam @Nonnull PagingContext pagingContext) { + final String urnType = _urnClass == null ? null : _urnClass.getSimpleName(); + return searchV2(input, aspectNames, filter, sortCriterion, preference, pagingContext, + _resourceLix.testSearchV2(urnType)); + } + + @Nonnull + private Task> searchV2(@QueryParam(PARAM_INPUT) @Nonnull String input, + @QueryParam(PARAM_ASPECTS) @Optional @Nullable String[] aspectNames, + @QueryParam(PARAM_FILTER) @Optional @Nullable Filter filter, + @QueryParam(PARAM_SORT) @Optional @Nullable SortCriterion sortCriterion, + @QueryParam(PARAM_PREFERENCE) @Optional @Nullable String preference, + @PagingContextParam @Nonnull PagingContext pagingContext, boolean isInternalModelsEnabled) { final Filter searchFilter = filter != null ? filter : QueryUtils.EMPTY_FILTER; final SearchResult searchResult = getSearchDAO().searchV2(input, searchFilter, sortCriterion, preference, pagingContext.getStart(), pagingContext.getCount()); return RestliUtils.toTask( - () -> getSearchQueryCollectionResult(searchResult, aspectNames)); + () -> getSearchQueryCollectionResult(searchResult, aspectNames, isInternalModelsEnabled)); } @Action(name = ACTION_AUTOCOMPLETE) @@ -156,14 +210,16 @@ public Task autocomplete(@ActionParam(PARAM_QUERY) @Nonnull * @return CollectionResult which contains: 1. aspect values fetched from MySQL DB, 2. Total count 3. Search result metadata. */ @Nonnull - public CollectionResult getSearchQueryCollectionResult(@Nonnull SearchResult searchResult, - @Nullable String[] aspectNames) { + private CollectionResult getSearchQueryCollectionResult(@Nonnull SearchResult searchResult, + @Nullable String[] aspectNames, boolean isInternalModelsEnabled) { final List matchedUrns = searchResult.getDocumentList() .stream() .map(d -> (URN) ModelUtils.getUrnFromDocument(d)) .collect(Collectors.toList()); - final Map urnValueMap = getInternalNonEmpty(matchedUrns, parseAspectsParam(aspectNames)); + final Map urnValueMap = + getInternalNonEmpty(matchedUrns, parseAspectsParam(aspectNames, isInternalModelsEnabled), + isInternalModelsEnabled); final List existingUrns = matchedUrns.stream().filter(urn -> urnValueMap.containsKey(urn)).collect(Collectors.toList()); return new CollectionResult<>( existingUrns.stream().map(urn -> urnValueMap.get(urn)).collect(Collectors.toList()), diff --git a/restli-resources/src/main/java/com/linkedin/metadata/restli/BaseSingleAspectEntityResource.java b/restli-resources/src/main/java/com/linkedin/metadata/restli/BaseSingleAspectEntityResource.java index 977e5096f..0e7e33cd5 100644 --- a/restli-resources/src/main/java/com/linkedin/metadata/restli/BaseSingleAspectEntityResource.java +++ b/restli-resources/src/main/java/com/linkedin/metadata/restli/BaseSingleAspectEntityResource.java @@ -22,6 +22,7 @@ import java.util.Set; import java.util.stream.Collectors; import javax.annotation.Nonnull; +import javax.annotation.Nullable; import static com.linkedin.metadata.dao.BaseReadDAO.*; @@ -36,6 +37,9 @@ * @param must be a valid aspect of the resource * @param must be a valid aspect union type * @param must be a valid snapshot type defined in com.linkedin.metadata.snapshot + * @param must be a valid internal snapshot type defined in com.linkedin.metadata.snapshot + * @param must be a valid internal aspect union type supported by the internal snapshot + * @param must be a valid asset type defined in com.linkedin.metadata.asset */ public abstract class BaseSingleAspectEntityResource< // @formatter:off @@ -44,9 +48,13 @@ public abstract class BaseSingleAspectEntityResource< URN extends Urn, ASPECT extends RecordTemplate, ASPECT_UNION extends UnionTemplate, - SNAPSHOT extends RecordTemplate> + SNAPSHOT extends RecordTemplate, + INTERNAL_SNAPSHOT extends RecordTemplate, + INTERNAL_ASPECT_UNION extends UnionTemplate, + ASSET extends RecordTemplate> // @formatter:on - extends BaseEntityResource { + extends + BaseEntityResource { private final Class _valueClass; private final Class _aspectClass; @@ -55,10 +63,11 @@ public abstract class BaseSingleAspectEntityResource< * Constructor. */ public BaseSingleAspectEntityResource(@Nonnull Class aspectClass, - @Nonnull Class aspectUnionClass, @Nonnull Class valueClass, - @Nonnull Class snapshotClass) { + @Nullable Class aspectUnionClass, @Nonnull Class valueClass, + @Nullable Class snapshotClass, @Nonnull Class internalSnapshotClass, + @Nonnull Class internalAspectUnionClass, @Nonnull Class assetClass) { - super(snapshotClass, aspectUnionClass); + super(snapshotClass, aspectUnionClass, internalSnapshotClass, internalAspectUnionClass, assetClass); this._aspectClass = aspectClass; this._valueClass = valueClass; } @@ -82,7 +91,7 @@ public BaseSingleAspectEntityResource(@Nonnull Class aspectClass, @Override @Nonnull protected Map getInternal(@Nonnull Collection urns, - @Nonnull Set> aspectClasses) { + @Nonnull Set> aspectClasses, boolean isInternalModelsEnabled) { return getUrnEntityMapInternal(urns); } @@ -94,7 +103,7 @@ protected Map getInternal(@Nonnull Collection urns, @Override @Nonnull protected Map getInternalNonEmpty(@Nonnull Collection urns, - @Nonnull Set> aspectClasses) { + @Nonnull Set> aspectClasses, boolean isInternalModelsEnabled) { return getUrnEntityMapInternal(urns); } @@ -135,7 +144,7 @@ private VALUE createPartialEntityFromAspect(@Nonnull ASPECT aspect) { /** * Throwing an exception with a `not implemented` error message as this method is only required - * by parent class {@link BaseEntityResource} method- {@link #getInternalNonEmpty(Collection, Set)}, + * by parent class {@link BaseEntityResource} method- {@link #getInternalNonEmpty(Collection, Set, boolean)} (Collection, Set)}, * which has been overridden here. */ @Override @@ -146,7 +155,7 @@ protected VALUE toValue(@Nonnull SNAPSHOT snapshot) { /** * Throwing an exception with a `not implemented` error message as this method is only required - * by parent class {@link BaseEntityResource} method- {@link #getInternalNonEmpty(Collection, Set)}, + * by parent class {@link BaseEntityResource} method- {@link #getInternalNonEmpty(Collection, Set, boolean)} (Collection, Set)}, * which has been overridden here. */ @Override diff --git a/restli-resources/src/main/java/com/linkedin/metadata/restli/BaseSingleAspectSearchableEntityResource.java b/restli-resources/src/main/java/com/linkedin/metadata/restli/BaseSingleAspectSearchableEntityResource.java index 111787edb..0bf390797 100644 --- a/restli-resources/src/main/java/com/linkedin/metadata/restli/BaseSingleAspectSearchableEntityResource.java +++ b/restli-resources/src/main/java/com/linkedin/metadata/restli/BaseSingleAspectSearchableEntityResource.java @@ -12,6 +12,7 @@ import java.util.Set; import java.util.stream.Collectors; import javax.annotation.Nonnull; +import javax.annotation.Nullable; import static com.linkedin.metadata.dao.BaseReadDAO.*; @@ -27,6 +28,9 @@ * @param must be a valid aspect union type supported by the snapshot * @param must be a valid snapshot type defined in com.linkedin.metadata.snapshot * @param must be a valid search document type defined in com.linkedin.metadata.search + * @param must be a valid internal snapshot type defined in com.linkedin.metadata.snapshot + * @param must be a valid internal aspect union type supported by the internal snapshot + * @param must be a valid asset type defined in com.linkedin.metadata.asset */ public abstract class BaseSingleAspectSearchableEntityResource< // @formatter:off @@ -36,9 +40,13 @@ public abstract class BaseSingleAspectSearchableEntityResource< ASPECT extends RecordTemplate, ASPECT_UNION extends UnionTemplate, SNAPSHOT extends RecordTemplate, - DOCUMENT extends RecordTemplate> + DOCUMENT extends RecordTemplate, + ASSET extends RecordTemplate, + INTERNAL_SNAPSHOT extends RecordTemplate, + INTERNAL_ASPECT_UNION extends UnionTemplate> // @formatter:on - extends BaseSearchableEntityResource { + extends + BaseSearchableEntityResource { private final Class _aspectClass; private final Class _valueClass; @@ -46,13 +54,12 @@ public abstract class BaseSingleAspectSearchableEntityResource< /** * Constructor. * */ - public BaseSingleAspectSearchableEntityResource( - @Nonnull Class aspectClass, - @Nonnull Class aspectUnionClass, - @Nonnull Class valueClass, - @Nonnull Class snapshotClass) { + public BaseSingleAspectSearchableEntityResource(@Nonnull Class aspectClass, + @Nullable Class aspectUnionClass, @Nonnull Class valueClass, + @Nullable Class snapshotClass, @Nonnull Class internalSnapshotClass, + @Nonnull Class internalAspectUnionClass, @Nonnull Class assetClass) { - super(snapshotClass, aspectUnionClass); + super(snapshotClass, aspectUnionClass, internalSnapshotClass, internalAspectUnionClass, assetClass); _aspectClass = aspectClass; _valueClass = valueClass; } @@ -75,9 +82,8 @@ public BaseSingleAspectSearchableEntityResource( * */ @Override @Nonnull - protected Map getInternal( - @Nonnull Collection urns, - @Nonnull Set> aspectClasses) { + protected Map getInternal(@Nonnull Collection urns, + @Nonnull Set> aspectClasses, boolean isInternalModelsEnabled) { // ignore the second parameter as it is not required for single aspect entities return getUrnEntityMapInternal(urns); } @@ -89,9 +95,8 @@ protected Map getInternal( * */ @Override @Nonnull - protected Map getInternalNonEmpty( - @Nonnull Collection urns, - @Nonnull Set> aspectClasses) { + protected Map getInternalNonEmpty(@Nonnull Collection urns, + @Nonnull Set> aspectClasses, boolean isInternalModelsEnabled) { // ignore the second parameter as it is not required for single aspect entities return getUrnEntityMapInternal(urns); } @@ -133,7 +138,7 @@ private VALUE createPartialEntityFromAspect(@Nonnull ASPECT aspect) { /** * Throwing an exception with a `not implemented` error message as this method is only required - * by parent class {@link BaseEntityResource} method- {@link #getInternalNonEmpty(Collection, Set)}, + * by parent class {@link BaseEntityResource} method- {@link #getInternalNonEmpty(Collection, Set, boolean)}, * which has been overridden here. * */ @Override @@ -144,7 +149,7 @@ protected VALUE toValue(@Nonnull SNAPSHOT snapshot) { /** * Throwing an exception with a `not implemented` error message as this method is only required - * by parent class {@link BaseEntityResource} method- {@link #getInternalNonEmpty(Collection, Set)}, + * by parent class {@link BaseEntityResource} method- {@link #getInternalNonEmpty(Collection, Set, boolean)} (Collection, Set)}, * which has been overridden here. * */ @Override diff --git a/restli-resources/src/main/java/com/linkedin/metadata/restli/RestliConstants.java b/restli-resources/src/main/java/com/linkedin/metadata/restli/RestliConstants.java index 6698e3f21..1126b5a95 100644 --- a/restli-resources/src/main/java/com/linkedin/metadata/restli/RestliConstants.java +++ b/restli-resources/src/main/java/com/linkedin/metadata/restli/RestliConstants.java @@ -20,8 +20,8 @@ private RestliConstants() { } public static final String ACTION_BROWSE = "browse"; public static final String ACTION_COUNT_AGGREGATE = "countAggregate"; public static final String ACTION_EMIT_NO_CHANGE_METADATA_AUDIT_EVENT = "emitNoChangeMetadataAuditEvent"; - public static final String ACTION_GET_BROWSE_PATHS = "getBrowsePaths"; public static final String ACTION_GET_ASSET = "getAsset"; + public static final String ACTION_GET_BROWSE_PATHS = "getBrowsePaths"; public static final String ACTION_GET_SNAPSHOT = "getSnapshot"; public static final String ACTION_INGEST = "ingest"; public static final String ACTION_INGEST_ASSET = "ingestAsset"; diff --git a/restli-resources/src/main/java/com/linkedin/metadata/restli/lix/DummyResourceLix.java b/restli-resources/src/main/java/com/linkedin/metadata/restli/lix/DummyResourceLix.java index 057685cb9..1dbad9f39 100644 --- a/restli-resources/src/main/java/com/linkedin/metadata/restli/lix/DummyResourceLix.java +++ b/restli-resources/src/main/java/com/linkedin/metadata/restli/lix/DummyResourceLix.java @@ -21,16 +21,6 @@ public boolean testBatchGetWithErrors(@Nullable String urn, @Nullable String typ return false; } - @Override - public boolean testIngest(@Nonnull String urn, @Nonnull String entityType, @Nullable String aspectName) { - return false; - } - - @Override - public boolean testIngestWithTracking(@Nonnull String urn, @Nonnull String entityType, @Nullable String aspectName) { - return false; - } - @Override public boolean testGetSnapshot(@Nullable String urn, @Nullable String entityType) { return false; diff --git a/restli-resources/src/main/java/com/linkedin/metadata/restli/lix/ResourceLix.java b/restli-resources/src/main/java/com/linkedin/metadata/restli/lix/ResourceLix.java index b9c315091..2b2ac9f29 100644 --- a/restli-resources/src/main/java/com/linkedin/metadata/restli/lix/ResourceLix.java +++ b/restli-resources/src/main/java/com/linkedin/metadata/restli/lix/ResourceLix.java @@ -34,24 +34,6 @@ public interface ResourceLix { */ boolean testBatchGetWithErrors(@Nullable String urn, @Nullable String type); - /** - * Experiment on the Ingest. - * @param urn urnString of the entity - * @param entityType type of the entity - * @param aspectName aspect FQCN of the urn - * @return enabling/not - */ - boolean testIngest(@Nonnull String urn, @Nonnull String entityType, @Nullable String aspectName); - - /** - * Experiment on the IngestWithTracking. - * @param urn urnString of the entity - * @param entityType type of the entity - * @param aspectName aspect FQCN of the urn - * @return enabling/not enabling/not - */ - boolean testIngestWithTracking(@Nonnull String urn, @Nonnull String entityType, @Nullable String aspectName); - /** * Experiment on the GetSnapshot. * @param urn urnString of the entity diff --git a/restli-resources/src/test/java/com/linkedin/metadata/restli/BaseAspectRoutingResourceTest.java b/restli-resources/src/test/java/com/linkedin/metadata/restli/BaseAspectRoutingResourceTest.java index 07e96a80a..87c276688 100644 --- a/restli-resources/src/test/java/com/linkedin/metadata/restli/BaseAspectRoutingResourceTest.java +++ b/restli-resources/src/test/java/com/linkedin/metadata/restli/BaseAspectRoutingResourceTest.java @@ -21,10 +21,13 @@ import com.linkedin.testing.AspectFoo; import com.linkedin.testing.EntityAspectUnion; import com.linkedin.testing.EntityAspectUnionArray; +import com.linkedin.testing.EntityAsset; import com.linkedin.testing.EntityDocument; import com.linkedin.testing.EntityKey; import com.linkedin.testing.EntitySnapshot; import com.linkedin.testing.EntityValue; +import com.linkedin.testing.InternalEntityAspectUnion; +import com.linkedin.testing.InternalEntitySnapshot; import com.linkedin.testing.urn.BazUrn; import com.linkedin.testing.urn.FooUrn; import java.net.URISyntaxException; @@ -59,15 +62,17 @@ public class BaseAspectRoutingResourceTest extends BaseEngineTest { class TestResource extends BaseAspectRoutingResource< // format - ComplexResourceKey, EntityValue, FooUrn, EntitySnapshot, EntityAspectUnion, EntityDocument> { + ComplexResourceKey, EntityValue, FooUrn, EntitySnapshot, EntityAspectUnion, EntityDocument, + InternalEntitySnapshot, InternalEntityAspectUnion, EntityAsset> { public TestResource() { - super(EntitySnapshot.class, EntityAspectUnion.class, FooUrn.class, EntityValue.class); + super(EntitySnapshot.class, EntityAspectUnion.class, FooUrn.class, EntityValue.class, InternalEntitySnapshot.class, + InternalEntityAspectUnion.class, EntityAsset.class); } @Nonnull @Override - protected BaseLocalDAO getLocalDAO() { + protected BaseLocalDAO getLocalDAO() { return _mockLocalDAO; } diff --git a/restli-resources/src/test/java/com/linkedin/metadata/restli/BaseBrowsableEntityResourceTest.java b/restli-resources/src/test/java/com/linkedin/metadata/restli/BaseBrowsableEntityResourceTest.java index 7a5d9e93a..b0c455a0f 100644 --- a/restli-resources/src/test/java/com/linkedin/metadata/restli/BaseBrowsableEntityResourceTest.java +++ b/restli-resources/src/test/java/com/linkedin/metadata/restli/BaseBrowsableEntityResourceTest.java @@ -14,10 +14,13 @@ import com.linkedin.restli.common.ComplexResourceKey; import com.linkedin.restli.common.EmptyRecord; import com.linkedin.testing.EntityAspectUnion; +import com.linkedin.testing.EntityAsset; import com.linkedin.testing.EntityDocument; import com.linkedin.testing.EntityKey; import com.linkedin.testing.EntitySnapshot; import com.linkedin.testing.EntityValue; +import com.linkedin.testing.InternalEntityAspectUnion; +import com.linkedin.testing.InternalEntitySnapshot; import java.net.URISyntaxException; import java.util.List; import javax.annotation.Nonnull; @@ -37,15 +40,17 @@ public class BaseBrowsableEntityResourceTest extends BaseEngineTest { class TestResource extends BaseBrowsableEntityResource< // format - ComplexResourceKey, EntityValue, Urn, EntitySnapshot, EntityAspectUnion, EntityDocument> { + ComplexResourceKey, EntityValue, Urn, EntitySnapshot, EntityAspectUnion, EntityDocument, + InternalEntitySnapshot, InternalEntityAspectUnion, EntityAsset> { public TestResource() { - super(EntitySnapshot.class, EntityAspectUnion.class); + super(EntitySnapshot.class, EntityAspectUnion.class, InternalEntitySnapshot.class, + InternalEntityAspectUnion.class, EntityAsset.class); } @Nonnull @Override - protected BaseLocalDAO getLocalDAO() { + protected BaseLocalDAO getLocalDAO() { throw new RuntimeException("Not implemented"); } diff --git a/restli-resources/src/test/java/com/linkedin/metadata/restli/BaseEntityResourceTest.java b/restli-resources/src/test/java/com/linkedin/metadata/restli/BaseEntityResourceTest.java index 733bbb06e..7c2e67dd7 100644 --- a/restli-resources/src/test/java/com/linkedin/metadata/restli/BaseEntityResourceTest.java +++ b/restli-resources/src/test/java/com/linkedin/metadata/restli/BaseEntityResourceTest.java @@ -24,6 +24,7 @@ import com.linkedin.metadata.query.IndexSortCriterion; import com.linkedin.metadata.query.MapMetadata; import com.linkedin.metadata.query.SortOrder; +import com.linkedin.metadata.restli.lix.ResourceLix; import com.linkedin.parseq.BaseEngineTest; import com.linkedin.restli.common.ComplexResourceKey; import com.linkedin.restli.common.EmptyRecord; @@ -36,12 +37,16 @@ import com.linkedin.testing.AspectAttributes; import com.linkedin.testing.AspectBar; import com.linkedin.testing.AspectFoo; +import com.linkedin.testing.AspectFooEvolved; import com.linkedin.testing.BarUrnArray; import com.linkedin.testing.EntityAspectUnion; import com.linkedin.testing.EntityAspectUnionArray; +import com.linkedin.testing.EntityAsset; import com.linkedin.testing.EntityKey; import com.linkedin.testing.EntitySnapshot; import com.linkedin.testing.EntityValue; +import com.linkedin.testing.InternalEntityAspectUnion; +import com.linkedin.testing.InternalEntitySnapshot; import com.linkedin.testing.localrelationship.AspectFooBar; import com.linkedin.testing.localrelationship.BelongsTo; import com.linkedin.testing.urn.BarUrn; @@ -57,6 +62,7 @@ import java.util.Set; import java.util.stream.Collectors; import javax.annotation.Nonnull; +import javax.annotation.Nullable; import org.testng.annotations.BeforeMethod; import org.testng.annotations.Test; @@ -68,18 +74,176 @@ public class BaseEntityResourceTest extends BaseEngineTest { - private BaseLocalDAO _mockLocalDAO; + private BaseLocalDAO _mockLocalDAO; private TestResource _resource = new TestResource(); + private TestInternalResource _internalResource = new TestInternalResource(); - class TestResource extends BaseEntityResource, EntityValue, FooUrn, EntitySnapshot, EntityAspectUnion> { + class TestResource extends + BaseEntityResource, EntityValue, FooUrn, EntitySnapshot, + EntityAspectUnion, InternalEntitySnapshot, InternalEntityAspectUnion, EntityAsset> { public TestResource() { - super(EntitySnapshot.class, EntityAspectUnion.class, FooUrn.class); + super(EntitySnapshot.class, EntityAspectUnion.class, FooUrn.class, InternalEntitySnapshot.class, + InternalEntityAspectUnion.class, EntityAsset.class); } @Nonnull @Override - protected BaseLocalDAO getLocalDAO() { + protected BaseLocalDAO getLocalDAO() { + return _mockLocalDAO; + } + + @Nonnull + @Override + protected FooUrn createUrnFromString(@Nonnull String urnString) { + try { + return FooUrn.createFromString(urnString); + } catch (URISyntaxException e) { + throw RestliUtils.badRequestException("Invalid URN: " + urnString); + } + } + + @Nonnull + @Override + protected FooUrn toUrn(@Nonnull ComplexResourceKey key) { + return makeFooUrn(key.getKey().getId().intValue()); + } + + @Nonnull + @Override + protected ComplexResourceKey toKey(@Nonnull FooUrn urn) { + return new ComplexResourceKey<>(new EntityKey().setId(urn.getIdAsLong()), new EmptyRecord()); + } + + @Nonnull + @Override + protected EntityValue toValue(@Nonnull EntitySnapshot snapshot) { + EntityValue value = new EntityValue(); + ModelUtils.getAspectsFromSnapshot(snapshot).forEach(a -> { + if (a instanceof AspectFoo) { + value.setFoo(AspectFoo.class.cast(a)); + } else if (a instanceof AspectBar) { + value.setBar(AspectBar.class.cast(a)); + } else if (a instanceof AspectAttributes) { + value.setAttributes(AspectAttributes.class.cast(a)); + } + }); + return value; + } + + @Nonnull + @Override + protected EntitySnapshot toSnapshot(@Nonnull EntityValue value, @Nonnull FooUrn urn) { + EntitySnapshot snapshot = new EntitySnapshot().setUrn(urn); + EntityAspectUnionArray aspects = new EntityAspectUnionArray(); + if (value.hasFoo()) { + aspects.add(ModelUtils.newAspectUnion(EntityAspectUnion.class, value.getFoo())); + } + if (value.hasBar()) { + aspects.add(ModelUtils.newAspectUnion(EntityAspectUnion.class, value.getBar())); + } + if (value.hasAttributes()) { + aspects.add(ModelUtils.newAspectUnion(EntityAspectUnion.class, value.getAttributes())); + } + + snapshot.setAspects(aspects); + return snapshot; + } + + @Override + public ResourceContext getContext() { + return mock(ResourceContext.class); + } + } + + class TestInternalResource extends + BaseEntityResource, EntityValue, FooUrn, EntitySnapshot, + EntityAspectUnion, InternalEntitySnapshot, InternalEntityAspectUnion, EntityAsset> { + + public TestInternalResource() { + super(EntitySnapshot.class, EntityAspectUnion.class, FooUrn.class, InternalEntitySnapshot.class, + InternalEntityAspectUnion.class, EntityAsset.class, new ResourceLix() { + + @Override + public boolean testGet(@Nonnull String urn, @Nonnull String entityType) { + return true; + } + + @Override + public boolean testBatchGet(@Nullable String urn, @Nullable String entityType) { + return true; + } + + @Override + public boolean testBatchGetWithErrors(@Nullable String urn, @Nullable String type) { + return false; + } + + @Override + public boolean testGetSnapshot(@Nullable String urn, @Nullable String entityType) { + return true; + } + + @Override + public boolean testBackfillLegacy(@Nullable String urn, @Nullable String entityType) { + return false; + } + + @Override + public boolean testBackfillWithUrns(@Nullable String urn, @Nullable String entityType) { + return false; + } + + @Override + public boolean testEmitNoChangeMetadataAuditEvent(@Nullable String urn, @Nullable String entityType) { + return false; + } + + @Override + public boolean testBackfillWithNewValue(@Nullable String urn, @Nullable String entityType) { + return false; + } + + @Override + public boolean testBackfillEntityTables(@Nullable String urn, @Nullable String entityType) { + return false; + } + + @Override + public boolean testBackfillRelationshipTables(@Nullable String urn, @Nullable String entityType) { + return false; + } + + @Override + public boolean testBackfill(@Nonnull String assetType, @Nonnull String mode) { + return false; + } + + @Override + public boolean testFilter(@Nonnull String assetType) { + return true; + } + + @Override + public boolean testGetAll(@Nullable String urnType) { + return false; + } + + @Override + public boolean testSearch(@Nullable String urnType) { + return false; + } + + @Override + public boolean testSearchV2(@Nullable String urnType) { + return false; + } + }); + } + + @Nonnull + @Override + protected BaseLocalDAO getLocalDAO() { return _mockLocalDAO; } @@ -169,6 +333,26 @@ public void testGet() { assertFalse(value.hasBar()); } + @Test + public void testInternalModelGet() { + FooUrn urn = makeFooUrn(1234); + AspectFoo foo = new AspectFoo().setValue("foo"); + AspectKey aspect1Key = new AspectKey<>(AspectFoo.class, urn, LATEST_VERSION); + AspectKey aspect2Key = new AspectKey<>(AspectBar.class, urn, LATEST_VERSION); + AspectKey aspect3Key = new AspectKey<>(AspectFooBar.class, urn, LATEST_VERSION); + AspectKey aspect4Key = new AspectKey<>(AspectAttributes.class, urn, LATEST_VERSION); + AspectKey aspect5Key = new AspectKey<>(AspectFooEvolved.class, urn, LATEST_VERSION); + when(_mockLocalDAO.exists(urn)).thenReturn(true); + when(_mockLocalDAO.get( + new HashSet<>(Arrays.asList(aspect1Key, aspect2Key, aspect3Key, aspect4Key, aspect5Key)))).thenReturn( + Collections.singletonMap(aspect1Key, Optional.of(foo))); + + EntityValue value = runAndWait(_internalResource.get(makeResourceKey(urn), null, true)); + + assertEquals(value.getFoo(), foo); + assertFalse(value.hasBar()); + } + @Test public void testGetUrnNotFound() { FooUrn urn = makeFooUrn(1234); @@ -266,6 +450,43 @@ public void testBatchGet() { assertFalse(keyValueMap.get(makeKey(2)).hasFoo()); } + @Test + public void testInternalModelBatchGet() { + FooUrn urn1 = makeFooUrn(1); + FooUrn urn2 = makeFooUrn(2); + AspectFoo foo = new AspectFoo().setValue("foo"); + AspectBar bar = new AspectBar().setValue("bar"); + + AspectKey aspectFooKey1 = new AspectKey<>(AspectFoo.class, urn1, LATEST_VERSION); + AspectKey aspectFooEvolvedKey1 = + new AspectKey<>(AspectFooEvolved.class, urn1, LATEST_VERSION); + AspectKey aspectBarKey1 = new AspectKey<>(AspectBar.class, urn1, LATEST_VERSION); + AspectKey aspectFooBarKey1 = new AspectKey<>(AspectFooBar.class, urn1, LATEST_VERSION); + AspectKey aspectAttKey1 = new AspectKey<>(AspectAttributes.class, urn1, LATEST_VERSION); + AspectKey aspectFooKey2 = new AspectKey<>(AspectFoo.class, urn2, LATEST_VERSION); + AspectKey aspectFooEvolvedKey2 = + new AspectKey<>(AspectFooEvolved.class, urn2, LATEST_VERSION); + AspectKey aspectBarKey2 = new AspectKey<>(AspectBar.class, urn2, LATEST_VERSION); + AspectKey aspectFooBarKey2 = new AspectKey<>(AspectFooBar.class, urn2, LATEST_VERSION); + AspectKey aspectAttKey2 = new AspectKey<>(AspectAttributes.class, urn2, LATEST_VERSION); + + when(_mockLocalDAO.get( + ImmutableSet.of(aspectFooBarKey1, aspectFooBarKey2, aspectFooKey1, aspectBarKey1, aspectFooKey2, aspectBarKey2, + aspectAttKey1, aspectAttKey2, aspectFooEvolvedKey1, aspectFooEvolvedKey2))).thenReturn( + ImmutableMap.of(aspectFooKey1, Optional.of(foo), aspectFooKey2, Optional.of(bar))); + + Map keyValueMap = runAndWait( + _internalResource.batchGet(ImmutableSet.of(makeResourceKey(urn1), makeResourceKey(urn2)), null)).entrySet() + .stream() + .collect(Collectors.toMap(e -> e.getKey().getKey(), e -> e.getValue())); + + assertEquals(keyValueMap.size(), 2); + assertEquals(keyValueMap.get(makeKey(1)).getFoo(), foo); + assertFalse(keyValueMap.get(makeKey(1)).hasBar()); + assertEquals(keyValueMap.get(makeKey(2)).getBar(), bar); + assertFalse(keyValueMap.get(makeKey(2)).hasFoo()); + } + @Test public void testBatchGetSpecificAspect() { FooUrn urn1 = makeFooUrn(1); @@ -482,6 +703,45 @@ public void testIngestWithTracking() { verifyNoMoreInteractions(_mockLocalDAO); } + @Test + public void testInternalModelIngest() { + FooUrn urn = makeFooUrn(1); + AspectFoo foo = new AspectFoo().setValue("foo"); + AspectBar bar = new AspectBar().setValue("bar"); + List aspects = Arrays.asList(ModelUtils.newAspectUnion(EntityAspectUnion.class, foo), + ModelUtils.newAspectUnion(EntityAspectUnion.class, bar)); + EntitySnapshot snapshot = ModelUtils.newSnapshot(EntitySnapshot.class, urn, aspects); + + runAndWait(_internalResource.ingest(snapshot)); + + verify(_mockLocalDAO, times(1)).add(eq(urn), eq(foo), any(), eq(null), eq(null)); + verify(_mockLocalDAO, times(1)).add(eq(urn), eq(bar), any(), eq(null), eq(null)); + verifyNoMoreInteractions(_mockLocalDAO); + } + + @Test + public void testInternalModelIngestWithTracking() { + FooUrn urn = makeFooUrn(1); + AspectFoo foo = new AspectFoo().setValue("foo"); + AspectBar bar = new AspectBar().setValue("bar"); + List aspects = Arrays.asList(ModelUtils.newAspectUnion(EntityAspectUnion.class, foo), + ModelUtils.newAspectUnion(EntityAspectUnion.class, bar)); + EntitySnapshot snapshot = ModelUtils.newSnapshot(EntitySnapshot.class, urn, aspects); + IngestionTrackingContext trackingContext = new IngestionTrackingContext(); + + runAndWait(_internalResource.ingestWithTracking(snapshot, trackingContext, null)); + + verify(_mockLocalDAO, times(1)).add(eq(urn), eq(foo), any(), eq(trackingContext), eq(null)); + verify(_mockLocalDAO, times(1)).add(eq(urn), eq(bar), any(), eq(trackingContext), eq(null)); + + IngestionParams ingestionParams = new IngestionParams().setIngestionMode(IngestionMode.LIVE); + runAndWait(_internalResource.ingestWithTracking(snapshot, trackingContext, ingestionParams)); + + verify(_mockLocalDAO, times(1)).add(eq(urn), eq(foo), any(), eq(trackingContext), eq(ingestionParams)); + verify(_mockLocalDAO, times(1)).add(eq(urn), eq(bar), any(), eq(trackingContext), eq(ingestionParams)); + verifyNoMoreInteractions(_mockLocalDAO); + } + @Test public void testSkipIngestAspect() { FooUrn urn = makeFooUrn(1); @@ -539,6 +799,36 @@ public void testGetSnapshotWithAllAspects() { assertEquals(aspects, ImmutableSet.of(foo, bar, fooBar, attributes)); } + @Test + public void testInternalModelGetSnapshotWithAllAspects() { + FooUrn urn = makeFooUrn(1); + AspectFoo foo = new AspectFoo().setValue("foo"); + AspectFooEvolved fooEvolved = new AspectFooEvolved().setValue("fooEvolved"); + AspectBar bar = new AspectBar().setValue("bar"); + AspectFooBar fooBar = new AspectFooBar().setBars(new BarUrnArray(new BarUrn(1))); + AspectAttributes attributes = new AspectAttributes().setAttributes(new StringArray("a")); + + AspectKey fooKey = new AspectKey<>(AspectFoo.class, urn, LATEST_VERSION); + AspectKey fooEvolvedKey = new AspectKey<>(AspectFooEvolved.class, urn, LATEST_VERSION); + AspectKey barKey = new AspectKey<>(AspectBar.class, urn, LATEST_VERSION); + AspectKey fooBarKey = new AspectKey<>(AspectFooBar.class, urn, LATEST_VERSION); + AspectKey attKey = new AspectKey<>(AspectAttributes.class, urn, LATEST_VERSION); + + Set> aspectKeys = + ImmutableSet.of(fooKey, fooEvolvedKey, barKey, fooBarKey, attKey); + when(_mockLocalDAO.get(aspectKeys)).thenReturn( + ImmutableMap.of(fooKey, Optional.of(foo), fooEvolvedKey, Optional.of(fooEvolved), barKey, Optional.of(bar), + fooBarKey, Optional.of(fooBar), attKey, Optional.of(attributes))); + + EntitySnapshot snapshot = runAndWait(_internalResource.getSnapshot(urn.toString(), null, true)); + + assertEquals(snapshot.getUrn(), urn); + + Set aspects = + snapshot.getAspects().stream().map(RecordUtils::getSelectedRecordTemplateFromUnion).collect(Collectors.toSet()); + assertEquals(aspects, ImmutableSet.of(foo, bar, fooBar, attributes)); + } + @Test public void testGetSnapshotWithInvalidUrn() { try { @@ -607,7 +897,7 @@ public void testBatchBackfill() { AspectBar bar1 = new AspectBar().setValue("bar1"); AspectBar bar2 = new AspectBar().setValue("bar2"); String[] aspects = new String[]{"com.linkedin.testing.AspectFoo", "com.linkedin.testing.AspectBar"}; - when(_mockLocalDAO.backfill(_resource.parseAspectsParam(aspects), ImmutableSet.of(urn1, urn2))).thenReturn( + when(_mockLocalDAO.backfill(_resource.parseAspectsParam(aspects, false), ImmutableSet.of(urn1, urn2))).thenReturn( ImmutableMap.of(urn1, ImmutableMap.of(AspectFoo.class, Optional.of(foo1), AspectBar.class, Optional.of(bar1)), urn2, ImmutableMap.of(AspectBar.class, Optional.of(bar2)))); @@ -638,7 +928,7 @@ public void testBackfillUsingSCSI() { AspectBar bar2 = new AspectBar().setValue("bar2"); String[] aspects = new String[]{"com.linkedin.testing.AspectFoo", "com.linkedin.testing.AspectBar"}; when( - _mockLocalDAO.backfill(BackfillMode.BACKFILL_ALL, _resource.parseAspectsParam(aspects), FooUrn.class, null, 10)) + _mockLocalDAO.backfill(BackfillMode.BACKFILL_ALL, _resource.parseAspectsParam(aspects, false), FooUrn.class, null, 10)) .thenReturn(ImmutableMap.of(urn1, ImmutableMap.of(AspectFoo.class, Optional.of(foo1), AspectBar.class, Optional.of(bar1)), urn2, ImmutableMap.of(AspectBar.class, Optional.of(bar2)))); @@ -668,7 +958,7 @@ public void testBackfillWithNewValue() { AspectBar bar1 = new AspectBar().setValue("bar1"); AspectBar bar2 = new AspectBar().setValue("bar2"); String[] aspects = new String[]{"com.linkedin.testing.AspectFoo", "com.linkedin.testing.AspectBar"}; - when(_mockLocalDAO.backfillWithNewValue(_resource.parseAspectsParam(aspects), ImmutableSet.of(urn1, urn2))) + when(_mockLocalDAO.backfillWithNewValue(_resource.parseAspectsParam(aspects, false), ImmutableSet.of(urn1, urn2))) .thenReturn( ImmutableMap.of(urn1, ImmutableMap.of(AspectFoo.class, Optional.of(foo1), AspectBar.class, Optional.of(bar1)), urn2, ImmutableMap.of(AspectBar.class, Optional.of(bar2))) @@ -700,7 +990,7 @@ public void testEmitNoChangeMetadataAuditEvent() { AspectBar bar1 = new AspectBar().setValue("bar1"); AspectBar bar2 = new AspectBar().setValue("bar2"); String[] aspects = new String[]{"com.linkedin.testing.AspectFoo", "com.linkedin.testing.AspectBar"}; - when(_mockLocalDAO.backfill(BackfillMode.BACKFILL_INCLUDING_LIVE_INDEX, _resource.parseAspectsParam(aspects), ImmutableSet.of(urn1, urn2))) + when(_mockLocalDAO.backfill(BackfillMode.BACKFILL_INCLUDING_LIVE_INDEX, _resource.parseAspectsParam(aspects, false), ImmutableSet.of(urn1, urn2))) .thenReturn( ImmutableMap.of(urn1, ImmutableMap.of(AspectFoo.class, Optional.of(foo1), AspectBar.class, Optional.of(bar1)), urn2, ImmutableMap.of(AspectBar.class, Optional.of(bar2))) @@ -733,7 +1023,7 @@ public void testEmitNoChangeMetadataAuditEventBootstrap() { AspectBar bar1 = new AspectBar().setValue("bar1"); AspectBar bar2 = new AspectBar().setValue("bar2"); String[] aspects = new String[]{"com.linkedin.testing.AspectFoo", "com.linkedin.testing.AspectBar"}; - when(_mockLocalDAO.backfill(BackfillMode.BACKFILL_ALL, _resource.parseAspectsParam(aspects), ImmutableSet.of(urn1, urn2))) + when(_mockLocalDAO.backfill(BackfillMode.BACKFILL_ALL, _resource.parseAspectsParam(aspects, false), ImmutableSet.of(urn1, urn2))) .thenReturn( ImmutableMap.of(urn1, ImmutableMap.of(AspectFoo.class, Optional.of(foo1), AspectBar.class, Optional.of(bar1)), urn2, ImmutableMap.of(AspectBar.class, Optional.of(bar2))) @@ -973,20 +1263,105 @@ public void testFilterFromIndexWithAspects() { assertEquals(actual4.getPageSize(), urnsListResult.getPageSize()); } + @Test + public void testInternalModelFilterFromIndexWithAspects() { + FooUrn urn1 = makeFooUrn(1); + FooUrn urn2 = makeFooUrn(2); + AspectFoo foo1 = new AspectFoo().setValue("val1"); + AspectFoo foo2 = new AspectFoo().setValue("val2"); + AspectBar bar1 = new AspectBar().setValue("val1"); + AspectBar bar2 = new AspectBar().setValue("val2"); + + UrnAspectEntry entry1 = new UrnAspectEntry<>(urn1, Arrays.asList(foo1, bar1)); + UrnAspectEntry entry2 = new UrnAspectEntry<>(urn2, Arrays.asList(foo2, bar2)); + + IndexCriterion criterion = new IndexCriterion().setAspect(AspectFoo.class.getCanonicalName()); + IndexCriterionArray criterionArray = new IndexCriterionArray(criterion); + IndexFilter indexFilter = new IndexFilter().setCriteria(criterionArray); + IndexSortCriterion indexSortCriterion = new IndexSortCriterion().setAspect(AspectFoo.class.getCanonicalName()) + .setOrder(SortOrder.DESCENDING); + String[] aspectNames = {ModelUtils.getAspectName(AspectFoo.class), ModelUtils.getAspectName(AspectBar.class)}; + + // case 1: aspect list is provided, null last urn + List> listResult1 = Arrays.asList(entry1, entry2); + + when(_mockLocalDAO.getAspects(ImmutableSet.of(AspectFoo.class, AspectBar.class), indexFilter, null, null, 2)) + .thenReturn(listResult1); + + List actual1 = + runAndWait(_resource.filter(indexFilter, aspectNames, null, new PagingContext(0, 2))); + + assertEquals(actual1.size(), 2); + assertEquals(actual1.get(0), new EntityValue().setFoo(foo1).setBar(bar1)); + assertEquals(actual1.get(1), new EntityValue().setFoo(foo2).setBar(bar2)); + + // case 2: null aspects is provided i.e. all aspects in the aspect union will be returned, non-null last urn + List> listResult2 = Collections.singletonList(entry2); + + when(_mockLocalDAO.getAspects( + ImmutableSet.of(AspectFoo.class, AspectBar.class, AspectFooEvolved.class, AspectFooBar.class, + AspectAttributes.class), indexFilter, null, urn1, 2)).thenReturn(listResult2); + + List actual2 = runAndWait( + _internalResource.filter(indexFilter, null, null, urn1.toString(), new PagingContext(0, 2).getCount())); + assertEquals(actual2.size(), 1); + assertEquals(actual2.get(0), new EntityValue().setFoo(foo2).setBar(bar2)); + + // case 3: non-null sort criterion is provided + List> listResult3 = Arrays.asList(entry2, entry1); + + when(_mockLocalDAO.getAspects(ImmutableSet.of(AspectFoo.class, AspectBar.class), indexFilter, indexSortCriterion, null, 2)) + .thenReturn(listResult3); + + List actual3 = + runAndWait(_resource.filter(indexFilter, indexSortCriterion, aspectNames, null, 2)); + + assertEquals(actual3.size(), 2); + assertEquals(actual3.get(0), new EntityValue().setFoo(foo2).setBar(bar2)); + assertEquals(actual3.get(1), new EntityValue().setFoo(foo1).setBar(bar1)); + + // case 4: offset pagination + ListResult> urnsListResult = ListResult.>builder() + .values(Arrays.asList(entry2, entry1)) + .metadata(null) + .nextStart(ListResult.INVALID_NEXT_START) + .havingMore(false) + .totalCount(2) + .totalPageCount(1) + .pageSize(2) + .build(); + + when(_mockLocalDAO.getAspects(ImmutableSet.of(AspectFoo.class, AspectBar.class), indexFilter, indexSortCriterion, 0, 2)) + .thenReturn(urnsListResult); + + ListResult actual4 = + runAndWait(_resource.filter(indexFilter, indexSortCriterion, aspectNames, new PagingContext(0, 2))); + + List actualValues = actual4.getValues(); + assertEquals(actualValues.size(), 2); + assertEquals(actualValues.get(0), new EntityValue().setFoo(foo2).setBar(bar2)); + assertEquals(actualValues.get(1), new EntityValue().setFoo(foo1).setBar(bar1)); + assertEquals(actual4.getNextStart(), urnsListResult.getNextStart()); + assertEquals(actual4.isHavingMore(), urnsListResult.isHavingMore()); + assertEquals(actual4.getTotalCount(), urnsListResult.getTotalCount()); + assertEquals(actual4.getTotalPageCount(), urnsListResult.getTotalPageCount()); + assertEquals(actual4.getPageSize(), urnsListResult.getPageSize()); + } + @Test public void testParseAspectsParam() { // Only 1 aspect Set> aspectClasses = - _resource.parseAspectsParam(new String[]{AspectFoo.class.getCanonicalName()}); + _resource.parseAspectsParam(new String[]{AspectFoo.class.getCanonicalName()}, false); assertEquals(aspectClasses.size(), 1); assertTrue(aspectClasses.contains(AspectFoo.class)); // No aspect - aspectClasses = _resource.parseAspectsParam(new String[]{}); + aspectClasses = _resource.parseAspectsParam(new String[]{}, false); assertEquals(aspectClasses.size(), 0); // All aspects - aspectClasses = _resource.parseAspectsParam(null); + aspectClasses = _resource.parseAspectsParam(null, false); assertEquals(aspectClasses.size(), 4); assertTrue(aspectClasses.contains(AspectFoo.class)); assertTrue(aspectClasses.contains(AspectBar.class)); @@ -994,6 +1369,28 @@ public void testParseAspectsParam() { assertTrue(aspectClasses.contains(AspectAttributes.class)); } + @Test + public void testInternalModelParseAspectsParam() { + // Only 1 aspect + Set> aspectClasses = + _resource.parseAspectsParam(new String[]{AspectFoo.class.getCanonicalName()}, true); + assertEquals(aspectClasses.size(), 1); + assertTrue(aspectClasses.contains(AspectFoo.class)); + + // No aspect + aspectClasses = _resource.parseAspectsParam(new String[]{}, true); + assertEquals(aspectClasses.size(), 0); + + // All aspects + aspectClasses = _resource.parseAspectsParam(null, true); + assertEquals(aspectClasses.size(), 5); + assertTrue(aspectClasses.contains(AspectFoo.class)); + assertTrue(aspectClasses.contains(AspectFooEvolved.class)); + assertTrue(aspectClasses.contains(AspectBar.class)); + assertTrue(aspectClasses.contains(AspectFooBar.class)); + assertTrue(aspectClasses.contains(AspectAttributes.class)); + } + @Test public void testCountAggregate() { FooUrn urn1 = makeFooUrn(1); @@ -1049,4 +1446,61 @@ public void testCountAggregateFilter() { assertEquals(actual.getMetadata().getLongMap(), new LongMap(mapResult)); } + + @Test + public void testIngestAsset() { + FooUrn urn = makeFooUrn(1); + EntityAsset asset = new EntityAsset(); + AspectFoo foo = new AspectFoo().setValue("foo"); + AspectBar bar = new AspectBar().setValue("bar"); + asset.setUrn(urn); + asset.setAspectFoo(foo); + asset.setAspectBar(bar); + IngestionTrackingContext trackingContext = new IngestionTrackingContext(); + + IngestionParams ingestionParams1 = new IngestionParams().setTestMode(true); + runAndWait(_internalResource.ingestAsset(asset, trackingContext, ingestionParams1)); + + verify(_mockLocalDAO, times(1)).add(eq(urn), eq(foo), any(), eq(trackingContext), eq(ingestionParams1)); + verify(_mockLocalDAO, times(1)).add(eq(urn), eq(bar), any(), eq(trackingContext), eq(ingestionParams1)); + + IngestionParams ingestionParams2 = new IngestionParams().setIngestionMode(IngestionMode.LIVE); + runAndWait(_internalResource.ingestAsset(asset, trackingContext, ingestionParams2)); + + verify(_mockLocalDAO, times(1)).add(eq(urn), eq(foo), any(), eq(trackingContext), eq(ingestionParams2)); + verify(_mockLocalDAO, times(1)).add(eq(urn), eq(bar), any(), eq(trackingContext), eq(ingestionParams2)); + verifyNoMoreInteractions(_mockLocalDAO); + } + + @Test + public void testGetAsset() { + FooUrn urn = makeFooUrn(1); + AspectFoo foo = new AspectFoo().setValue("foo"); + AspectBar bar = new AspectBar().setValue("bar"); + AspectFooEvolved fooEvolved = new AspectFooEvolved().setValue("fooEvolved"); + AspectFooBar fooBar = new AspectFooBar().setBars(new BarUrnArray(new BarUrn(1))); + AspectAttributes attributes = new AspectAttributes().setAttributes(new StringArray("a")); + + AspectKey fooKey = new AspectKey<>(AspectFoo.class, urn, LATEST_VERSION); + AspectKey fooEvolvedKey = new AspectKey<>(AspectFooEvolved.class, urn, LATEST_VERSION); + AspectKey barKey = new AspectKey<>(AspectBar.class, urn, LATEST_VERSION); + AspectKey fooBarKey = new AspectKey<>(AspectFooBar.class, urn, LATEST_VERSION); + AspectKey attKey = new AspectKey<>(AspectAttributes.class, urn, LATEST_VERSION); + + Set> aspectKeys = + ImmutableSet.of(fooKey, fooEvolvedKey, barKey, fooBarKey, attKey); + when(_mockLocalDAO.get(aspectKeys)).thenReturn( + ImmutableMap.of(fooKey, Optional.of(foo), fooEvolvedKey, Optional.of(fooEvolved), barKey, Optional.of(bar), + fooBarKey, Optional.of(fooBar), attKey, Optional.of(attributes))); + + EntityAsset asset = runAndWait(_resource.getAsset(urn.toString(), null)); + + assertEquals(asset.getUrn(), urn); + + assertEquals(asset.getAspectFoo(), foo); + assertEquals(asset.getAspectFooEvolved(), fooEvolved); + assertEquals(asset.getAspectBar(), bar); + assertEquals(asset.getAspectFooBar(), fooBar); + assertEquals(asset.getAspectAttributes(), attributes); + } } diff --git a/restli-resources/src/test/java/com/linkedin/metadata/restli/BaseEntitySimpleKeyResourceTest.java b/restli-resources/src/test/java/com/linkedin/metadata/restli/BaseEntitySimpleKeyResourceTest.java index f9f35dffe..4c25bdacf 100644 --- a/restli-resources/src/test/java/com/linkedin/metadata/restli/BaseEntitySimpleKeyResourceTest.java +++ b/restli-resources/src/test/java/com/linkedin/metadata/restli/BaseEntitySimpleKeyResourceTest.java @@ -9,6 +9,7 @@ import com.linkedin.metadata.dao.BaseLocalDAO; import com.linkedin.metadata.dao.utils.ModelUtils; import com.linkedin.metadata.dao.utils.RecordUtils; +import com.linkedin.metadata.restli.lix.ResourceLix; import com.linkedin.parseq.BaseEngineTest; import com.linkedin.restli.common.HttpStatus; import com.linkedin.restli.server.ResourceContext; @@ -16,11 +17,15 @@ import com.linkedin.testing.AspectAttributes; import com.linkedin.testing.AspectBar; import com.linkedin.testing.AspectFoo; +import com.linkedin.testing.AspectFooEvolved; import com.linkedin.testing.BarUrnArray; import com.linkedin.testing.EntityAspectUnion; import com.linkedin.testing.EntityAspectUnionArray; +import com.linkedin.testing.EntityAsset; import com.linkedin.testing.EntitySnapshot; import com.linkedin.testing.EntityValue; +import com.linkedin.testing.InternalEntityAspectUnion; +import com.linkedin.testing.InternalEntitySnapshot; import com.linkedin.testing.localrelationship.AspectFooBar; import com.linkedin.testing.urn.BarUrn; import java.net.URISyntaxException; @@ -33,6 +38,7 @@ import java.util.Set; import java.util.stream.Collectors; import javax.annotation.Nonnull; +import javax.annotation.Nullable; import org.testng.annotations.BeforeMethod; import org.testng.annotations.Test; @@ -44,8 +50,9 @@ public class BaseEntitySimpleKeyResourceTest extends BaseEngineTest { - private BaseLocalDAO _mockLocalDAO; + private BaseLocalDAO _mockLocalDAO; private TestResource _resource = new TestResource(); + private TestInternalResource _internalResource = new TestInternalResource(); @SuppressWarnings("unchecked") @BeforeMethod @@ -73,6 +80,28 @@ public void testGet() { assertFalse(value.hasBar()); } + @Test + public void testInternalModelGet() { + long id = 1234; + Urn urn = makeUrn(id); + AspectFoo foo = new AspectFoo().setValue("foo"); + AspectKey aspect1Key = new AspectKey<>(AspectFoo.class, urn, LATEST_VERSION); + AspectKey aspect2Key = new AspectKey<>(AspectBar.class, urn, LATEST_VERSION); + AspectKey aspect3Key = new AspectKey<>(AspectFooBar.class, urn, LATEST_VERSION); + AspectKey aspect4Key = new AspectKey<>(AspectAttributes.class, urn, LATEST_VERSION); + AspectKey aspect5Key = new AspectKey<>(AspectFooEvolved.class, urn, LATEST_VERSION); + + when(_mockLocalDAO.exists(urn)).thenReturn(true); + when(_mockLocalDAO.get( + new HashSet<>(Arrays.asList(aspect1Key, aspect2Key, aspect3Key, aspect4Key, aspect5Key)))).thenReturn( + Collections.singletonMap(aspect1Key, Optional.of(foo))); + + EntityValue value = runAndWait(_resource.get(id, null, true)); + + assertEquals(value.getFoo(), foo); + assertFalse(value.hasBar()); + } + @Test public void testGetUrnNotFound() { long id = 1234; @@ -177,6 +206,43 @@ public void testBatchGet() { assertFalse(keyValueMap.get(id2).hasFoo()); } + @Test + public void testInternalModelBatchGet() { + long id1 = 1; + Urn urn1 = makeUrn(id1); + long id2 = 2; + Urn urn2 = makeUrn(id2); + AspectFoo foo = new AspectFoo().setValue("foo"); + AspectBar bar = new AspectBar().setValue("bar"); + + AspectKey aspectFooKey1 = new AspectKey<>(AspectFoo.class, urn1, LATEST_VERSION); + AspectKey aspectFooEvolvedKey1 = new AspectKey<>(AspectFooEvolved.class, urn1, LATEST_VERSION); + AspectKey aspectBarKey1 = new AspectKey<>(AspectBar.class, urn1, LATEST_VERSION); + AspectKey aspectFooKey2 = new AspectKey<>(AspectFoo.class, urn2, LATEST_VERSION); + AspectKey aspectFooEvolvedKey2 = new AspectKey<>(AspectFooEvolved.class, urn2, LATEST_VERSION); + AspectKey aspectBarKey2 = new AspectKey<>(AspectBar.class, urn2, LATEST_VERSION); + AspectKey aspectFooBarKey1 = new AspectKey<>(AspectFooBar.class, urn1, LATEST_VERSION); + AspectKey aspectFooBarKey2 = new AspectKey<>(AspectFooBar.class, urn2, LATEST_VERSION); + AspectKey aspectAttKey1 = new AspectKey<>(AspectAttributes.class, urn1, LATEST_VERSION); + AspectKey aspectAttKey2 = new AspectKey<>(AspectAttributes.class, urn2, LATEST_VERSION); + + when(_mockLocalDAO.get( + ImmutableSet.of(aspectFooKey1, aspectFooEvolvedKey1, aspectBarKey1, aspectAttKey1, aspectFooKey2, + aspectFooEvolvedKey2, aspectBarKey2, aspectAttKey2, aspectFooBarKey1, aspectFooBarKey2))).thenReturn( + ImmutableMap.of(aspectFooKey1, Optional.of(foo), aspectFooKey2, Optional.of(bar))); + + Map keyValueMap = runAndWait(_internalResource.batchGet(ImmutableSet.of(id1, id2), null)) + .entrySet() + .stream() + .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue)); + + assertEquals(keyValueMap.size(), 2); + assertEquals(keyValueMap.get(id1).getFoo(), foo); + assertFalse(keyValueMap.get(id1).hasBar()); + assertEquals(keyValueMap.get(id2).getBar(), bar); + assertFalse(keyValueMap.get(id2).hasFoo()); + } + @Test public void testBatchGetSpecificAspect() { long id1 = 1; @@ -248,6 +314,34 @@ public void testGetSnapshotWithAllAspects() { assertEquals(aspects, ImmutableSet.of(foo, bar, fooBar, att)); } + @Test + public void testInternalModelGetSnapshotWithAllAspects() { + Urn urn = makeUrn(1); + AspectFoo foo = new AspectFoo().setValue("foo"); + AspectFooEvolved fooEvolved = new AspectFooEvolved().setValue("fooEvolved"); + AspectBar bar = new AspectBar().setValue("bar"); + AspectFooBar fooBar = new AspectFooBar().setBars(new BarUrnArray(new BarUrn(1))); + AspectAttributes att = new AspectAttributes().setAttributes(new StringArray("a")); + AspectKey fooKey = new AspectKey<>(AspectFoo.class, urn, LATEST_VERSION); + AspectKey fooEvolvedKey = new AspectKey<>(AspectFooEvolved.class, urn, LATEST_VERSION); + AspectKey barKey = new AspectKey<>(AspectBar.class, urn, LATEST_VERSION); + AspectKey fooBarKey = new AspectKey<>(AspectFooBar.class, urn, LATEST_VERSION); + AspectKey attKey = new AspectKey<>(AspectAttributes.class, urn, LATEST_VERSION); + Set> aspectKeys = + ImmutableSet.of(fooKey, fooEvolvedKey, barKey, fooBarKey, attKey); + when(_mockLocalDAO.get(aspectKeys)).thenReturn( + ImmutableMap.of(fooKey, Optional.of(foo), fooEvolvedKey, Optional.of(fooEvolved), barKey, Optional.of(bar), + fooBarKey, Optional.of(fooBar), attKey, Optional.of(att))); + + EntitySnapshot snapshot = runAndWait(_resource.getSnapshot(urn.toString(), null, true)); + + assertEquals(snapshot.getUrn(), urn); + + Set aspects = + snapshot.getAspects().stream().map(RecordUtils::getSelectedRecordTemplateFromUnion).collect(Collectors.toSet()); + assertEquals(aspects, ImmutableSet.of(foo, bar, fooBar, att)); + } + @Test public void testGetSnapshotWithInvalidUrn() { try { @@ -305,16 +399,173 @@ public void testBackfillWithInvalidUrn() { /** * Test class for {@link BaseEntityResource}. * */ - private class TestResource extends BaseEntityResource< - Long, EntityValue, Urn, EntitySnapshot, EntityAspectUnion> { + private class TestResource extends + BaseEntityResource { TestResource() { - super(EntitySnapshot.class, EntityAspectUnion.class); + super(EntitySnapshot.class, EntityAspectUnion.class, InternalEntitySnapshot.class, + InternalEntityAspectUnion.class, EntityAsset.class); } + + + @Override + @Nonnull + protected BaseLocalDAO getLocalDAO() { + return _mockLocalDAO; + } + + @Override + @Nonnull + protected Urn createUrnFromString(@Nonnull String urnString) { + try { + return Urn.createFromString(urnString); + } catch (URISyntaxException e) { + throw RestliUtils.badRequestException("Invalid URN: " + urnString); + } + } + + @Override + @Nonnull + protected Urn toUrn(@Nonnull Long key) { + return makeUrn(key); + } + + @Nonnull + @Override + protected Long toKey(@Nonnull Urn urn) { + return urn.getIdAsLong(); + } + + @Override + @Nonnull + protected EntityValue toValue(@Nonnull EntitySnapshot snapshot) { + EntityValue value = new EntityValue(); + ModelUtils.getAspectsFromSnapshot(snapshot).forEach(a -> { + if (a instanceof AspectFoo) { + value.setFoo((AspectFoo) a); + } else if (a instanceof AspectBar) { + value.setBar((AspectBar) a); + } + }); + return value; + } + + @Override + @Nonnull + protected EntitySnapshot toSnapshot(@Nonnull EntityValue value, @Nonnull Urn urn) { + EntitySnapshot snapshot = new EntitySnapshot().setUrn(urn); + EntityAspectUnionArray aspects = new EntityAspectUnionArray(); + if (value.hasFoo()) { + aspects.add(ModelUtils.newAspectUnion(EntityAspectUnion.class, value.getFoo())); + } + if (value.hasBar()) { + aspects.add(ModelUtils.newAspectUnion(EntityAspectUnion.class, value.getBar())); + } + + snapshot.setAspects(aspects); + return snapshot; + } + + @Override + public ResourceContext getContext() { + return mock(ResourceContext.class); + } + } + + /** + * Test class for {@link BaseEntityResource}. + * */ + private class TestInternalResource extends + BaseEntityResource { + + TestInternalResource() { + super(EntitySnapshot.class, EntityAspectUnion.class, InternalEntitySnapshot.class, + InternalEntityAspectUnion.class, EntityAsset.class, new ResourceLix() { + @Override + public boolean testGet(@Nonnull String urn, @Nonnull String entityType) { + return false; + } + + @Override + public boolean testBatchGet(@Nullable String urn, @Nullable String entityType) { + return true; + } + + @Override + public boolean testBatchGetWithErrors(@Nullable String urn, @Nullable String type) { + return false; + } + + @Override + public boolean testGetSnapshot(@Nullable String urn, @Nullable String entityType) { + return false; + } + + @Override + public boolean testBackfillLegacy(@Nullable String urn, @Nullable String entityType) { + return false; + } + + @Override + public boolean testBackfillWithUrns(@Nullable String urn, @Nullable String entityType) { + return false; + } + + @Override + public boolean testEmitNoChangeMetadataAuditEvent(@Nullable String urn, @Nullable String entityType) { + return false; + } + + @Override + public boolean testBackfillWithNewValue(@Nullable String urn, @Nullable String entityType) { + return false; + } + + @Override + public boolean testBackfillEntityTables(@Nullable String urn, @Nullable String entityType) { + return false; + } + + @Override + public boolean testBackfillRelationshipTables(@Nullable String urn, @Nullable String entityType) { + return false; + } + + @Override + public boolean testBackfill(@Nonnull String assetType, @Nonnull String mode) { + return false; + } + + @Override + public boolean testFilter(@Nonnull String assetType) { + return false; + } + + @Override + public boolean testGetAll(@Nullable String urnType) { + return false; + } + + @Override + public boolean testSearch(@Nullable String urnType) { + return false; + } + + @Override + public boolean testSearchV2(@Nullable String urnType) { + return false; + } + }); + } + + + @Override @Nonnull - protected BaseLocalDAO getLocalDAO() { + protected BaseLocalDAO getLocalDAO() { return _mockLocalDAO; } diff --git a/restli-resources/src/test/java/com/linkedin/metadata/restli/BaseSearchableEntityResourceTest.java b/restli-resources/src/test/java/com/linkedin/metadata/restli/BaseSearchableEntityResourceTest.java index 3ce02aba6..aaa786656 100644 --- a/restli-resources/src/test/java/com/linkedin/metadata/restli/BaseSearchableEntityResourceTest.java +++ b/restli-resources/src/test/java/com/linkedin/metadata/restli/BaseSearchableEntityResourceTest.java @@ -28,10 +28,13 @@ import com.linkedin.testing.AspectFoo; import com.linkedin.testing.EntityAspectUnion; import com.linkedin.testing.EntityAspectUnionArray; +import com.linkedin.testing.EntityAsset; import com.linkedin.testing.EntityDocument; import com.linkedin.testing.EntityKey; import com.linkedin.testing.EntitySnapshot; import com.linkedin.testing.EntityValue; +import com.linkedin.testing.InternalEntityAspectUnion; +import com.linkedin.testing.InternalEntitySnapshot; import java.net.URISyntaxException; import java.util.Arrays; import java.util.List; @@ -48,21 +51,23 @@ public class BaseSearchableEntityResourceTest extends BaseEngineTest { - private BaseLocalDAO _mockLocalDAO; + private BaseLocalDAO _mockLocalDAO; private BaseSearchDAO _mockSearchDAO; private TestResource _resource = new TestResource(); class TestResource extends BaseSearchableEntityResource< // format - ComplexResourceKey, EntityValue, Urn, EntitySnapshot, EntityAspectUnion, EntityDocument> { + ComplexResourceKey, EntityValue, Urn, EntitySnapshot, EntityAspectUnion, EntityDocument, + InternalEntitySnapshot, InternalEntityAspectUnion, EntityAsset> { public TestResource() { - super(EntitySnapshot.class, EntityAspectUnion.class); + super(EntitySnapshot.class, EntityAspectUnion.class, InternalEntitySnapshot.class, + InternalEntityAspectUnion.class, EntityAsset.class); } @Nonnull @Override - protected BaseLocalDAO getLocalDAO() { + protected BaseLocalDAO getLocalDAO() { return _mockLocalDAO; } diff --git a/restli-resources/src/test/java/com/linkedin/metadata/restli/BaseSearchableEntitySimpleKeyResourceTest.java b/restli-resources/src/test/java/com/linkedin/metadata/restli/BaseSearchableEntitySimpleKeyResourceTest.java index dff297301..4157e2b85 100644 --- a/restli-resources/src/test/java/com/linkedin/metadata/restli/BaseSearchableEntitySimpleKeyResourceTest.java +++ b/restli-resources/src/test/java/com/linkedin/metadata/restli/BaseSearchableEntitySimpleKeyResourceTest.java @@ -28,9 +28,12 @@ import com.linkedin.testing.AspectFoo; import com.linkedin.testing.EntityAspectUnion; import com.linkedin.testing.EntityAspectUnionArray; +import com.linkedin.testing.EntityAsset; import com.linkedin.testing.EntityDocument; import com.linkedin.testing.EntitySnapshot; import com.linkedin.testing.EntityValue; +import com.linkedin.testing.InternalEntityAspectUnion; +import com.linkedin.testing.InternalEntitySnapshot; import java.net.URISyntaxException; import java.util.Arrays; import java.util.List; @@ -46,7 +49,7 @@ public class BaseSearchableEntitySimpleKeyResourceTest extends BaseEngineTest { - private BaseLocalDAO _mockLocalDAO; + private BaseLocalDAO _mockLocalDAO; private BaseSearchDAO _mockSearchDAO; private TestResource _resource = new TestResource(); @@ -178,16 +181,18 @@ private AutoCompleteResult makeAutoCompleteResult(String query, List sug /** * Test resource class for {@link BaseSearchableEntityResource}. * */ - private class TestResource extends BaseSearchableEntityResource< - Long, EntityValue, Urn, EntitySnapshot, EntityAspectUnion, EntityDocument> { + private class TestResource extends + BaseSearchableEntityResource { TestResource() { - super(EntitySnapshot.class, EntityAspectUnion.class); + super(EntitySnapshot.class, EntityAspectUnion.class, InternalEntitySnapshot.class, + InternalEntityAspectUnion.class, EntityAsset.class); } @Override @Nonnull - protected BaseLocalDAO getLocalDAO() { + protected BaseLocalDAO getLocalDAO() { return _mockLocalDAO; } diff --git a/restli-resources/src/test/java/com/linkedin/metadata/restli/BaseSingleAspectEntitySimpleKeyResourceTest.java b/restli-resources/src/test/java/com/linkedin/metadata/restli/BaseSingleAspectEntitySimpleKeyResourceTest.java index ff7270554..1247a2c80 100644 --- a/restli-resources/src/test/java/com/linkedin/metadata/restli/BaseSingleAspectEntitySimpleKeyResourceTest.java +++ b/restli-resources/src/test/java/com/linkedin/metadata/restli/BaseSingleAspectEntitySimpleKeyResourceTest.java @@ -18,6 +18,9 @@ import com.linkedin.restli.server.ResourceContext; import com.linkedin.restli.server.RestLiServiceException; import com.linkedin.testing.AspectBar; +import com.linkedin.testing.EntityAsset; +import com.linkedin.testing.singleaspectentity.InternalEntityAspectUnion; +import com.linkedin.testing.singleaspectentity.InternalEntitySnapshot; import com.linkedin.testing.singleaspectentity.EntityAspectUnion; import com.linkedin.testing.singleaspectentity.EntitySnapshot; import com.linkedin.testing.singleaspectentity.EntityValue; @@ -44,7 +47,7 @@ public class BaseSingleAspectEntitySimpleKeyResourceTest extends BaseEngineTest { - private BaseLocalDAO _mockLocalDao; + private BaseLocalDAO _mockLocalDao; private TestResource _resource = new TestResource(); @SuppressWarnings("unchecked") @@ -218,15 +221,17 @@ private ExtraInfo makeExtraInfo(Urn urn, Long version, AuditStamp audit) { * Test implementation of BaseSingleAspectEntitySimpleKeyResource. * */ private class TestResource extends - BaseSingleAspectEntityResource { + BaseSingleAspectEntityResource { TestResource() { - super(AspectBar.class, EntityAspectUnion.class, EntityValue.class, EntitySnapshot.class); + super(AspectBar.class, EntityAspectUnion.class, EntityValue.class, EntitySnapshot.class, + InternalEntitySnapshot.class, InternalEntityAspectUnion.class, EntityAsset.class); } @Override @Nonnull - protected BaseLocalDAO getLocalDAO() { + protected BaseLocalDAO getLocalDAO() { return _mockLocalDao; } diff --git a/restli-resources/src/test/java/com/linkedin/metadata/restli/BaseSingleAspectEntitySimpleKeyVersionedSubResourceTest.java b/restli-resources/src/test/java/com/linkedin/metadata/restli/BaseSingleAspectEntitySimpleKeyVersionedSubResourceTest.java index 1a95a3ad7..72237fd6f 100644 --- a/restli-resources/src/test/java/com/linkedin/metadata/restli/BaseSingleAspectEntitySimpleKeyVersionedSubResourceTest.java +++ b/restli-resources/src/test/java/com/linkedin/metadata/restli/BaseSingleAspectEntitySimpleKeyVersionedSubResourceTest.java @@ -12,6 +12,9 @@ import com.linkedin.restli.server.RestLiServiceException; import com.linkedin.restli.server.annotations.RestLiCollection; import com.linkedin.testing.AspectBar; +import com.linkedin.testing.EntityAsset; +import com.linkedin.testing.InternalEntityAspectUnion; +import com.linkedin.testing.InternalEntitySnapshot; import com.linkedin.testing.singleaspectentity.EntityAspectUnion; import com.linkedin.testing.singleaspectentity.EntitySnapshot; import com.linkedin.testing.singleaspectentity.EntityValue; @@ -30,7 +33,7 @@ public class BaseSingleAspectEntitySimpleKeyVersionedSubResourceTest extends BaseEngineTest { - private BaseLocalDAO _mockLocalDao; + private BaseLocalDAO _mockLocalDao; private TestVersionedSubResource _resource = new TestVersionedSubResource(); @SuppressWarnings("unchecked") @@ -108,7 +111,7 @@ public void testGetAllWithMetadata() throws URISyntaxException { * */ @RestLiCollection(name = "versions", namespace = "com.linkedin.testing", parent = TestResource.class, keyName = "versionId") private class TestVersionedSubResource extends BaseSingleAspectEntityVersionedSubResource< - EntityValue, SingleAspectEntityUrn, AspectBar, EntityAspectUnion> { + EntityValue, SingleAspectEntityUrn, AspectBar, InternalEntityAspectUnion> { TestVersionedSubResource() { super(AspectBar.class, EntityValue.class); @@ -116,7 +119,7 @@ private class TestVersionedSubResource extends BaseSingleAspectEntityVersionedSu @Override @Nonnull - protected BaseLocalDAO getLocalDAO() { + protected BaseLocalDAO getLocalDAO() { return _mockLocalDao; } @@ -151,16 +154,18 @@ public ResourceContext getContext() { * Test resource class for BaseSingleAspectEntitySimpleKeyResource. * */ @RestLiCollection(name = "SingleAspectEntity", namespace = "com.linkedin.testing", keyName = "testId") - private class TestResource extends BaseSingleAspectEntityResource< - Long, EntityValue, SingleAspectEntityUrn, AspectBar, EntityAspectUnion, EntitySnapshot> { + private class TestResource extends + BaseSingleAspectEntityResource { TestResource() { - super(AspectBar.class, EntityAspectUnion.class, EntityValue.class, EntitySnapshot.class); + super(AspectBar.class, EntityAspectUnion.class, EntityValue.class, EntitySnapshot.class, + InternalEntitySnapshot.class, InternalEntityAspectUnion.class, EntityAsset.class); } @Override @Nonnull - protected BaseLocalDAO getLocalDAO() { + protected BaseLocalDAO getLocalDAO() { return _mockLocalDao; } diff --git a/restli-resources/src/test/java/com/linkedin/metadata/restli/BaseSingleAspectSearchableEntitySimpleKeyResourceTest.java b/restli-resources/src/test/java/com/linkedin/metadata/restli/BaseSingleAspectSearchableEntitySimpleKeyResourceTest.java index 2608df80e..b73b52ba1 100644 --- a/restli-resources/src/test/java/com/linkedin/metadata/restli/BaseSingleAspectSearchableEntitySimpleKeyResourceTest.java +++ b/restli-resources/src/test/java/com/linkedin/metadata/restli/BaseSingleAspectSearchableEntitySimpleKeyResourceTest.java @@ -17,6 +17,9 @@ import com.linkedin.restli.server.CollectionResult; import com.linkedin.restli.server.PagingContext; import com.linkedin.testing.AspectBar; +import com.linkedin.testing.EntityAsset; +import com.linkedin.testing.InternalEntityAspectUnion; +import com.linkedin.testing.InternalEntitySnapshot; import com.linkedin.testing.singleaspectentity.EntityAspectUnion; import com.linkedin.testing.singleaspectentity.EntityDocument; import com.linkedin.testing.singleaspectentity.EntitySnapshot; @@ -37,7 +40,7 @@ public class BaseSingleAspectSearchableEntitySimpleKeyResourceTest extends BaseEngineTest { - private BaseLocalDAO _mockLocalDao; + private BaseLocalDAO _mockLocalDao; private BaseSearchDAO _mockSearchDao; private TestResource _resource = new TestResource(); @@ -143,11 +146,13 @@ public void testGetAll() throws URISyntaxException { /** * Test class for BaseSingleAspectSearchableEntitySimpleKeyResource. * */ - private class TestResource extends BaseSingleAspectSearchableEntityResource { + private class TestResource extends + BaseSingleAspectSearchableEntityResource { TestResource() { - super(AspectBar.class, EntityAspectUnion.class, EntityValue.class, EntitySnapshot.class); + super(AspectBar.class, EntityAspectUnion.class, EntityValue.class, EntitySnapshot.class, + InternalEntitySnapshot.class, InternalEntityAspectUnion.class, EntityAsset.class); } @Override @@ -158,7 +163,7 @@ protected BaseSearchDAO getSearchDAO() { @Override @Nonnull - protected BaseLocalDAO getLocalDAO() { + protected BaseLocalDAO getLocalDAO() { return _mockLocalDao; } diff --git a/testing/test-models/src/main/pegasus/com/linkedin/testing/singleaspectentity/InternalEntitySnapshot.pdl b/testing/test-models/src/main/pegasus/com/linkedin/testing/singleaspectentity/InternalEntitySnapshot.pdl index 1c11bb98d..4e67e35b7 100644 --- a/testing/test-models/src/main/pegasus/com/linkedin/testing/singleaspectentity/InternalEntitySnapshot.pdl +++ b/testing/test-models/src/main/pegasus/com/linkedin/testing/singleaspectentity/InternalEntitySnapshot.pdl @@ -16,4 +16,4 @@ record InternalEntitySnapshot { * For unit tests */ aspects: array[InternalEntityAspectUnion] -} +} \ No newline at end of file diff --git a/version.properties b/version.properties index 6f785b137..d951ca33a 100644 --- a/version.properties +++ b/version.properties @@ -1 +1 @@ -version=0.5.* +version=0.6.*