bound relationship type in LocalRelationshipUpdates (#343)

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

@@ -16,9 +16,9 @@ public abstract class BaseLocalRelationshipBuilder<ASPECT extends RecordTemplate
   private final Class<ASPECT> _aspectClass;
 
   @Value
-  public static class LocalRelationshipUpdates {
-    List<? extends RecordTemplate> relationships;
-    Class<? extends RecordTemplate>[] relationshipClasses;
+  public static class LocalRelationshipUpdates<RELATIONSHIP extends RecordTemplate> {
+    List<RELATIONSHIP> relationships;
+    Class<RELATIONSHIP> relationshipClass;
     BaseGraphWriterDAO.RemovalOption removalOption;
   }
 

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

@@ -51,7 +51,7 @@ public <ASPECT extends RecordTemplate> void processLocalRelationshipUpdates(@Non
       @Nonnull List<LocalRelationshipUpdates> relationshipUpdates) {
     for (LocalRelationshipUpdates relationshipUpdate : relationshipUpdates) {
       if (relationshipUpdate.getRelationships().isEmpty()) {
-        clearRelationshipsByEntity(urn, relationshipUpdate.getRelationshipClasses(),
+        clearRelationshipsByEntity(urn, relationshipUpdate.getRelationshipClass(),
             relationshipUpdate.getRemovalOption());
       } else {
         addRelationships(relationshipUpdate.getRelationships(), relationshipUpdate.getRemovalOption());
@@ -62,32 +62,26 @@ public <ASPECT extends RecordTemplate> void processLocalRelationshipUpdates(@Non
   /**
    * 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
+   * @param relationshipClass relationship that needs to be cleared
    */
   public void clearRelationshipsByEntity(@Nonnull Urn urn,
-      @Nonnull Class<? extends RecordTemplate>[] relationshipClasses, @Nonnull RemovalOption removalOption) {
+      @Nonnull Class<? extends RecordTemplate> relationshipClass, @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();
+    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();
   }
 
   @Override

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

@@ -2713,7 +2713,7 @@ public void testBackfillLocalRelationshipsFromEntityTables() throws URISyntaxExc
       BarUrnArray sources = new BarUrnArray();
       for (int i = 0; i < results.size(); i++) {
         try {
-          RecordTemplate relationship = relationshipUpdates.get(0).getRelationships().get(i);
+          RecordTemplate relationship = (RecordTemplate) relationshipUpdates.get(0).getRelationships().get(i);
           Urn source = (Urn) relationship.getClass().getMethod("getSource").invoke(relationship);
           Urn dest = (Urn) relationship.getClass().getMethod("getDestination").invoke(relationship);
           assertEquals(dest.toString(), "urn:li:foo:1");

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

@@ -197,7 +197,7 @@ public void testClearRelationshipsByEntityUrn() throws URISyntaxException {
     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},
+    _localRelationshipWriterDAO.clearRelationshipsByEntity(barUrn, PairsWith.class,
         BaseGraphWriterDAO.RemovalOption.REMOVE_ALL_EDGES_FROM_SOURCE);
 
     // After processing verification
@@ -211,7 +211,7 @@ public void testClearRelationshipsByEntityUrn() throws URISyntaxException {
     _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},
+    _localRelationshipWriterDAO.clearRelationshipsByEntity(fooUrn, PairsWith.class,
         BaseGraphWriterDAO.RemovalOption.REMOVE_ALL_EDGES_TO_DESTINATION);
 
     // After processing verification

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

@@ -25,7 +25,7 @@ public <URN extends Urn> List<LocalRelationshipUpdates> buildRelationships(@Nonn
     }
 
     LocalRelationshipUpdates localRelationshipUpdates =
-        new LocalRelationshipUpdates(belongsToRelationships, new Class[]{BelongsTo.class},
+        new LocalRelationshipUpdates(belongsToRelationships, 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,7 +26,7 @@ public <URN extends Urn> List<LocalRelationshipUpdates> buildRelationships(@Nonn
     }
 
     LocalRelationshipUpdates localRelationshipUpdates = new LocalRelationshipUpdates(pairsWithRelationships,
-        new Class[]{PairsWith.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,7 @@ public <URN extends Urn> List<LocalRelationshipUpdates> buildRelationships(@Nonn
     }
 
     LocalRelationshipUpdates localRelationshipUpdates =
-        new LocalRelationshipUpdates(reportsToRelationships, new Class[]{ReportsTo.class},
+        new LocalRelationshipUpdates(reportsToRelationships, 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,7 +26,7 @@ public <URN extends Urn> List<LocalRelationshipUpdates> buildRelationships(@Nonn
     }
 
     LocalRelationshipUpdates localRelationshipUpdates = new LocalRelationshipUpdates(versionOfRelationships,
-        new Class[]{VersionOf.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, new Class[]{BelongsTo.class},
+    LocalRelationshipUpdates updates = new LocalRelationshipUpdates(belongsTos, BelongsTo.class,
         BaseGraphWriterDAO.RemovalOption.REMOVE_ALL_EDGES_FROM_SOURCE);
     List<LocalRelationshipUpdates> relationships = Collections.singletonList(updates);