Skip to content

Commit 402f1e6

Browse files
jphuijphui
authored
(feat) Add deleteWithReturn (#497)
1 parent 2de35c0 commit 402f1e6

File tree

2 files changed

+193
-15
lines changed

2 files changed

+193
-15
lines changed

‎dao-api/src/main/java/com/linkedin/metadata/dao/BaseLocalDAO.java‎

Lines changed: 48 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,7 @@
4949
import com.linkedin.metadata.query.IndexSortCriterion;
5050
import java.sql.Timestamp;
5151
import java.time.Clock;
52+
import java.util.Collection;
5253
import java.util.Collections;
5354
import java.util.ArrayList;
5455
import java.util.HashMap;
@@ -57,6 +58,7 @@
5758
import java.util.LinkedList;
5859
import java.util.List;
5960
import java.util.Map;
61+
import java.util.Objects;
6062
import java.util.Optional;
6163
import java.util.Set;
6264
import java.util.UUID;
@@ -659,7 +661,10 @@ private <ASPECT extends RecordTemplate> AddResult<ASPECT> aspectUpdateHelper(URN
659661

660662
private <ASPECT extends RecordTemplate> ASPECT_UNION unwrapAddResultToUnion(URN urn, AddResult<ASPECT> result,
661663
@Nonnull AuditStamp auditStamp, @Nullable IngestionTrackingContext trackingContext) {
664+
// handle post-update hooks and emit MAE + return the newValue
662665
ASPECT rawResult = unwrapAddResult(urn, result, auditStamp, trackingContext);
666+
667+
// package it into a union
663668
return ModelUtils.newEntityUnion(_aspectUnionClass, rawResult);
664669
}
665670

@@ -817,53 +822,65 @@ public <ASPECT extends RecordTemplate> ASPECT add(@Nonnull URN urn, AspectUpdate
817822
* <p>The new aspect will have an automatically assigned version number, which is guaranteed to be positive and
818823
* monotonically increasing. Older versions of aspect will be purged automatically based on the retention setting.
819824
*
820-
* <p>Note that we do not support Post-update hooks while soft deleting an aspect
825+
* <p>Note that we do not currently support pre- or post- update hooks while soft deleting an aspect.
821826
*
822827
* @param urn urn the URN for the entity the aspects are attached to
823828
* @param aspectClasses Aspect Classes of the aspects being deleted, must be supported aspect types in {@code ASPECT_UNION}
824829
* Because Aspect Classes must be unique for a given Entity, we use a set to avoid duplicates.
825830
* @param auditStamp the audit stamp of this action
826831
* @param maxTransactionRetry maximum number of transaction retries before throwing an exception
832+
* @return a collection of the deleted aspects (their value before deletion), each wrapped in an instance of {@link ASPECT_UNION}
833+
* If an aspect is already null or deleted, the return collection will not contain it whatsoever (it's "null").
827834
*/
828-
public void deleteMany(@Nonnull URN urn,
835+
@Nonnull
836+
public Collection<ASPECT_UNION> deleteMany(@Nonnull URN urn,
829837
@Nonnull Set<Class<? extends RecordTemplate>> aspectClasses,
830838
@Nonnull AuditStamp auditStamp,
831839
int maxTransactionRetry) {
832-
deleteMany(urn, aspectClasses, auditStamp, maxTransactionRetry, null);
840+
return deleteMany(urn, aspectClasses, auditStamp, maxTransactionRetry, null);
833841
}
834842

835843
/**
836844
* Similar to {@link #deleteMany(Urn, Set, AuditStamp, int)} but uses the default maximum transaction retry.
837845
*/
838846
@Nonnull
839-
public void deleteMany(@Nonnull URN urn, @Nonnull Set<Class<? extends RecordTemplate>> aspectClasses,
847+
public Collection<ASPECT_UNION> deleteMany(
848+
@Nonnull URN urn, @Nonnull Set<Class<? extends RecordTemplate>> aspectClasses,
840849
@Nonnull AuditStamp auditStamp) {
841-
deleteMany(urn, aspectClasses, auditStamp, DEFAULT_MAX_TRANSACTION_RETRY);
850+
return deleteMany(urn, aspectClasses, auditStamp, DEFAULT_MAX_TRANSACTION_RETRY);
842851
}
843852

844853
/**
845854
* Same as above {@link #deleteMany(Urn, Set, AuditStamp)} but with tracking context.
846855
*/
847856
@Nonnull
848-
public void deleteMany(@Nonnull URN urn, @Nonnull Set<Class<? extends RecordTemplate>> aspectClasses,
857+
public Collection<ASPECT_UNION> deleteMany(
858+
@Nonnull URN urn, @Nonnull Set<Class<? extends RecordTemplate>> aspectClasses,
849859
@Nonnull AuditStamp auditStamp, @Nullable IngestionTrackingContext trackingContext) {
850-
deleteMany(urn, aspectClasses, auditStamp, DEFAULT_MAX_TRANSACTION_RETRY, trackingContext);
860+
return deleteMany(urn, aspectClasses, auditStamp, DEFAULT_MAX_TRANSACTION_RETRY, trackingContext);
851861
}
852862

853863
/**
854864
* Same as {@link #deleteMany(Urn, Set, AuditStamp, int)} but with tracking context.
855865
*/
856-
public void deleteMany(@Nonnull URN urn,
866+
@Nonnull
867+
public Collection<ASPECT_UNION> deleteMany(@Nonnull URN urn,
857868
@Nonnull Set<Class<? extends RecordTemplate>> aspectClasses,
858869
@Nonnull AuditStamp auditStamp,
859870
int maxTransactionRetry,
860871
@Nullable IngestionTrackingContext trackingContext) {
861872

862873
// entire delete operation should be atomic
863-
runInTransactionWithRetry(() -> {
864-
aspectClasses.forEach(x -> delete(urn, x, auditStamp, maxTransactionRetry, trackingContext));
865-
return null;
866-
}, maxTransactionRetry);
874+
final Collection<RecordTemplate> results = runInTransactionWithRetry(() -> aspectClasses.stream()
875+
.map(x -> deleteWithReturn(urn, x, auditStamp, maxTransactionRetry, trackingContext))
876+
.collect(Collectors.toList()), maxTransactionRetry);
877+
878+
// package into ASPECT_UNION, this is logic performed in unwrapAddResultToUnion()
879+
// Aspect Union members are not allowed to be 'null' by convention (see ModelUtils' call tracing), so we must
880+
// filter out all nulls from the results.
881+
return results.stream()
882+
.filter(Objects::nonNull)
883+
.map(x -> ModelUtils.newEntityUnion(_aspectUnionClass, x)).collect(Collectors.toList());
867884
}
868885

869886
/**
@@ -890,16 +907,34 @@ public <ASPECT extends RecordTemplate> void delete(@Nonnull URN urn, @Nonnull Cl
890907
*/
891908
public <ASPECT extends RecordTemplate> void delete(@Nonnull URN urn, @Nonnull Class<ASPECT> aspectClass,
892909
@Nonnull AuditStamp auditStamp, int maxTransactionRetry, @Nullable IngestionTrackingContext trackingContext) {
910+
deleteWithReturn(urn, aspectClass, auditStamp, maxTransactionRetry, trackingContext);
911+
}
912+
913+
/**
914+
* Deletes the latest version of an aspect for an entity and returns the ***old value***.
915+
*
916+
* <p>Note that if the aspect is already null or deleted, the return value will be null. Mechanistically, upon a normal
917+
* "LIVE" ingestion, the deletion operation is skipped altogether.
918+
*/
919+
@Nullable
920+
public <ASPECT extends RecordTemplate> ASPECT deleteWithReturn(@Nonnull URN urn, @Nonnull Class<ASPECT> aspectClass,
921+
@Nonnull AuditStamp auditStamp, int maxTransactionRetry, @Nullable IngestionTrackingContext trackingContext) {
893922

894923
checkValidAspect(aspectClass);
895924

896-
runInTransactionWithRetry(() -> {
925+
final AddResult<ASPECT> result = runInTransactionWithRetry(() -> {
897926
final AspectEntry<ASPECT> latest = getLatest(urn, aspectClass, false);
898927
final IngestionParams ingestionParams = new IngestionParams().setIngestionMode(IngestionMode.LIVE);
899928
return addCommon(urn, latest, null, aspectClass, auditStamp, new DefaultEqualityTester<>(), trackingContext, ingestionParams);
900929
}, maxTransactionRetry);
901930

931+
return result.getOldValue();
932+
902933
// TODO: add support for sending MAE for soft deleted aspects
934+
// FY25H2 Note: When performing an Aspect UPDATE, unwrapAddResultToUnion() is called, which emits MAE and does post-update hooks.
935+
// When doing similar for DELETE, we should end up doing something similar, but specific to deletion.
936+
// We *could* modify the existing unwrapAddResultToUnion() to account for both cases and just reuse it completely,
937+
// but this might be confusing, so it might be best to make a deletion-specific version of that method.
903938
}
904939

905940
/**

‎dao-impl/ebean-dao/src/test/java/com/linkedin/metadata/dao/EbeanLocalDAOTest.java‎

Lines changed: 145 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -90,6 +90,7 @@
9090
import java.time.Instant;
9191
import java.util.ArrayList;
9292
import java.util.Arrays;
93+
import java.util.Collection;
9394
import java.util.Collections;
9495
import java.util.HashMap;
9596
import java.util.HashSet;
@@ -2633,8 +2634,38 @@ public void testRemoveRelationshipsDuringAspectSoftDeletion() throws URISyntaxEx
26332634
public void testDeleteManyWithRelationshipRemoval() throws URISyntaxException {
26342635
FooUrn fooUrn = makeFooUrn(1);
26352636
EbeanLocalDAO<EntityAspectUnion, FooUrn> fooDao = createDao(FooUrn.class);
2637+
// necessary flag to prevent removal of existing same-type relationships in "another aspect"
2638+
fooDao.setUseAspectColumnForRelationshipRemoval(true);
26362639

2637-
setupAspectsAndRelationships(fooUrn, fooDao);
2640+
EbeanLocalDAO<EntityAspectUnion, BarUrn> barDao = createDao(BarUrn.class);
2641+
2642+
// add an aspect (AspectFooBar) which includes BelongsTo relationships and ReportsTo relationships
2643+
BarUrn barUrn1 = BarUrn.createFromString("urn:li:bar:1");
2644+
BelongsToV2 belongsTo1 = new BelongsToV2().setDestination(BelongsToV2.Destination.create(barUrn1.toString()));
2645+
BarUrn barUrn2 = BarUrn.createFromString("urn:li:bar:2");
2646+
BelongsToV2 belongsTo2 = new BelongsToV2().setDestination(BelongsToV2.Destination.create(barUrn2.toString()));
2647+
BarUrn barUrn3 = BarUrn.createFromString("urn:li:bar:3");
2648+
BelongsToV2 belongsTo3 = new BelongsToV2().setDestination(BelongsToV2.Destination.create(barUrn3.toString()));
2649+
BelongsToV2Array belongsToArray = new BelongsToV2Array(belongsTo1, belongsTo2, belongsTo3);
2650+
ReportsTo reportsTo = new ReportsTo().setSource(fooUrn).setDestination(barUrn1);
2651+
ReportsToArray reportsToArray = new ReportsToArray(reportsTo);
2652+
AspectFooBar aspectFooBar = new AspectFooBar()
2653+
.setBars(new BarUrnArray(barUrn1, barUrn2, barUrn3)).setBelongsTos(belongsToArray).setReportsTos(reportsToArray);
2654+
AuditStamp auditStamp = makeAuditStamp("foo", System.currentTimeMillis());
2655+
2656+
fooDao.add(fooUrn, aspectFooBar, auditStamp);
2657+
barDao.add(barUrn1, new AspectFoo().setValue("1"), auditStamp);
2658+
barDao.add(barUrn2, new AspectFoo().setValue("2"), auditStamp);
2659+
barDao.add(barUrn3, new AspectFoo().setValue("3"), auditStamp);
2660+
2661+
// add an aspect (AspectFooBaz) which includes BelongsTo relationships
2662+
BarUrn barUrn4 = BarUrn.createFromString("urn:li:bar:4");
2663+
BelongsToV2 belongsTo4 = new BelongsToV2().setDestination(BelongsToV2.Destination.create(barUrn4.toString()));
2664+
BelongsToV2Array belongsToArray2 = new BelongsToV2Array(belongsTo4);
2665+
AspectFooBaz aspectFooBaz = new AspectFooBaz().setBars(new BarUrnArray(barUrn4)).setBelongsTos(belongsToArray2);
2666+
2667+
fooDao.add(fooUrn, aspectFooBaz, auditStamp);
2668+
barDao.add(barUrn4, new AspectFoo().setValue("4"), auditStamp);
26382669

26392670
// Verify local relationships and entities are added.
26402671
EbeanLocalRelationshipQueryDAO ebeanLocalRelationshipQueryDAO = new EbeanLocalRelationshipQueryDAO(_server);
@@ -2658,7 +2689,24 @@ public void testDeleteManyWithRelationshipRemoval() throws URISyntaxException {
26582689
assertEquals(aspects.size(), 1);
26592690

26602691
// soft delete the AspectFooBar and AspectFooBaz aspects
2661-
fooDao.deleteMany(fooUrn, new HashSet<>(Arrays.asList(AspectFooBar.class, AspectFooBaz.class)), _dummyAuditStamp);
2692+
Collection<EntityAspectUnion> deletedAspects =
2693+
fooDao.deleteMany(fooUrn, new HashSet<>(Arrays.asList(AspectFooBar.class, AspectFooBaz.class)), _dummyAuditStamp);
2694+
2695+
assertEquals(deletedAspects.size(), 2);
2696+
2697+
// check that the AspectFooBar content returned matches the pre-deletion content
2698+
Optional<EntityAspectUnion> aspectFooBarDeleted = deletedAspects.stream()
2699+
.filter(EntityAspectUnion::isAspectFooBar)
2700+
.findFirst();
2701+
assertTrue(aspectFooBarDeleted.isPresent());
2702+
assertEquals(aspectFooBarDeleted.get().getAspectFooBar(), aspectFooBar);
2703+
2704+
// check that the AspectFooBaz content returned matches the pre-deletion content
2705+
Optional<EntityAspectUnion> aspectFooBazDeleted = deletedAspects.stream()
2706+
.filter(EntityAspectUnion::isAspectFooBaz)
2707+
.findFirst();
2708+
assertTrue(aspectFooBazDeleted.isPresent());
2709+
assertEquals(aspectFooBazDeleted.get().getAspectFooBaz(), aspectFooBaz);
26622710

26632711
// check that the belongsTo relationships 1, 2, 3, and 4 were soft deleted
26642712
resultBelongsTos = ebeanLocalRelationshipQueryDAO.findRelationships(FooSnapshot.class, EMPTY_FILTER, BarSnapshot.class,
@@ -2682,6 +2730,101 @@ public void testDeleteManyWithRelationshipRemoval() throws URISyntaxException {
26822730
assertFalse(optionalAspect2.isPresent());
26832731
}
26842732

2733+
@Test
2734+
public void testDeleteWithReturnOnNonexistentAsset() {
2735+
EbeanLocalDAO<EntityAspectUnion, FooUrn> dao = createDao(FooUrn.class);
2736+
FooUrn urn = makeFooUrn(1);
2737+
2738+
AspectFoo foo = dao.deleteWithReturn(urn, AspectFoo.class, _dummyAuditStamp, 3, null);
2739+
assertNull(foo);
2740+
}
2741+
2742+
@Test
2743+
public void testDeleteWithReturnOnNullAspect() {
2744+
EbeanLocalDAO<EntityAspectUnion, FooUrn> dao = createDao(FooUrn.class);
2745+
FooUrn urn = makeFooUrn(1);
2746+
2747+
// add aspect so the row exists in the entity table, but the column for other aspects will be empty
2748+
AspectFoo v0 = new AspectFoo().setValue("foo");
2749+
dao.add(urn, v0, _dummyAuditStamp);
2750+
2751+
// attempt to delete an aspect that doesn't exist
2752+
AspectBar foo = dao.deleteWithReturn(urn, AspectBar.class, _dummyAuditStamp, 3, null);
2753+
assertNull(foo);
2754+
}
2755+
2756+
@Test
2757+
public void testDeleteWithReturnOnAlreadyDeletedAspect() {
2758+
EbeanLocalDAO<EntityAspectUnion, FooUrn> dao = createDao(FooUrn.class);
2759+
FooUrn urn = makeFooUrn(1);
2760+
AspectFoo v0 = new AspectFoo().setValue("foo");
2761+
dao.add(urn, v0, _dummyAuditStamp);
2762+
AspectFoo foo = dao.deleteWithReturn(urn, AspectFoo.class, _dummyAuditStamp, 3, null);
2763+
2764+
// make sure that the content matches the original
2765+
assertEquals(foo, v0);
2766+
2767+
// attempt to delete an aspect that has already been deleted
2768+
AspectFoo fooAgain = dao.deleteWithReturn(urn, AspectFoo.class, _dummyAuditStamp, 3, null);
2769+
assertNull(fooAgain);
2770+
}
2771+
2772+
@Test
2773+
public void testDeleteManyOnNonexistentAsset() {
2774+
EbeanLocalDAO<EntityAspectUnion, FooUrn> dao = createDao(FooUrn.class);
2775+
FooUrn urn = makeFooUrn(1);
2776+
2777+
Collection<EntityAspectUnion> deletionResults =
2778+
dao.deleteMany(urn, new HashSet<>(Collections.singletonList(AspectFoo.class)), _dummyAuditStamp, 3, null);
2779+
2780+
// make sure return collection is empty
2781+
assertEquals(deletionResults.size(), 0);
2782+
}
2783+
2784+
@Test
2785+
public void testDeleteManyOnNullAspect() {
2786+
EbeanLocalDAO<EntityAspectUnion, FooUrn> dao = createDao(FooUrn.class);
2787+
FooUrn urn = makeFooUrn(1);
2788+
2789+
// add aspect so the row exists in the entity table, but the column for other aspects will be empty
2790+
AspectFoo v0 = new AspectFoo().setValue("foo");
2791+
dao.add(urn, v0, _dummyAuditStamp);
2792+
2793+
// attempt to delete an aspect that doesn't exist
2794+
Collection<EntityAspectUnion> deletionResults =
2795+
dao.deleteMany(urn, new HashSet<>(Collections.singletonList(AspectBar.class)), _dummyAuditStamp, 3, null);
2796+
2797+
// make sure return collection is empty
2798+
assertEquals(deletionResults.size(), 0);
2799+
}
2800+
2801+
@Test
2802+
public void testDeleteManyOnAlreadyDeletedAspect() {
2803+
EbeanLocalDAO<EntityAspectUnion, FooUrn> dao = createDao(FooUrn.class);
2804+
FooUrn urn = makeFooUrn(1);
2805+
AspectFoo v0 = new AspectFoo().setValue("foo");
2806+
dao.add(urn, v0, _dummyAuditStamp);
2807+
2808+
// delete the aspect
2809+
Collection<EntityAspectUnion> deletionResults =
2810+
dao.deleteMany(urn, new HashSet<>(Collections.singletonList(AspectFoo.class)), _dummyAuditStamp, 3, null);
2811+
assertEquals(deletionResults.size(), 1);
2812+
2813+
// make sure that the content matches the original
2814+
Optional<EntityAspectUnion> aspectFooDeleted = deletionResults.stream()
2815+
.filter(EntityAspectUnion::isAspectFoo)
2816+
.findFirst();
2817+
assertTrue(aspectFooDeleted.isPresent());
2818+
assertEquals(aspectFooDeleted.get().getAspectFoo(), v0);
2819+
2820+
// attempt to delete an aspect that has already been deleted
2821+
Collection<EntityAspectUnion> deletionResultsAgain =
2822+
dao.deleteMany(urn, new HashSet<>(Collections.singletonList(AspectFoo.class)), _dummyAuditStamp, 3, null);
2823+
2824+
// make sure return collection is empty
2825+
assertEquals(deletionResultsAgain.size(), 0);
2826+
}
2827+
26852828
@Test
26862829
public void testGetWithExtraInfoMultipleKeys() {
26872830
EbeanLocalDAO<EntityAspectUnion, FooUrn> dao = createDao(FooUrn.class);

0 commit comments

Comments
 (0)