feat: Support deletion in EbeanGenericLocalDAO and BaseEntityAgnostic…

…AspectResource (#454)
linkedin · Oct 23, 2024 · b762dac · b762dac
1 parent ea9061a
commit b762dac
Show file tree

Hide file tree

Showing 6 changed files with 147 additions and 6 deletions.
diff --git a/dao-api/src/main/java/com/linkedin/metadata/dao/GenericLocalDAO.java b/dao-api/src/main/java/com/linkedin/metadata/dao/GenericLocalDAO.java
@@ -54,4 +54,13 @@ void save(@Nonnull Urn urn, @Nonnull Class aspectClass, @Nonnull String metadata
    */
   Map<Urn, Map<Class<? extends RecordTemplate>, Optional<? extends RecordTemplate>>> backfill(@Nonnull BackfillMode mode,
       @Nonnull Map<Urn, Set<Class<? extends RecordTemplate>>> urnToAspect);
+
+  /**
+   * Delete the metadata from database.
+   *
+   * @param urn The identifier of the entity which the metadata is associated with.
+   * @param aspectClass The aspect class for the metadata.
+   * @param auditStamp audit stamp containing information on who and when the metadata is deleted.
+   */
+  void delete(@Nonnull Urn urn, @Nonnull Class aspectClass, @Nonnull AuditStamp auditStamp);
 }
diff --git a/dao-impl/ebean-dao/src/main/java/com/linkedin/metadata/dao/EbeanGenericLocalDAO.java b/dao-impl/ebean-dao/src/main/java/com/linkedin/metadata/dao/EbeanGenericLocalDAO.java
@@ -87,13 +87,25 @@ public void setEqualityTesters(Map<Class, GenericEqualityTester> equalityTesters
    */
   public void save(@Nonnull Urn urn, @Nonnull Class aspectClass, @Nonnull String metadata, @Nonnull AuditStamp auditStamp,
       @Nullable IngestionTrackingContext trackingContext, @Nullable IngestionMode ingestionMode) {
+    saveCommon(urn, aspectClass, metadata, auditStamp, trackingContext, ingestionMode);
+  }
+
+  /* a common helper method for saving metadata */
+  void saveCommon(@Nonnull Urn urn, @Nonnull Class aspectClass, @Nullable String metadata, @Nonnull AuditStamp auditStamp,
+      @Nullable IngestionTrackingContext trackingContext, @Nullable IngestionMode ingestionMode) {
     runInTransactionWithRetry(() -> {
       final Optional<GenericLocalDAO.MetadataWithExtraInfo> latest = queryLatest(urn, aspectClass);
-      RecordTemplate newValue = toRecordTemplate(aspectClass, metadata);
+
+      RecordTemplate newValue = null;
+      if (metadata != null) {
+        newValue = toRecordTemplate(aspectClass, metadata);
+      }
 
       if (!latest.isPresent()) {
         saveLatest(urn, aspectClass, newValue, null, auditStamp, null);
-        _producer.produceAspectSpecificMetadataAuditEvent(urn, null, newValue, auditStamp, trackingContext, ingestionMode);
+        if (!shouldSkipMAEUpdate(newValue)) {
+          _producer.produceAspectSpecificMetadataAuditEvent(urn, null, newValue, auditStamp, trackingContext, ingestionMode);
+        }
       } else {
         RecordTemplate currentValue = toRecordTemplate(aspectClass, latest.get().getAspect());
         final AuditStamp oldAuditStamp = latest.get().getExtraInfo() == null ? null : latest.get().getExtraInfo().getAudit();
@@ -107,15 +119,25 @@ public void save(@Nonnull Urn urn, @Nonnull Class aspectClass, @Nonnull String m
         }
 
         // Skip update if current value and new value are equal.
-        if (!areEqual(currentValue, newValue, _equalityTesters.get(aspectClass))) {
-          saveLatest(urn, aspectClass, newValue, currentValue, auditStamp, latest.get().getExtraInfo().getAudit());
+        // currentValue is always not null in this case
+        if (newValue != null && areEqual(currentValue, newValue, _equalityTesters.get(aspectClass))) {
+          return null;
+        }
+        saveLatest(urn, aspectClass, newValue, currentValue, auditStamp, latest.get().getExtraInfo().getAudit());
+
+        if (!shouldSkipMAEUpdate(newValue)) {
           _producer.produceAspectSpecificMetadataAuditEvent(urn, currentValue, newValue, auditStamp, trackingContext, ingestionMode);
         }
       }
       return null;
     }, 5);
   }
 
+  private boolean shouldSkipMAEUpdate(@Nullable RecordTemplate newValue) {
+    // do not send MAE for null new value (deletion), to keep the same behavior as in BaseLocalDao
+    return newValue == null;
+  }
+
   /**
    * Query the latest metadata from database.
    * @param urn The identifier of the entity which the metadata is associated with.
@@ -128,7 +150,7 @@ public Optional<GenericLocalDAO.MetadataWithExtraInfo> queryLatest(@Nonnull Urn
     final PrimaryKey key = new PrimaryKey(urn.toString(), aspectName, LATEST_VERSION);
     EbeanMetadataAspect metadata = _server.find(EbeanMetadataAspect.class, key);
 
-    if (metadata == null || metadata.getMetadata() == null) {
+    if (metadata == null || metadata.getMetadata() == null || DELETED_VALUE.equals(metadata.getMetadata())) {
       return Optional.empty();
     }
 
@@ -170,6 +192,11 @@ public Map<Urn, Map<Class<? extends RecordTemplate>, Optional<? extends RecordTe
     return urnToAspects;
   }
 
+  @Override
+  public void delete(@Nonnull Urn urn, @Nonnull Class aspectClass, @Nonnull AuditStamp auditStamp) {
+    saveCommon(urn, aspectClass, null, auditStamp, null, null);
+  }
+
   /**
    * Emits backfill MAE for an aspect of an entity depending on the backfill mode.
    *
@@ -195,7 +222,7 @@ private void backfill(@Nonnull BackfillMode mode, @Nonnull RecordTemplate aspect
   /**
    * Save metadata into database.
    */
-  private void saveLatest(@Nonnull Urn urn, @Nonnull Class aspectClass, @Nonnull RecordTemplate newValue,
+  private void saveLatest(@Nonnull Urn urn, @Nonnull Class aspectClass, @Nullable RecordTemplate newValue,
       @Nullable RecordTemplate currentValue, @Nonnull AuditStamp newAuditStamp, @Nullable AuditStamp currentAuditStamp) {
 
     // Save oldValue as the largest version + 1

diff --git a/dao-impl/ebean-dao/src/test/java/com/linkedin/metadata/dao/EbeanGenericLocalDAOTest.java b/dao-impl/ebean-dao/src/test/java/com/linkedin/metadata/dao/EbeanGenericLocalDAOTest.java
@@ -239,4 +239,81 @@ public void testBackfill() throws URISyntaxException {
     assertEquals(backfillResults.size(), 1);
     assertEquals(backfillResults.get(fooUrn).get(AspectFoo.class).get(), aspectFoo);
   }
+
+  @Test
+  public void testDelete() throws URISyntaxException {
+    FooUrn fooUrn = FooUrn.createFromString("urn:li:foo:1");
+    AspectFoo aspectFoo = new AspectFoo().setValue("foo");
+
+    _genericLocalDAO.save(fooUrn, AspectFoo.class, RecordUtils.toJsonString(aspectFoo),
+        makeAuditStamp("tester"), null, null);
+
+    verify(_producer, times(1)).produceAspectSpecificMetadataAuditEvent(eq(fooUrn),
+        eq(null), eq(aspectFoo), eq(makeAuditStamp("tester")), eq(null), eq(null));
+
+    Optional<GenericLocalDAO.MetadataWithExtraInfo> metadata = _genericLocalDAO.queryLatest(fooUrn, AspectFoo.class);
+
+    // {"value":"foo"} is inserted later so it is the latest metadata.
+    assertTrue(metadata.isPresent());
+    assertEquals(metadata.get().getAspect(), RecordUtils.toJsonString(aspectFoo));
+
+    reset(_producer);
+
+    // Delete the record and verify it is deleted.
+    _genericLocalDAO.delete(fooUrn, AspectFoo.class, makeAuditStamp("tester"));
+
+    metadata = _genericLocalDAO.queryLatest(fooUrn, AspectFoo.class);
+    assertFalse(metadata.isPresent());
+
+    // does not produce MAE for deletion
+    verify(_producer, times(0)).produceAspectSpecificMetadataAuditEvent(eq(fooUrn),
+        any(), any(), any(), any(), any());
+    verifyNoMoreInteractions(_producer);
+  }
+
+  @Test
+  public void testDeleteVoid() throws URISyntaxException {
+    FooUrn fooUrn = FooUrn.createFromString("urn:li:foo:1");
+
+    Optional<GenericLocalDAO.MetadataWithExtraInfo> metadata = _genericLocalDAO.queryLatest(fooUrn, AspectFoo.class);
+
+    // no record is returned.
+    assertFalse(metadata.isPresent());
+
+    // Delete the record and verify no record is returned.
+    _genericLocalDAO.delete(fooUrn, AspectFoo.class, makeAuditStamp("tester"));
+
+    metadata = _genericLocalDAO.queryLatest(fooUrn, AspectFoo.class);
+    assertFalse(metadata.isPresent());
+
+    // does not produce MAE for deletion
+    verify(_producer, times(0)).produceAspectSpecificMetadataAuditEvent(eq(fooUrn),
+        any(), any(), any(), any(), any());
+    verifyNoMoreInteractions(_producer);
+  }
+
+  @Test
+  public void testBackfillAfterDelete() throws URISyntaxException {
+    FooUrn fooUrn = FooUrn.createFromString("urn:li:foo:1");
+    AspectFoo aspectFoo = new AspectFoo().setValue("foo");
+
+    _genericLocalDAO.save(fooUrn, AspectFoo.class, RecordUtils.toJsonString(aspectFoo),
+        makeAuditStamp("tester"), null, null);
+
+    Map<Urn, Set<Class<? extends RecordTemplate>>> aspects = Collections.singletonMap(fooUrn, Collections.singleton(AspectFoo.class));
+
+    Map<Urn, Map<Class<? extends RecordTemplate>, Optional<? extends RecordTemplate>>> backfillResults
+        = _genericLocalDAO.backfill(BackfillMode.BACKFILL_ALL, aspects);
+
+    assertEquals(backfillResults.size(), 1);
+    assertEquals(backfillResults.get(fooUrn).get(AspectFoo.class).get(), aspectFoo);
+
+
+    // verify no aspect will be backfilled after deletion
+    _genericLocalDAO.delete(fooUrn, AspectFoo.class, makeAuditStamp("tester"));
+
+    backfillResults = _genericLocalDAO.backfill(BackfillMode.BACKFILL_ALL, aspects);
+    assertEquals(backfillResults.size(), 1);
+    assertEquals(backfillResults.get(fooUrn).size(), 0);
+  }
 }
diff --git a/...esources/src/main/java/com/linkedin/metadata/restli/BaseEntityAgnosticAspectResource.java b/...esources/src/main/java/com/linkedin/metadata/restli/BaseEntityAgnosticAspectResource.java
@@ -131,4 +131,22 @@ public Task<BackfillResult> backfill(
       throw new RestLiServiceException(HttpStatus.S_400_BAD_REQUEST, String.format("Urn %s is malformed.", urn));
     }
   }
+
+  @Action(name = ACTION_DELETE)
+  @Nonnull
+  public Task<Void> delete(
+      @ActionParam(PARAM_URN) @Nonnull String urn,
+      @ActionParam(PARAM_ASPECT_CLASS) @Nonnull String aspectClass) {
+    final AuditStamp auditStamp = getAuditor().requestAuditStamp(getContext().getRawRequestContext());
+
+    try {
+      Class clazz = this.getClass().getClassLoader().loadClass(aspectClass);
+      genericLocalDAO().delete(Urn.createFromCharSequence(urn), clazz, auditStamp);
+      return Task.value(null);
+    } catch (ClassNotFoundException e) {
+      throw new RestLiServiceException(HttpStatus.S_400_BAD_REQUEST, String.format("No such class %s.", aspectClass));
+    } catch (URISyntaxException e) {
+      throw new RestLiServiceException(HttpStatus.S_400_BAD_REQUEST, String.format("Urn %s is malformed.", urn));
+    }
+  }
 }
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
@@ -31,6 +31,7 @@ private RestliConstants() { }
   public static final String ACTION_RAW_INGEST_ASSET = "rawIngestAsset";
   public static final String ACTION_LIST_URNS_FROM_INDEX = "listUrnsFromIndex";
   public static final String ACTION_LIST_URNS = "listUrns";
+  public static final String ACTION_DELETE = "delete";
   public static final String PARAM_INPUT = "input";
   public static final String PARAM_ASPECTS = "aspects";
   public static final String PARAM_ASPECT = "aspect";

diff --git a/...rces/src/test/java/com/linkedin/metadata/restli/BaseEntityAgnosticAspectResourceTest.java b/...rces/src/test/java/com/linkedin/metadata/restli/BaseEntityAgnosticAspectResourceTest.java
@@ -86,4 +86,13 @@ public void testBackfill() {
 
     verifyNoMoreInteractions(_mockLocalDAO);
   }
+
+  @Test
+  public void testDelete() {
+    runAndWait(_resource.delete(ENTITY_URN.toString(), AspectFoo.class.getCanonicalName()));
+
+    verify(_mockLocalDAO, times(1)).delete(eq(ENTITY_URN), eq(AspectFoo.class), any(AuditStamp.class));
+
+    verifyNoMoreInteractions(_mockLocalDAO);
+  }
 }