make unit test use v2 + fix removeRelationships logic

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

@@ -653,7 +653,7 @@ protected <ASPECT extends RecordTemplate> long saveLatest(@Nonnull URN urn, @Non
       List<RecordTemplate> relationships = extractRelationshipsFromAspect(oldValue).stream()
           .flatMap(List::stream)
           .collect(Collectors.toList());
-      _localRelationshipWriterDAO.removeRelationships(relationships);
+      _localRelationshipWriterDAO.removeRelationshipsV2(relationships, urn);
     // Otherwise, add any local relationships that are derived from the aspect.
     } else {
       addRelationshipsIfAny(urn, newValue, aspectClass, isTestMode);

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

@@ -120,10 +120,14 @@ public <RELATIONSHIP extends RecordTemplate> void addRelationships(@Nonnull List
 
   @Override
   public <RELATIONSHIP extends RecordTemplate> void removeRelationships(@Nonnull List<RELATIONSHIP> relationships) {
+    removeRelationshipsV2(relationships, null);
+  }
+
+  public <RELATIONSHIP extends RecordTemplate> void removeRelationshipsV2(@Nonnull List<RELATIONSHIP> relationships, @Nullable Urn sourceUrn) {
     for (RELATIONSHIP relationship : relationships) {
       _server.createSqlUpdate(SQLStatementUtils.deleteLocalRelationshipSQL(SQLSchemaUtils.getRelationshipTableName(relationship),
               RemovalOption.REMOVE_ALL_EDGES_FROM_SOURCE_TO_DESTINATION))
-          .setParameter(CommonColumnName.SOURCE, getSourceUrnFromRelationship(relationship).toString())
+          .setParameter(CommonColumnName.SOURCE, sourceUrn != null ? sourceUrn.toString() : getSourceUrnFromRelationship(relationship).toString())
           .setParameter(CommonColumnName.DESTINATION, getDestinationUrnFromRelationship(relationship).toString())
           .execute();
     }

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

@@ -64,7 +64,8 @@
 import com.linkedin.testing.MixedRecord;
 import com.linkedin.testing.localrelationship.AspectFooBar;
 import com.linkedin.testing.localrelationship.BelongsTo;
-import com.linkedin.testing.localrelationship.BelongsToArray;
+import com.linkedin.testing.localrelationship.BelongsToV2;
+import com.linkedin.testing.localrelationship.BelongsToV2Array;
 import com.linkedin.testing.localrelationship.ReportsTo;
 import com.linkedin.testing.localrelationship.ReportsToArray;
 import com.linkedin.testing.urn.BarUrn;
@@ -2548,12 +2549,12 @@ public void testRemoveRelationshipsDuringAspectSoftDeletion() throws URISyntaxEx
     // add an aspect (AspectFooBar) which includes BelongsTo relationships and ReportsTo relationships
     FooUrn fooUrn = makeFooUrn(1);
     BarUrn barUrn1 = BarUrn.createFromString("urn:li:bar:1");
-    BelongsTo belongsTo1 = new BelongsTo().setSource(fooUrn).setDestination(barUrn1);
+    BelongsToV2 belongsTo1 = new BelongsToV2().setDestination(BelongsToV2.Destination.create(barUrn1.toString()));
     BarUrn barUrn2 = BarUrn.createFromString("urn:li:bar:2");
-    BelongsTo belongsTo2 = new BelongsTo().setSource(fooUrn).setDestination(barUrn2);
+    BelongsToV2 belongsTo2 = new BelongsToV2().setDestination(BelongsToV2.Destination.create(barUrn2.toString()));
     BarUrn barUrn3 = BarUrn.createFromString("urn:li:bar:3");
-    BelongsTo belongsTo3 = new BelongsTo().setSource(fooUrn).setDestination(barUrn3);
-    BelongsToArray belongsToArray = new BelongsToArray(belongsTo1, belongsTo2, belongsTo3);
+    BelongsToV2 belongsTo3 = new BelongsToV2().setDestination(BelongsToV2.Destination.create(barUrn3.toString()));
+    BelongsToV2Array belongsToArray = new BelongsToV2Array(belongsTo1, belongsTo2, belongsTo3);
     ReportsTo reportsTo = new ReportsTo().setSource(fooUrn).setDestination(barUrn1);
     ReportsToArray reportsToArray = new ReportsToArray(reportsTo);
     AspectFooBar aspectFooBar = new AspectFooBar()
@@ -2569,9 +2570,9 @@ public void testRemoveRelationshipsDuringAspectSoftDeletion() throws URISyntaxEx
     EbeanLocalRelationshipQueryDAO ebeanLocalRelationshipQueryDAO = new EbeanLocalRelationshipQueryDAO(_server);
     ebeanLocalRelationshipQueryDAO.setSchemaConfig(_schemaConfig);
 
-    List<BelongsTo> resultBelongsTos =
+    List<BelongsToV2> resultBelongsTos =
         ebeanLocalRelationshipQueryDAO.findRelationships(FooSnapshot.class, EMPTY_FILTER, BarSnapshot.class,
-            EMPTY_FILTER, BelongsTo.class, OUTGOING_FILTER, 0, 10);
+            EMPTY_FILTER, BelongsToV2.class, OUTGOING_FILTER, 0, 10);
 
     assertEquals(resultBelongsTos.size(), 3);
 
@@ -2591,7 +2592,7 @@ public void testRemoveRelationshipsDuringAspectSoftDeletion() throws URISyntaxEx
 
     // check that the belongsTo relationships 1, 2, & 3 were soft deleted
     resultBelongsTos = ebeanLocalRelationshipQueryDAO.findRelationships(FooSnapshot.class, EMPTY_FILTER, BarSnapshot.class,
-            EMPTY_FILTER, BelongsTo.class, OUTGOING_FILTER, 0, 10);
+            EMPTY_FILTER, BelongsToV2.class, OUTGOING_FILTER, 0, 10);
 
     // check that the reportsTo relationship was soft deleted
     resultReportsTos =
@@ -3132,6 +3133,7 @@ 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);
@@ -3141,12 +3143,12 @@ public void testAddRelationshipsFromAspect() throws URISyntaxException {
 
     FooUrn fooUrn = makeFooUrn(1);
     BarUrn barUrn1 = BarUrn.createFromString("urn:li:bar:1");
-    BelongsTo belongsTo1 = new BelongsTo().setSource(fooUrn).setDestination(barUrn1);
+    BelongsToV2 belongsTo1 = new BelongsToV2().setDestination(BelongsToV2.Destination.create(barUrn1.toString()));
     BarUrn barUrn2 = BarUrn.createFromString("urn:li:bar:2");
-    BelongsTo belongsTo2 = new BelongsTo().setSource(fooUrn).setDestination(barUrn2);
+    BelongsToV2 belongsTo2 = new BelongsToV2().setDestination(BelongsToV2.Destination.create(barUrn2.toString()));
     BarUrn barUrn3 = BarUrn.createFromString("urn:li:bar:3");
-    BelongsTo belongsTo3 = new BelongsTo().setSource(fooUrn).setDestination(barUrn3);
-    BelongsToArray belongsToArray = new BelongsToArray(belongsTo1, belongsTo2, belongsTo3);
+    BelongsToV2 belongsTo3 = new BelongsToV2().setDestination(BelongsToV2.Destination.create(barUrn3.toString()));
+    BelongsToV2Array belongsToArray = new BelongsToV2Array(belongsTo1, belongsTo2, belongsTo3);
     AspectFooBar aspectFooBar = new AspectFooBar().setBars(new BarUrnArray(barUrn1, barUrn2, barUrn3)).setBelongsTos(belongsToArray);
     AuditStamp auditStamp = makeAuditStamp("foo", System.currentTimeMillis());
 
@@ -3159,9 +3161,9 @@ public void testAddRelationshipsFromAspect() throws URISyntaxException {
     EbeanLocalRelationshipQueryDAO ebeanLocalRelationshipQueryDAO = new EbeanLocalRelationshipQueryDAO(_server);
     ebeanLocalRelationshipQueryDAO.setSchemaConfig(_schemaConfig);
 
-    List<BelongsTo> relationships =
+    List<BelongsToV2> relationships =
         ebeanLocalRelationshipQueryDAO.findRelationships(FooSnapshot.class, EMPTY_FILTER, BarSnapshot.class,
-            EMPTY_FILTER, BelongsTo.class, OUTGOING_FILTER, 0, 10);
+            EMPTY_FILTER, BelongsToV2.class, OUTGOING_FILTER, 0, 10);
 
     AspectKey<FooUrn, AspectFooBar> key = new AspectKey<>(AspectFooBar.class, fooUrn, 0L);
     List<EbeanMetadataAspect> aspects = fooDao.batchGetHelper(Collections.singletonList(key), 1, 0);

dao-impl/ebean-dao/src/test/resources/ebean-local-dao-create-all-with-non-dollar-virtual-column-names.sql

@@ -6,6 +6,7 @@ DROP TABLE IF EXISTS metadata_aspect;
 DROP TABLE IF EXISTS metadata_id;
 DROP TABLE IF EXISTS metadata_index;
 DROP TABLE IF EXISTS metadata_relationship_belongsto;
+DROP TABLE IF EXISTS metadata_relationship_belongstov2;
 
 -- initialize foo entity table
 CREATE TABLE IF NOT EXISTS metadata_entity_foo (
@@ -56,6 +57,19 @@ CREATE TABLE IF NOT EXISTS metadata_relationship_belongsto (
     PRIMARY KEY (id)
 );
 
+CREATE TABLE IF NOT EXISTS metadata_relationship_belongstov2 (
+    id BIGINT NOT NULL AUTO_INCREMENT,
+    metadata LONGTEXT NOT NULL,
+    source VARCHAR(1000) NOT NULL,
+    source_type VARCHAR(100) NOT NULL,
+    destination VARCHAR(1000) NOT NULL,
+    destination_type VARCHAR(100) NOT NULL,
+    lastmodifiedon TIMESTAMP NOT NULL,
+    lastmodifiedby VARCHAR(255) NOT NULL,
+    deleted_ts DATETIME(6) DEFAULT NULL,
+    PRIMARY KEY (id)
+);
+
 CREATE TABLE IF NOT EXISTS metadata_relationship_reportsto (
     id BIGINT NOT NULL AUTO_INCREMENT,
     metadata LONGTEXT NOT NULL,

dao-impl/ebean-dao/src/test/resources/ebean-local-dao-create-all.sql

@@ -6,6 +6,7 @@ DROP TABLE IF EXISTS metadata_aspect;
 DROP TABLE IF EXISTS metadata_id;
 DROP TABLE IF EXISTS metadata_index;
 DROP TABLE IF EXISTS metadata_relationship_belongsto;
+DROP TABLE IF EXISTS metadata_relationship_belongstov2;
 
 -- initialize foo entity table
 CREATE TABLE IF NOT EXISTS metadata_entity_foo (
@@ -56,6 +57,19 @@ CREATE TABLE IF NOT EXISTS metadata_relationship_belongsto (
     PRIMARY KEY (id)
 );
 
+CREATE TABLE IF NOT EXISTS metadata_relationship_belongstov2 (
+    id BIGINT NOT NULL AUTO_INCREMENT,
+    metadata LONGTEXT NOT NULL,
+    source VARCHAR(1000) NOT NULL,
+    source_type VARCHAR(100) NOT NULL,
+    destination VARCHAR(1000) NOT NULL,
+    destination_type VARCHAR(100) NOT NULL,
+    lastmodifiedon TIMESTAMP NOT NULL,
+    lastmodifiedby VARCHAR(255) NOT NULL,
+    deleted_ts DATETIME(6) DEFAULT NULL,
+    PRIMARY KEY (id)
+);
+
 CREATE TABLE IF NOT EXISTS metadata_relationship_reportsto (
     id BIGINT NOT NULL AUTO_INCREMENT,
     metadata LONGTEXT NOT NULL,

dao-impl/ebean-dao/src/test/resources/gma-create-all.sql

@@ -44,6 +44,19 @@ CREATE TABLE IF NOT EXISTS metadata_relationship_belongsto (
     PRIMARY KEY (id)
 );
 
+CREATE TABLE IF NOT EXISTS metadata_relationship_belongstov2 (
+    id BIGINT NOT NULL AUTO_INCREMENT,
+    metadata JSON NOT NULL,
+    source VARCHAR(500) NOT NULL,
+    source_type VARCHAR(100) NOT NULL,
+    destination VARCHAR(500) NOT NULL,
+    destination_type VARCHAR(100) NOT NULL,
+    lastmodifiedon DATETIME(6) NOT NULL,
+    lastmodifiedby VARCHAR(255) NOT NULL,
+    deleted_ts DATETIME(6) DEFAULT NULL,
+    PRIMARY KEY (id)
+);
+
 CREATE TABLE IF NOT EXISTS metadata_relationship_reportsto (
     id BIGINT NOT NULL AUTO_INCREMENT,
     metadata JSON NOT NULL,

dao-impl/ebean-dao/src/test/resources/gma-drop-all.sql

@@ -6,4 +6,6 @@ drop table if exists metadata_index;
 
 drop table if exists metadata_relationship_belongsto;
 
+drop table if exists metadata_relationship_belongstov2;
+
 drop table if exists metadata_relationship_reportsto;

testing/test-models/src/main/pegasus/com/linkedin/testing/localrelationship/AspectFooBar.pdl

@@ -5,6 +5,6 @@ import com.linkedin.testing.BarUrn
 @gma.aspect.column.name = "aspectfoobar"
 record AspectFooBar {
   bars: array[BarUrn]
-  belongsTos: optional array[BelongsTo]
+  belongsTos: optional array[BelongsToV2]
   reportsTos: optional array[ReportsTo]
 }

testing/test-models/src/main/pegasus/com/linkedin/testing/localrelationship/AspectFooBarBaz.pdl

@@ -5,5 +5,5 @@ import com.linkedin.testing.BarUrn
 @gma.aspect.column.name = "aspectfoobarbaz"
 record AspectFooBarBaz {
   bars: array[BarUrn]
-  belongsTos: optional array[BelongsTo]
+  belongsTos: optional array[BelongsToV2]
 }

testing/test-models/src/main/pegasus/com/linkedin/testing/localrelationship/BelongsToV2.pdl

@@ -0,0 +1,8 @@
+namespace com.linkedin.testing.localrelationship
+
+@gma.model = "RELATIONSHIP"
+record BelongsToV2 {
+  destination: union [
+    string
+  ]
+}