From 2e11f80612c96cf8f683b91be3c08a8279b102fd Mon Sep 17 00:00:00 2001
From: Yanbo Zhao <ybzhao.p3k@gmail.com>
Date: Thu, 7 Nov 2024 12:34:36 -0800
Subject: [PATCH] use util

---
 .../linkedin/metadata/dao/EbeanLocalRelationshipWriterDAO.java  | 2 +-
 .../test/java/com/linkedin/metadata/dao/EbeanLocalDAOTest.java  | 1 -
 2 files changed, 1 insertion(+), 2 deletions(-)

diff --git a/dao-impl/ebean-dao/src/main/java/com/linkedin/metadata/dao/EbeanLocalRelationshipWriterDAO.java b/dao-impl/ebean-dao/src/main/java/com/linkedin/metadata/dao/EbeanLocalRelationshipWriterDAO.java
index 6fa3a06a0..5df4bebdb 100644
--- a/dao-impl/ebean-dao/src/main/java/com/linkedin/metadata/dao/EbeanLocalRelationshipWriterDAO.java
+++ b/dao-impl/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
index d207d50e5..3a2ba7edf 100644
--- 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);