(feat) support clear relationship by source entity urn (#342)

* support clear relationship by source entity urn

* including RemovalOption support

dao-api/src/main/java/com/linkedin/metadata/dao/builder/BaseLocalRelationshipBuilder.java

@@ -18,6 +18,7 @@ public abstract class BaseLocalRelationshipBuilder<ASPECT extends RecordTemplate
   @Value
   public static class LocalRelationshipUpdates {
     List<? extends RecordTemplate> relationships;
+    Class<? extends RecordTemplate>[] relationshipClasses;
     BaseGraphWriterDAO.RemovalOption removalOption;
   }
 

dao-impl/ebean-dao/src/main/java/com/linkedin/metadata/dao/EbeanLocalAccess.java

@@ -169,17 +169,14 @@ public <ASPECT extends RecordTemplate> int addWithOptimisticLocking(
   }
 
   @Override
-  public <ASPECT extends RecordTemplate> List<LocalRelationshipUpdates>
-  addRelationships(@Nonnull URN urn, @Nonnull ASPECT aspect, @Nonnull Class<ASPECT> aspectClass) {
+  public <ASPECT extends RecordTemplate> List<LocalRelationshipUpdates> addRelationships(@Nonnull URN urn,
+      @Nonnull ASPECT aspect, @Nonnull Class<ASPECT> aspectClass) {
     if (_localRelationshipBuilderRegistry != null && _localRelationshipBuilderRegistry.isRegistered(aspectClass)) {
       List<LocalRelationshipUpdates> localRelationshipUpdates =
           _localRelationshipBuilderRegistry.getLocalRelationshipBuilder(aspect).buildRelationships(urn, aspect);
-
-      _localRelationshipWriterDAO.processLocalRelationshipUpdates(localRelationshipUpdates);
-
+      _localRelationshipWriterDAO.processLocalRelationshipUpdates(urn, localRelationshipUpdates);
       return localRelationshipUpdates;
     }
-
     return new ArrayList<>();
   }
 

dao-impl/ebean-dao/src/main/java/com/linkedin/metadata/dao/EbeanLocalRelationshipWriterDAO.java

@@ -47,11 +47,46 @@ public EbeanLocalRelationshipWriterDAO(EbeanServer server) {
    * @param relationshipUpdates Updates to local relationship tables.
    */
   @Transactional
-  public <ASPECT extends RecordTemplate> void processLocalRelationshipUpdates(
+  public <ASPECT extends RecordTemplate> void processLocalRelationshipUpdates(@Nonnull Urn urn,
       @Nonnull List<LocalRelationshipUpdates> relationshipUpdates) {
-
     for (LocalRelationshipUpdates relationshipUpdate : relationshipUpdates) {
-      addRelationships(relationshipUpdate.getRelationships(), relationshipUpdate.getRemovalOption());
+      if (relationshipUpdate.getRelationships().isEmpty()) {
+        clearRelationshipsByEntity(urn, relationshipUpdate.getRelationshipClasses(),
+            relationshipUpdate.getRemovalOption());
+      } else {
+        addRelationships(relationshipUpdate.getRelationships(), relationshipUpdate.getRemovalOption());
+      }
+    }
+  }
+
+  /**
+   * This method is to serve for the purpose to clear all the relationships from a source entity urn.
+   * @param urn entity urn could be either source or destination, depends on the RemovalOption
+   * @param relationshipClasses relationship that needs to be cleared
+   */
+  public void clearRelationshipsByEntity(@Nonnull Urn urn,
+      @Nonnull Class<? extends RecordTemplate>[] relationshipClasses, @Nonnull RemovalOption removalOption) {
+    if (removalOption == RemovalOption.REMOVE_NONE
+        || removalOption == RemovalOption.REMOVE_ALL_EDGES_FROM_SOURCE_TO_DESTINATION) {
+      // this method is to handle the case of adding empty relationship list to clear relationships of an entity urn
+      // REMOVE_NONE and REMOVE_ALL_EDGES_FROM_SOURCE_TO_DESTINATION won't apply for this case.
+      return;
+    }
+    if (relationshipClasses.length == 0) {
+      // if no relationship supported relationship classes are declared, then there's no relationship tables to delete.
+      return;
+    }
+    for (Class<? extends RecordTemplate> relationshipClass : relationshipClasses) {
+      RelationshipValidator.validateRelationshipSchema(relationshipClass);
+      SqlUpdate deletionSQL = _server.createSqlUpdate(
+          SQLStatementUtils.deleteLocaRelationshipSQL(SQLSchemaUtils.getRelationshipTableName(relationshipClass),
+              removalOption));
+      if (removalOption == RemovalOption.REMOVE_ALL_EDGES_FROM_SOURCE) {
+        deletionSQL.setParameter(CommonColumnName.SOURCE, urn.toString());
+      } else if (removalOption == RemovalOption.REMOVE_ALL_EDGES_TO_DESTINATION) {
+        deletionSQL.setParameter(CommonColumnName.DESTINATION, urn.toString());
+      }
+      deletionSQL.execute();
     }
   }
 

dao-impl/ebean-dao/src/test/java/com/linkedin/metadata/dao/localrelationship/EbeanLocalRelationshipWriterDAOTest.java

@@ -2,6 +2,7 @@
 
 import com.google.common.io.Resources;
 import com.linkedin.metadata.dao.EbeanLocalRelationshipWriterDAO;
+import com.linkedin.metadata.dao.internal.BaseGraphWriterDAO;
 import com.linkedin.metadata.dao.localrelationship.builder.BelongsToLocalRelationshipBuilder;
 import com.linkedin.metadata.dao.localrelationship.builder.PairsWithLocalRelationshipBuilder;
 import com.linkedin.metadata.dao.localrelationship.builder.ReportsToLocalRelationshipBuilder;
@@ -10,6 +11,7 @@
 import com.linkedin.metadata.dao.builder.BaseLocalRelationshipBuilder.LocalRelationshipUpdates;
 import com.linkedin.testing.BarUrnArray;
 import com.linkedin.testing.localrelationship.AspectFooBar;
+import com.linkedin.testing.localrelationship.PairsWith;
 import com.linkedin.testing.urn.BarUrn;
 import com.linkedin.testing.urn.FooUrn;
 import io.ebean.Ebean;
@@ -54,7 +56,7 @@ public void testAddRelationshipWithRemoveAllEdgesToDestination() throws URISynta
     List<SqlRow> before = _server.createSqlQuery("select * from metadata_relationship_belongsto where source='urn:li:bar:000'").findList();
     assertEquals(before.size(), 1);
 
-    _localRelationshipWriterDAO.processLocalRelationshipUpdates(updates);
+    _localRelationshipWriterDAO.processLocalRelationshipUpdates(FooUrn.createFromString("urn:li:foo:123"), updates);
 
     // After processing verification
     List<SqlRow> all = _server.createSqlQuery("select * from metadata_relationship_belongsto").findList();
@@ -89,7 +91,7 @@ public void testAddRelationshipWithRemoveNone() throws URISyntaxException {
     List<SqlRow> before = _server.createSqlQuery("select * from metadata_relationship_reportsto where source='urn:li:bar:000'").findList();
     assertEquals(before.size(), 1);
 
-    _localRelationshipWriterDAO.processLocalRelationshipUpdates(updates);
+    _localRelationshipWriterDAO.processLocalRelationshipUpdates(FooUrn.createFromString("urn:li:foo:123"), updates);
 
     // After processing verification
     List<SqlRow> after = _server.createSqlQuery("select * from metadata_relationship_reportsto where destination='urn:li:foo:123'").findList();
@@ -121,7 +123,7 @@ public void testAddRelationshipWithRemoveAllEdgesFromSourceToDestination() throw
     List<SqlRow> before = _server.createSqlQuery("select * from metadata_relationship_pairswith").findList();
     assertEquals(before.size(), 3);
 
-    _localRelationshipWriterDAO.processLocalRelationshipUpdates(updates);
+    _localRelationshipWriterDAO.processLocalRelationshipUpdates(FooUrn.createFromString("urn:li:foo:123"), updates);
 
     // After processing verification
     List<SqlRow> all = _server.createSqlQuery("select * from metadata_relationship_pairswith").findList();
@@ -158,7 +160,7 @@ public void testAddRelationshipWithRemoveAllEdgesFromSource() throws URISyntaxEx
     List<SqlRow> before = _server.createSqlQuery("select * from metadata_relationship_versionof").findList();
     assertEquals(before.size(), 3);
 
-    _localRelationshipWriterDAO.processLocalRelationshipUpdates(updates);
+    _localRelationshipWriterDAO.processLocalRelationshipUpdates(FooUrn.createFromString("urn:li:foo:123"), updates);
 
     // After processing verification
     List<SqlRow> all = _server.createSqlQuery("select * from metadata_relationship_versionof").findList();
@@ -179,6 +181,47 @@ public void testAddRelationshipWithRemoveAllEdgesFromSource() throws URISyntaxEx
     _server.execute(Ebean.createSqlUpdate("truncate metadata_relationship_versionof"));
   }
 
+
+  @Test
+  public void testClearRelationshipsByEntityUrn() throws URISyntaxException {
+    _server.execute(Ebean.createSqlUpdate(insertRelationships("metadata_relationship_pairswith", "urn:li:bar:123",
+        "bar", "urn:li:foo:123", "foo")));
+
+    _server.execute(Ebean.createSqlUpdate(insertRelationships("metadata_relationship_pairswith", "urn:li:bar:123",
+        "bar", "urn:li:foo:456", "foo")));
+
+    BarUrn barUrn = BarUrn.createFromString("urn:li:bar:123");
+    FooUrn fooUrn = FooUrn.createFromString("urn:li:foo:123");
+
+    // Before processing
+    List<SqlRow> before = _server.createSqlQuery("select * from metadata_relationship_pairswith where deleted_ts is null").findList();
+    assertEquals(before.size(), 2);
+
+    _localRelationshipWriterDAO.clearRelationshipsByEntity(barUrn, new Class[]{PairsWith.class},
+        BaseGraphWriterDAO.RemovalOption.REMOVE_ALL_EDGES_FROM_SOURCE);
+
+    // After processing verification
+    List<SqlRow> all = _server.createSqlQuery("select * from metadata_relationship_pairswith where deleted_ts is null").findList();
+    assertEquals(all.size(), 0); // Total number of edges is 0
+
+
+    _server.execute(Ebean.createSqlUpdate(insertRelationships("metadata_relationship_pairswith", "urn:li:bar:123",
+        "bar", "urn:li:foo:123", "foo")));
+
+    _server.execute(Ebean.createSqlUpdate(insertRelationships("metadata_relationship_pairswith", "urn:li:bar:123",
+        "bar", "urn:li:foo:456", "foo")));
+
+    _localRelationshipWriterDAO.clearRelationshipsByEntity(fooUrn, new Class[]{PairsWith.class},
+        BaseGraphWriterDAO.RemovalOption.REMOVE_ALL_EDGES_TO_DESTINATION);
+
+    // After processing verification
+    all = _server.createSqlQuery("select * from metadata_relationship_pairswith where deleted_ts is null").findList();
+    assertEquals(all.size(), 1); // Total number of edges is 1
+
+    // Clean up
+    _server.execute(Ebean.createSqlUpdate("truncate metadata_relationship_pairswith"));
+  }
+
   private String insertRelationships(String table, String sourceUrn, String sourceType, String destinationUrn, String destinationType) {
     String insertTemplate = "INSERT INTO %s (metadata, source, source_type, destination, destination_type, lastmodifiedon, lastmodifiedby)"
         + " VALUES ('{\"metadata\": true}', '%s', '%s', '%s', '%s', '1970-01-01 00:00:01', 'unknown')";

dao-impl/ebean-dao/src/test/java/com/linkedin/metadata/dao/localrelationship/builder/BelongsToLocalRelationshipBuilder.java

@@ -24,8 +24,9 @@ public <URN extends Urn> List<LocalRelationshipUpdates> buildRelationships(@Nonn
       belongsToRelationships.add(new BelongsTo().setSource(barUrn).setDestination(urn));
     }
 
-    LocalRelationshipUpdates localRelationshipUpdates = new LocalRelationshipUpdates(belongsToRelationships,
-        BaseGraphWriterDAO.RemovalOption.REMOVE_ALL_EDGES_TO_DESTINATION);
+    LocalRelationshipUpdates localRelationshipUpdates =
+        new LocalRelationshipUpdates(belongsToRelationships, new Class[]{BelongsTo.class},
+            BaseGraphWriterDAO.RemovalOption.REMOVE_ALL_EDGES_TO_DESTINATION);
 
     return Collections.singletonList(localRelationshipUpdates);
   }

dao-impl/ebean-dao/src/test/java/com/linkedin/metadata/dao/localrelationship/builder/PairsWithLocalRelationshipBuilder.java

@@ -26,6 +26,7 @@ public <URN extends Urn> List<LocalRelationshipUpdates> buildRelationships(@Nonn
     }
 
     LocalRelationshipUpdates localRelationshipUpdates = new LocalRelationshipUpdates(pairsWithRelationships,
+        new Class[]{PairsWith.class},
         BaseGraphWriterDAO.RemovalOption.REMOVE_ALL_EDGES_FROM_SOURCE_TO_DESTINATION);
 
     return Collections.singletonList(localRelationshipUpdates);

dao-impl/ebean-dao/src/test/java/com/linkedin/metadata/dao/localrelationship/builder/ReportsToLocalRelationshipBuilder.java

@@ -26,7 +26,8 @@ public <URN extends Urn> List<LocalRelationshipUpdates> buildRelationships(@Nonn
     }
 
     LocalRelationshipUpdates localRelationshipUpdates =
-        new LocalRelationshipUpdates(reportsToRelationships, BaseGraphWriterDAO.RemovalOption.REMOVE_NONE);
+        new LocalRelationshipUpdates(reportsToRelationships, new Class[]{ReportsTo.class},
+            BaseGraphWriterDAO.RemovalOption.REMOVE_NONE);
 
     return Collections.singletonList(localRelationshipUpdates);
   }

dao-impl/ebean-dao/src/test/java/com/linkedin/metadata/dao/localrelationship/builder/VersionOfLocalRelationshipBuilder.java

@@ -26,6 +26,7 @@ public <URN extends Urn> List<LocalRelationshipUpdates> buildRelationships(@Nonn
     }
 
     LocalRelationshipUpdates localRelationshipUpdates = new LocalRelationshipUpdates(versionOfRelationships,
+        new Class[]{VersionOf.class},
         BaseGraphWriterDAO.RemovalOption.REMOVE_ALL_EDGES_FROM_SOURCE);
 
     return Collections.singletonList(localRelationshipUpdates);

restli-resources/src/test/java/com/linkedin/metadata/restli/BaseEntityResourceTest.java

@@ -776,7 +776,7 @@ public void testBackfillRelationshipTables() {
     BelongsTo belongsTo = new BelongsTo().setSource(fooUrn).setDestination(barUrn);
     List<BelongsTo> belongsTos = Collections.singletonList(belongsTo);
 
-    LocalRelationshipUpdates updates = new LocalRelationshipUpdates(belongsTos,
+    LocalRelationshipUpdates updates = new LocalRelationshipUpdates(belongsTos, new Class[]{BelongsTo.class},
         BaseGraphWriterDAO.RemovalOption.REMOVE_ALL_EDGES_FROM_SOURCE);
     List<LocalRelationshipUpdates> relationships = Collections.singletonList(updates);