use util

linkedin · Nov 7, 2024 · 2e11f80 · 2e11f80
1 parent 6eaf478
commit 2e11f80
Show file tree

Hide file tree

Showing 2 changed files with 1 addition and 2 deletions.
diff --git a/...pl/ebean-dao/src/main/java/com/linkedin/metadata/dao/EbeanLocalRelationshipWriterDAO.java b/...pl/ebean-dao/src/main/java/com/linkedin/metadata/dao/EbeanLocalRelationshipWriterDAO.java
@@ -127,7 +127,7 @@ public <RELATIONSHIP extends RecordTemplate> void removeRelationshipsV2(@Nonnull
     for (RELATIONSHIP relationship : relationships) {
       _server.createSqlUpdate(SQLStatementUtils.deleteLocalRelationshipSQL(SQLSchemaUtils.getRelationshipTableName(relationship),
               RemovalOption.REMOVE_ALL_EDGES_FROM_SOURCE_TO_DESTINATION))
-          .setParameter(CommonColumnName.SOURCE, sourceUrn != null ? sourceUrn.toString() : getSourceUrnFromRelationship(relationship).toString())
+          .setParameter(CommonColumnName.SOURCE, GraphUtils.getSourceUrnBasedOnRelationshipVersion(relationship, sourceUrn))
           .setParameter(CommonColumnName.DESTINATION, getDestinationUrnFromRelationship(relationship).toString())
           .execute();
     }

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
@@ -3133,7 +3133,6 @@ public void testAddWithLocalRelationshipBuilder() throws URISyntaxException {
     assertEquals(aspects.size(), 1);
   }
 
-  // TODO: fix this
   @Test
   public void testAddRelationshipsFromAspect() throws URISyntaxException {
     EbeanLocalDAO<EntityAspectUnion, FooUrn> fooDao = createDao(FooUrn.class);