fix(relationship): ignore 'pegasus.' prefix when comparing aspect FQCN (#557)

Co-authored-by: Justin Donn <jdonn@jdonn-mn3766.linkedin.biz>

Author: jsdonn

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

@@ -218,7 +218,15 @@ private <ASPECT extends RecordTemplate> void removeRelationshipsBySource(@Nonnul
     SqlUpdate deletionSQL = _server.createSqlUpdate(SQLStatementUtils.deleteLocalRelationshipSQL(tableName, _useAspectColumnForRelationshipRemoval));
     deletionSQL.setParameter(CommonColumnName.SOURCE, source.toString());
     if (_useAspectColumnForRelationshipRemoval) {
-      deletionSQL.setParameter(CommonColumnName.ASPECT, aspectClass.getCanonicalName());
+      // treat "pegasus.com.linkedin..." and "com.linkedin..." as equivalent
+      String aspectClassFQCN = aspectClass.getCanonicalName();
+      String pegasusPrefix = "pegasus.";
+      // normalize the aspect FQCN first by removing any 'pegasus.' prefix
+      if (aspectClassFQCN.startsWith(pegasusPrefix)) {
+        aspectClassFQCN = aspectClassFQCN.substring(pegasusPrefix.length());
+      }
+      deletionSQL.setParameter(CommonColumnName.ASPECT, aspectClassFQCN); // WHERE aspect = "com.linkedin..."
+      deletionSQL.setParameter("pegasus_" + CommonColumnName.ASPECT, pegasusPrefix + aspectClassFQCN); // OR aspect = "pegasus.com.linkedin..."
     }
       batchCount = 0;
       while (batchCount < MAX_BATCHES) {

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

@@ -141,7 +141,7 @@ public class SQLStatementUtils {
       + "WHERE source = :source AND deleted_ts IS NULL";
 
   private static final String DELETE_BY_SOURCE_AND_ASPECT = "UPDATE %s SET deleted_ts=NOW() "
-      + "WHERE source = :source AND aspect = :aspect AND deleted_ts IS NULL";
+      + "WHERE source = :source AND (aspect = :aspect OR aspect = :pegasus_aspect) AND deleted_ts IS NULL";
 
   /**
    *  Filter query has pagination params in the existing APIs. To accommodate this, we use subquery to include total result counts in the query response.

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

@@ -318,6 +318,37 @@ public void testRemoveRelationshipsDifferentAspect() throws URISyntaxException {
     _server.execute(Ebean.createSqlUpdate("truncate metadata_relationship_pairswith"));
   }
 
+  @Test
+  public void testRemoveRelationshipsSameAspectDifferentNamespace() throws URISyntaxException {
+    if (!_useAspectColumnForRelationshipRemoval) {
+      return; // this test doesn't apply to this case
+    }
+    _localRelationshipWriterDAO.setUseAspectColumnForRelationshipRemoval(_useAspectColumnForRelationshipRemoval);
+
+    BarUrn barUrn = BarUrn.createFromString("urn:li:bar:123");
+    FooUrn fooUrn123 = FooUrn.createFromString("urn:li:foo:123");
+    FooUrn fooUrn456 = FooUrn.createFromString("urn:li:foo:456");
+    _server.execute(Ebean.createSqlUpdate(insertRelationships("metadata_relationship_pairswith", barUrn.toString(),
+        "bar", fooUrn123.toString(), "foo", AspectFooBar.class.getCanonicalName())));
+
+    _server.execute(Ebean.createSqlUpdate(insertRelationships("metadata_relationship_pairswith", barUrn.toString(),
+        "bar", fooUrn456.toString(), "foo", "pegasus." + AspectFooBar.class.getCanonicalName())));
+
+    // Before processing
+    List<SqlRow> before = _server.createSqlQuery("select * from metadata_relationship_pairswith where deleted_ts is null").findList();
+    assertEquals(before.size(), 2);
+
+    PairsWith pairsWith = new PairsWith().setSource(barUrn).setDestination(fooUrn123);
+    _localRelationshipWriterDAO.removeRelationships(barUrn, AspectFooBar.class, Collections.singletonList(pairsWith));
+
+    // After processing verification - both relationships should have been deleted.
+    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
+
+    // Clean up
+    _server.execute(Ebean.createSqlUpdate("truncate metadata_relationship_pairswith"));
+  }
+
   private String insertRelationships(String table, String sourceUrn, String sourceType, String destinationUrn, String destinationType, String aspect) {
     String insertWithAspectTemplate = "INSERT INTO %s (metadata, source, source_type, destination, destination_type, lastmodifiedon, lastmodifiedby, aspect)"
         + " VALUES ('{\"metadata\": true}', '%s', '%s', '%s', '%s', CURRENT_TIMESTAMP, 'unknown', '%s')";