From 2f4a4fdf1fc93708d33d924de65e75f34f932f61 Mon Sep 17 00:00:00 2001 From: Jonathan Hui Date: Tue, 11 Feb 2025 22:58:39 -0800 Subject: [PATCH] (feat) Add deleteWithReturn --- .../linkedin/metadata/dao/BaseLocalDAO.java | 49 +++++++--- .../metadata/dao/EbeanLocalDAOTest.java | 91 ++++++++++++++++++- 2 files changed, 125 insertions(+), 15 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 ad87c1058..1b27f46be 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 @@ -49,6 +49,7 @@ import com.linkedin.metadata.query.IndexSortCriterion; import java.sql.Timestamp; import java.time.Clock; +import java.util.Collection; import java.util.Collections; import java.util.ArrayList; import java.util.HashMap; @@ -659,7 +660,10 @@ private AddResult aspectUpdateHelper(URN private ASPECT_UNION unwrapAddResultToUnion(URN urn, AddResult result, @Nonnull AuditStamp auditStamp, @Nullable IngestionTrackingContext trackingContext) { + // handle post-update hooks and emit MAE + return the newValue ASPECT rawResult = unwrapAddResult(urn, result, auditStamp, trackingContext); + + // package it into a union return ModelUtils.newEntityUnion(_aspectUnionClass, rawResult); } @@ -817,53 +821,58 @@ public ASPECT add(@Nonnull URN urn, AspectUpdate *

The new aspect will have an automatically assigned version number, which is guaranteed to be positive and * monotonically increasing. Older versions of aspect will be purged automatically based on the retention setting. * - *

Note that we do not support Post-update hooks while soft deleting an aspect + *

Note that we do not currently support pre- or post- update hooks while soft deleting an aspect. * * @param urn urn the URN for the entity the aspects are attached to * @param aspectClasses Aspect Classes of the aspects being deleted, must be supported aspect types in {@code ASPECT_UNION} * Because Aspect Classes must be unique for a given Entity, we use a set to avoid duplicates. * @param auditStamp the audit stamp of this action * @param maxTransactionRetry maximum number of transaction retries before throwing an exception + * @return a collection of the deleted aspects (their value before deletion), each wrapped in an instance of {@link ASPECT_UNION} */ - public void deleteMany(@Nonnull URN urn, + public Collection deleteMany(@Nonnull URN urn, @Nonnull Set> aspectClasses, @Nonnull AuditStamp auditStamp, int maxTransactionRetry) { - deleteMany(urn, aspectClasses, auditStamp, maxTransactionRetry, null); + return deleteMany(urn, aspectClasses, auditStamp, maxTransactionRetry, null); } /** * Similar to {@link #deleteMany(Urn, Set, AuditStamp, int)} but uses the default maximum transaction retry. */ @Nonnull - public void deleteMany(@Nonnull URN urn, @Nonnull Set> aspectClasses, + public Collection deleteMany( + @Nonnull URN urn, @Nonnull Set> aspectClasses, @Nonnull AuditStamp auditStamp) { - deleteMany(urn, aspectClasses, auditStamp, DEFAULT_MAX_TRANSACTION_RETRY); + return deleteMany(urn, aspectClasses, auditStamp, DEFAULT_MAX_TRANSACTION_RETRY); } /** * Same as above {@link #deleteMany(Urn, Set, AuditStamp)} but with tracking context. */ @Nonnull - public void deleteMany(@Nonnull URN urn, @Nonnull Set> aspectClasses, + public Collection deleteMany( + @Nonnull URN urn, @Nonnull Set> aspectClasses, @Nonnull AuditStamp auditStamp, @Nullable IngestionTrackingContext trackingContext) { - deleteMany(urn, aspectClasses, auditStamp, DEFAULT_MAX_TRANSACTION_RETRY, trackingContext); + return deleteMany(urn, aspectClasses, auditStamp, DEFAULT_MAX_TRANSACTION_RETRY, trackingContext); } /** * Same as {@link #deleteMany(Urn, Set, AuditStamp, int)} but with tracking context. */ - public void deleteMany(@Nonnull URN urn, + public Collection deleteMany(@Nonnull URN urn, @Nonnull Set> aspectClasses, @Nonnull AuditStamp auditStamp, int maxTransactionRetry, @Nullable IngestionTrackingContext trackingContext) { // entire delete operation should be atomic - runInTransactionWithRetry(() -> { - aspectClasses.forEach(x -> delete(urn, x, auditStamp, maxTransactionRetry, trackingContext)); - return null; - }, maxTransactionRetry); + final Collection results = runInTransactionWithRetry(() -> aspectClasses.stream() + .map(x -> deleteWithReturn(urn, x, auditStamp, maxTransactionRetry, trackingContext)) + .collect(Collectors.toList()), maxTransactionRetry); + + // package into ASPECT_UNION, this is logic performed in unwrapAddResultToUnion() + return results.stream().map(x -> ModelUtils.newEntityUnion(_aspectUnionClass, x)).collect(Collectors.toList()); } /** @@ -890,16 +899,30 @@ public void delete(@Nonnull URN urn, @Nonnull Cl */ public void delete(@Nonnull URN urn, @Nonnull Class aspectClass, @Nonnull AuditStamp auditStamp, int maxTransactionRetry, @Nullable IngestionTrackingContext trackingContext) { + deleteWithReturn(urn, aspectClass, auditStamp, maxTransactionRetry, trackingContext); + } + + /** + * Deletes the latest version of an aspect for an entity and returns the ***old value***. + */ + public ASPECT deleteWithReturn(@Nonnull URN urn, @Nonnull Class aspectClass, + @Nonnull AuditStamp auditStamp, int maxTransactionRetry, @Nullable IngestionTrackingContext trackingContext) { checkValidAspect(aspectClass); - runInTransactionWithRetry(() -> { + final AddResult result = runInTransactionWithRetry(() -> { final AspectEntry latest = getLatest(urn, aspectClass, false); final IngestionParams ingestionParams = new IngestionParams().setIngestionMode(IngestionMode.LIVE); return addCommon(urn, latest, null, aspectClass, auditStamp, new DefaultEqualityTester<>(), trackingContext, ingestionParams); }, maxTransactionRetry); + return result.getOldValue(); + // TODO: add support for sending MAE for soft deleted aspects + // FY25H2 Note: When performing an Aspect UPDATE, unwrapAddResultToUnion() is called, which emits MAE and does post-update hooks. + // When doing similar for DELETE, we should end up doing something similar, but specific to deletion. + // We *could* modify the existing unwrapAddResultToUnion() to account for both cases and just reuse it completely, + // but this might be confusing, so it might be best to make a deletion-specific version of that method. } /** diff --git a/dao-impl/ebean-dao/src/test/java/com/linkedin/metadata/dao/EbeanLocalDAOTest.java b/dao-impl/ebean-dao/src/test/java/com/linkedin/metadata/dao/EbeanLocalDAOTest.java index d9d4692ad..985af08f5 100644 --- a/dao-impl/ebean-dao/src/test/java/com/linkedin/metadata/dao/EbeanLocalDAOTest.java +++ b/dao-impl/ebean-dao/src/test/java/com/linkedin/metadata/dao/EbeanLocalDAOTest.java @@ -90,6 +90,7 @@ import java.time.Instant; import java.util.ArrayList; import java.util.Arrays; +import java.util.Collection; import java.util.Collections; import java.util.HashMap; import java.util.HashSet; @@ -2633,8 +2634,38 @@ public void testRemoveRelationshipsDuringAspectSoftDeletion() throws URISyntaxEx public void testDeleteManyWithRelationshipRemoval() throws URISyntaxException { FooUrn fooUrn = makeFooUrn(1); EbeanLocalDAO fooDao = createDao(FooUrn.class); + // necessary flag to prevent removal of existing same-type relationships in "another aspect" + fooDao.setUseAspectColumnForRelationshipRemoval(true); - setupAspectsAndRelationships(fooUrn, fooDao); + EbeanLocalDAO barDao = createDao(BarUrn.class); + + // add an aspect (AspectFooBar) which includes BelongsTo relationships and ReportsTo relationships + BarUrn barUrn1 = BarUrn.createFromString("urn:li:bar:1"); + BelongsToV2 belongsTo1 = new BelongsToV2().setDestination(BelongsToV2.Destination.create(barUrn1.toString())); + BarUrn barUrn2 = BarUrn.createFromString("urn:li:bar:2"); + BelongsToV2 belongsTo2 = new BelongsToV2().setDestination(BelongsToV2.Destination.create(barUrn2.toString())); + BarUrn barUrn3 = BarUrn.createFromString("urn:li:bar:3"); + BelongsToV2 belongsTo3 = new BelongsToV2().setDestination(BelongsToV2.Destination.create(barUrn3.toString())); + BelongsToV2Array belongsToArray = new BelongsToV2Array(belongsTo1, belongsTo2, belongsTo3); + ReportsTo reportsTo = new ReportsTo().setSource(fooUrn).setDestination(barUrn1); + ReportsToArray reportsToArray = new ReportsToArray(reportsTo); + AspectFooBar aspectFooBar = new AspectFooBar() + .setBars(new BarUrnArray(barUrn1, barUrn2, barUrn3)).setBelongsTos(belongsToArray).setReportsTos(reportsToArray); + AuditStamp auditStamp = makeAuditStamp("foo", System.currentTimeMillis()); + + fooDao.add(fooUrn, aspectFooBar, auditStamp); + barDao.add(barUrn1, new AspectFoo().setValue("1"), auditStamp); + barDao.add(barUrn2, new AspectFoo().setValue("2"), auditStamp); + barDao.add(barUrn3, new AspectFoo().setValue("3"), auditStamp); + + // add an aspect (AspectFooBaz) which includes BelongsTo relationships + BarUrn barUrn4 = BarUrn.createFromString("urn:li:bar:4"); + BelongsToV2 belongsTo4 = new BelongsToV2().setDestination(BelongsToV2.Destination.create(barUrn4.toString())); + BelongsToV2Array belongsToArray2 = new BelongsToV2Array(belongsTo4); + AspectFooBaz aspectFooBaz = new AspectFooBaz().setBars(new BarUrnArray(barUrn4)).setBelongsTos(belongsToArray2); + + fooDao.add(fooUrn, aspectFooBaz, auditStamp); + barDao.add(barUrn4, new AspectFoo().setValue("4"), auditStamp); // Verify local relationships and entities are added. EbeanLocalRelationshipQueryDAO ebeanLocalRelationshipQueryDAO = new EbeanLocalRelationshipQueryDAO(_server); @@ -2658,7 +2689,24 @@ public void testDeleteManyWithRelationshipRemoval() throws URISyntaxException { assertEquals(aspects.size(), 1); // soft delete the AspectFooBar and AspectFooBaz aspects - fooDao.deleteMany(fooUrn, new HashSet<>(Arrays.asList(AspectFooBar.class, AspectFooBaz.class)), _dummyAuditStamp); + Collection deletedAspects = + fooDao.deleteMany(fooUrn, new HashSet<>(Arrays.asList(AspectFooBar.class, AspectFooBaz.class)), _dummyAuditStamp); + + assertEquals(deletedAspects.size(), 2); + + // check that the AspectFooBar content returned matches the pre-deletion content + Optional aspectFooBarDeleted = deletedAspects.stream() + .filter(aspect -> aspect.getAspectFooBar() != null) + .findFirst(); + assertTrue(aspectFooBarDeleted.isPresent()); + assertEquals(aspectFooBarDeleted.get().getAspectFooBar(), aspectFooBar); + + // check that the AspectFooBaz content returned matches the pre-deletion content + Optional aspectFooBazDeleted = deletedAspects.stream() + .filter(aspect -> aspect.getAspectFooBaz() != null) + .findFirst(); + assertTrue(aspectFooBazDeleted.isPresent()); + assertEquals(aspectFooBazDeleted.get().getAspectFooBaz(), aspectFooBaz); // check that the belongsTo relationships 1, 2, 3, and 4 were soft deleted resultBelongsTos = ebeanLocalRelationshipQueryDAO.findRelationships(FooSnapshot.class, EMPTY_FILTER, BarSnapshot.class, @@ -2682,6 +2730,45 @@ public void testDeleteManyWithRelationshipRemoval() throws URISyntaxException { assertFalse(optionalAspect2.isPresent()); } + @Test + public void testDeleteWithReturnOnNonexistentAsset() { + EbeanLocalDAO dao = createDao(FooUrn.class); + FooUrn urn = makeFooUrn(1); + + AspectFoo foo = dao.deleteWithReturn(urn, AspectFoo.class, _dummyAuditStamp, 3, null); + assertNull(foo); + } + + @Test + public void testDeleteWithReturnOnNullAspect() { + EbeanLocalDAO dao = createDao(FooUrn.class); + FooUrn urn = makeFooUrn(1); + + // add aspect so the row exists in the entity table, but the column for other aspects will be empty + AspectFoo v0 = new AspectFoo().setValue("foo"); + dao.add(urn, v0, _dummyAuditStamp); + + // attempt to delete an aspect that doesn't exist + AspectBaz foo = dao.deleteWithReturn(urn, AspectBaz.class, _dummyAuditStamp, 3, null); + assertNull(foo); + } + + @Test + public void testDeleteWithReturnOnAlreadyDeletedAspect() { + EbeanLocalDAO dao = createDao(FooUrn.class); + FooUrn urn = makeFooUrn(1); + AspectFoo v0 = new AspectFoo().setValue("foo"); + dao.add(urn, v0, _dummyAuditStamp); + AspectFoo foo = dao.deleteWithReturn(urn, AspectFoo.class, _dummyAuditStamp, 3, null); + + // make sure that the content matches the original + assertEquals(foo, v0); + + // attempt to delete an aspect that has already been deleted + AspectFoo fooAgain = dao.deleteWithReturn(urn, AspectFoo.class, _dummyAuditStamp, 3, null); + assertNull(fooAgain); + } + @Test public void testGetWithExtraInfoMultipleKeys() { EbeanLocalDAO dao = createDao(FooUrn.class);