From 827d2a0f16fd161ce82776c3d07e27c54aa573b2 Mon Sep 17 00:00:00 2001 From: Yang Yang Date: Thu, 18 Jan 2024 15:13:39 -0800 Subject: [PATCH] (feat) support clear relationship by source entity urn (#342) * support clear relationship by source entity urn * including RemovalOption support --- .../builder/BaseLocalRelationshipBuilder.java | 1 + .../metadata/dao/EbeanLocalAccess.java | 9 ++-- .../dao/EbeanLocalRelationshipWriterDAO.java | 41 +++++++++++++-- .../EbeanLocalRelationshipWriterDAOTest.java | 51 +++++++++++++++++-- .../BelongsToLocalRelationshipBuilder.java | 5 +- .../PairsWithLocalRelationshipBuilder.java | 1 + .../ReportsToLocalRelationshipBuilder.java | 3 +- .../VersionOfLocalRelationshipBuilder.java | 1 + .../restli/BaseEntityResourceTest.java | 2 +- 9 files changed, 97 insertions(+), 17 deletions(-) diff --git a/dao-api/src/main/java/com/linkedin/metadata/dao/builder/BaseLocalRelationshipBuilder.java b/dao-api/src/main/java/com/linkedin/metadata/dao/builder/BaseLocalRelationshipBuilder.java index a4444453e..8eb2fe0b6 100644 --- a/dao-api/src/main/java/com/linkedin/metadata/dao/builder/BaseLocalRelationshipBuilder.java +++ b/dao-api/src/main/java/com/linkedin/metadata/dao/builder/BaseLocalRelationshipBuilder.java @@ -18,6 +18,7 @@ public abstract class BaseLocalRelationshipBuilder relationships; + Class[] relationshipClasses; BaseGraphWriterDAO.RemovalOption removalOption; } diff --git a/dao-impl/ebean-dao/src/main/java/com/linkedin/metadata/dao/EbeanLocalAccess.java b/dao-impl/ebean-dao/src/main/java/com/linkedin/metadata/dao/EbeanLocalAccess.java index 2bcef3f03..646b06fc0 100644 --- a/dao-impl/ebean-dao/src/main/java/com/linkedin/metadata/dao/EbeanLocalAccess.java +++ b/dao-impl/ebean-dao/src/main/java/com/linkedin/metadata/dao/EbeanLocalAccess.java @@ -169,17 +169,14 @@ public int addWithOptimisticLocking( } @Override - public List - addRelationships(@Nonnull URN urn, @Nonnull ASPECT aspect, @Nonnull Class aspectClass) { + public List addRelationships(@Nonnull URN urn, + @Nonnull ASPECT aspect, @Nonnull Class aspectClass) { if (_localRelationshipBuilderRegistry != null && _localRelationshipBuilderRegistry.isRegistered(aspectClass)) { List localRelationshipUpdates = _localRelationshipBuilderRegistry.getLocalRelationshipBuilder(aspect).buildRelationships(urn, aspect); - - _localRelationshipWriterDAO.processLocalRelationshipUpdates(localRelationshipUpdates); - + _localRelationshipWriterDAO.processLocalRelationshipUpdates(urn, localRelationshipUpdates); return localRelationshipUpdates; } - return new ArrayList<>(); } 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 45fb3cb25..68d28b992 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 @@ -47,11 +47,46 @@ public EbeanLocalRelationshipWriterDAO(EbeanServer server) { * @param relationshipUpdates Updates to local relationship tables. */ @Transactional - public void processLocalRelationshipUpdates( + public void processLocalRelationshipUpdates(@Nonnull Urn urn, @Nonnull List relationshipUpdates) { - for (LocalRelationshipUpdates relationshipUpdate : relationshipUpdates) { - addRelationships(relationshipUpdate.getRelationships(), relationshipUpdate.getRemovalOption()); + if (relationshipUpdate.getRelationships().isEmpty()) { + clearRelationshipsByEntity(urn, relationshipUpdate.getRelationshipClasses(), + relationshipUpdate.getRemovalOption()); + } else { + addRelationships(relationshipUpdate.getRelationships(), relationshipUpdate.getRemovalOption()); + } + } + } + + /** + * 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 + */ + public void clearRelationshipsByEntity(@Nonnull Urn urn, + @Nonnull Class[] relationshipClasses, @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 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(); } } diff --git a/dao-impl/ebean-dao/src/test/java/com/linkedin/metadata/dao/localrelationship/EbeanLocalRelationshipWriterDAOTest.java b/dao-impl/ebean-dao/src/test/java/com/linkedin/metadata/dao/localrelationship/EbeanLocalRelationshipWriterDAOTest.java index b6324cdb9..7e484cd89 100644 --- a/dao-impl/ebean-dao/src/test/java/com/linkedin/metadata/dao/localrelationship/EbeanLocalRelationshipWriterDAOTest.java +++ b/dao-impl/ebean-dao/src/test/java/com/linkedin/metadata/dao/localrelationship/EbeanLocalRelationshipWriterDAOTest.java @@ -2,6 +2,7 @@ import com.google.common.io.Resources; import com.linkedin.metadata.dao.EbeanLocalRelationshipWriterDAO; +import com.linkedin.metadata.dao.internal.BaseGraphWriterDAO; import com.linkedin.metadata.dao.localrelationship.builder.BelongsToLocalRelationshipBuilder; import com.linkedin.metadata.dao.localrelationship.builder.PairsWithLocalRelationshipBuilder; import com.linkedin.metadata.dao.localrelationship.builder.ReportsToLocalRelationshipBuilder; @@ -10,6 +11,7 @@ import com.linkedin.metadata.dao.builder.BaseLocalRelationshipBuilder.LocalRelationshipUpdates; import com.linkedin.testing.BarUrnArray; import com.linkedin.testing.localrelationship.AspectFooBar; +import com.linkedin.testing.localrelationship.PairsWith; import com.linkedin.testing.urn.BarUrn; import com.linkedin.testing.urn.FooUrn; import io.ebean.Ebean; @@ -54,7 +56,7 @@ public void testAddRelationshipWithRemoveAllEdgesToDestination() throws URISynta List before = _server.createSqlQuery("select * from metadata_relationship_belongsto where source='urn:li:bar:000'").findList(); assertEquals(before.size(), 1); - _localRelationshipWriterDAO.processLocalRelationshipUpdates(updates); + _localRelationshipWriterDAO.processLocalRelationshipUpdates(FooUrn.createFromString("urn:li:foo:123"), updates); // After processing verification List all = _server.createSqlQuery("select * from metadata_relationship_belongsto").findList(); @@ -89,7 +91,7 @@ public void testAddRelationshipWithRemoveNone() throws URISyntaxException { List before = _server.createSqlQuery("select * from metadata_relationship_reportsto where source='urn:li:bar:000'").findList(); assertEquals(before.size(), 1); - _localRelationshipWriterDAO.processLocalRelationshipUpdates(updates); + _localRelationshipWriterDAO.processLocalRelationshipUpdates(FooUrn.createFromString("urn:li:foo:123"), updates); // After processing verification List after = _server.createSqlQuery("select * from metadata_relationship_reportsto where destination='urn:li:foo:123'").findList(); @@ -121,7 +123,7 @@ public void testAddRelationshipWithRemoveAllEdgesFromSourceToDestination() throw List before = _server.createSqlQuery("select * from metadata_relationship_pairswith").findList(); assertEquals(before.size(), 3); - _localRelationshipWriterDAO.processLocalRelationshipUpdates(updates); + _localRelationshipWriterDAO.processLocalRelationshipUpdates(FooUrn.createFromString("urn:li:foo:123"), updates); // After processing verification List all = _server.createSqlQuery("select * from metadata_relationship_pairswith").findList(); @@ -158,7 +160,7 @@ public void testAddRelationshipWithRemoveAllEdgesFromSource() throws URISyntaxEx List before = _server.createSqlQuery("select * from metadata_relationship_versionof").findList(); assertEquals(before.size(), 3); - _localRelationshipWriterDAO.processLocalRelationshipUpdates(updates); + _localRelationshipWriterDAO.processLocalRelationshipUpdates(FooUrn.createFromString("urn:li:foo:123"), updates); // After processing verification List all = _server.createSqlQuery("select * from metadata_relationship_versionof").findList(); @@ -179,6 +181,47 @@ public void testAddRelationshipWithRemoveAllEdgesFromSource() throws URISyntaxEx _server.execute(Ebean.createSqlUpdate("truncate metadata_relationship_versionof")); } + + @Test + public void testClearRelationshipsByEntityUrn() throws URISyntaxException { + _server.execute(Ebean.createSqlUpdate(insertRelationships("metadata_relationship_pairswith", "urn:li:bar:123", + "bar", "urn:li:foo:123", "foo"))); + + _server.execute(Ebean.createSqlUpdate(insertRelationships("metadata_relationship_pairswith", "urn:li:bar:123", + "bar", "urn:li:foo:456", "foo"))); + + BarUrn barUrn = BarUrn.createFromString("urn:li:bar:123"); + FooUrn fooUrn = FooUrn.createFromString("urn:li:foo:123"); + + // Before processing + List 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}, + BaseGraphWriterDAO.RemovalOption.REMOVE_ALL_EDGES_FROM_SOURCE); + + // After processing verification + List all = _server.createSqlQuery("select * from metadata_relationship_pairswith where deleted_ts is null").findList(); + assertEquals(all.size(), 0); // Total number of edges is 0 + + + _server.execute(Ebean.createSqlUpdate(insertRelationships("metadata_relationship_pairswith", "urn:li:bar:123", + "bar", "urn:li:foo:123", "foo"))); + + _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}, + BaseGraphWriterDAO.RemovalOption.REMOVE_ALL_EDGES_TO_DESTINATION); + + // After processing verification + all = _server.createSqlQuery("select * from metadata_relationship_pairswith where deleted_ts is null").findList(); + assertEquals(all.size(), 1); // Total number of edges is 1 + + // Clean up + _server.execute(Ebean.createSqlUpdate("truncate metadata_relationship_pairswith")); + } + private String insertRelationships(String table, String sourceUrn, String sourceType, String destinationUrn, String destinationType) { String insertTemplate = "INSERT INTO %s (metadata, source, source_type, destination, destination_type, lastmodifiedon, lastmodifiedby)" + " VALUES ('{\"metadata\": true}', '%s', '%s', '%s', '%s', '1970-01-01 00:00:01', 'unknown')"; diff --git a/dao-impl/ebean-dao/src/test/java/com/linkedin/metadata/dao/localrelationship/builder/BelongsToLocalRelationshipBuilder.java b/dao-impl/ebean-dao/src/test/java/com/linkedin/metadata/dao/localrelationship/builder/BelongsToLocalRelationshipBuilder.java index 1aa4bd6fd..c7e465df2 100644 --- a/dao-impl/ebean-dao/src/test/java/com/linkedin/metadata/dao/localrelationship/builder/BelongsToLocalRelationshipBuilder.java +++ b/dao-impl/ebean-dao/src/test/java/com/linkedin/metadata/dao/localrelationship/builder/BelongsToLocalRelationshipBuilder.java @@ -24,8 +24,9 @@ public List buildRelationships(@Nonn belongsToRelationships.add(new BelongsTo().setSource(barUrn).setDestination(urn)); } - LocalRelationshipUpdates localRelationshipUpdates = new LocalRelationshipUpdates(belongsToRelationships, - BaseGraphWriterDAO.RemovalOption.REMOVE_ALL_EDGES_TO_DESTINATION); + LocalRelationshipUpdates localRelationshipUpdates = + new LocalRelationshipUpdates(belongsToRelationships, new Class[]{BelongsTo.class}, + BaseGraphWriterDAO.RemovalOption.REMOVE_ALL_EDGES_TO_DESTINATION); return Collections.singletonList(localRelationshipUpdates); } diff --git a/dao-impl/ebean-dao/src/test/java/com/linkedin/metadata/dao/localrelationship/builder/PairsWithLocalRelationshipBuilder.java b/dao-impl/ebean-dao/src/test/java/com/linkedin/metadata/dao/localrelationship/builder/PairsWithLocalRelationshipBuilder.java index 9ce5a6a66..c06f31d06 100644 --- a/dao-impl/ebean-dao/src/test/java/com/linkedin/metadata/dao/localrelationship/builder/PairsWithLocalRelationshipBuilder.java +++ b/dao-impl/ebean-dao/src/test/java/com/linkedin/metadata/dao/localrelationship/builder/PairsWithLocalRelationshipBuilder.java @@ -26,6 +26,7 @@ public List buildRelationships(@Nonn } LocalRelationshipUpdates localRelationshipUpdates = new LocalRelationshipUpdates(pairsWithRelationships, + new Class[]{PairsWith.class}, BaseGraphWriterDAO.RemovalOption.REMOVE_ALL_EDGES_FROM_SOURCE_TO_DESTINATION); return Collections.singletonList(localRelationshipUpdates); diff --git a/dao-impl/ebean-dao/src/test/java/com/linkedin/metadata/dao/localrelationship/builder/ReportsToLocalRelationshipBuilder.java b/dao-impl/ebean-dao/src/test/java/com/linkedin/metadata/dao/localrelationship/builder/ReportsToLocalRelationshipBuilder.java index 5d3c154f1..595cacf67 100644 --- a/dao-impl/ebean-dao/src/test/java/com/linkedin/metadata/dao/localrelationship/builder/ReportsToLocalRelationshipBuilder.java +++ b/dao-impl/ebean-dao/src/test/java/com/linkedin/metadata/dao/localrelationship/builder/ReportsToLocalRelationshipBuilder.java @@ -26,7 +26,8 @@ public List buildRelationships(@Nonn } LocalRelationshipUpdates localRelationshipUpdates = - new LocalRelationshipUpdates(reportsToRelationships, BaseGraphWriterDAO.RemovalOption.REMOVE_NONE); + new LocalRelationshipUpdates(reportsToRelationships, new Class[]{ReportsTo.class}, + BaseGraphWriterDAO.RemovalOption.REMOVE_NONE); return Collections.singletonList(localRelationshipUpdates); } diff --git a/dao-impl/ebean-dao/src/test/java/com/linkedin/metadata/dao/localrelationship/builder/VersionOfLocalRelationshipBuilder.java b/dao-impl/ebean-dao/src/test/java/com/linkedin/metadata/dao/localrelationship/builder/VersionOfLocalRelationshipBuilder.java index 2c73d1b25..368a7a8cf 100644 --- a/dao-impl/ebean-dao/src/test/java/com/linkedin/metadata/dao/localrelationship/builder/VersionOfLocalRelationshipBuilder.java +++ b/dao-impl/ebean-dao/src/test/java/com/linkedin/metadata/dao/localrelationship/builder/VersionOfLocalRelationshipBuilder.java @@ -26,6 +26,7 @@ public List buildRelationships(@Nonn } LocalRelationshipUpdates localRelationshipUpdates = new LocalRelationshipUpdates(versionOfRelationships, + new Class[]{VersionOf.class}, BaseGraphWriterDAO.RemovalOption.REMOVE_ALL_EDGES_FROM_SOURCE); return Collections.singletonList(localRelationshipUpdates); diff --git a/restli-resources/src/test/java/com/linkedin/metadata/restli/BaseEntityResourceTest.java b/restli-resources/src/test/java/com/linkedin/metadata/restli/BaseEntityResourceTest.java index cac534302..9052498f5 100644 --- a/restli-resources/src/test/java/com/linkedin/metadata/restli/BaseEntityResourceTest.java +++ b/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 belongsTos = Collections.singletonList(belongsTo); - LocalRelationshipUpdates updates = new LocalRelationshipUpdates(belongsTos, + LocalRelationshipUpdates updates = new LocalRelationshipUpdates(belongsTos, new Class[]{BelongsTo.class}, BaseGraphWriterDAO.RemovalOption.REMOVE_ALL_EDGES_FROM_SOURCE); List relationships = Collections.singletonList(updates);