diff --git a/build.gradle b/build.gradle index 235456d31..8efa37917 100644 --- a/build.gradle +++ b/build.gradle @@ -61,8 +61,6 @@ project.ext.externalDependency = [ 'mockito': 'org.mockito:mockito-core:4.11.0', 'mockitoInline': 'org.mockito:mockito-inline:3.11.2', 'mysql': 'mysql:mysql-connector-java:8.0.29', - 'neo4jHarness': 'org.neo4j.test:neo4j-harness:3.4.11', - 'neo4jJavaDriver': 'org.neo4j.driver:neo4j-java-driver:4.0.0', 'parseqTest': 'com.linkedin.parseq:parseq-test-api:5.1.20', 'postgresql': 'org.postgresql:postgresql:42.2.14', 'reflections': 'org.reflections:reflections:0.9.11', 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 2ed6dba68..6266ca971 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 @@ -5,7 +5,8 @@ import com.linkedin.metadata.dao.internal.BaseGraphWriterDAO; import java.util.List; import javax.annotation.Nonnull; -import lombok.Value; +import lombok.AllArgsConstructor; +import lombok.Data; /** @@ -15,7 +16,8 @@ public abstract class BaseLocalRelationshipBuilder _aspectClass; - @Value + @AllArgsConstructor + @Data public static class LocalRelationshipUpdates { List relationships; Class relationshipClass; diff --git a/dao-api/src/main/java/com/linkedin/metadata/dao/internal/BaseGraphWriterDAO.java b/dao-api/src/main/java/com/linkedin/metadata/dao/internal/BaseGraphWriterDAO.java index e0dc9155d..7894801f9 100644 --- a/dao-api/src/main/java/com/linkedin/metadata/dao/internal/BaseGraphWriterDAO.java +++ b/dao-api/src/main/java/com/linkedin/metadata/dao/internal/BaseGraphWriterDAO.java @@ -14,6 +14,7 @@ */ public abstract class BaseGraphWriterDAO { + // These removal options are deprecated; all relationship ingestion uses REMOVE_ALL_EDGES_FROM_SOURCE. public enum RemovalOption { REMOVE_NONE, REMOVE_ALL_EDGES_FROM_SOURCE, diff --git a/dao-api/src/main/java/com/linkedin/metadata/dao/utils/GraphUtils.java b/dao-api/src/main/java/com/linkedin/metadata/dao/utils/GraphUtils.java index 61c40d6f5..9c092c13e 100644 --- a/dao-api/src/main/java/com/linkedin/metadata/dao/utils/GraphUtils.java +++ b/dao-api/src/main/java/com/linkedin/metadata/dao/utils/GraphUtils.java @@ -2,7 +2,6 @@ import com.linkedin.common.urn.Urn; import com.linkedin.data.template.RecordTemplate; -import com.linkedin.metadata.dao.internal.BaseGraphWriterDAO; import java.util.List; import javax.annotation.Nonnull; import javax.annotation.Nullable; @@ -11,37 +10,34 @@ public class GraphUtils { - private static final String SOURCE = "source"; private GraphUtils() { // Util class } /** - * Check if a group relationship shares the same source urn, destination urn or both based on the remove option. + * Check if a group relationship shares the same source urn. * @param relationships list of relationships - * @param removalOption removal option to specify which relationships to be removed - * @param sourceField name of the source field - * @param destinationField name of the destination field - * @param urn source urn to compare. Optional for V1. Needed for V2. + * @param assetUrn source urn to compare. Optional for V1. Needed for V2. */ - public static void checkSameUrn(@Nonnull final List relationships, - @Nonnull final BaseGraphWriterDAO.RemovalOption removalOption, @Nonnull final String sourceField, - @Nonnull final String destinationField, @Nullable Urn urn) { - + public static void checkSameSourceUrn(@Nonnull final List relationships, @Nullable Urn assetUrn) { if (relationships.isEmpty()) { return; } - - final Urn sourceUrn = getSourceUrnBasedOnRelationshipVersion(relationships.get(0), urn); - final Urn destinationUrn = getDestinationUrnFromRelationship(relationships.get(0)); - - if (removalOption == BaseGraphWriterDAO.RemovalOption.REMOVE_ALL_EDGES_FROM_SOURCE) { - checkSameUrn(relationships, sourceField, sourceUrn); - } else if (removalOption == BaseGraphWriterDAO.RemovalOption.REMOVE_ALL_EDGES_TO_DESTINATION) { - checkSameUrn(relationships, destinationField, destinationUrn); - } else if (removalOption == BaseGraphWriterDAO.RemovalOption.REMOVE_ALL_EDGES_FROM_SOURCE_TO_DESTINATION) { - checkSameUrn(relationships, sourceField, sourceUrn); - checkSameUrn(relationships, destinationField, destinationUrn); + for (RecordTemplate relationship : relationships) { + if (ModelUtils.isRelationshipInV2(relationship.schema())) { + if (assetUrn == null) { + throw new IllegalArgumentException("Something went wrong. The asset urn is missing which is required during " + + "ingestion of a model 2.0 relationship. Relationship model: " + relationship); + } + // Skip source urn check for V2 relationships since they don't have source field + } else { + Urn compare = assetUrn == null ? getSourceUrnFromRelationship(relationships.get(0)) : assetUrn; + Urn source = getSourceUrnFromRelationship(relationship); + if (!compare.equals(source)) { + throw new IllegalArgumentException( + String.format("Relationships have different source urns. Urn being compared to: %s, Relationship source: %s", compare, source)); + } + } } } @@ -67,24 +63,4 @@ public static Urn getSourceUrnBasedOnRelat } return sourceUrn; } - - public static void checkSameUrn(@Nonnull final List relationships, - @Nonnull final BaseGraphWriterDAO.RemovalOption removalOption, @Nonnull final String sourceField, - @Nonnull final String destinationField) { - checkSameUrn(relationships, removalOption, sourceField, destinationField, null); - } - - private static void checkSameUrn(@Nonnull List records, @Nonnull String field, - @Nonnull Urn compare) { - for (RecordTemplate relation : records) { - if (ModelUtils.isRelationshipInV2(relation.schema()) && field.equals(SOURCE)) { - // Skip source urn check for V2 relationships since they don't have source field - // ToDo: enhance the source check for V2 relationships - return; - } - if (!compare.equals(ModelUtils.getUrnFromRelationship(relation, field))) { - throw new IllegalArgumentException("Records have different " + field + " urn"); - } - } - } } diff --git a/dao-api/src/test/java/com/linkedin/metadata/dao/utils/GraphUtilsTest.java b/dao-api/src/test/java/com/linkedin/metadata/dao/utils/GraphUtilsTest.java index 2750e1f08..27510d71b 100644 --- a/dao-api/src/test/java/com/linkedin/metadata/dao/utils/GraphUtilsTest.java +++ b/dao-api/src/test/java/com/linkedin/metadata/dao/utils/GraphUtilsTest.java @@ -1,7 +1,6 @@ package com.linkedin.metadata.dao.utils; import com.linkedin.data.template.RecordTemplate; -import com.linkedin.metadata.dao.internal.BaseGraphWriterDAO; import com.linkedin.testing.RelationshipBar; import com.linkedin.testing.RelationshipFoo; import com.linkedin.testing.RelationshipV2Bar; @@ -19,15 +18,14 @@ public class GraphUtilsTest { @Test - public void testCheckSameUrnWithEmptyRelationships() { + public void testCheckSameSourceUrnWithEmptyRelationships() { List relationships = Collections.emptyList(); - GraphUtils.checkSameUrn(relationships, BaseGraphWriterDAO.RemovalOption.REMOVE_ALL_EDGES_FROM_SOURCE, "source", "destination"); - GraphUtils.checkSameUrn(relationships, BaseGraphWriterDAO.RemovalOption.REMOVE_ALL_EDGES_TO_DESTINATION, "source", "destination", new BarUrn(1)); + GraphUtils.checkSameSourceUrn(relationships, null); // No exception should be thrown } @Test - public void testCheckSameUrnWithSameSourceUrn() { + public void testCheckSameSourceUrnWithSameSourceUrn() throws URISyntaxException { // Test cases for relationship V1 RelationshipFoo relationship; try { @@ -37,41 +35,46 @@ public void testCheckSameUrnWithSameSourceUrn() { } List relationships = Lists.newArrayList(relationship, relationship); + // when the assetUrn is not provided for relationship v1, use the first relationship's source to compare instead try { - GraphUtils.checkSameUrn(relationships, BaseGraphWriterDAO.RemovalOption.REMOVE_ALL_EDGES_FROM_SOURCE, "source", "destination"); + GraphUtils.checkSameSourceUrn(relationships, null); } catch (IllegalArgumentException e) { fail("Expected no IllegalArgumentException to be thrown, but got: " + e.getMessage()); } - // when urn is provided for relationship V1, the provided urn should be ignored + // when the assetUrn is provided for relationship V1, check all relationships to see if they have the same source try { - GraphUtils.checkSameUrn(relationships, BaseGraphWriterDAO.RemovalOption.REMOVE_ALL_EDGES_FROM_SOURCE, "source", "destination", new BarUrn(10)); + GraphUtils.checkSameSourceUrn(relationships, new FooUrn(1)); } catch (IllegalArgumentException e) { fail("Expected no IllegalArgumentException to be thrown, but got: " + e.getMessage()); } + // given an assetUrn that is not the same as the source urns of the relationships, throw an exception + try { + GraphUtils.checkSameSourceUrn(relationships, new BarUrn(10)); + fail("Expected an IllegalArgumentException to be thrown, but it wasn't"); + } catch (IllegalArgumentException e) { + // do nothing + } + // Test cases for relationship V2 RelationshipV2Bar relationshipV2 = mockRelationshipV2Bar(new BarUrn(2)); List relationshipsV2 = Lists.newArrayList(relationshipV2, relationshipV2); // when urn is provided for relationship V2, the check should pass try { - GraphUtils.checkSameUrn(relationshipsV2, BaseGraphWriterDAO.RemovalOption.REMOVE_ALL_EDGES_FROM_SOURCE, "source", "destination", new BarUrn(2)); + GraphUtils.checkSameSourceUrn(relationshipsV2, new BarUrn(2)); } catch (IllegalArgumentException e) { fail("Expected no IllegalArgumentException to be thrown, but got: " + e.getMessage()); } // when urn is not provided for relationship V2, it should throw IllegalArgumentException assertThrows(IllegalArgumentException.class, - () -> GraphUtils.checkSameUrn( - relationshipsV2, - BaseGraphWriterDAO.RemovalOption.REMOVE_ALL_EDGES_FROM_SOURCE, - "source", - "destination") + () -> GraphUtils.checkSameSourceUrn(relationshipsV2, null) ); } @Test - public void testCheckSameUrnWithDifferentSourceUrn() { - // Test cases for relationship V1 + public void testCheckSameSourceUrnWithDifferentSourceUrn() { + // Test cases for relationship V1 (not applicable to relationship v2) RecordTemplate relationship1; RecordTemplate relationship2; try { @@ -83,75 +86,17 @@ public void testCheckSameUrnWithDifferentSourceUrn() { List relationships = Lists.newArrayList(relationship1, relationship2); assertThrows(IllegalArgumentException.class, - () -> GraphUtils.checkSameUrn( - relationships, - BaseGraphWriterDAO.RemovalOption.REMOVE_ALL_EDGES_FROM_SOURCE, - "source", - "destination") + () -> GraphUtils.checkSameSourceUrn(relationships, null) ); // when urn is provided for relationship V1, the provided urn should be ignored assertThrows(IllegalArgumentException.class, - () -> GraphUtils.checkSameUrn( - relationships, - BaseGraphWriterDAO.RemovalOption.REMOVE_ALL_EDGES_FROM_SOURCE, - "source", - "destination", - new BarUrn(10)) - ); - - // ToDo: add test cases for V2. Right now it check if a list of relationships have the same source urn. - } - - @Test - public void testCheckSameUrnWithSameDestinationUrn() { - // Test cases for relationship V1 - RelationshipFoo relationship1; - try { - relationship1 = mockRelationshipFoo(new FooUrn(1), new BarUrn(2)); - } catch (URISyntaxException e) { - throw new RuntimeException(e); - } - List relationships1 = Lists.newArrayList(relationship1, relationship1); - - try { - GraphUtils.checkSameUrn(relationships1, BaseGraphWriterDAO.RemovalOption.REMOVE_ALL_EDGES_TO_DESTINATION, "source", - "destination"); - } catch (IllegalArgumentException e) { - fail("Expected no IllegalArgumentException to be thrown, but got: " + e.getMessage()); - } - - // when urn is provided for relationship V1, the provided urn should be ignored - try { - GraphUtils.checkSameUrn(relationships1, BaseGraphWriterDAO.RemovalOption.REMOVE_ALL_EDGES_TO_DESTINATION, "source", - "destination", new BarUrn(10)); - } catch (IllegalArgumentException e) { - fail("Expected no IllegalArgumentException to be thrown, but got: " + e.getMessage()); - } - - // Test cases for relationship V2 - RelationshipV2Bar relationship2 = mockRelationshipV2Bar(new BarUrn(2)); - List relationships2 = Lists.newArrayList(relationship2, relationship2); - - // throws exception if V2 relationships without source urn provided - assertThrows(IllegalArgumentException.class, - () -> GraphUtils.checkSameUrn( - relationships2, - BaseGraphWriterDAO.RemovalOption.REMOVE_ALL_EDGES_TO_DESTINATION, - "source", - "destination") + () -> GraphUtils.checkSameSourceUrn(relationships, new BarUrn(10)) ); - - try { - GraphUtils.checkSameUrn(relationships2, BaseGraphWriterDAO.RemovalOption.REMOVE_ALL_EDGES_TO_DESTINATION, "source", - "destination", new BarUrn(10)); - } catch (IllegalArgumentException e) { - fail("Expected no IllegalArgumentException to be thrown, but got: " + e.getMessage()); - } } @Test - public void testCheckSameUrnWithDifferentDestinationUrn() { - // Test cases for relationship V1 + public void testCheckSameSourceUrnWithDifferentDestinationUrn() { + // Test cases for relationship V1 (already tested v2 case in testCheckSameSourceUrnWithSameSourceUrn) RelationshipFoo relationship1; RelationshipBar relationship2; try { @@ -163,41 +108,11 @@ public void testCheckSameUrnWithDifferentDestinationUrn() { List relationships1 = Lists.newArrayList(relationship1, relationship2); assertThrows(IllegalArgumentException.class, - () -> GraphUtils.checkSameUrn( - relationships1, - BaseGraphWriterDAO.RemovalOption.REMOVE_ALL_EDGES_TO_DESTINATION, - "source", - "destination") - ); - - assertThrows(IllegalArgumentException.class, - () -> GraphUtils.checkSameUrn( - relationships1, - BaseGraphWriterDAO.RemovalOption.REMOVE_ALL_EDGES_TO_DESTINATION, - "source", - "destination", - new BarUrn(10)) - ); - - // Test cases for relationship V2 - RelationshipV2Bar relationship3 = mockRelationshipV2Bar(new BarUrn(3)); - RelationshipV2Bar relationship4 = mockRelationshipV2Bar(new BarUrn(4)); - List relationships2 = Lists.newArrayList(relationship3, relationship4); - assertThrows(IllegalArgumentException.class, - () -> GraphUtils.checkSameUrn( - relationships2, - BaseGraphWriterDAO.RemovalOption.REMOVE_ALL_EDGES_TO_DESTINATION, - "source", - "destination") + () -> GraphUtils.checkSameSourceUrn(relationships1, null) ); assertThrows(IllegalArgumentException.class, - () -> GraphUtils.checkSameUrn( - relationships2, - BaseGraphWriterDAO.RemovalOption.REMOVE_ALL_EDGES_TO_DESTINATION, - "source", - "destination", - new BarUrn(10)) + () -> GraphUtils.checkSameSourceUrn(relationships1, new BarUrn(10)) ); } diff --git a/dao-api/src/test/java/com/linkedin/metadata/dao/utils/ModelUtilsTest.java b/dao-api/src/test/java/com/linkedin/metadata/dao/utils/ModelUtilsTest.java index 459e42a58..c2b5153c7 100644 --- a/dao-api/src/test/java/com/linkedin/metadata/dao/utils/ModelUtilsTest.java +++ b/dao-api/src/test/java/com/linkedin/metadata/dao/utils/ModelUtilsTest.java @@ -57,6 +57,7 @@ import com.linkedin.testing.SnapshotUnionWithEntitySnapshotOptionalFields; import com.linkedin.testing.TyperefPizzaAspect; import com.linkedin.testing.localrelationship.AspectFooBar; +import com.linkedin.testing.localrelationship.AspectFooBaz; import com.linkedin.testing.localrelationship.AspectFooBarBaz; import com.linkedin.testing.namingedgecase.EntityValueNamingEdgeCase; import com.linkedin.testing.namingedgecase.InternalEntityAspectUnionNamingEdgeCase; @@ -108,7 +109,7 @@ public void testGetValidAspectTypes() { Set> validTypes = ModelUtils.getValidAspectTypes(EntityAspectUnion.class); assertEquals(validTypes, - ImmutableSet.of(AspectFoo.class, AspectBar.class, AspectFooBar.class, AspectFooBarBaz.class, AspectAttributes.class)); + ImmutableSet.of(AspectFoo.class, AspectBar.class, AspectFooBar.class, AspectFooBaz.class, AspectFooBarBaz.class, AspectAttributes.class)); } @Test diff --git a/dao-impl/ebean-dao/src/main/java/com/linkedin/metadata/dao/EbeanLocalDAO.java b/dao-impl/ebean-dao/src/main/java/com/linkedin/metadata/dao/EbeanLocalDAO.java index 7bb9271d1..b047438e6 100644 --- a/dao-impl/ebean-dao/src/main/java/com/linkedin/metadata/dao/EbeanLocalDAO.java +++ b/dao-impl/ebean-dao/src/main/java/com/linkedin/metadata/dao/EbeanLocalDAO.java @@ -880,6 +880,8 @@ protected void insert(@Nonnull URN urn, @Nullabl /** * If the aspect is associated with at least one relationship, upsert the relationship into the corresponding local * relationship table. Associated means that the aspect has a registered relationship build or it includes a relationship field. + * It will first try to find a registered relationship builder; if one doesn't exist or returns no relationship updates, + * try to find relationships from the aspect itself. * @param urn Urn of the metadata update * @param aspect Aspect of the metadata update * @param aspectClass Aspect class of the metadata update @@ -892,20 +894,24 @@ public List return Collections.emptyList(); } List localRelationshipUpdates = Collections.emptyList(); - if (_relationshipSource == RelationshipSource.ASPECT_METADATA) { + // Try to get relationships using relationship builders first. If there is not a relationship builder registered + // for the aspect class, try to get relationships from the aspect metadata instead. After most relationship models + // are using Model 2.0, switch the priority i.e. try to get the relationship from the aspect first before falling back + // on relationship builders. + // TODO: fix the gap where users can define new relationships in the aspect while still using graph builders to extract existing relationships + if (_localRelationshipBuilderRegistry != null && _localRelationshipBuilderRegistry.isRegistered(aspectClass)) { + localRelationshipUpdates = _localRelationshipBuilderRegistry.getLocalRelationshipBuilder(aspect).buildRelationships(urn, aspect); + // default all relationship updates to use REMOVE_ALL_EDGES_FROM_SOURCE + localRelationshipUpdates.forEach(update -> update.setRemovalOption(BaseGraphWriterDAO.RemovalOption.REMOVE_ALL_EDGES_FROM_SOURCE)); + } + // If no relationship updates were found using relationship builders, try to get them via the aspect. + if (localRelationshipUpdates.isEmpty()) { Map, Set> allRelationships = EBeanDAOUtils.extractRelationshipsFromAspect(aspect); localRelationshipUpdates = allRelationships.entrySet().stream() .filter(entry -> !entry.getValue().isEmpty()) // ensure at least 1 relationship in sublist to avoid index out of bounds .map(entry -> new LocalRelationshipUpdates( Arrays.asList(entry.getValue().toArray()), entry.getKey(), BaseGraphWriterDAO.RemovalOption.REMOVE_ALL_EDGES_FROM_SOURCE)) .collect(Collectors.toList()); - } else if (_relationshipSource == RelationshipSource.RELATIONSHIP_BUILDERS) { - if (_localRelationshipBuilderRegistry != null && _localRelationshipBuilderRegistry.isRegistered(aspectClass)) { - localRelationshipUpdates = _localRelationshipBuilderRegistry.getLocalRelationshipBuilder(aspect).buildRelationships(urn, aspect); - } - } else { - throw new UnsupportedOperationException("Please ensure that the RelationshipSource enum is properly set using " - + "setRelationshipSource method."); } _localRelationshipWriterDAO.processLocalRelationshipUpdates(urn, localRelationshipUpdates, isTestMode); return localRelationshipUpdates; 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 16e3702e5..7edfb5561 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 @@ -49,73 +49,59 @@ public EbeanLocalRelationshipWriterDAO(EbeanServer server) { */ @Transactional public void processLocalRelationshipUpdates(@Nonnull Urn urn, - @Nonnull List relationshipUpdates, @Nonnull boolean isTestMode) { + @Nonnull List relationshipUpdates, boolean isTestMode) { for (LocalRelationshipUpdates relationshipUpdate : relationshipUpdates) { if (relationshipUpdate.getRelationships().isEmpty()) { - clearRelationshipsByEntity(urn, relationshipUpdate.getRelationshipClass(), - relationshipUpdate.getRemovalOption(), isTestMode); + clearRelationshipsByEntity(urn, relationshipUpdate.getRelationshipClass(), isTestMode); } else { - addRelationships(relationshipUpdate.getRelationships(), relationshipUpdate.getRemovalOption(), isTestMode, urn); + addRelationships(relationshipUpdate.getRelationships(), isTestMode, urn); } } } /** - * This method is to serve for the purpose to clear all the relationships from a source entity urn. + * This method clears all the relationships from a source entity urn using REMOVE_ALL_EDGES_FROM_SOURCE. * @param urn entity urn could be either source or destination, depends on the RemovalOption * @param relationshipClass relationship that needs to be cleared - * @param removalOption removal option to specify which relationships to be removed * @param isTestMode whether to use test schema */ public void clearRelationshipsByEntity(@Nonnull Urn urn, - @Nonnull Class relationshipClass, @Nonnull RemovalOption removalOption, - @Nonnull boolean isTestMode) { - 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; - } - RelationshipValidator.validateRelationshipSchema(relationshipClass); + @Nonnull Class relationshipClass, boolean isTestMode) { + RelationshipValidator.validateRelationshipSchema(relationshipClass, isRelationshipInV2(relationshipClass)); SqlUpdate deletionSQL = _server.createSqlUpdate(SQLStatementUtils.deleteLocalRelationshipSQL( isTestMode ? SQLSchemaUtils.getTestRelationshipTableName(relationshipClass) - : 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()); - } + : SQLSchemaUtils.getRelationshipTableName(relationshipClass), RemovalOption.REMOVE_ALL_EDGES_FROM_SOURCE)); + deletionSQL.setParameter(CommonColumnName.SOURCE, urn.toString()); deletionSQL.execute(); } /** - * Persist the given list of relationships to the local relationship tables. + * Persist the given list of relationships to the local relationship using REMOVE_ALL_EDGES_FROM_SOURCE. * @param relationships the list of relationships to be persisted - * @param removalOption whether to remove existing relationship of the same type * @param isTestMode whether to use test schema * @param urn Urn of the entity to update relationships. * For Relationship V1: Optional, can be source or destination urn. * For Relationship V2: Required, is the source urn. */ public void addRelationships(@Nonnull List relationships, - @Nonnull RemovalOption removalOption, @Nonnull boolean isTestMode, @Nullable Urn urn) { + boolean isTestMode, @Nullable Urn urn) { // split relationships by relationship type Map> relationshipGroupMap = relationships.stream() .collect(Collectors.groupingBy(relationship -> relationship.getClass().getCanonicalName())); // validate if all relationship groups have valid urns - relationshipGroupMap.values().forEach(relationshipGroup - -> GraphUtils.checkSameUrn(relationshipGroup, removalOption, CommonColumnName.SOURCE, CommonColumnName.DESTINATION, urn)); + relationshipGroupMap.values().forEach(relationshipGroup -> GraphUtils.checkSameSourceUrn(relationshipGroup, urn)); relationshipGroupMap.values().forEach(relationshipGroup -> { - addRelationshipGroup(relationshipGroup, removalOption, isTestMode, urn); + addRelationshipGroup(relationshipGroup, isTestMode, urn); }); } + // This method only supports Relationship 1.0 (i.e. source present in model) ingestion using graph builders. @Override public void addRelationships(@Nonnull List relationships, - @Nonnull RemovalOption removalOption, @Nonnull boolean isTestMode) { - addRelationships(relationships, removalOption, isTestMode, null); + @Nonnull RemovalOption removalOption, boolean isTestMode) { + addRelationships(relationships, isTestMode, null); } @Override @@ -146,23 +132,22 @@ public void removeEntities(@Nonnull List urns) { /** * Add the given list of relationships to the local relationship tables. * @param relationshipGroup the list of relationships to be persisted - * @param removalOption whether to remove existing relationship of the same type * @param isTestMode whether to use test schema * @param urn the source urn to be used for the relationships. Optional for Relationship V1. * Needed for Relationship V2 because source is not included in the relationshipV2 metadata. */ private void addRelationshipGroup(@Nonnull final List relationshipGroup, - @Nonnull RemovalOption removalOption, @Nonnull boolean isTestMode, @Nullable Urn urn) { + boolean isTestMode, @Nullable Urn urn) { if (relationshipGroup.size() == 0) { return; } RELATIONSHIP firstRelationship = relationshipGroup.get(0); - RelationshipValidator.validateRelationshipSchema(firstRelationship.getClass()); + RelationshipValidator.validateRelationshipSchema(firstRelationship.getClass(), isRelationshipInV2(firstRelationship.getClass())); - // Process remove option to delete some local relationships if needed before adding new relationships. - processRemovalOption(isTestMode ? SQLSchemaUtils.getTestRelationshipTableName(firstRelationship) - : SQLSchemaUtils.getRelationshipTableName(firstRelationship), firstRelationship, removalOption, urn); + // Remove some local relationships if needed before adding new relationships using REMOVE_ALL_EDGES_FROM_SOURCE. + removeRelationshipsBySource(isTestMode ? SQLSchemaUtils.getTestRelationshipTableName(firstRelationship) + : SQLSchemaUtils.getRelationshipTableName(firstRelationship), firstRelationship, urn); long now = Instant.now().toEpochMilli(); @@ -191,30 +176,14 @@ private void addRelationshipGroup(@Nonnull * Process the relationship removal in the DB tableName based on the removal option. * @param tableName the table name of the relationship * @param relationship the relationship to be removed - * @param removalOption the removal option * @param urn the source urn to be used for the relationships. Optional for Relationship V1. * Needed for Relationship V2 because source is not included in the relationshipV2 metadata. */ - private void processRemovalOption(@Nonnull String tableName, - @Nonnull RELATIONSHIP relationship, @Nonnull RemovalOption removalOption, @Nullable Urn urn) { - - if (removalOption == RemovalOption.REMOVE_NONE) { - return; - } - - SqlUpdate deletionSQL = _server.createSqlUpdate(SQLStatementUtils.deleteLocalRelationshipSQL(tableName, removalOption)); + private void removeRelationshipsBySource(@Nonnull String tableName, + @Nonnull RELATIONSHIP relationship, @Nullable Urn urn) { + SqlUpdate deletionSQL = _server.createSqlUpdate(SQLStatementUtils.deleteLocalRelationshipSQL(tableName, RemovalOption.REMOVE_ALL_EDGES_FROM_SOURCE)); Urn source = GraphUtils.getSourceUrnBasedOnRelationshipVersion(relationship, urn); - Urn destination = getDestinationUrnFromRelationship(relationship); - - if (removalOption == RemovalOption.REMOVE_ALL_EDGES_FROM_SOURCE_TO_DESTINATION) { - deletionSQL.setParameter(CommonColumnName.DESTINATION, destination.toString()); - deletionSQL.setParameter(CommonColumnName.SOURCE, source.toString()); - } else if (removalOption == RemovalOption.REMOVE_ALL_EDGES_FROM_SOURCE) { - deletionSQL.setParameter(CommonColumnName.SOURCE, source.toString()); - } else if (removalOption == RemovalOption.REMOVE_ALL_EDGES_TO_DESTINATION) { - deletionSQL.setParameter(CommonColumnName.DESTINATION, destination.toString()); - } - + deletionSQL.setParameter(CommonColumnName.SOURCE, source.toString()); deletionSQL.execute(); } } diff --git a/dao-impl/ebean-dao/src/main/java/com/linkedin/metadata/dao/utils/SQLStatementUtils.java b/dao-impl/ebean-dao/src/main/java/com/linkedin/metadata/dao/utils/SQLStatementUtils.java index 9f0323fc7..28a2782f4 100644 --- a/dao-impl/ebean-dao/src/main/java/com/linkedin/metadata/dao/utils/SQLStatementUtils.java +++ b/dao-impl/ebean-dao/src/main/java/com/linkedin/metadata/dao/utils/SQLStatementUtils.java @@ -342,11 +342,10 @@ public static String deleteLocalRelationshipSQL(final String tableName, final Ba return String.format(DELETE_BY_SOURCE, tableName); } else if (removalOption == BaseGraphWriterDAO.RemovalOption.REMOVE_ALL_EDGES_FROM_SOURCE_TO_DESTINATION) { return String.format(DELETE_BY_SOURCE_AND_DESTINATION, tableName); - } else if (removalOption == BaseGraphWriterDAO.RemovalOption.REMOVE_ALL_EDGES_TO_DESTINATION) { - return String.format(DELETE_BY_DESTINATION, tableName); } - - throw new IllegalArgumentException(String.format("Removal option %s is not valid.", removalOption)); + throw new IllegalArgumentException(String.format("Relationships can only be removed using either REMOVE_ALL_EDGES_FROM_SOURCE " + + "when inserting new relationships or REMOVE_ALL_EDGES_FROM_SOURCE_TO_DESTINATION when soft deleting an aspect " + + "from which relationships are derived from. Table name: %s, removal option: %s", tableName, removalOption)); } /** 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 3a2ba7edf..4c41dcb90 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 @@ -63,6 +63,7 @@ import com.linkedin.testing.FooSnapshot; import com.linkedin.testing.MixedRecord; import com.linkedin.testing.localrelationship.AspectFooBar; +import com.linkedin.testing.localrelationship.AspectFooBaz; import com.linkedin.testing.localrelationship.BelongsTo; import com.linkedin.testing.localrelationship.BelongsToV2; import com.linkedin.testing.localrelationship.BelongsToV2Array; @@ -3138,8 +3139,6 @@ public void testAddRelationshipsFromAspect() throws URISyntaxException { EbeanLocalDAO fooDao = createDao(FooUrn.class); EbeanLocalDAO barDao = createDao(BarUrn.class); - fooDao.setRelationshipSource(EbeanLocalDAO.RelationshipSource.ASPECT_METADATA); - FooUrn fooUrn = makeFooUrn(1); BarUrn barUrn1 = BarUrn.createFromString("urn:li:bar:1"); BelongsToV2 belongsTo1 = new BelongsToV2().setDestination(BelongsToV2.Destination.create(barUrn1.toString())); @@ -3148,7 +3147,7 @@ public void testAddRelationshipsFromAspect() throws URISyntaxException { BarUrn barUrn3 = BarUrn.createFromString("urn:li:bar:3"); 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); + AspectFooBar aspectFooBar = new AspectFooBar().setBelongsTos(belongsToArray); AuditStamp auditStamp = makeAuditStamp("foo", System.currentTimeMillis()); fooDao.add(fooUrn, aspectFooBar, auditStamp); @@ -3171,6 +3170,125 @@ public void testAddRelationshipsFromAspect() throws URISyntaxException { assertEquals(aspects.size(), 1); } + @Test + public void testAddRelationshipsV2WithRelationshipBuilders() throws URISyntaxException { + EbeanLocalDAO fooDao = createDao(FooUrn.class); + EbeanLocalDAO barDao = createDao(BarUrn.class); + + // create AspectFooBaz without setting relationship-type fields + FooUrn fooUrn = makeFooUrn(1); + BarUrn barUrn1 = BarUrn.createFromString("urn:li:bar:1"); + BarUrn barUrn2 = BarUrn.createFromString("urn:li:bar:2"); + BarUrn barUrn3 = BarUrn.createFromString("urn:li:bar:3"); + AspectFooBaz aspectFooBaz = new AspectFooBaz().setBars(new BarUrnArray(barUrn1, barUrn2, barUrn3)); + AuditStamp auditStamp = makeAuditStamp("foo", System.currentTimeMillis()); + + fooDao.setLocalRelationshipBuilderRegistry(new SampleLocalRelationshipRegistryImpl()); + + fooDao.add(fooUrn, aspectFooBaz, auditStamp); + barDao.add(barUrn1, new AspectFoo().setValue("1"), auditStamp); + barDao.add(barUrn2, new AspectFoo().setValue("2"), auditStamp); + barDao.add(barUrn3, new AspectFoo().setValue("3"), auditStamp); + + // Verify local relationships and entity are added. + EbeanLocalRelationshipQueryDAO ebeanLocalRelationshipQueryDAO = new EbeanLocalRelationshipQueryDAO(_server); + ebeanLocalRelationshipQueryDAO.setSchemaConfig(_schemaConfig); + + List relationships = + ebeanLocalRelationshipQueryDAO.findRelationships(FooSnapshot.class, EMPTY_FILTER, BarSnapshot.class, + EMPTY_FILTER, BelongsToV2.class, OUTGOING_FILTER, 0, 10); + + AspectKey key = new AspectKey<>(AspectFooBaz.class, fooUrn, 0L); + List aspects = fooDao.batchGetHelper(Collections.singletonList(key), 1, 0); + + assertEquals(relationships.size(), 3); + assertEquals(aspects.size(), 1); + } + + @Test + public void testAddRelationshipsV2WithRegisteredButEmptyRelationshipBuilders() throws URISyntaxException { + EbeanLocalDAO fooDao = createDao(FooUrn.class); + EbeanLocalDAO barDao = createDao(BarUrn.class); + + // create AspectFooBaz with null bar array and non-null relationship-type fields. + // the relationship builder will return an empty list of relationship updates so the DAO should try + // to extract relationships from aspect metadata + FooUrn fooUrn = makeFooUrn(1); + BarUrn barUrn1 = BarUrn.createFromString("urn:li:bar:1"); + BelongsToV2 belongsTo1 = new BelongsToV2().setDestination(BelongsToV2.Destination.create(barUrn1.toString())); + BarUrn barUrn2 = BarUrn.createFromString("urn:li:bar:2"); + BelongsToV2 belongsTo2 = new BelongsToV2().setDestination(BelongsToV2.Destination.create(barUrn2.toString())); + BarUrn barUrn3 = BarUrn.createFromString("urn:li:bar:3"); + BelongsToV2 belongsTo3 = new BelongsToV2().setDestination(BelongsToV2.Destination.create(barUrn3.toString())); + BelongsToV2Array belongsToArray = new BelongsToV2Array(belongsTo1, belongsTo2, belongsTo3); + AspectFooBaz aspectFooBaz = new AspectFooBaz().setBelongsTos(belongsToArray); + AuditStamp auditStamp = makeAuditStamp("foo", System.currentTimeMillis()); + + fooDao.setLocalRelationshipBuilderRegistry(new SampleLocalRelationshipRegistryImpl()); + + fooDao.add(fooUrn, aspectFooBaz, auditStamp); + barDao.add(barUrn1, new AspectFoo().setValue("1"), auditStamp); + barDao.add(barUrn2, new AspectFoo().setValue("2"), auditStamp); + barDao.add(barUrn3, new AspectFoo().setValue("3"), auditStamp); + + // Verify local relationships and entity are added. + EbeanLocalRelationshipQueryDAO ebeanLocalRelationshipQueryDAO = new EbeanLocalRelationshipQueryDAO(_server); + ebeanLocalRelationshipQueryDAO.setSchemaConfig(_schemaConfig); + + List relationships = + ebeanLocalRelationshipQueryDAO.findRelationships(FooSnapshot.class, EMPTY_FILTER, BarSnapshot.class, + EMPTY_FILTER, BelongsToV2.class, OUTGOING_FILTER, 0, 10); + + AspectKey key = new AspectKey<>(AspectFooBaz.class, fooUrn, 0L); + List aspects = fooDao.batchGetHelper(Collections.singletonList(key), 1, 0); + + assertEquals(relationships.size(), 3); + assertEquals(aspects.size(), 1); + } + + @Test + public void testAddRelationshipsV2DefaultToRelationshipBuilders() throws URISyntaxException { + EbeanLocalDAO fooDao = createDao(FooUrn.class); + EbeanLocalDAO barDao = createDao(BarUrn.class); + + // create AspectFooBaz with non-null bar array and non-null relationship-type fields but with different values. + // bars: 1, 2, 3 -RelationshipBuilder-> foo1->bar1, foo1->bar2, foo1->bar3 + // belongsTos: foo1->bar4 + // the DAO should default to using the relationships from the relationship builders over the relationships + // extracted from the aspect. + FooUrn fooUrn = makeFooUrn(1); + BarUrn barUrn1 = BarUrn.createFromString("urn:li:bar:1"); + BarUrn barUrn2 = BarUrn.createFromString("urn:li:bar:2"); + BarUrn barUrn3 = BarUrn.createFromString("urn:li:bar:3"); + BarUrn barUrn4 = BarUrn.createFromString("urn:li:bar:4"); + BelongsToV2 belongsTo4 = new BelongsToV2().setDestination(BelongsToV2.Destination.create(barUrn4.toString())); + BelongsToV2Array belongsToArray = new BelongsToV2Array(belongsTo4); + AspectFooBaz aspectFooBaz = new AspectFooBaz().setBars(new BarUrnArray(barUrn1, barUrn2, barUrn3)).setBelongsTos(belongsToArray); + AuditStamp auditStamp = makeAuditStamp("foo", System.currentTimeMillis()); + + fooDao.setLocalRelationshipBuilderRegistry(new SampleLocalRelationshipRegistryImpl()); + + fooDao.add(fooUrn, aspectFooBaz, auditStamp); + barDao.add(barUrn1, new AspectFoo().setValue("1"), auditStamp); + barDao.add(barUrn2, new AspectFoo().setValue("2"), auditStamp); + barDao.add(barUrn3, new AspectFoo().setValue("3"), auditStamp); + + // Verify local relationships and entity are added. + EbeanLocalRelationshipQueryDAO ebeanLocalRelationshipQueryDAO = new EbeanLocalRelationshipQueryDAO(_server); + ebeanLocalRelationshipQueryDAO.setSchemaConfig(_schemaConfig); + + List relationships = + ebeanLocalRelationshipQueryDAO.findRelationships(FooSnapshot.class, EMPTY_FILTER, BarSnapshot.class, + EMPTY_FILTER, BelongsToV2.class, OUTGOING_FILTER, 0, 10); + + AspectKey key = new AspectKey<>(AspectFooBaz.class, fooUrn, 0L); + List aspects = fooDao.batchGetHelper(Collections.singletonList(key), 1, 0); + + assertEquals(relationships.size(), 3); + assertFalse(relationships.contains(belongsTo4)); + assertEquals(aspects.size(), 1); + } + @Test public void testNewSchemaFilterByArray() { if (_schemaConfig == SchemaConfig.NEW_SCHEMA_ONLY) { diff --git a/dao-impl/ebean-dao/src/test/java/com/linkedin/metadata/dao/localrelationship/EbeanLocalRelationshipQueryDAOTest.java b/dao-impl/ebean-dao/src/test/java/com/linkedin/metadata/dao/localrelationship/EbeanLocalRelationshipQueryDAOTest.java index 9659a3a08..0ab2dd552 100644 --- a/dao-impl/ebean-dao/src/test/java/com/linkedin/metadata/dao/localrelationship/EbeanLocalRelationshipQueryDAOTest.java +++ b/dao-impl/ebean-dao/src/test/java/com/linkedin/metadata/dao/localrelationship/EbeanLocalRelationshipQueryDAOTest.java @@ -1,5 +1,6 @@ package com.linkedin.metadata.dao.localrelationship; +import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableSet; import com.google.common.io.Resources; import com.linkedin.common.AuditStamp; @@ -265,13 +266,11 @@ public void testFindOneRelationshipWithFilter(EbeanLocalDAO.SchemaConfig schemaC ConsumeFrom sparkConsumeFromHdfs = new ConsumeFrom().setSource(spark).setDestination(hdfs).setEnvironment(EnvorinmentType.OFFLINE); _localRelationshipWriterDAO.addRelationship(sparkConsumeFromHdfs, false); - // Add Samza consume-from kafka relationship + // Add Samza consume-from kafka and Samza consume-from restli relationships ConsumeFrom samzaConsumeFromKafka = new ConsumeFrom().setSource(samza).setDestination(kafka).setEnvironment(EnvorinmentType.NEARLINE); - _localRelationshipWriterDAO.addRelationship(samzaConsumeFromKafka, false); - - // Add Samza consume-from restli relationship ConsumeFrom samzaConsumeFromRestli = new ConsumeFrom().setSource(samza).setDestination(restli).setEnvironment(EnvorinmentType.ONLINE); - _localRelationshipWriterDAO.addRelationship(samzaConsumeFromRestli, false); + + _localRelationshipWriterDAO.addRelationships(ImmutableList.of(samzaConsumeFromKafka, samzaConsumeFromRestli), false); // Find all consume-from relationship for Samza. LocalRelationshipCriterion filterUrnCriterion = EBeanDAOUtils.buildRelationshipFieldCriterion( @@ -414,13 +413,11 @@ public void testFindOneRelationshipWithFilterWithEntityUrn(EbeanLocalDAO.SchemaC ConsumeFrom sparkConsumeFromHdfs = new ConsumeFrom().setSource(spark).setDestination(hdfs).setEnvironment(EnvorinmentType.OFFLINE); _localRelationshipWriterDAO.addRelationship(sparkConsumeFromHdfs, false); - // Add Samza consume-from kafka relationship + // Add Samza consume-from kafka and Samza consume-from restli relationships ConsumeFrom samzaConsumeFromKafka = new ConsumeFrom().setSource(samza).setDestination(kafka).setEnvironment(EnvorinmentType.NEARLINE); - _localRelationshipWriterDAO.addRelationship(samzaConsumeFromKafka, false); - - // Add Samza consume-from restli relationship ConsumeFrom samzaConsumeFromRestli = new ConsumeFrom().setSource(samza).setDestination(restli).setEnvironment(EnvorinmentType.ONLINE); - _localRelationshipWriterDAO.addRelationship(samzaConsumeFromRestli, false); + + _localRelationshipWriterDAO.addRelationships(ImmutableList.of(samzaConsumeFromRestli, samzaConsumeFromKafka), false); // Find all consume-from relationship for Samza. LocalRelationshipCriterion filterUrnCriterion = EBeanDAOUtils.buildRelationshipFieldCriterion( @@ -777,8 +774,7 @@ public void testFindEntitiesOneHopAwayOutgoingDirection() throws Exception { // Add Alice belongs to MIT and Stanford. BelongsTo aliceBelongsToMit = new BelongsTo().setSource(alice).setDestination(mit); BelongsTo aliceBelongsToStanford = new BelongsTo().setSource(alice).setDestination(stanford); - _localRelationshipWriterDAO.addRelationship(aliceBelongsToMit, false); - _localRelationshipWriterDAO.addRelationship(aliceBelongsToStanford, false); + _localRelationshipWriterDAO.addRelationships(ImmutableList.of(aliceBelongsToStanford, aliceBelongsToMit), false); // Add Bob belongs to Stanford. BelongsTo bobBelongsToStandford = new BelongsTo().setSource(bob).setDestination(stanford); 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 6ac2f17ec..0cb3567a3 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,13 +2,12 @@ import com.google.common.io.Resources; import com.linkedin.metadata.dao.EbeanLocalRelationshipWriterDAO; +import com.linkedin.metadata.dao.builder.BaseLocalRelationshipBuilder.LocalRelationshipUpdates; 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; import com.linkedin.metadata.dao.localrelationship.builder.VersionOfLocalRelationshipBuilder; import com.linkedin.metadata.dao.utils.EmbeddedMariaInstance; -import com.linkedin.metadata.dao.builder.BaseLocalRelationshipBuilder.LocalRelationshipUpdates; import com.linkedin.testing.BarUrnArray; import com.linkedin.testing.RelationshipV2Bar; import com.linkedin.testing.localrelationship.AspectFooBar; @@ -42,46 +41,9 @@ public void init() throws IOException { _localRelationshipWriterDAO = new EbeanLocalRelationshipWriterDAO(_server); } - @Test - public void testAddRelationshipWithRemoveAllEdgesToDestination() throws URISyntaxException { - _server.execute(Ebean.createSqlUpdate(insertRelationships("metadata_relationship_belongsto", "urn:li:foo:123", - "foo", "urn:li:bar:000", "bar"))); - - AspectFooBar aspectFooBar = new AspectFooBar().setBars(new BarUrnArray( - BarUrn.createFromString("urn:li:bar:123"), - BarUrn.createFromString("urn:li:bar:456"), - BarUrn.createFromString("urn:li:bar:789"))); - - List updates = new BelongsToLocalRelationshipBuilder(AspectFooBar.class) - .buildRelationships(FooUrn.createFromString("urn:li:foo:123"), aspectFooBar); - - // Before processing - List before = _server.createSqlQuery("select * from metadata_relationship_belongsto where destination='urn:li:bar:000'").findList(); - assertEquals(before.size(), 1); - - _localRelationshipWriterDAO.processLocalRelationshipUpdates(FooUrn.createFromString("urn:li:foo:123"), updates, false); - - // After processing verification - List all = _server.createSqlQuery("select * from metadata_relationship_belongsto").findList(); - assertEquals(all.size(), 4); // Total number of edges is 4 - - List softDeleted = _server.createSqlQuery("select * from metadata_relationship_belongsto where deleted_ts IS NOT NULL").findList(); - assertEquals(softDeleted.size(), 1); // 1 soft deleted edge - - List newEdges = _server.createSqlQuery( - "select * from metadata_relationship_belongsto where source='urn:li:foo:123' and deleted_ts IS NULL").findList(); - - assertEquals(newEdges.size(), 3); // 3 new edges added. - - // Clean up - _server.execute(Ebean.createSqlUpdate("truncate metadata_relationship_belongsto")); - } - @Test public void testAddRelationshipWithRemoveNone() throws URISyntaxException { - _server.execute(Ebean.createSqlUpdate(insertRelationships("metadata_relationship_reportsto", "urn:li:bar:000", - "bar", "urn:li:foo:123", "foo"))); - + // All relationship ingestion should default to REMOVE_ALL_EDGES_FROM_SOURCE. REMOVE_NONE is no longer supported. AspectFooBar aspectFooBar = new AspectFooBar().setBars(new BarUrnArray( BarUrn.createFromString("urn:li:bar:123"), BarUrn.createFromString("urn:li:bar:456"), @@ -90,93 +52,29 @@ public void testAddRelationshipWithRemoveNone() throws URISyntaxException { List updates = new ReportsToLocalRelationshipBuilder(AspectFooBar.class) .buildRelationships(FooUrn.createFromString("urn:li:foo:123"), aspectFooBar); - // Before processing - List before = _server.createSqlQuery("select * from metadata_relationship_reportsto where source='urn:li:bar:000'").findList(); - assertEquals(before.size(), 1); - - _localRelationshipWriterDAO.processLocalRelationshipUpdates(FooUrn.createFromString("urn:li:foo:123"), updates, false); - - // After processing verification - List after = _server.createSqlQuery("select * from metadata_relationship_reportsto where destination='urn:li:foo:123'").findList(); - assertEquals(after.size(), 4); - List edges = _server.createSqlQuery("select * from metadata_relationship_reportsto where source='urn:li:bar:000'").findList(); - assertEquals(edges.size(), 1); - - // Clean up - _server.execute(Ebean.createSqlUpdate("truncate metadata_relationship_reportsto")); - } - - @Test - public void testAddRelationshipWithRemoveNoneInTestMode() throws URISyntaxException { - _server.execute(Ebean.createSqlUpdate( - insertRelationships("metadata_relationship_reportsto_test", "urn:li:bar:000", "bar", "urn:li:foo:123", "foo"))); - - AspectFooBar aspectFooBar = new AspectFooBar().setBars( - new BarUrnArray(BarUrn.createFromString("urn:li:bar:123"), BarUrn.createFromString("urn:li:bar:456"), - BarUrn.createFromString("urn:li:bar:789"))); - - List updates = - new ReportsToLocalRelationshipBuilder(AspectFooBar.class).buildRelationships( - FooUrn.createFromString("urn:li:foo:123"), aspectFooBar); - - // Before processing - List beforeTest = - _server.createSqlQuery("select * from metadata_relationship_reportsto_test where source='urn:li:bar:000'") - .findList(); - assertEquals(beforeTest.size(), 1); - - _localRelationshipWriterDAO.processLocalRelationshipUpdates(FooUrn.createFromString("urn:li:foo:123"), updates, - true); - - // After processing verification - List afterTest = - _server.createSqlQuery("select * from metadata_relationship_reportsto_test where destination='urn:li:foo:123'") - .findList(); - assertEquals(afterTest.size(), 4); - List edgesTest = - _server.createSqlQuery("select * from metadata_relationship_reportsto_test where source='urn:li:bar:000'") - .findList(); - assertEquals(edgesTest.size(), 1); - - // Clean up - _server.execute(Ebean.createSqlUpdate("truncate metadata_relationship_reportsto_test")); + try { + _localRelationshipWriterDAO.processLocalRelationshipUpdates(FooUrn.createFromString("urn:li:foo:123"), updates, false); + fail("This test should throw an exception since only REMOVE_ALL_EDGES_FROM_SOURCE is supported"); + } catch (IllegalArgumentException e) { + // do nothing + } } @Test public void testAddRelationshipWithRemoveAllEdgesFromSourceToDestination() 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:123", "foo"))); - - _server.execute(Ebean.createSqlUpdate(insertRelationships("metadata_relationship_pairswith", "urn:li:bar:000", - "bar", "urn:li:foo:123", "foo"))); - + // All relationship ingestion should default to REMOVE_ALL_EDGES_FROM_SOURCE. REMOVE_ALL_EDGES_FROM_SOURCE_TO_DESTINATION is no longer supported. AspectFooBar aspectFooBar = new AspectFooBar().setBars(new BarUrnArray(BarUrn.createFromString("urn:li:bar:123"))); List updates = new PairsWithLocalRelationshipBuilder(AspectFooBar.class) .buildRelationships(FooUrn.createFromString("urn:li:foo:123"), aspectFooBar); - // Before processing - List before = _server.createSqlQuery("select * from metadata_relationship_pairswith").findList(); - assertEquals(before.size(), 3); - - _localRelationshipWriterDAO.processLocalRelationshipUpdates(FooUrn.createFromString("urn:li:foo:123"), updates, false); - - // After processing verification - List all = _server.createSqlQuery("select * from metadata_relationship_pairswith").findList(); - assertEquals(all.size(), 4); // Total number of edges is 4 - - List softDeleted = _server.createSqlQuery("select * from metadata_relationship_pairswith where deleted_ts IS NOT NULL").findList(); - assertEquals(softDeleted.size(), 2); // 2 edges are soft-deleted. - - List oldEdge = _server.createSqlQuery( - "select * from metadata_relationship_pairswith where source='urn:li:bar:000' and deleted_ts IS NULL").findList(); - assertEquals(oldEdge.size(), 1); // 1 old edge untouched. - - // Clean up - _server.execute(Ebean.createSqlUpdate("truncate metadata_relationship_pairswith")); + try { + _localRelationshipWriterDAO.processLocalRelationshipUpdates(FooUrn.createFromString("urn:li:foo:123"), updates, + false); + fail("This test should throw an exception since only REMOVE_ALL_EDGES_FROM_SOURCE is supported"); + } catch (IllegalArgumentException e) { + // do nothing + } } @Test @@ -206,7 +104,7 @@ public void testAddRelationshipWithRemoveAllEdgesFromSource() throws URISyntaxEx List before = _server.createSqlQuery("select * from metadata_relationship_versionof").findList(); assertEquals(before.size(), 3); - _localRelationshipWriterDAO.processLocalRelationshipUpdates(FooUrn.createFromString("urn:li:foo:123"), updates, false); + _localRelationshipWriterDAO.processLocalRelationshipUpdates(BarUrn.createFromString("urn:li:bar:123"), updates, false); // After processing verification // now the relationship table should have the following relationships: @@ -319,27 +217,12 @@ public void testClearRelationshipsByEntityUrn() throws URISyntaxException { List before = _server.createSqlQuery("select * from metadata_relationship_pairswith where deleted_ts is null").findList(); assertEquals(before.size(), 2); - _localRelationshipWriterDAO.clearRelationshipsByEntity(barUrn, PairsWith.class, - BaseGraphWriterDAO.RemovalOption.REMOVE_ALL_EDGES_FROM_SOURCE, false); + _localRelationshipWriterDAO.clearRelationshipsByEntity(barUrn, PairsWith.class, false); // 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, PairsWith.class, - BaseGraphWriterDAO.RemovalOption.REMOVE_ALL_EDGES_TO_DESTINATION, false); - - // 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")); } diff --git a/dao-impl/ebean-dao/src/test/java/com/linkedin/metadata/dao/localrelationship/SampleLocalRelationshipRegistryImpl.java b/dao-impl/ebean-dao/src/test/java/com/linkedin/metadata/dao/localrelationship/SampleLocalRelationshipRegistryImpl.java index 6e09da72a..446728a1a 100644 --- a/dao-impl/ebean-dao/src/test/java/com/linkedin/metadata/dao/localrelationship/SampleLocalRelationshipRegistryImpl.java +++ b/dao-impl/ebean-dao/src/test/java/com/linkedin/metadata/dao/localrelationship/SampleLocalRelationshipRegistryImpl.java @@ -5,7 +5,9 @@ import com.linkedin.metadata.dao.builder.BaseLocalRelationshipBuilder; import com.linkedin.metadata.dao.builder.LocalRelationshipBuilderRegistry; import com.linkedin.metadata.dao.localrelationship.builder.BelongsToLocalRelationshipBuilder; +import com.linkedin.metadata.dao.localrelationship.builder.BelongsToV2LocalRelationshipBuilder; import com.linkedin.testing.localrelationship.AspectFooBar; +import com.linkedin.testing.localrelationship.AspectFooBaz; import java.util.Map; import javax.annotation.Nonnull; import javax.annotation.Nullable; @@ -13,7 +15,10 @@ public class SampleLocalRelationshipRegistryImpl implements LocalRelationshipBuilderRegistry { private Map, BaseLocalRelationshipBuilder> builders = - new ImmutableMap.Builder().put(AspectFooBar.class, new BelongsToLocalRelationshipBuilder(AspectFooBar.class)).build(); + new ImmutableMap.Builder() + .put(AspectFooBar.class, new BelongsToLocalRelationshipBuilder(AspectFooBar.class)) + .put(AspectFooBaz.class, new BelongsToV2LocalRelationshipBuilder(AspectFooBaz.class)) + .build(); @Nullable @Override diff --git a/dao-impl/ebean-dao/src/test/java/com/linkedin/metadata/dao/localrelationship/builder/BelongsToV2LocalRelationshipBuilder.java b/dao-impl/ebean-dao/src/test/java/com/linkedin/metadata/dao/localrelationship/builder/BelongsToV2LocalRelationshipBuilder.java new file mode 100644 index 000000000..a7160b42a --- /dev/null +++ b/dao-impl/ebean-dao/src/test/java/com/linkedin/metadata/dao/localrelationship/builder/BelongsToV2LocalRelationshipBuilder.java @@ -0,0 +1,37 @@ +package com.linkedin.metadata.dao.localrelationship.builder; + +import com.linkedin.common.urn.Urn; +import com.linkedin.metadata.dao.builder.BaseLocalRelationshipBuilder; +import com.linkedin.metadata.dao.internal.BaseGraphWriterDAO; +import com.linkedin.testing.localrelationship.AspectFooBaz; +import com.linkedin.testing.localrelationship.BelongsToV2; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import javax.annotation.Nonnull; + + +public class BelongsToV2LocalRelationshipBuilder extends BaseLocalRelationshipBuilder { + public BelongsToV2LocalRelationshipBuilder(@Nonnull Class aspectFooBazClass) { + super(aspectFooBazClass); + } + + @Nonnull + @Override + public List buildRelationships(@Nonnull URN urn, + @Nonnull AspectFooBaz aspectFooBaz) { + List belongsToRelationships = new ArrayList<>(); + if (!aspectFooBaz.hasBars()) { + return Collections.emptyList(); + } + for (Urn barUrn : aspectFooBaz.getBars()) { + belongsToRelationships.add(new BelongsToV2().setDestination(BelongsToV2.Destination.create(barUrn.toString()))); + } + + BaseLocalRelationshipBuilder.LocalRelationshipUpdates localRelationshipUpdates = + new BaseLocalRelationshipBuilder.LocalRelationshipUpdates(belongsToRelationships, BelongsToV2.class, + BaseGraphWriterDAO.RemovalOption.REMOVE_ALL_EDGES_FROM_SOURCE); + + return Collections.singletonList(localRelationshipUpdates); + } +} diff --git a/dao-impl/ebean-dao/src/test/resources/ebean-local-access-create-all-with-non-dollar-virtual-column-names.sql b/dao-impl/ebean-dao/src/test/resources/ebean-local-access-create-all-with-non-dollar-virtual-column-names.sql index 955d7e6b8..edd209bbe 100644 --- a/dao-impl/ebean-dao/src/test/resources/ebean-local-access-create-all-with-non-dollar-virtual-column-names.sql +++ b/dao-impl/ebean-dao/src/test/resources/ebean-local-access-create-all-with-non-dollar-virtual-column-names.sql @@ -81,6 +81,9 @@ ALTER TABLE metadata_entity_foo ADD a_aspectbar JSON; -- add foobar aspect to foo entity ALTER TABLE metadata_entity_foo ADD a_aspectfoobar JSON; +-- add foobaz aspect to foo entity +ALTER TABLE metadata_entity_foo ADD a_aspectfoobaz JSON; + -- add foo aspect to burger entity ALTER TABLE metadata_entity_burger ADD a_aspectfoo JSON; diff --git a/dao-impl/ebean-dao/src/test/resources/ebean-local-access-create-all.sql b/dao-impl/ebean-dao/src/test/resources/ebean-local-access-create-all.sql index 3b5c51f7a..c19bc9e75 100644 --- a/dao-impl/ebean-dao/src/test/resources/ebean-local-access-create-all.sql +++ b/dao-impl/ebean-dao/src/test/resources/ebean-local-access-create-all.sql @@ -81,6 +81,9 @@ ALTER TABLE metadata_entity_foo ADD a_aspectbar JSON; -- add foobar aspect to foo entity ALTER TABLE metadata_entity_foo ADD a_aspectfoobar JSON; +-- add foobaz aspect to foo entity +ALTER TABLE metadata_entity_foo ADD a_aspectfoobaz JSON; + -- add foo aspect to burger entity ALTER TABLE metadata_entity_burger ADD a_aspectfoo JSON; diff --git a/dao-impl/ebean-dao/src/test/resources/ebean-local-dao-create-all-with-non-dollar-virtual-column-names.sql b/dao-impl/ebean-dao/src/test/resources/ebean-local-dao-create-all-with-non-dollar-virtual-column-names.sql index acb76acd3..aecda0a66 100644 --- a/dao-impl/ebean-dao/src/test/resources/ebean-local-dao-create-all-with-non-dollar-virtual-column-names.sql +++ b/dao-impl/ebean-dao/src/test/resources/ebean-local-dao-create-all-with-non-dollar-virtual-column-names.sql @@ -119,6 +119,9 @@ ALTER TABLE metadata_entity_foo_test ADD a_aspectbar JSON; -- add foobar aspect to foo entity ALTER TABLE metadata_entity_foo ADD a_aspectfoobar JSON; +-- add foobaz aspect to foo entity +ALTER TABLE metadata_entity_foo ADD a_aspectfoobaz JSON; + -- add foobar aspect to foo entity ALTER TABLE metadata_entity_foo ADD a_aspectfoobarbaz JSON; diff --git a/dao-impl/ebean-dao/src/test/resources/ebean-local-dao-create-all.sql b/dao-impl/ebean-dao/src/test/resources/ebean-local-dao-create-all.sql index 3bf5b19e1..17722f5ef 100644 --- a/dao-impl/ebean-dao/src/test/resources/ebean-local-dao-create-all.sql +++ b/dao-impl/ebean-dao/src/test/resources/ebean-local-dao-create-all.sql @@ -119,6 +119,9 @@ ALTER TABLE metadata_entity_foo_test ADD a_aspectbar JSON; -- add foobar aspect to foo entity ALTER TABLE metadata_entity_foo ADD a_aspectfoobar JSON; +-- add foobaz aspect to foo entity +ALTER TABLE metadata_entity_foo ADD a_aspectfoobaz JSON; + -- add foobar aspect to foo entity ALTER TABLE metadata_entity_foo ADD a_aspectfoobarbaz JSON; diff --git a/dao-impl/neo4j-dao/build.gradle b/dao-impl/neo4j-dao/build.gradle deleted file mode 100644 index d7022b044..000000000 --- a/dao-impl/neo4j-dao/build.gradle +++ /dev/null @@ -1,14 +0,0 @@ -apply plugin: 'java' - -apply from: "$rootDir/gradle/java-publishing.gradle" - -dependencies { - compile project(':dao-api') - compile externalDependency.neo4jJavaDriver - - compileOnly externalDependency.lombok - annotationProcessor externalDependency.lombok - - testCompile project(':testing:test-models') - testCompile externalDependency.neo4jHarness -} diff --git a/dao-impl/neo4j-dao/src/main/java/com/linkedin/metadata/dao/Neo4jQueryDAO.java b/dao-impl/neo4j-dao/src/main/java/com/linkedin/metadata/dao/Neo4jQueryDAO.java deleted file mode 100644 index 3257f879e..000000000 --- a/dao-impl/neo4j-dao/src/main/java/com/linkedin/metadata/dao/Neo4jQueryDAO.java +++ /dev/null @@ -1,365 +0,0 @@ -package com.linkedin.metadata.dao; - -import com.linkedin.data.template.RecordTemplate; -import com.linkedin.metadata.dao.utils.Statement; -import com.linkedin.metadata.query.Filter; -import com.linkedin.metadata.query.RelationshipDirection; -import com.linkedin.metadata.query.RelationshipFilter; -import com.linkedin.metadata.validator.EntityValidator; -import com.linkedin.metadata.validator.RelationshipValidator; -import java.util.ArrayList; -import java.util.HashMap; -import java.util.List; -import java.util.Map; -import java.util.function.Function; -import java.util.stream.StreamSupport; -import javax.annotation.Nonnull; -import javax.annotation.Nullable; -import org.javatuples.Triplet; -import org.neo4j.driver.Driver; -import org.neo4j.driver.Record; -import org.neo4j.driver.Session; -import org.neo4j.driver.Value; -import org.neo4j.driver.types.Path; - -import static com.linkedin.metadata.dao.Neo4jUtil.*; - - -/** - * An Neo4j implementation of {@link BaseQueryDAO}. - */ -public class Neo4jQueryDAO extends BaseQueryDAO { - - private final Driver _driver; - - public Neo4jQueryDAO(@Nonnull Driver driver) { - this._driver = driver; - } - - @Nonnull - @Override - public List findEntities(@Nonnull Class entityClass, - @Nonnull Filter filter, int offset, int count) { - EntityValidator.validateEntitySchema(entityClass); - - return findEntities(entityClass, findNodes(entityClass, filter, offset, count)); - } - - @Nonnull - @Override - public List findEntities(@Nonnull Class entityClass, - @Nonnull Statement queryStatement) { - EntityValidator.validateEntitySchema(entityClass); - - return runQuery(queryStatement, record -> nodeRecordToEntity(entityClass, record)); - } - - @Nonnull - @Override - public List findMixedTypesEntities(@Nonnull Statement queryStatement) { - return runQuery(queryStatement, this::nodeRecordToEntity); - } - - /** - * Similar to other free form APIs, such as findEntities, findRelationships, findMixedTypesEntities, etc. - * findPaths should be used when there is a specific need to query the graph DB and no existing APIs could be used. - * - * @return A list of paths, each of which should be [Node1, Edge1, Node2, Edge2, ....] - */ - @Nonnull - public List> findPaths(@Nonnull Statement queryStatement) { - return runQuery(queryStatement, this::pathRecordToPathList); - } - - @Nonnull - @Override - public List findEntities( - @Nullable Class sourceEntityClass, @Nonnull Filter sourceEntityFilter, - @Nullable Class destinationEntityClass, @Nonnull Filter destinationEntityFilter, - @Nonnull Class relationshipType, @Nonnull RelationshipFilter relationshipFilter, int minHops, - int maxHops, int offset, int count) { - - if (sourceEntityClass != null) { - EntityValidator.validateEntitySchema(sourceEntityClass); - } - if (destinationEntityClass != null) { - EntityValidator.validateEntitySchema(destinationEntityClass); - } - RelationshipValidator.validateRelationshipSchema(relationshipType); - - final String srcType = getTypeOrEmptyString(sourceEntityClass); - final String srcCriteria = filterToCriteria(sourceEntityFilter); - final String destType = getTypeOrEmptyString(destinationEntityClass); - final String destCriteria = filterToCriteria(destinationEntityFilter); - final String edgeType = getType(relationshipType); - final String edgeCriteria = criterionToString(relationshipFilter.getCriteria()); - - final RelationshipDirection relationshipDirection = relationshipFilter.getDirection(); - - String matchTemplate = "MATCH (src%s %s)-[r:%s*%d..%d %s]-(dest%s %s) RETURN dest"; - if (relationshipDirection == RelationshipDirection.INCOMING) { - matchTemplate = "MATCH (src%s %s)<-[r:%s*%d..%d %s]-(dest%s %s) RETURN dest"; - } else if (relationshipDirection == RelationshipDirection.OUTGOING) { - matchTemplate = "MATCH (src%s %s)-[r:%s*%d..%d %s]->(dest%s %s) RETURN dest"; - } - - final String statementString = - String.format(matchTemplate, srcType, srcCriteria, edgeType, minHops, maxHops, edgeCriteria, destType, - destCriteria); - - final Statement statement = buildStatement(statementString, offset, count); - - return runQuery(statement, this::nodeRecordToEntity); - } - - @Nonnull - @Override - public List findEntities( - @Nullable Class sourceEntityClass, @Nonnull Filter sourceEntityFilter, - @Nonnull List, RelationshipFilter, Class>> traversePaths, int offset, - int count) { - if (sourceEntityClass != null) { - EntityValidator.validateEntitySchema(sourceEntityClass); - } - - final String srcType = getTypeOrEmptyString(sourceEntityClass); - final String srcCriteria = filterToCriteria(sourceEntityFilter); - - StringBuilder matchTemplate = new StringBuilder().append("MATCH (src%s %s)"); - int pathCounter = 0; - - // for each triplet, construct substatement via relationship type + relationship filer + inter entity - for (Triplet, RelationshipFilter, Class> path : traversePaths) { - - pathCounter++; // Cannot use the same relationship variable 'r' or 'dest' for multiple patterns - - final String edgeType = getTypeOrEmptyString(path.getValue0()); - final String edgeCriteria = path.getValue1() == null ? "" : criterionToString(path.getValue1().getCriteria()); - final RelationshipDirection relationshipDirection = - path.getValue1() == null ? RelationshipDirection.UNDIRECTED : path.getValue1().getDirection(); - final String destType = getTypeOrEmptyString(path.getValue2()); - - String subTemplate = "-[r%d%s %s]-(dest%d%s)"; - if (relationshipDirection == RelationshipDirection.INCOMING) { - subTemplate = "<-[r%d%s %s]-(dest%d%s)"; - } else if (relationshipDirection == RelationshipDirection.OUTGOING) { - subTemplate = "-[r%d%s %s]->(dest%d%s)"; - } - String subStatementString = - String.format(subTemplate, pathCounter, edgeType, edgeCriteria, pathCounter, destType); - - matchTemplate.append(subStatementString); - } - - // last INTER_ENTITY will be the Destination Entity - String lastEntity = String.format("dest%d", pathCounter); - matchTemplate.append("RETURN ").append(lastEntity); - - final String statementString = String.format(matchTemplate.toString(), srcType, srcCriteria); - final Statement statement = buildStatement(statementString, offset, count); - - return runQuery(statement, this::nodeRecordToEntity); - } - - @Nonnull - @Override - public List findRelationships( - @Nullable Class sourceEntityClass, @Nonnull Filter sourceEntityFilter, - @Nullable Class destinationEntityClass, @Nonnull Filter destinationEnityFilter, - @Nonnull Class relationshipType, @Nonnull Filter relationshipFilter, int offset, int count) { - if (sourceEntityClass != null) { - EntityValidator.validateEntitySchema(sourceEntityClass); - } - if (destinationEntityClass != null) { - EntityValidator.validateEntitySchema(destinationEntityClass); - } - RelationshipValidator.validateRelationshipSchema(relationshipType); - - return runQuery(findEdges(sourceEntityClass, sourceEntityFilter, destinationEntityClass, destinationEnityFilter, - relationshipType, relationshipFilter, offset, count), record -> edgeRecordToRelationship(relationshipType, record)); - } - - @Nonnull - @Override - public List findRelationships( - @Nonnull Class relationshipClass, @Nonnull Statement queryStatement) { - RelationshipValidator.validateRelationshipSchema(relationshipClass); - - return runQuery(queryStatement, record -> edgeRecordToRelationship(relationshipClass, record)); - } - - @Nonnull - @Override - public List findMixedTypesRelationships(@Nonnull Statement queryStatement) { - return runQuery(queryStatement, this::edgeRecordToRelationship); - } - - @Nonnull - public - List> findPaths( - @Nullable Class sourceEntityClass, @Nonnull Filter sourceEntityFilter, - @Nullable Class destinationEntityClass, @Nonnull Filter destinationEntityFilter, - @Nonnull Class relationshipType, @Nonnull RelationshipFilter relationshipFilter, - int minHops, int maxHops, int offset, int count) { - - if (sourceEntityClass != null) { - EntityValidator.validateEntitySchema(sourceEntityClass); - } - if (destinationEntityClass != null) { - EntityValidator.validateEntitySchema(destinationEntityClass); - } - RelationshipValidator.validateRelationshipSchema(relationshipType); - - final String srcType = getTypeOrEmptyString(sourceEntityClass); - final String srcCriteria = filterToCriteria(sourceEntityFilter); - final String destType = getTypeOrEmptyString(destinationEntityClass); - final String destCriteria = filterToCriteria(destinationEntityFilter); - final String edgeType = getType(relationshipType); - final String edgeCriteria = criterionToString(relationshipFilter.getCriteria()); - - final RelationshipDirection relationshipDirection = relationshipFilter.getDirection(); - - String matchTemplate = "MATCH p=(src%s %s)-[r:%s*%d..%d %s]-(dest%s %s) RETURN p"; - if (relationshipDirection == RelationshipDirection.INCOMING) { - matchTemplate = "MATCH p=(src%s %s)<-[r:%s*%d..%d %s]-(dest%s %s) RETURN p"; - } else if (relationshipDirection == RelationshipDirection.OUTGOING) { - matchTemplate = "MATCH p=(src%s %s)-[r:%s*%d..%d %s]->(dest%s %s) RETURN p"; - } - - final String statementString = - String.format(matchTemplate, srcType, srcCriteria, edgeType, minHops, maxHops, edgeCriteria, destType, - destCriteria); - - final Statement statement = buildStatement(statementString, "length(p), dest.urn", offset, count); - - return runQuery(statement, this::pathRecordToPathList); - } - - /** - * Runs a query statement with parameters and return StatementResult. - * - * @param statement a statement with parameters to be executed - * @param mapperFunction lambda to transform query result - * @return list of elements in the query result - */ - @Nonnull - private List runQuery(@Nonnull Statement statement, @Nonnull Function mapperFunction) { - try (final Session session = _driver.session()) { - return session.run(statement.getCommandText(), statement.getParams()).list(mapperFunction); - } - } - - - /** - * Runs a free-form Cypher query. - * - * @param query Cypher query to be executed - * @return query result as a list of {@link Record} - */ - @Nonnull - public List runFreeFormQuery(@Nonnull String query) { - try (final Session session = _driver.session()) { - return session.run(query).list(); - } - } - - @Nonnull - private Statement findNodes(@Nonnull Class entityClass, - @Nonnull Filter filter, int offset, int count) { - final String nodeType = getType(entityClass); - - final String findTemplate = "MATCH (node:%s %s) RETURN node"; - final String criteria = filterToCriteria(filter); - final String statement = String.format(findTemplate, nodeType, criteria); - - return buildStatement(statement, "node.urn", offset, count); - } - - @Nonnull - private Statement findEdges( - @Nullable Class sourceEntityClass, @Nonnull Filter sourceEnityFilter, - @Nullable Class destinationEnityClass, @Nonnull Filter destinationFilter, - @Nonnull Class relationshipClass, @Nonnull Filter relationshipFilter, int offset, int count) { - - final String srcType = sourceEntityClass == null ? "" : ":" + getType(sourceEntityClass); - final String srcCriteria = filterToCriteria(sourceEnityFilter); - final String destType = destinationEnityClass == null ? "" : ":" + getType(destinationEnityClass); - final String destCriteria = filterToCriteria(destinationFilter); - final String edgeType = getType(relationshipClass); - final String edgeCriteria = filterToCriteria(relationshipFilter); - - final String findTemplate = "MATCH (src%s %s)-[r:%s %s]->(dest%s %s) RETURN src, r, dest"; - final String statement = - String.format(findTemplate, srcType, srcCriteria, edgeType, edgeCriteria, destType, destCriteria); - - return buildStatement(statement, "src.urn", offset, count); - } - - @Nonnull - private Statement buildStatement(@Nonnull String statement, int offset, int count) { - return buildStatement(statement, null, offset, count); - } - - @Nonnull - private Statement buildStatement(@Nonnull String statement, @Nullable String orderBy, int offset, int count) { - String orderStatement = statement; - if (orderBy != null) { - orderStatement += " ORDER BY " + orderBy; - } - - final Map params = new HashMap<>(); - if (offset > 0) { - orderStatement += " SKIP $offset"; - params.put("offset", offset); - } - if (count >= 0) { - orderStatement += " LIMIT $count"; - params.put("count", count); - } - - return new Statement(orderStatement, params); - } - - @Nonnull - ENTITY nodeRecordToEntity(@Nonnull Class entityClass, - @Nonnull Record nodeRecord) { - return nodeToEntity(entityClass, nodeRecord.values().get(0).asNode()); - } - - @Nonnull - private List pathRecordToPathList(@Nonnull Record pathRecord) { - final Path path = pathRecord.values().get(0).asPath(); - final List pathList = new ArrayList<>(); - - StreamSupport.stream(path.spliterator(), false) - .map(Neo4jUtil::pathSegmentToRecordList) - .forEach(segment -> { - if (pathList.isEmpty()) { - pathList.add(segment.get(0)); - } - pathList.add(segment.get(1)); - pathList.add(segment.get(2)); - }); - - return pathList; - } - - @Nonnull - private RecordTemplate nodeRecordToEntity(@Nonnull Record nodeRecord) { - return nodeToEntity(nodeRecord.values().get(0).asNode()); - } - - @Nonnull - private RELATIONSHIP edgeRecordToRelationship( - @Nonnull Class relationshipClass, @Nonnull Record edgeRecord) { - final List values = edgeRecord.values(); - return edgeToRelationship(relationshipClass, values.get(0).asNode(), values.get(2).asNode(), - values.get(1).asRelationship()); - } - - @Nonnull - private RecordTemplate edgeRecordToRelationship(@Nonnull Record edgeRecord) { - final List values = edgeRecord.values(); - return edgeToRelationship(values.get(0).asNode(), values.get(2).asNode(), values.get(1).asRelationship()); - } -} diff --git a/dao-impl/neo4j-dao/src/main/java/com/linkedin/metadata/dao/Neo4jUtil.java b/dao-impl/neo4j-dao/src/main/java/com/linkedin/metadata/dao/Neo4jUtil.java deleted file mode 100644 index f61e35b41..000000000 --- a/dao-impl/neo4j-dao/src/main/java/com/linkedin/metadata/dao/Neo4jUtil.java +++ /dev/null @@ -1,282 +0,0 @@ -package com.linkedin.metadata.dao; - -import com.linkedin.data.DataMap; -import com.linkedin.data.template.RecordTemplate; -import com.linkedin.metadata.dao.utils.RecordUtils; -import com.linkedin.metadata.query.Condition; -import com.linkedin.metadata.query.CriterionArray; -import com.linkedin.metadata.query.Filter; -import com.linkedin.metadata.query.RelationshipDirection; -import com.linkedin.metadata.query.RelationshipFilter; -import java.util.Arrays; -import java.util.Collections; -import java.util.HashMap; -import java.util.List; -import java.util.Map; -import java.util.StringJoiner; -import javax.annotation.Nonnull; -import javax.annotation.Nullable; -import org.apache.commons.lang3.ClassUtils; -import org.neo4j.driver.types.Node; -import org.neo4j.driver.types.Path; -import org.neo4j.driver.types.Relationship; - -import static com.linkedin.metadata.dao.utils.QueryUtils.*; - - -public class Neo4jUtil { - - public static final String URN_FIELD = "urn"; - public static final String SOURCE_FIELD = "source"; - public static final String DESTINATION_FIELD = "destination"; - - private Neo4jUtil() { - // Util class - } - - /** - * Converts ENTITY to node (field:value map). - * - * @param entity ENTITY defined in models - * @return unmodifiable field value map - */ - @Nonnull - public static Map entityToNode(@Nonnull ENTITY entity) { - final Map fields = new HashMap<>(); - - // put all field values - entity.data().forEach((k, v) -> fields.put(k, toValueObject(v))); - - return fields; - } - - /** - * Converts RELATIONSHIP to edge (field:value map), excluding source and destination. - * - * @param relationship RELATIONSHIP defined in models - * @return unmodifiable field value map - */ - @Nonnull - public static Map relationshipToEdge( - @Nonnull RELATIONSHIP relationship) { - final Map fields = new HashMap<>(); - - // put all field values except source and destination - relationship.data().forEach((k, v) -> { - if (!SOURCE_FIELD.equals(k) && !DESTINATION_FIELD.equals(k)) { - fields.put(k, toValueObject(v)); - } - }); - - return Collections.unmodifiableMap(fields); - } - - /** - * Converts RELATIONSHIP to cypher matching criteria, excluding source and destination, e.g. {key: "value"}. - * - * @param relationship RELATIONSHIP defined in models - * @return Criteria String, or "" if no additional fields in relationship - */ - @Nonnull - public static String relationshipToCriteria( - @Nonnull RELATIONSHIP relationship) { - final StringJoiner joiner = new StringJoiner(",", "{", "}"); - - // put all field values except source and destination - relationship.data().forEach((k, v) -> { - if (!SOURCE_FIELD.equals(k) && !DESTINATION_FIELD.equals(k)) { - joiner.add(toCriterionString(k, v)); - } - }); - - return joiner.length() <= 2 ? "" : joiner.toString(); - } - - // Returns self if primitive type, otherwise, return toString() - @Nonnull - private static Object toValueObject(@Nonnull Object obj) { - if (ClassUtils.isPrimitiveOrWrapper(obj.getClass())) { - return obj; - } - - return obj.toString(); - } - - // Returns "key:value" String, if value is not primitive, then use toString() and double quote it - @Nonnull - private static String toCriterionString(@Nonnull String key, @Nonnull Object value) { - if (ClassUtils.isPrimitiveOrWrapper(value.getClass())) { - return key + ":" + value; - } - - return key + ":\"" + value.toString() + "\""; - } - - /** - * Converts {@link Filter} to neo4j query criteria, filter criterion condition requires to be EQUAL. - * - * @param filter Query Filter - * @return Neo4j criteria string - */ - @Nonnull - public static String filterToCriteria(@Nonnull Filter filter) { - return criterionToString(filter.getCriteria()); - } - - /** - * Converts {@link CriterionArray} to neo4j query string. - * - * @param criterionArray CriterionArray in a Filter - * @return Neo4j criteria string - */ - @Nonnull - public static String criterionToString(@Nonnull CriterionArray criterionArray) { - if (!criterionArray.stream().allMatch(criterion -> Condition.EQUAL.equals(criterion.getCondition()))) { - throw new RuntimeException("Neo4j query filter only support EQUAL condition " + criterionArray); - } - - final StringJoiner joiner = new StringJoiner(",", "{", "}"); - - criterionArray.forEach(criterion -> joiner.add(toCriterionString(criterion.getField(), criterion.getValue()))); - - return joiner.length() <= 2 ? "" : joiner.toString(); - } - - /** - * Converts node (field:value map) to ENTITY. - * - * @param entityClass Class of Entity - * @param node Neo4j Node of entityClass type - * @return ENTITY - */ - @Nonnull - public static ENTITY nodeToEntity(@Nonnull Class entityClass, - @Nonnull Node node) { - return RecordUtils.toRecordTemplate(entityClass, new DataMap(node.asMap())); - } - - /** - * Converts node (field:value map) to ENTITY RecordTemplate. - * - * @param node Neo4j Node of entityClass type - * @return RecordTemplate - */ - @Nonnull - public static RecordTemplate nodeToEntity(@Nonnull Node node) { - - final String className = node.labels().iterator().next(); - return RecordUtils.toRecordTemplate(className, new DataMap(node.asMap())); - } - - /** - * Converts path segment (field:value map) list of {@link RecordTemplate}s of nodes and edges. - * - * @param segment the segment of a path containing nodes and edges - */ - @Nonnull - public static List pathSegmentToRecordList(@Nonnull Path.Segment segment) { - final Node startNode = segment.start(); - final Node endNode = segment.end(); - final Relationship edge = segment.relationship(); - - return Arrays.asList( - nodeToEntity(startNode), - edgeToRelationship(startNode, endNode, edge), - nodeToEntity(endNode) - ); - } - - /** - * Converts edge (source-relationship->destination) to RELATIONSHIP. - * - * @param relationshipClass Class of RELATIONSHIP - * @param source Neo4j source Node - * @param destination Neo4j destination Node - * @param relationship Neo4j relationship - * @return ENTITY - */ - @Nonnull - public static RELATIONSHIP edgeToRelationship( - @Nonnull Class relationshipClass, @Nonnull Node source, @Nonnull Node destination, - @Nonnull Relationship relationship) { - - final DataMap dataMap = relationshipDataMap(source, destination, relationship); - return RecordUtils.toRecordTemplate(relationshipClass, dataMap); - } - - /** - * Converts edge (source-relationship->destination) to RELATIONSHIP RecordTemplate. - * - * @param source Neo4j source Node - * @param destination Neo4j destination Node - * @param relationship Neo4j relationship - * @return ENTITY RecordTemplate - */ - @Nonnull - public static RecordTemplate edgeToRelationship(@Nonnull Node source, @Nonnull Node destination, - @Nonnull Relationship relationship) { - - final String className = relationship.type(); - final DataMap dataMap = relationshipDataMap(source, destination, relationship); - return RecordUtils.toRecordTemplate(className, dataMap); - } - - @Nonnull - private static DataMap relationshipDataMap(@Nonnull Node source, @Nonnull Node destination, - @Nonnull Relationship relationship) { - - final DataMap dataMap = new DataMap(relationship.asMap()); - dataMap.put(SOURCE_FIELD, source.get(URN_FIELD).asString()); - dataMap.put(DESTINATION_FIELD, destination.get(URN_FIELD).asString()); - return dataMap; - } - - // Gets the Node/Edge type from an Entity/Relationship, using the backtick-quoted FQCN - @Nonnull - public static String getType(@Nullable RecordTemplate record) { - return record == null ? "" : getType(record.getClass()); - } - - // Gets the Node/Edge type from an Entity/Relationship class, return empty string if null - @Nonnull - public static String getTypeOrEmptyString(@Nullable Class recordClass) { - return recordClass == null ? "" : ":" + getType(recordClass); - } - - // Gets the Node/Edge type from an Entity/Relationship class, using the backtick-quoted FQCN - @Nonnull - public static String getType(@Nonnull Class recordClass) { - return new StringBuilder("`").append(recordClass.getCanonicalName()).append("`").toString(); - } - - /** - * Create {@link RelationshipFilter} using filter and relationship direction. - * @deprecated Use `newRelationshipFilter` in {@link com.linkedin.metadata.dao.utils.QueryUtils} instead. - * - * @param filter {@link Filter} filter - * @param relationshipDirection {@link RelationshipDirection} relationship direction - * @return RelationshipFilter - */ - @Nonnull - @Deprecated - public static RelationshipFilter createRelationshipFilter(@Nonnull Filter filter, - @Nonnull RelationshipDirection relationshipDirection) { - return newRelationshipFilter(filter, relationshipDirection); - } - - /** - * Create {@link RelationshipFilter} using filter conditions and relationship direction. - * @deprecated Use `newRelationshipFilter` in {@link com.linkedin.metadata.dao.utils.QueryUtils} instead. - * - * @param field field to create a filter on - * @param value field value to be filtered - * @param relationshipDirection {@link RelationshipDirection} relationship direction - * @return RelationshipFilter - */ - @Nonnull - @Deprecated - public static RelationshipFilter createRelationshipFilter(@Nonnull String field, @Nonnull String value, - @Nonnull RelationshipDirection relationshipDirection) { - return newRelationshipFilter(field, value, relationshipDirection); - } -} diff --git a/dao-impl/neo4j-dao/src/main/java/com/linkedin/metadata/dao/internal/Neo4jGraphWriterDAO.java b/dao-impl/neo4j-dao/src/main/java/com/linkedin/metadata/dao/internal/Neo4jGraphWriterDAO.java deleted file mode 100644 index aea4c5eee..000000000 --- a/dao-impl/neo4j-dao/src/main/java/com/linkedin/metadata/dao/internal/Neo4jGraphWriterDAO.java +++ /dev/null @@ -1,209 +0,0 @@ -package com.linkedin.metadata.dao.internal; - -import com.linkedin.common.urn.Urn; -import com.linkedin.data.template.RecordTemplate; -import com.linkedin.metadata.dao.utils.GraphUtils; -import java.util.ArrayList; -import java.util.HashSet; -import java.util.List; -import java.util.Set; -import javax.annotation.Nonnull; -import lombok.extern.slf4j.Slf4j; -import org.neo4j.driver.Driver; -import org.neo4j.driver.Query; -import org.neo4j.driver.SessionConfig; - -import static com.linkedin.metadata.dao.Neo4jUtil.*; - - -/** - * An Neo4j implementation of {@link BaseGraphWriterDAO}. - */ -@Slf4j -public class Neo4jGraphWriterDAO extends BaseGraphWriterDAO { - /** - * Event listening interface to consume certain neo4j events and report metrics to some specific metric recording - * framework. - * - *

This allows for recording lower level metrics than just recording how long these method calls take; there is - * other overhead associated with the methods in this class, and these callbacks attempt to more accurately reflect - * how long neo4j transactions are taking. - */ - public interface MetricListener { - /** - * Event when entities are successfully added to the neo4j graph. - * - * @param entityCount how many entities were added in this transaction - * @param updateTimeMs how long the update took in total (across all retries) - * @param retries how many retries were needed before the update was successful (0 means first attempt was a - * success) - */ - void onEntitiesAdded(int entityCount, long updateTimeMs, int retries); - - /** - * Event when relationships are successfully added to the neo4j graph. - * - * @param relationshipCount how many relationships were added in this transaction - * @param updateTimeMs how long the update took in total (across all retries) - * @param retries how many retries were needed before the update was successful (0 means first attempt was a - * success) - */ - void onRelationshipsAdded(int relationshipCount, long updateTimeMs, int retries); - - /** - * Event when entities are successfully removed from the neo4j graph. - * - * @param entityCount how many entities were removed in this transaction - * @param updateTimeMs how long the update took in total (across all retries) - * @param retries how many retries were needed before the update was successful (0 means first attempt was a - * success) - */ - void onEntitiesRemoved(int entityCount, long updateTimeMs, int retries); - - /** - * Event when relationships are successfully removed from the neo4j graph. - * - * @param relationshipCount how many relationships were added in this transaction - * @param updateTimeMs how long the update took in total (across all retries) - * @param retries how many retries were needed before the update was successful (0 means first attempt was a - * success) - */ - void onRelationshipsRemoved(int relationshipCount, long updateTimeMs, int retries); - } - - private static final class DelegateMetricListener implements MetricListener { - private final Set _metricListeners = new HashSet<>(); - - void addMetricListener(@Nonnull MetricListener metricListener) { - _metricListeners.add(metricListener); - } - - @Override - public void onEntitiesAdded(int entityCount, long updateTimeMs, int retries) { - for (MetricListener m : _metricListeners) { - m.onEntitiesAdded(entityCount, updateTimeMs, retries); - } - } - - @Override - public void onRelationshipsAdded(int relationshipCount, long updateTimeMs, int retries) { - for (MetricListener m : _metricListeners) { - m.onRelationshipsAdded(relationshipCount, updateTimeMs, retries); - } - } - - @Override - public void onEntitiesRemoved(int entityCount, long updateTimeMs, int retries) { - for (MetricListener m : _metricListeners) { - m.onEntitiesRemoved(entityCount, updateTimeMs, retries); - } - } - - @Override - public void onRelationshipsRemoved(int relationshipCount, long updateTimeMs, int retries) { - for (MetricListener m : _metricListeners) { - m.onRelationshipsRemoved(relationshipCount, updateTimeMs, retries); - } - } - } - - private DelegateMetricListener _metricListener = new DelegateMetricListener(); - private final Neo4jQueriesTransformer _queriesTransformer; - private final Neo4jQueryExecutor _queryExecutor; - - private Neo4jGraphWriterDAO(@Nonnull Neo4jQueriesTransformer queriesTransformer, - @Nonnull Neo4jQueryExecutor queryExecutor) { - _queriesTransformer = queriesTransformer; - _queryExecutor = queryExecutor; - } - - public Neo4jGraphWriterDAO(@Nonnull Driver driver) { - this(new Neo4jQueriesTransformer(), new Neo4jQueryExecutor(driver)); - } - - /* Should only be used for testing */ - public Neo4jGraphWriterDAO(@Nonnull Driver driver, @Nonnull Set> allEntities) { - this(new Neo4jQueriesTransformer(allEntities), new Neo4jQueryExecutor(driver)); - } - - /** - * WARNING: Do NOT use this! This is not tested yet. - * Multi-DB support comes with Neo4j 4+. - * Although DAO works with Neo4j 4+, we can't bump Neo4j test harness to 4+ to test this because it needs Java 11 - * And Java 11 build is blocked by ES7 migration. - */ - public Neo4jGraphWriterDAO(@Nonnull Driver driver, @Nonnull String databaseName) { - this(new Neo4jQueriesTransformer(), new Neo4jQueryExecutor(driver, SessionConfig.forDatabase(databaseName))); - } - - public void addMetricListener(@Nonnull MetricListener metricListener) { - _metricListener.addMetricListener(metricListener); - } - - @Override - public void addEntities(@Nonnull List entities) { - final List list = new ArrayList<>(); - - for (ENTITY entity : entities) { - list.add(_queriesTransformer.addEntityQuery(entity)); - } - - final Neo4jQueryResult result = _queryExecutor.execute(list); - log.trace("Added {} entities over {} retries, which took {} millis", entities.size(), result.getTookMs(), - result.getRetries()); - _metricListener.onEntitiesAdded(entities.size(), result.getTookMs(), result.getRetries()); - } - - @Override - public void removeEntities(@Nonnull List urns) { - final List list = new ArrayList<>(); - for (URN urn : urns) { - list.add(_queriesTransformer.removeEntityQuery(urn)); - } - - final Neo4jQueryResult result = _queryExecutor.execute(list); - log.trace("Removed {} entities over {} retries, which took {} millis", urns.size(), result.getTookMs(), - result.getRetries()); - _metricListener.onEntitiesRemoved(urns.size(), result.getTookMs(), result.getRetries()); - } - - @Override - public void addRelationships(@Nonnull List relationships, - @Nonnull RemovalOption removalOption, boolean isTestMode) { - if (relationships.isEmpty()) { - return; - } - - final List list = new ArrayList<>(); - - _queriesTransformer.relationshipRemovalOptionQuery(relationships.get(0), removalOption).ifPresent(list::add); - - GraphUtils.checkSameUrn(relationships, removalOption, SOURCE_FIELD, DESTINATION_FIELD); - - for (RELATIONSHIP relationship : relationships) { - list.add(_queriesTransformer.addRelationshipQuery(relationship)); - } - - final Neo4jQueryResult result = _queryExecutor.execute(list); - log.trace("Added {} relationships over {} retries, which took {} millis", relationships.size(), result.getTookMs(), - result.getRetries()); - _metricListener.onRelationshipsAdded(relationships.size(), result.getTookMs(), result.getRetries()); - } - - @Override - public void removeRelationships(@Nonnull List relationships) { - if (relationships.isEmpty()) { - return; - } - - final List list = new ArrayList<>(); - for (RELATIONSHIP relationship : relationships) { - list.add(_queriesTransformer.removeEdge(relationship)); - } - - final Neo4jQueryResult result = _queryExecutor.execute(list); - log.trace("Removed {} relationships over {} retries, which took {} millis", relationships.size(), - result.getTookMs(), result.getRetries()); - _metricListener.onRelationshipsRemoved(relationships.size(), result.getTookMs(), result.getRetries()); - } -} diff --git a/dao-impl/neo4j-dao/src/main/java/com/linkedin/metadata/dao/internal/Neo4jQueriesTransformer.java b/dao-impl/neo4j-dao/src/main/java/com/linkedin/metadata/dao/internal/Neo4jQueriesTransformer.java deleted file mode 100644 index cb0f9eaa0..000000000 --- a/dao-impl/neo4j-dao/src/main/java/com/linkedin/metadata/dao/internal/Neo4jQueriesTransformer.java +++ /dev/null @@ -1,195 +0,0 @@ -package com.linkedin.metadata.dao.internal; - -import com.linkedin.common.urn.Urn; -import com.linkedin.data.template.RecordTemplate; -import com.linkedin.metadata.validator.EntityValidator; -import com.linkedin.metadata.validator.RelationshipValidator; -import java.util.HashMap; -import java.util.Map; -import java.util.Optional; -import java.util.Set; -import javax.annotation.Nonnull; -import org.neo4j.driver.Query; - -import static com.linkedin.metadata.dao.Neo4jUtil.*; -import static com.linkedin.metadata.dao.utils.ModelUtils.*; - - -/** - * Can transform GMA entities and relationships into Neo4j queries for upserting. - * - *

This separates out transformation logic from query execution logic ({@link Neo4jGraphWriterDAO}). - */ -public final class Neo4jQueriesTransformer { - private static final Map DEFAULT_URN_TO_ENTITY_MAP = buildUrnToEntityMap(getAllEntities()); - private final Map _urnToEntityMap; - - public Neo4jQueriesTransformer() { - this(DEFAULT_URN_TO_ENTITY_MAP); - } - - /** - * For use in unit testing. - */ - public Neo4jQueriesTransformer(@Nonnull Set> entitiesSet) { - this(buildUrnToEntityMap(entitiesSet)); - } - - private Neo4jQueriesTransformer(@Nonnull Map urnToEntityMap) { - _urnToEntityMap = urnToEntityMap; - } - - @Nonnull - private static Map buildUrnToEntityMap(@Nonnull Set> entitiesSet) { - Map map = new HashMap<>(); - for (Class entity : entitiesSet) { - if (map.put(getEntityTypeFromUrnClass(urnClassForEntity(entity)), getType(entity)) != null) { - throw new IllegalStateException("Duplicate key"); - } - } - return map; - } - - @Nonnull - private Object toPropertyValue(@Nonnull Object obj) { - if (obj instanceof Urn) { - return obj.toString(); - } - return obj; - } - - // visible for testing - @Nonnull - String getNodeType(@Nonnull Urn urn) { - return ":" + _urnToEntityMap.getOrDefault(urn.getEntityType(), "UNKNOWN"); - } - - @Nonnull - private Query buildQuery(@Nonnull String queryTemplate, @Nonnull Map params) { - for (Map.Entry entry : params.entrySet()) { - String k = entry.getKey(); - Object v = entry.getValue(); - params.put(k, toPropertyValue(v)); - } - return new Query(queryTemplate, params); - } - - @Nonnull - public Query addEntityQuery(@Nonnull RecordTemplate entity) { - EntityValidator.validateEntitySchema(entity.getClass()); - - final Urn urn = getUrnFromEntity(entity); - final String nodeType = getNodeType(urn); - - // Use += to ensure this doesn't override the node but merges in the new properties to allow for partial updates. - final String mergeTemplate = "MERGE (node%s {urn: $urn}) SET node += $properties RETURN node"; - final String statement = String.format(mergeTemplate, nodeType); - - final Map params = new HashMap<>(); - params.put("urn", urn.toString()); - final Map props = entityToNode(entity); - props.remove("urn"); // no need to set twice (this is implied by MERGE), and they can be quite long. - params.put("properties", props); - - return buildQuery(statement, params); - } - - @Nonnull - public Query removeEntityQuery(@Nonnull Urn urn) { - // also delete any relationship going to or from it - final String nodeType = getNodeType(urn); - - final String matchTemplate = "MATCH (node%s {urn: $urn}) DETACH DELETE node"; - final String statement = String.format(matchTemplate, nodeType); - - final Map params = new HashMap<>(); - params.put("urn", urn.toString()); - - return buildQuery(statement, params); - } - - @Nonnull - public Optional relationshipRemovalOptionQuery(@Nonnull RecordTemplate relationship, - BaseGraphWriterDAO.RemovalOption removalOption) { - // remove existing edges according to RemovalOption - final Urn source0Urn = getSourceUrnFromRelationship(relationship); - final Urn destination0Urn = getDestinationUrnFromRelationship(relationship); - final String relationType = getType(relationship); - - final String sourceType = getNodeType(source0Urn); - final String destinationType = getNodeType(destination0Urn); - - final Map params = new HashMap<>(); - - if (removalOption == BaseGraphWriterDAO.RemovalOption.REMOVE_ALL_EDGES_FROM_SOURCE) { - final String removeTemplate = "MATCH (source%s {urn: $urn})-[relation:%s]->() DELETE relation"; - final String statement = String.format(removeTemplate, sourceType, relationType); - - params.put("urn", source0Urn.toString()); - - return Optional.of(buildQuery(statement, params)); - } else if (removalOption == BaseGraphWriterDAO.RemovalOption.REMOVE_ALL_EDGES_TO_DESTINATION) { - final String removeTemplate = "MATCH ()-[relation:%s]->(destination%s {urn: $urn}) DELETE relation"; - final String statement = String.format(removeTemplate, relationType, destinationType); - - params.put("urn", destination0Urn.toString()); - - return Optional.of(buildQuery(statement, params)); - } else if (removalOption == BaseGraphWriterDAO.RemovalOption.REMOVE_ALL_EDGES_FROM_SOURCE_TO_DESTINATION) { - final String removeTemplate = - "MATCH (source%s {urn: $sourceUrn})-[relation:%s]->(destination%s {urn: $destinationUrn}) DELETE relation"; - final String statement = String.format(removeTemplate, sourceType, relationType, destinationType); - - params.put("sourceUrn", source0Urn.toString()); - params.put("destinationUrn", destination0Urn.toString()); - - return Optional.of(buildQuery(statement, params)); - } - - return Optional.empty(); - } - - @Nonnull - public Query addRelationshipQuery(@Nonnull RecordTemplate relationship) { - RelationshipValidator.validateRelationshipSchema(relationship.getClass()); - final Urn srcUrn = getSourceUrnFromRelationship(relationship); - final Urn destUrn = getDestinationUrnFromRelationship(relationship); - final String sourceNodeType = getNodeType(srcUrn); - final String destinationNodeType = getNodeType(destUrn); - - // Add/Update relationship. Use MERGE on nodes to prevent needing to have separate queries to create them. - final String mergeRelationshipTemplate = - "MERGE (source%s {urn: $sourceUrn}) " + "MERGE (destination%s {urn: $destinationUrn}) " - + "MERGE (source)-[r:%s]->(destination) SET r += $properties"; - final String statement = - String.format(mergeRelationshipTemplate, sourceNodeType, destinationNodeType, getType(relationship)); - - final Map paramsMerge = new HashMap<>(); - paramsMerge.put("sourceUrn", srcUrn.toString()); - paramsMerge.put("destinationUrn", destUrn.toString()); - paramsMerge.put("properties", relationshipToEdge(relationship)); - - return new Query(statement, paramsMerge); - } - - @Nonnull - public Query removeEdge(@Nonnull RecordTemplate relationship) { - final Urn sourceUrn = getSourceUrnFromRelationship(relationship); - final Urn destinationUrn = getDestinationUrnFromRelationship(relationship); - - final String sourceType = getNodeType(sourceUrn); - final String destinationType = getNodeType(destinationUrn); - - final String removeMatchTemplate = - "MATCH (source%s {urn: $sourceUrn})-[relation:%s %s]->(destination%s {urn: $destinationUrn}) DELETE relation"; - final String criteria = relationshipToCriteria(relationship); - final String statement = - String.format(removeMatchTemplate, sourceType, getType(relationship), criteria, destinationType); - - final Map params = new HashMap<>(); - params.put("sourceUrn", sourceUrn.toString()); - params.put("destinationUrn", destinationUrn.toString()); - - return buildQuery(statement, params); - } -} diff --git a/dao-impl/neo4j-dao/src/main/java/com/linkedin/metadata/dao/internal/Neo4jQueryExecutor.java b/dao-impl/neo4j-dao/src/main/java/com/linkedin/metadata/dao/internal/Neo4jQueryExecutor.java deleted file mode 100644 index a4ebbb726..000000000 --- a/dao-impl/neo4j-dao/src/main/java/com/linkedin/metadata/dao/internal/Neo4jQueryExecutor.java +++ /dev/null @@ -1,64 +0,0 @@ -package com.linkedin.metadata.dao.internal; - -import com.linkedin.metadata.dao.exception.RetryLimitReached; -import java.util.List; -import javax.annotation.Nonnull; -import org.apache.commons.lang.time.StopWatch; -import org.neo4j.driver.Driver; -import org.neo4j.driver.Query; -import org.neo4j.driver.Session; -import org.neo4j.driver.SessionConfig; -import org.neo4j.driver.exceptions.Neo4jException; - - -public final class Neo4jQueryExecutor { - private static final int MAX_TRANSACTION_RETRY = 3; - private final Driver _driver; - private final SessionConfig _sessionConfig; - - public Neo4jQueryExecutor(@Nonnull Driver driver, @Nonnull SessionConfig sessionConfig) { - _driver = driver; - _sessionConfig = sessionConfig; - } - - public Neo4jQueryExecutor(@Nonnull Driver driver) { - this(driver, SessionConfig.defaultConfig()); - } - - /** - * Executes a list of queries with parameters in one transaction. - * - * @param queries List of queries with parameters to be executed in order - */ - @Nonnull - public Neo4jQueryResult execute(@Nonnull List queries) { - int retry = 0; - final StopWatch stopWatch = new StopWatch(); - stopWatch.start(); - Exception lastException; - try (final Session session = _driver.session(_sessionConfig)) { - do { - try { - session.writeTransaction(tx -> { - for (Query query : queries) { - tx.run(query); - } - return null; - }); - lastException = null; - break; - } catch (Neo4jException e) { - lastException = e; - } - } while (++retry <= MAX_TRANSACTION_RETRY); - } - - if (lastException != null) { - throw new RetryLimitReached( - "Failed to execute Neo4j write transaction after " + MAX_TRANSACTION_RETRY + " retries", lastException); - } - - stopWatch.stop(); - return Neo4jQueryResult.builder().tookMs(stopWatch.getTime()).retries(retry).build(); - } -} diff --git a/dao-impl/neo4j-dao/src/main/java/com/linkedin/metadata/dao/internal/Neo4jQueryResult.java b/dao-impl/neo4j-dao/src/main/java/com/linkedin/metadata/dao/internal/Neo4jQueryResult.java deleted file mode 100644 index ffe058ec1..000000000 --- a/dao-impl/neo4j-dao/src/main/java/com/linkedin/metadata/dao/internal/Neo4jQueryResult.java +++ /dev/null @@ -1,12 +0,0 @@ -package com.linkedin.metadata.dao.internal; - -import lombok.Builder; -import lombok.Data; - - -@Builder -@Data -public final class Neo4jQueryResult { - private final long tookMs; - private final int retries; -} diff --git a/dao-impl/neo4j-dao/src/test/java/com/linkedin/metadata/dao/Neo4jQueryDAOTest.java b/dao-impl/neo4j-dao/src/test/java/com/linkedin/metadata/dao/Neo4jQueryDAOTest.java deleted file mode 100644 index 36f4114d7..000000000 --- a/dao-impl/neo4j-dao/src/test/java/com/linkedin/metadata/dao/Neo4jQueryDAOTest.java +++ /dev/null @@ -1,623 +0,0 @@ -package com.linkedin.metadata.dao; - -import com.linkedin.data.template.RecordTemplate; -import com.linkedin.metadata.dao.internal.Neo4jGraphWriterDAO; -import com.linkedin.metadata.dao.utils.Statement; -import com.linkedin.metadata.query.Filter; -import com.linkedin.metadata.query.RelationshipDirection; -import com.linkedin.metadata.query.RelationshipFilter; -import com.linkedin.testing.EntityBar; -import com.linkedin.testing.EntityBaz; -import com.linkedin.testing.EntityFoo; -import com.linkedin.testing.RelationshipBar; -import com.linkedin.testing.RelationshipFoo; -import com.linkedin.testing.TestUtils; -import com.linkedin.testing.urn.BarUrn; -import com.linkedin.testing.urn.BazUrn; -import com.linkedin.testing.urn.FooUrn; -import java.util.ArrayList; -import java.util.Arrays; -import java.util.Collections; -import java.util.HashMap; -import java.util.List; -import java.util.Map; -import java.util.stream.Collectors; -import org.javatuples.Triplet; -import org.neo4j.driver.Driver; -import org.neo4j.driver.GraphDatabase; -import org.neo4j.driver.Record; -import org.testng.annotations.AfterMethod; -import org.testng.annotations.BeforeMethod; -import org.testng.annotations.Test; - -import static com.linkedin.metadata.dao.utils.QueryUtils.*; -import static com.linkedin.testing.TestUtils.*; -import static org.testng.Assert.*; - - -public class Neo4jQueryDAOTest { - - private Neo4jTestServerBuilder _serverBuilder; - private Neo4jQueryDAO _dao; - private Neo4jGraphWriterDAO _writer; - - @BeforeMethod - public void init() { - _serverBuilder = new Neo4jTestServerBuilder(); - _serverBuilder.newServer(); - - final Driver driver = GraphDatabase.driver(_serverBuilder.boltURI()); - _dao = new Neo4jQueryDAO(driver); - _writer = new Neo4jGraphWriterDAO(driver, TestUtils.getAllTestEntities()); - } - - @AfterMethod - public void tearDown() { - _serverBuilder.shutdown(); - } - - @Test - public void testFindEntityByUrn() throws Exception { - FooUrn urn = makeFooUrn(1); - EntityFoo entity = new EntityFoo().setUrn(urn).setValue("foo"); - - _writer.addEntity(entity); - - Filter filter = newFilter("urn", urn.toString()); - List found = _dao.findEntities(EntityFoo.class, filter, -1, -1); - - assertEquals(found.size(), 1); - assertEquals(found.get(0), entity); - } - - @Test - public void testFindEntityByAttribute() throws Exception { - FooUrn urn1 = makeFooUrn(1); - EntityFoo entity1 = new EntityFoo().setUrn(urn1).setValue("foo"); - _writer.addEntity(entity1); - - FooUrn urn2 = makeFooUrn(2); - EntityFoo entity2 = new EntityFoo().setUrn(urn2).setValue("foo"); - _writer.addEntity(entity2); - - // find by removed - Filter filter1 = newFilter("value", "foo"); - List found = _dao.findEntities(EntityFoo.class, filter1, -1, -1); - assertEquals(found, Arrays.asList(entity1, entity2)); - - // find with limit - found = _dao.findEntities(EntityFoo.class, filter1, -1, 1); - assertEquals(found, Collections.singletonList(entity1)); - - // find with offset - found = _dao.findEntities(EntityFoo.class, filter1, 1, -1); - assertEquals(found, Collections.singletonList(entity2)); - } - - @Test - public void testFindEntityByQuery() throws Exception { - FooUrn urn1 = makeFooUrn(1); - EntityFoo entity1 = new EntityFoo().setUrn(urn1).setValue("foo"); - _writer.addEntity(entity1); - - Statement statement = new Statement("MATCH (n {value:\"foo\"}) RETURN n ORDER BY n.urn", Collections.emptyMap()); - List found = _dao.findEntities(EntityFoo.class, statement); - assertEquals(found.size(), 1); - assertEquals(found.get(0), entity1); - - FooUrn urn2 = makeFooUrn(2); - EntityFoo entity2 = new EntityFoo().setUrn(urn2).setValue("foo"); - _writer.addEntity(entity2); - - found = _dao.findEntities(EntityFoo.class, statement); - assertEquals(found.size(), 2); - assertEquals(found.get(0), entity1); - assertEquals(found.get(1), entity2); - } - - @Test - public void testFindEntityWithOneRelationship() throws Exception { - // Test interface 1 & 3 - FooUrn urn1 = makeFooUrn(1); - EntityFoo entity1 = new EntityFoo().setUrn(urn1).setValue("foo1"); - _writer.addEntity(entity1); - - FooUrn urn2 = makeFooUrn(2); - EntityFoo entity2 = new EntityFoo().setUrn(urn2).setValue("foo2"); - _writer.addEntity(entity2); - - FooUrn urn3 = makeFooUrn(3); - EntityFoo entity3 = new EntityFoo().setUrn(urn3).setValue("foo3"); - _writer.addEntity(entity3); - - BarUrn urn4 = makeBarUrn(4); - EntityBar entity4 = new EntityBar().setUrn(urn4).setValue("bar4"); - _writer.addEntity(entity4); - - BarUrn urn5 = makeBarUrn(5); - EntityBar entity5 = new EntityBar().setUrn(urn5).setValue("bar5"); - _writer.addEntity(entity5); - - // create relationship urn1 -(r:Foo)-> urn2 -(r:Foo)-> urn3 (example use case: ReportTo list) - // also relationship urn1 -(r:Foo)-> urn4, and urn1 -(r:Bar)-> urn5 - RelationshipFoo relationshipFoo1To2 = new RelationshipFoo().setSource(urn1).setDestination(urn2); - _writer.addRelationship(relationshipFoo1To2, false); - - RelationshipFoo relationshipFoo2o3 = new RelationshipFoo().setSource(urn2).setDestination(urn3).setType("apa"); - _writer.addRelationship(relationshipFoo2o3, false); - - RelationshipFoo relationshipFoo1o4 = new RelationshipFoo().setSource(urn1).setDestination(urn4); - _writer.addRelationship(relationshipFoo1o4, false); - - RelationshipBar relationshipBar1o5 = new RelationshipBar().setSource(urn1).setDestination(urn5); - _writer.addRelationship(relationshipBar1o5, false); - - // test source filter with urn - Filter sourceFilter = newFilter("urn", urn2.toString()); - - // use case: the direct reportee to me (urn2) - List resultIncoming = - _dao.findEntities(EntityFoo.class, sourceFilter, EntityFoo.class, EMPTY_FILTER, RelationshipFoo.class, - newRelationshipFilter(EMPTY_FILTER, RelationshipDirection.INCOMING), 0, 10); - assertEquals(resultIncoming.size(), 1); - assertEquals(resultIncoming.get(0), entity1); - - // use case: the manager I (urn2) am report to - List resultOutgoing = - _dao.findEntities(EntityFoo.class, sourceFilter, EntityFoo.class, EMPTY_FILTER, RelationshipFoo.class, - newRelationshipFilter(EMPTY_FILTER, RelationshipDirection.OUTGOING), 0, 10); - assertEquals(resultOutgoing.size(), 1); - assertEquals(resultOutgoing.get(0), entity3); - - // use case: give me my friends at one degree/hop - List resultUndirected = - _dao.findEntities(EntityFoo.class, sourceFilter, EntityFoo.class, EMPTY_FILTER, RelationshipFoo.class, - newRelationshipFilter(EMPTY_FILTER, RelationshipDirection.UNDIRECTED), 0, 10); - assertEquals(resultUndirected.size(), 2); - assertEquals(resultUndirected.get(0), entity1); - assertEquals(resultUndirected.get(1), entity3); - - // Test: if dest entity type is not specified, from urn1, it will return a mixed collection of entities like urn2 and urn4 - // urn5 should not be returned because it is based on RelationshipBar. - sourceFilter = newFilter("urn", urn1.toString()); - List resultNullDest = - _dao.findEntities(EntityFoo.class, sourceFilter, null, EMPTY_FILTER, RelationshipFoo.class, - newRelationshipFilter(EMPTY_FILTER, RelationshipDirection.OUTGOING), 0, 10); - assertEquals(resultNullDest.size(), 2); - assertEquals(resultNullDest.get(0), entity2); - assertEquals(resultNullDest.get(1), entity4); - - // Test: source filter on other attributes - sourceFilter = newFilter("value", "foo2"); - List result = - _dao.findEntities(EntityFoo.class, sourceFilter, EntityFoo.class, EMPTY_FILTER, RelationshipFoo.class, - newRelationshipFilter(EMPTY_FILTER, RelationshipDirection.INCOMING), 0, 10); - assertEquals(result, resultIncoming); - - //Test relationship filters on the attributes - Filter filter = newFilter("type", "apa"); - List resultWithRelationshipFilter = - _dao.findEntities(EntityFoo.class, sourceFilter, EntityFoo.class, EMPTY_FILTER, RelationshipFoo.class, - newRelationshipFilter(filter, RelationshipDirection.OUTGOING), 0, 10); - assertEquals(resultWithRelationshipFilter.size(), 1); - assertEquals(resultWithRelationshipFilter.get(0), entity3); - - //Test a wrong value for relationship filters - Filter relationshipFilterWrongValue = newFilter("type", "wrongValue"); - List resultWithRelationshipFilter2 = - _dao.findEntities(EntityFoo.class, sourceFilter, EntityFoo.class, EMPTY_FILTER, RelationshipFoo.class, - newRelationshipFilter(relationshipFilterWrongValue, RelationshipDirection.OUTGOING), 0, 10); - assertEquals(resultWithRelationshipFilter2.size(), 0); - } - - @Test - public void testFindEntitiesMultiHops() throws Exception { - // Test interface 5 - FooUrn urn1 = makeFooUrn(1); - EntityFoo entity1 = new EntityFoo().setUrn(urn1).setValue("foo1"); - _writer.addEntity(entity1); - - FooUrn urn2 = makeFooUrn(2); - EntityFoo entity2 = new EntityFoo().setUrn(urn2).setValue("foo2"); - _writer.addEntity(entity2); - - FooUrn urn3 = makeFooUrn(3); - EntityFoo entity3 = new EntityFoo().setUrn(urn3).setValue("foo3"); - _writer.addEntity(entity3); - - // create relationship urn1 -> urn2 -> urn3 (use case: ReportTo list) - RelationshipFoo relationshipFoo1To2 = new RelationshipFoo().setSource(urn1).setDestination(urn2); - _writer.addRelationship(relationshipFoo1To2, false); - - RelationshipFoo relationshipFoo2o3 = new RelationshipFoo().setSource(urn2).setDestination(urn3); - _writer.addRelationship(relationshipFoo2o3, false); - - // From urn1, get result with one hop, such as direct manager - Filter sourceFilter = newFilter("urn", urn1.toString()); - List result = - _dao.findEntities(EntityFoo.class, sourceFilter, EntityFoo.class, EMPTY_FILTER, RelationshipFoo.class, - newRelationshipFilter(EMPTY_FILTER, RelationshipDirection.OUTGOING), 1, 1, 0, 10); - assertEquals(result.size(), 1); - assertEquals(result.get(0), entity2); - - // get result with 2 hops, two managers - List result2 = - _dao.findEntities(EntityFoo.class, sourceFilter, EntityFoo.class, EMPTY_FILTER, RelationshipFoo.class, - newRelationshipFilter(EMPTY_FILTER, RelationshipDirection.OUTGOING), 1, 2, 0, 10); - assertEquals(result2.size(), 2); - assertEquals(result2.get(0), entity2); - assertEquals(result2.get(1), entity3); - - // get result with >= 3 hops, until end of the chain - List result3 = - _dao.findEntities(EntityFoo.class, sourceFilter, EntityFoo.class, EMPTY_FILTER, RelationshipFoo.class, - newRelationshipFilter(EMPTY_FILTER, RelationshipDirection.OUTGOING), 1, 3, 0, 10); - assertEquals(result3.size(), 2); - assertEquals(result3.get(0), entity2); - assertEquals(result3.get(1), entity3); - - // test dest filter, filter the list of result and only get urn3 back - Filter destFilter = newFilter("urn", urn3.toString()); - List result4 = - _dao.findEntities(EntityFoo.class, sourceFilter, EntityFoo.class, destFilter, RelationshipFoo.class, - newRelationshipFilter(EMPTY_FILTER, RelationshipDirection.OUTGOING), 1, 3, 0, 10); - assertEquals(result4.size(), 1); - assertEquals(result4.get(0), entity3); - - // test relationship filter: TODO: for all the interfaces add relationship filter testing - - // corner cases 1: minHops set to 3, get no result because the relationship chain reaches the end - List result5 = - _dao.findEntities(EntityFoo.class, sourceFilter, EntityFoo.class, EMPTY_FILTER, RelationshipFoo.class, - newRelationshipFilter(EMPTY_FILTER, RelationshipDirection.OUTGOING), 3, 6, 0, 10); - assertEquals(result5.size(), 0); - - // corner cases 2: minHops < maxHops, return no result - List result6 = - _dao.findEntities(EntityFoo.class, sourceFilter, EntityFoo.class, EMPTY_FILTER, RelationshipFoo.class, - newRelationshipFilter(EMPTY_FILTER, RelationshipDirection.OUTGOING), 6, 0, 0, 10); - assertEquals(result6.size(), 0); - - // test the relationship directions - // From urn2, get result with one hop, such as direct manager. - // NOTE: without direction specified in the matchTemplate, urn1 could end up in the result too - Filter sourceFilter2 = newFilter("urn", urn2.toString()); - List result21 = - _dao.findEntities(EntityFoo.class, sourceFilter2, EntityFoo.class, EMPTY_FILTER, RelationshipFoo.class, - newRelationshipFilter(EMPTY_FILTER, RelationshipDirection.OUTGOING), 1, 1, 0, 10); - assertEquals(result21.size(), 1); - assertEquals(result21.get(0), entity3); - - // get result with 2 hops, 1 manager - List result22 = - _dao.findEntities(EntityFoo.class, sourceFilter2, EntityFoo.class, EMPTY_FILTER, RelationshipFoo.class, - newRelationshipFilter(EMPTY_FILTER, RelationshipDirection.OUTGOING), 1, 2, 0, 10); - assertEquals(result22.size(), 1); - assertEquals(result22.get(0), entity3); - - // let's see what we get if we use interface 1, it should return urn3 only - List resultInterface1 = - _dao.findEntities(EntityFoo.class, sourceFilter2, null, EMPTY_FILTER, RelationshipFoo.class, - newRelationshipFilter(EMPTY_FILTER, RelationshipDirection.OUTGOING), 0, 10); - assertEquals(resultInterface1.size(), 1); - assertEquals(resultInterface1.get(0), entity3); - - // test INCOMING direction, use case such as who report to URN3 - Filter sourceFilter3 = newFilter("urn", urn3.toString()); - List resultIncoming = - _dao.findEntities(EntityFoo.class, sourceFilter3, EntityFoo.class, EMPTY_FILTER, RelationshipFoo.class, - newRelationshipFilter(EMPTY_FILTER, RelationshipDirection.INCOMING), 1, 2, 0, 10); - assertEquals(resultIncoming.size(), 2); - assertEquals(resultIncoming.get(0), entity2); - assertEquals(resultIncoming.get(1), entity1); - } - - @Test - public void testFindEntitiesViaTraversePathes() throws Exception { - // Test interface 4 - FooUrn urn1 = makeFooUrn(1); - EntityFoo entity1 = new EntityFoo().setUrn(urn1).setValue("CorpGroup1"); - _writer.addEntity(entity1); - - BarUrn urn2 = makeBarUrn(2); - EntityBar entity2 = new EntityBar().setUrn(urn2).setValue("CorpUser2"); - _writer.addEntity(entity2); - - BazUrn urn3 = makeBazUrn(3); - EntityBaz entity3 = new EntityBaz().setUrn(urn3).setValue("Dataset3"); - _writer.addEntity(entity3); - - // create relationship urn1 -> urn2 with RelationshipFoo (ex: CorpGroup1 HasMember CorpUser2) - RelationshipFoo relationshipFoo1To2 = new RelationshipFoo().setSource(urn1).setDestination(urn2); - _writer.addRelationship(relationshipFoo1To2, false); - - // create relationship urn2 -> urn3 with RelationshipBar (ex: Dataset3 is OwnedBy CorpUser2) - RelationshipBar relationshipFoo2o3 = new RelationshipBar().setSource(urn3).setDestination(urn2); - _writer.addRelationship(relationshipFoo2o3, false); - - // test source filter with urn - Filter sourceFilter = newFilter("urn", urn1.toString()); - - // use case: return all the datasets owned by corpgroup1 - List paths = new ArrayList(); - paths.add( - Triplet.with(RelationshipFoo.class, newRelationshipFilter(EMPTY_FILTER, RelationshipDirection.OUTGOING), - EntityBar.class)); - paths.add( - Triplet.with(RelationshipBar.class, newRelationshipFilter(EMPTY_FILTER, RelationshipDirection.INCOMING), - EntityBaz.class)); - List result = _dao.findEntities(EntityFoo.class, sourceFilter, paths, 0, 10); - assertEquals(result.size(), 1); - assertEquals(result.get(0), entity3); - - // use case: return all the datasets owned by corpgroup1 - List paths2 = new ArrayList(); - paths2.add( - Triplet.with(RelationshipFoo.class, newRelationshipFilter(EMPTY_FILTER, RelationshipDirection.UNDIRECTED), - EntityBar.class)); - paths2.add( - Triplet.with(RelationshipBar.class, newRelationshipFilter(EMPTY_FILTER, RelationshipDirection.UNDIRECTED), - EntityBaz.class)); - List result2 = _dao.findEntities(EntityFoo.class, sourceFilter, paths2, 0, 10); - assertEquals(result2, result); - - // add another user & dataset - BarUrn urn4 = makeBarUrn(4); - EntityBar entity4 = new EntityBar().setUrn(urn4).setValue("CorpUser4"); - _writer.addEntity(entity4); - - BazUrn urn5 = makeBazUrn(5); - EntityBaz entity5 = new EntityBaz().setUrn(urn5).setValue("Dataset5"); - _writer.addEntity(entity5); - - // create relationship urn4 -> urn5 with RelationshipBar - RelationshipBar relationshipFoo4o5 = new RelationshipBar().setSource(urn5).setDestination(urn4); - _writer.addRelationship(relationshipFoo4o5, false); - - // create relationship urn1 -> urn4 with RelationshipFoo - RelationshipFoo relationshipFoo1To4 = new RelationshipFoo().setSource(urn1).setDestination(urn4); - _writer.addRelationship(relationshipFoo1To4, false); - - List paths3 = new ArrayList(); - paths3.add( - Triplet.with(RelationshipFoo.class, newRelationshipFilter(EMPTY_FILTER, RelationshipDirection.OUTGOING), - EntityBar.class)); - paths3.add( - Triplet.with(RelationshipBar.class, newRelationshipFilter(EMPTY_FILTER, RelationshipDirection.INCOMING), - EntityBaz.class)); - List result3 = _dao.findEntities(EntityFoo.class, sourceFilter, paths3, 0, 10); - assertEquals(result3.size(), 2); - assertEquals(result3.get(0), entity5); - assertEquals(result3.get(1), entity3); - - // test nulls - List paths4 = new ArrayList(); - paths4.add(Triplet.with(null, null, null)); - List result4 = _dao.findEntities(EntityFoo.class, sourceFilter, paths4, 0, 10); - assertEquals(result4.size(), 2); - assertEquals(result4.get(0), entity4); - assertEquals(result4.get(1), entity2); - - // test partial nulls with entity - List paths5 = new ArrayList(); - paths5.add( - Triplet.with(null, newRelationshipFilter(EMPTY_FILTER, RelationshipDirection.OUTGOING), EntityBar.class)); - List result5 = _dao.findEntities(EntityFoo.class, sourceFilter, paths5, 0, 10); - assertEquals(result5.size(), 2); - assertEquals(result5.get(0), entity4); - assertEquals(result5.get(1), entity2); - - // test partial nulls with relationship - List paths6 = new ArrayList(); - paths6.add(Triplet.with(null, null, EntityBar.class)); - List result6 = _dao.findEntities(EntityFoo.class, sourceFilter, paths6, 0, 10); - assertEquals(result6.size(), 2); - assertEquals(result6.get(0), entity4); - assertEquals(result6.get(1), entity2); - } - - @Test - public void testFindRelationship() throws Exception { - FooUrn urn1 = makeFooUrn(1); - EntityFoo entity1 = new EntityFoo().setUrn(urn1).setValue("foo"); - _writer.addEntity(entity1); - - BarUrn urn2 = makeBarUrn(2); - EntityBar entity2 = new EntityBar().setUrn(urn2).setValue("bar"); - _writer.addEntity(entity2); - - // create relationship urn1 -> urn2 - RelationshipFoo relationship = new RelationshipFoo().setSource(urn1).setDestination(urn2); - _writer.addRelationship(relationship, false); - - // find by source - Filter filter1 = newFilter("urn", urn1.toString()); - List found = - _dao.findRelationshipsFromSource(null, filter1, RelationshipFoo.class, EMPTY_FILTER, -1, -1); - assertEquals(found, Collections.singletonList(relationship)); - - // find by destination - Filter filter2 = newFilter("urn", urn2.toString()); - found = - _dao.findRelationshipsFromDestination(EntityBar.class, filter2, RelationshipFoo.class, EMPTY_FILTER, -1, -1); - assertEquals(found, Collections.singletonList(relationship)); - - // find by source and destination - found = _dao.findRelationships(null, filter1, null, filter2, RelationshipFoo.class, EMPTY_FILTER, 0, 10); - assertEquals(found, Collections.singletonList(relationship)); - } - - @Test - public void testFindRelationshipByQuery() throws Exception { - FooUrn urn1 = makeFooUrn(1); - EntityFoo entity1 = new EntityFoo().setUrn(urn1).setValue("foo"); - _writer.addEntity(entity1); - - BarUrn urn2 = makeBarUrn(2); - EntityBar entity2 = new EntityBar().setUrn(urn2).setValue("bar"); - _writer.addEntity(entity2); - - RelationshipFoo relationship = new RelationshipFoo().setSource(urn1).setDestination(urn2); - _writer.addRelationship(relationship, false); - - // with type - Map params = new HashMap<>(); - params.put("src", urn1.toString()); - Statement statement1 = new Statement("MATCH (src {urn:$src})-[r]->(dest) RETURN src, r, dest", params); - List found = _dao.findRelationships(RelationshipFoo.class, statement1); - assertEquals(found, Collections.singletonList(relationship)); - - // without type - params.put("dest", urn2.toString()); - Statement statement2 = new Statement("MATCH (src {urn:$src})-[r]->(dest {urn:$dest}) RETURN src, r, dest", params); - List foundNoType = _dao.findMixedTypesRelationships(statement2); - assertEquals(foundNoType, Collections.singletonList(relationship)); - } - - @Test - public void testFindNodesInPath() throws Exception { - FooUrn urn1 = makeFooUrn(1); - EntityFoo entity1 = new EntityFoo().setUrn(urn1).setValue("fooDirector"); - _writer.addEntity(entity1); - - FooUrn urn2 = makeFooUrn(2); - EntityFoo entity2 = new EntityFoo().setUrn(urn2).setValue("fooManager1"); - _writer.addEntity(entity2); - - FooUrn urn3 = makeFooUrn(3); - EntityFoo entity3 = new EntityFoo().setUrn(urn3).setValue("fooManager2"); - _writer.addEntity(entity3); - - FooUrn urn4 = makeFooUrn(4); - EntityFoo entity4 = new EntityFoo().setUrn(urn3).setValue("fooReport1ofManager1"); - _writer.addEntity(entity4); - - FooUrn urn5 = makeFooUrn(5); - EntityFoo entity5 = new EntityFoo().setUrn(urn3).setValue("fooReport1ofManager2"); - _writer.addEntity(entity5); - - FooUrn urn6 = makeFooUrn(6); - EntityFoo entity6 = new EntityFoo().setUrn(urn3).setValue("fooReport2ofManager2"); - _writer.addEntity(entity6); - - // Create relationships - simulate reportsto use case - _writer.addRelationship(new RelationshipFoo().setSource(urn6).setDestination(urn3), false); - _writer.addRelationship(new RelationshipFoo().setSource(urn5).setDestination(urn3), false); - _writer.addRelationship(new RelationshipFoo().setSource(urn4).setDestination(urn2), false); - _writer.addRelationship(new RelationshipFoo().setSource(urn3).setDestination(urn1), false); - _writer.addRelationship(new RelationshipFoo().setSource(urn2).setDestination(urn1), false); - - // Get reports roll-up - 2 levels - Filter sourceFilter = newFilter("urn", urn1.toString()); - RelationshipFilter relationshipFilter = newRelationshipFilter(EMPTY_FILTER, RelationshipDirection.INCOMING); - List> paths = _dao.findPaths(EntityFoo.class, sourceFilter, null, - EMPTY_FILTER, RelationshipFoo.class, relationshipFilter, 1, 2, -1, -1); - assertEquals(paths.size(), 5); - assertEquals(paths.stream().filter(l -> l.size() == 3).collect(Collectors.toList()).size(), 2); - assertEquals(paths.stream().filter(l -> l.size() == 5).collect(Collectors.toList()).size(), 3); - - // Get reports roll-up - 1 level - paths = _dao.findPaths(EntityFoo.class, sourceFilter, null, - EMPTY_FILTER, RelationshipFoo.class, relationshipFilter, 1, 1, -1, -1); - assertEquals(paths.size(), 2); - assertEquals(paths.stream().filter(l -> l.size() == 3).collect(Collectors.toList()).size(), 2); - assertEquals(paths.stream().filter(l -> l.size() == 5).collect(Collectors.toList()).size(), 0); - } - - @Test - public void testFindPaths() throws Exception { - BarUrn srcUrn = makeBarUrn(0); - BarUrn des1Urn = makeBarUrn(1); - BarUrn des2Urn = makeBarUrn(2); - - FooUrn srcField1Urn = makeFooUrn(1); - FooUrn srcField2Urn = makeFooUrn(2); - FooUrn des1Field1Urn = makeFooUrn(11); - FooUrn des1Field2Urn = makeFooUrn(12); - FooUrn des2Field1Urn = makeFooUrn(21); - - EntityBar src = new EntityBar().setUrn(srcUrn); - EntityBar des1 = new EntityBar().setUrn(des1Urn); - EntityBar des2 = new EntityBar().setUrn(des2Urn); - EntityFoo srcField1 = new EntityFoo().setUrn(srcField1Urn); - EntityFoo srcField2 = new EntityFoo().setUrn(srcField2Urn); - EntityFoo des1Field1 = new EntityFoo().setUrn(des1Field1Urn); - EntityFoo des1Field2 = new EntityFoo().setUrn(des1Field2Urn); - EntityFoo des2Field1 = new EntityFoo().setUrn(des2Field1Urn); - - _writer.addEntity(src); - _writer.addEntity(des1); - _writer.addEntity(des2); - _writer.addEntity(srcField1); - _writer.addEntity(srcField2); - _writer.addEntity(des1Field1); - _writer.addEntity(des1Field2); - _writer.addEntity(des2Field1); - - String commandText = "MATCH p=(n1 {urn:$src})-" - + "[r1]->(n2)-[r2:`com.linkedin.testing.RelationshipFoo`*0..100]->(n3)<-[r3]-() return p"; - Map params = new HashMap<>(); - params.put("src", srcUrn.toString()); - Statement statement = new Statement(commandText, params); - - // Test one path - createBarRelationship(srcUrn, srcField1Urn); - createBarRelationship(des1Urn, des1Field1Urn); - createFooRelationship(srcField1Urn, des1Field1Urn); - - List> paths = _dao.findPaths(statement); - assertEquals(paths.size(), 1); - - List path = paths.get(0); - assertEquals(path.size(), 7); - - assertEquals(path.get(0), src); - assertEquals(path.get(2), srcField1); - assertEquals(path.get(4), des1Field1); - assertEquals(path.get(6), des1); - - // Test multiple paths - createBarRelationship(des1Urn, des1Field2Urn); - createBarRelationship(des2Urn, des2Field1Urn); - createBarRelationship(srcUrn, srcField2Urn); - createFooRelationship(srcField2Urn, des1Field2Urn); - createFooRelationship(srcField2Urn, des2Field1Urn); - - paths = _dao.findPaths(statement); - - assertEquals(paths.size(), 3); - paths.forEach(p -> assertEquals(p.size(), 7)); - } - - @Test - public void testRunFreeFormQuery() throws Exception { - FooUrn urn1 = makeFooUrn(1); - FooUrn urn2 = makeFooUrn(2); - EntityFoo entity1 = new EntityFoo().setUrn(urn1).setValue("foo"); - EntityFoo entity2 = new EntityFoo().setUrn(urn2).setValue("foo"); - _writer.addEntity(entity1); - _writer.addEntity(entity2); - - String cypherQuery = "MATCH (n {value:\"foo\"}) RETURN n ORDER BY n.urn"; - List result = _dao.runFreeFormQuery(cypherQuery); - List nodes = result.stream() - .map(record -> _dao.nodeRecordToEntity(EntityFoo.class, record)) - .collect(Collectors.toList()); - assertEquals(nodes.size(), 2); - assertEquals(nodes.get(0), entity1); - assertEquals(nodes.get(1), entity2); - - cypherQuery = "MATCH (n {value:\"foo\"}) RETURN count(n)"; - result = _dao.runFreeFormQuery(cypherQuery); - assertEquals(result.size(), 1); - assertEquals(result.get(0).values().get(0).asInt(), 2); - } - - private void createFooRelationship(FooUrn f1, FooUrn f2) throws Exception { - _writer.addRelationship(new RelationshipFoo().setSource(f1).setDestination(f2), false); - } - - private void createBarRelationship(BarUrn d1, FooUrn f1) throws Exception { - _writer.addRelationship(new RelationshipBar().setSource(d1).setDestination(f1), false); - } -} diff --git a/dao-impl/neo4j-dao/src/test/java/com/linkedin/metadata/dao/Neo4jTestServerBuilder.java b/dao-impl/neo4j-dao/src/test/java/com/linkedin/metadata/dao/Neo4jTestServerBuilder.java deleted file mode 100644 index 84fd1ae18..000000000 --- a/dao-impl/neo4j-dao/src/test/java/com/linkedin/metadata/dao/Neo4jTestServerBuilder.java +++ /dev/null @@ -1,66 +0,0 @@ -package com.linkedin.metadata.dao; - -import java.io.File; -import java.net.URI; -import org.neo4j.graphdb.GraphDatabaseService; -import org.neo4j.harness.ServerControls; -import org.neo4j.harness.TestServerBuilder; -import org.neo4j.harness.TestServerBuilders; - - -public class Neo4jTestServerBuilder { - - private final TestServerBuilder builder; - private ServerControls controls; - - private Neo4jTestServerBuilder(TestServerBuilder builder) { - this.builder = builder; - } - - public Neo4jTestServerBuilder() { - this(TestServerBuilders.newInProcessBuilder()); - } - - public Neo4jTestServerBuilder(File workingDirectory) { - this(TestServerBuilders.newInProcessBuilder(workingDirectory)); - } - - public ServerControls newServer() { - if (controls == null) { - controls = builder.newServer(); - } - return controls; - } - - public void shutdown() { - if (controls != null) { - controls.close(); - controls = null; - } - } - - public URI boltURI() { - if (controls == null) { - throw new IllegalStateException("Cannot access instance URI."); - } - return controls.boltURI(); - } - - public URI httpURI() { - if (controls == null) { - throw new IllegalStateException("Cannot access instance URI."); - } - return controls.httpURI(); - } - - public URI httpsURI() { - if (controls == null) { - throw new IllegalStateException("Cannot access instance URI."); - } - return controls.httpURI(); - } - - public GraphDatabaseService getGraphDatabaseService() { - return controls.graph(); - } -} diff --git a/dao-impl/neo4j-dao/src/test/java/com/linkedin/metadata/dao/Neo4jUtilTest.java b/dao-impl/neo4j-dao/src/test/java/com/linkedin/metadata/dao/Neo4jUtilTest.java deleted file mode 100644 index eaf1e284a..000000000 --- a/dao-impl/neo4j-dao/src/test/java/com/linkedin/metadata/dao/Neo4jUtilTest.java +++ /dev/null @@ -1,214 +0,0 @@ -package com.linkedin.metadata.dao; - -import com.linkedin.common.urn.Urn; -import com.linkedin.data.template.DataTemplateUtil; -import com.linkedin.data.template.RecordTemplate; -import com.linkedin.metadata.query.Condition; -import com.linkedin.metadata.query.Criterion; -import com.linkedin.metadata.query.CriterionArray; -import com.linkedin.metadata.query.Filter; -import com.linkedin.testing.EntityBar; -import com.linkedin.testing.EntityFoo; -import com.linkedin.testing.RelationshipFoo; -import com.linkedin.testing.urn.BarUrn; -import com.linkedin.testing.urn.FooUrn; -import java.util.Collections; -import java.util.HashMap; -import java.util.List; -import java.util.Map; -import org.neo4j.driver.Value; -import org.neo4j.driver.internal.InternalNode; -import org.neo4j.driver.internal.InternalPath; -import org.neo4j.driver.internal.InternalRelationship; -import org.neo4j.driver.internal.value.IntegerValue; -import org.neo4j.driver.internal.value.StringValue; -import org.neo4j.driver.types.Node; -import org.neo4j.driver.types.Path; -import org.neo4j.driver.types.Relationship; -import org.testng.annotations.Test; - -import static com.linkedin.metadata.dao.Neo4jUtil.*; -import static com.linkedin.metadata.dao.utils.ModelUtils.*; -import static com.linkedin.metadata.dao.utils.QueryUtils.*; -import static com.linkedin.testing.TestUtils.*; -import static org.testng.Assert.*; - - -public class Neo4jUtilTest { - - @Test - public void testEntityToNode() { - FooUrn urn = makeFooUrn(1); - EntityFoo entity = new EntityFoo().setUrn(urn).setValue("foo"); - - Map actual = entityToNode(entity); - - Map expected = new HashMap<>(); - expected.put("urn", urn.toString()); - expected.put("value", "foo"); - - assertEquals(actual, expected); - } - - @Test - public void testRelationshipToEdge() { - Urn urn = makeUrn(1); - RelationshipFoo relationship = new RelationshipFoo().setSource(urn).setDestination(urn); - - Map fields = relationshipToEdge(relationship); - - assertTrue(fields.isEmpty()); - } - - @Test - public void testRelationshipToCriteria() { - Urn urn = makeUrn(1); - RelationshipFoo relationship = new RelationshipFoo().setSource(urn).setDestination(urn); - - String criteria = relationshipToCriteria(relationship); - - assertTrue(criteria.isEmpty()); - - relationship.data().put("bar", urn.toString()); - relationship.data().put("foo", 2); - - criteria = relationshipToCriteria(relationship); - - assertEquals(criteria, "{bar:\"urn:li:testing:1\",foo:2}"); - } - - @Test - public void testFilterToCriteria() { - Filter filter = newFilter("a", "b"); - - assertEquals(filterToCriteria(filter), "{a:\"b\"}"); - } - - @Test(expectedExceptions = RuntimeException.class) - public void testFilterToCriteriaNonEqual() { - Filter filter = new Filter().setCriteria(new CriterionArray()); - filter.getCriteria().add(new Criterion().setField("a").setValue("b").setCondition(Condition.CONTAIN)); - - filterToCriteria(filter); - fail("no exception"); - } - - @Test - public void testNodeToEntity() { - FooUrn urn = makeFooUrn(1); - Node node = new InternalNode(2, Collections.singletonList(EntityFoo.class.getCanonicalName()), - Collections.singletonMap("urn", new StringValue(urn.toString()))); - - // with type class - EntityFoo entity1 = nodeToEntity(EntityFoo.class, node); - - assertEquals(entity1, new EntityFoo().setUrn(urn)); - - // without type class - RecordTemplate entity2 = nodeToEntity(node); - - assertEquals(entity2.getClass(), EntityFoo.class); - assertEquals(entity2, new EntityFoo().setUrn(urn)); - } - - @Test - public void testEdgeToRelationship() { - Urn urn1 = makeUrn(1); - Node node1 = new InternalNode(3, Collections.singletonList(EntityFoo.class.getCanonicalName()), - Collections.singletonMap("urn", new StringValue(urn1.toString()))); - - Urn urn2 = makeUrn(2); - Node node2 = new InternalNode(4, Collections.singletonList(EntityBar.class.getCanonicalName()), - Collections.singletonMap("urn", new StringValue(urn2.toString()))); - - Relationship edge = new InternalRelationship(5, 3, 4, RelationshipFoo.class.getCanonicalName()); - - // with type class - RelationshipFoo relationship1 = edgeToRelationship(RelationshipFoo.class, node1, node2, edge); - - RelationshipFoo expected = new RelationshipFoo().setSource(urn1).setDestination(urn2); - assertEquals(relationship1, expected); - - // without type class - RecordTemplate relationship2 = edgeToRelationship(node1, node2, edge); - - assertEquals(relationship2.getClass(), RelationshipFoo.class); - assertEquals(relationship2, expected); - } - - @Test - public void testEdgeToRelationshipWithoutType() { - FooUrn fooUrn = makeFooUrn(1); - Node fooNode = new InternalNode(6, Collections.singletonList("foo"), - Collections.singletonMap("urn", new StringValue(fooUrn.toString()))); - - BarUrn barUrn = makeBarUrn(1); - Node barNode = new InternalNode(7, Collections.singletonList("bar"), - Collections.singletonMap("urn", new StringValue(barUrn.toString()))); - - Relationship edge = new InternalRelationship(8, 6, 7, RelationshipFoo.class.getCanonicalName()); - - RelationshipFoo relationship = edgeToRelationship(RelationshipFoo.class, fooNode, barNode, edge); - - assertEquals(relationship, new RelationshipFoo().setSource(fooUrn).setDestination(barUrn)); - } - - @Test - public void testGetUrn() { - FooUrn urn = makeFooUrn(1); - EntityFoo entity = new EntityFoo().setUrn(urn); - RelationshipFoo relationship = new RelationshipFoo().setSource(urn).setDestination(urn); - - assertEquals(getUrnFromEntity(entity), urn); - assertEquals(getSourceUrnFromRelationship(relationship), urn); - assertEquals(getDestinationUrnFromRelationship(relationship), urn); - } - - @Test - public void testGetType() { - assertEquals(getType(new EntityFoo()), "`com.linkedin.testing.EntityFoo`"); - assertEquals(getType(EntityBar.class), "`com.linkedin.testing.EntityBar`"); - } - - @Test - public void testGetTypeOrEmptyString() { - assertEquals(getTypeOrEmptyString(EntityBar.class), ":`com.linkedin.testing.EntityBar`"); - assertEquals(getTypeOrEmptyString(null), ""); - } - - @Test - public void testPathSegmentToRecordList() { - FooUrn fooUrn = makeFooUrn(1); - Node fooNode = new InternalNode(0, Collections.singletonList("com.linkedin.testing.EntityFoo"), - Collections.singletonMap("urn", new StringValue(fooUrn.toString()))); - - BarUrn barUrn = makeBarUrn(1); - Node barNode = new InternalNode(1, Collections.singletonList("com.linkedin.testing.EntityBar"), - Collections.singletonMap("urn", new StringValue(barUrn.toString()))); - - Relationship edge = new InternalRelationship(2, 0, 1, RelationshipFoo.class.getCanonicalName(), - new HashMap() { - { - put("intField", new IntegerValue(42)); - put("type", new StringValue("dummyType")); - } - }); - - Path path = new InternalPath(fooNode, edge, barNode); - - List pathList = pathSegmentToRecordList(path.iterator().next()); - - assertTrue(pathList.get(0) instanceof EntityFoo); - assertTrue(pathList.get(1) instanceof RelationshipFoo); - assertTrue(pathList.get(2) instanceof EntityBar); - - assertEquals(pathList.get(0), new EntityFoo().setUrn(fooUrn)); - assertTrue(DataTemplateUtil.areEqual(pathList.get(1), - new RelationshipFoo() - .setIntField(42) - .setSource(fooUrn) - .setDestination(barUrn) - .setType("dummyType"))); - assertEquals(pathList.get(2), new EntityBar().setUrn(barUrn)); - } -} diff --git a/dao-impl/neo4j-dao/src/test/java/com/linkedin/metadata/dao/internal/Neo4jGraphWriterDAOTest.java b/dao-impl/neo4j-dao/src/test/java/com/linkedin/metadata/dao/internal/Neo4jGraphWriterDAOTest.java deleted file mode 100644 index adf6c58e4..000000000 --- a/dao-impl/neo4j-dao/src/test/java/com/linkedin/metadata/dao/internal/Neo4jGraphWriterDAOTest.java +++ /dev/null @@ -1,412 +0,0 @@ -package com.linkedin.metadata.dao.internal; - -import com.linkedin.common.urn.Urn; -import com.linkedin.metadata.dao.BaseQueryDAO; -import com.linkedin.metadata.dao.Neo4jQueryDAO; -import com.linkedin.metadata.dao.Neo4jTestServerBuilder; -import com.linkedin.metadata.query.Criterion; -import com.linkedin.metadata.query.CriterionArray; -import com.linkedin.metadata.query.Filter; -import com.linkedin.testing.RelationshipFoo; -import com.linkedin.testing.EntityFoo; -import com.linkedin.testing.EntityBar; -import com.linkedin.testing.TestUtils; -import com.linkedin.testing.urn.BarUrn; -import com.linkedin.testing.urn.FooUrn; -import java.util.Arrays; -import java.util.Collections; -import java.util.List; -import java.util.Map; -import java.util.Optional; -import javax.annotation.Nonnull; -import org.neo4j.driver.Driver; -import org.neo4j.driver.GraphDatabase; -import org.testng.annotations.AfterMethod; -import org.testng.annotations.BeforeMethod; -import org.testng.annotations.Test; - -import static com.linkedin.metadata.dao.Neo4jUtil.*; -import static com.linkedin.metadata.dao.internal.BaseGraphWriterDAO.RemovalOption.*; -import static com.linkedin.testing.TestUtils.*; -import static org.testng.Assert.*; - - -public class Neo4jGraphWriterDAOTest { - - private Neo4jTestServerBuilder _serverBuilder; - private Neo4jGraphWriterDAO _dao; - private Neo4jTestHelper _helper; - private BaseQueryDAO _queryDao; - private TestMetricListener _testMetricListener; - - private static class TestMetricListener implements Neo4jGraphWriterDAO.MetricListener { - int entitiesAdded = 0; - int entityAddedEvents = 0; - int entitiesRemoved = 0; - int entityRemovedEvents = 0; - int relationshipsAdded = 0; - int relationshipAddedEvents = 0; - int relationshipsRemoved = 0; - int relationshipRemovedEvents = 0; - - @Override - public void onEntitiesAdded(int entityCount, long updateTimeMs, int retries) { - entityAddedEvents++; - entitiesAdded += entityCount; - } - - @Override - public void onRelationshipsAdded(int relationshipCount, long updateTimeMs, int retries) { - relationshipAddedEvents++; - relationshipsAdded += relationshipCount; - } - - @Override - public void onEntitiesRemoved(int entityCount, long updateTimeMs, int retries) { - entityRemovedEvents++; - entitiesRemoved += entityCount; - } - - @Override - public void onRelationshipsRemoved(int relationshipCount, long updateTimeMs, int retries) { - relationshipRemovedEvents++; - relationshipsRemoved += relationshipCount; - } - } - - @BeforeMethod - public void init() { - _serverBuilder = new Neo4jTestServerBuilder(); - _serverBuilder.newServer(); - _testMetricListener = new TestMetricListener(); - final Driver driver = GraphDatabase.driver(_serverBuilder.boltURI()); - _dao = new Neo4jGraphWriterDAO(driver, TestUtils.getAllTestEntities()); - _helper = new Neo4jTestHelper(driver, TestUtils.getAllTestEntities()); - _queryDao = new Neo4jQueryDAO(driver); - _dao.addMetricListener(_testMetricListener); - } - - @AfterMethod - public void tearDown() { - _serverBuilder.shutdown(); - } - - @Test - public void testAddRemoveEntity() throws Exception { - FooUrn urn = makeFooUrn(1); - EntityFoo entity = new EntityFoo().setUrn(urn).setValue("foo"); - - _dao.addEntity(entity); - Optional> node = _helper.getNode(urn); - assertEntityFoo(node.get(), entity); - assertEquals(_testMetricListener.entitiesAdded, 1); - assertEquals(_testMetricListener.entityAddedEvents, 1); - - _dao.removeEntity(urn); - node = _helper.getNode(urn); - assertFalse(node.isPresent()); - assertEquals(_testMetricListener.entitiesRemoved, 1); - assertEquals(_testMetricListener.entityRemovedEvents, 1); - } - - @Test - public void testPartialUpdateEntity() throws Exception { - FooUrn urn = makeFooUrn(1); - EntityFoo entity = new EntityFoo().setUrn(urn); - - _dao.addEntity(entity); - Optional> node = _helper.getNode(urn); - assertEntityFoo(node.get(), entity); - - // add value for optional field - EntityFoo entity2 = new EntityFoo().setUrn(urn).setValue("IamTheSameEntity"); - _dao.addEntity(entity2); - node = _helper.getNode(urn); - assertEquals(_helper.getAllNodes(urn).size(), 1); - assertEntityFoo(node.get(), entity2); - - // change value for optional field - EntityFoo entity3 = new EntityFoo().setUrn(urn).setValue("ChangeValue"); - _dao.addEntity(entity3); - node = _helper.getNode(urn); - assertEquals(_helper.getAllNodes(urn).size(), 1); - assertEntityFoo(node.get(), entity3); - } - - @Test - public void testAddRemoveEntities() throws Exception { - EntityFoo entity1 = new EntityFoo().setUrn(makeFooUrn(1)).setValue("foo"); - EntityFoo entity2 = new EntityFoo().setUrn(makeFooUrn(2)).setValue("bar"); - EntityFoo entity3 = new EntityFoo().setUrn(makeFooUrn(3)).setValue("baz"); - List entities = Arrays.asList(entity1, entity2, entity3); - - _dao.addEntities(entities); - assertEntityFoo(_helper.getNode(entity1.getUrn()).get(), entity1); - assertEntityFoo(_helper.getNode(entity2.getUrn()).get(), entity2); - assertEntityFoo(_helper.getNode(entity3.getUrn()).get(), entity3); - assertEquals(_testMetricListener.entitiesAdded, 3); - assertEquals(_testMetricListener.entityAddedEvents, 1); - - _dao.removeEntities(Arrays.asList(entity1.getUrn(), entity3.getUrn())); - assertFalse(_helper.getNode(entity1.getUrn()).isPresent()); - assertTrue(_helper.getNode(entity2.getUrn()).isPresent()); - assertFalse(_helper.getNode(entity3.getUrn()).isPresent()); - assertEquals(_testMetricListener.entitiesRemoved, 2); - assertEquals(_testMetricListener.entityRemovedEvents, 1); - } - - @Test - public void testAddRelationshipNodeNonExist() throws Exception { - FooUrn urn1 = makeFooUrn(1); - BarUrn urn2 = makeBarUrn(2); - RelationshipFoo relationship = new RelationshipFoo().setSource(urn1).setDestination(urn2); - - _dao.addRelationship(relationship, REMOVE_NONE, false); - - assertRelationshipFoo(_helper.getEdges(relationship), 1); - assertEntityFoo(_helper.getNode(urn1).get(), new EntityFoo().setUrn(urn1)); - assertEntityBar(_helper.getNode(urn2).get(), new EntityBar().setUrn(urn2)); - assertEquals(_testMetricListener.relationshipsAdded, 1); - assertEquals(_testMetricListener.relationshipAddedEvents, 1); - } - - @Test - public void testPartialUpdateEntityCreatedByRelationship() throws Exception { - FooUrn urn1 = makeFooUrn(1); - FooUrn urn2 = makeFooUrn(2); - RelationshipFoo relationship = new RelationshipFoo().setSource(urn1).setDestination(urn2); - - _dao.addRelationship(relationship, REMOVE_NONE, false); - - // Check if adding an entity with same urn and with label creates a new node - _dao.addEntity(new EntityFoo().setUrn(urn1)); - assertEquals(_helper.getAllNodes(urn1).size(), 1); - } - - @Test - public void testAddRemoveRelationships() throws Exception { - // Add entity1 - FooUrn urn1 = makeFooUrn(1); - EntityFoo entity1 = new EntityFoo().setUrn(urn1).setValue("foo"); - _dao.addEntity(entity1); - assertEntityFoo(_helper.getNode(urn1).get(), entity1); - - // Add entity2 - BarUrn urn2 = makeBarUrn(2); - EntityBar entity2 = new EntityBar().setUrn(urn2).setValue("bar"); - _dao.addEntity(entity2); - assertEntityBar(_helper.getNode(urn2).get(), entity2); - - // add relationship1 (urn1 -> urn2) - RelationshipFoo relationship1 = new RelationshipFoo().setSource(urn1).setDestination(urn2); - _dao.addRelationship(relationship1, REMOVE_NONE, false); - assertRelationshipFoo(_helper.getEdges(relationship1), 1); - - // add relationship1 again - _dao.addRelationship(relationship1, false); - assertRelationshipFoo(_helper.getEdges(relationship1), 1); - - // add relationship2 (urn1 -> urn3) - Urn urn3 = makeUrn(3); - RelationshipFoo relationship2 = new RelationshipFoo().setSource(urn1).setDestination(urn3); - _dao.addRelationship(relationship2, false); - assertRelationshipFoo(_helper.getEdgesFromSource(urn1, RelationshipFoo.class), 2); - - // remove relationship1 - _dao.removeRelationship(relationship1); - assertRelationshipFoo(_helper.getEdges(relationship1), 0); - - // remove relationship1 & relationship2 - _dao.removeRelationships(Arrays.asList(relationship1, relationship2)); - assertRelationshipFoo(_helper.getEdgesFromSource(urn1, RelationshipFoo.class), 0); - - - assertEquals(_testMetricListener.relationshipsAdded, 3); - assertEquals(_testMetricListener.relationshipAddedEvents, 3); - - assertEquals(_testMetricListener.relationshipsRemoved, 3); - assertEquals(_testMetricListener.relationshipRemovedEvents, 2); - } - - @Test - public void testAddRelationshipRemoveAll() throws Exception { - // Add entity1 - FooUrn urn1 = makeFooUrn(1); - EntityFoo entity1 = new EntityFoo().setUrn(urn1).setValue("foo"); - _dao.addEntity(entity1); - assertEntityFoo(_helper.getNode(urn1).get(), entity1); - - // Add entity2 - BarUrn urn2 = makeBarUrn(2); - EntityBar entity2 = new EntityBar().setUrn(urn2).setValue("bar"); - _dao.addEntity(entity2); - assertEntityBar(_helper.getNode(urn2).get(), entity2); - - // add relationship1 (urn1 -> urn2) - RelationshipFoo relationship1 = new RelationshipFoo().setSource(urn1).setDestination(urn2); - _dao.addRelationship(relationship1, REMOVE_NONE, false); - assertRelationshipFoo(_helper.getEdges(relationship1), 1); - - // add relationship2 (urn1 -> urn3), removeAll from source - Urn urn3 = makeUrn(3); - RelationshipFoo relationship2 = new RelationshipFoo().setSource(urn1).setDestination(urn3); - _dao.addRelationship(relationship2, REMOVE_ALL_EDGES_FROM_SOURCE, false); - assertRelationshipFoo(_helper.getEdgesFromSource(urn1, RelationshipFoo.class), 1); - - // add relationship3 (urn4 -> urn3), removeAll from destination - Urn urn4 = makeUrn(4); - RelationshipFoo relationship3 = new RelationshipFoo().setSource(urn4).setDestination(urn3); - _dao.addRelationship(relationship3, REMOVE_ALL_EDGES_TO_DESTINATION, false); - assertRelationshipFoo(_helper.getEdgesFromSource(urn1, RelationshipFoo.class), 0); - assertRelationshipFoo(_helper.getEdgesFromSource(urn4, RelationshipFoo.class), 1); - - // add relationship3 again without removal - _dao.addRelationship(relationship3, false); - assertRelationshipFoo(_helper.getEdgesFromSource(urn4, RelationshipFoo.class), 1); - - // add relationship3 again, removeAll from source & destination - _dao.addRelationship(relationship3, REMOVE_ALL_EDGES_FROM_SOURCE_TO_DESTINATION, false); - assertRelationshipFoo(_helper.getEdgesFromSource(urn1, RelationshipFoo.class), 0); - assertRelationshipFoo(_helper.getEdgesFromSource(urn4, RelationshipFoo.class), 1); - } - - @Test - public void upsertNodeAddNewProperty() throws Exception { - // given - final FooUrn urn = makeFooUrn(1); - final EntityFoo initialEntity = new EntityFoo().setUrn(urn); - final EntityFoo updatedEntity = new EntityFoo().setUrn(urn).setValue("updated"); - - // when - _dao.addEntity(initialEntity); - _dao.addEntity(updatedEntity); - - // then - assertEntityFoo(_helper.getNode(urn).get(), updatedEntity); - } - - @Test - public void upsertEdgeAddNewProperty() throws Exception { - // given - final EntityFoo foo = new EntityFoo().setUrn(makeFooUrn(1)); - final EntityBar bar = new EntityBar().setUrn(makeBarUrn(2)).setValue("bar"); - _dao.addEntity(foo); - _dao.addEntity(bar); - - final RelationshipFoo initialRelationship = - new RelationshipFoo().setSource(foo.getUrn()).setDestination(bar.getUrn()); - final RelationshipFoo updatedRelationship = - new RelationshipFoo().setSource(foo.getUrn()).setDestination(bar.getUrn()).setType("test"); - _dao.addRelationship(initialRelationship, false); - - // when - _dao.addRelationship(updatedRelationship, false); - - // then - assertEquals(_queryDao.findRelationships(EntityFoo.class, - new Filter().setCriteria(new CriterionArray(new Criterion().setField("urn").setValue(foo.getUrn().toString()))), - EntityBar.class, - new Filter().setCriteria(new CriterionArray(new Criterion().setField("urn").setValue(bar.getUrn().toString()))), - RelationshipFoo.class, new Filter().setCriteria(new CriterionArray()), 0, 10), - Collections.singletonList(updatedRelationship)); - } - - @Test - public void upsertNodeChangeProperty() throws Exception { - // given - final FooUrn urn = makeFooUrn(1); - final EntityFoo initialEntity = new EntityFoo().setUrn(urn).setValue("before"); - final EntityFoo updatedEntity = new EntityFoo().setUrn(urn).setValue("after"); - _dao.addEntity(initialEntity); - - // when - _dao.addEntity(updatedEntity); - - // then - assertEntityFoo(_helper.getNode(urn).get(), updatedEntity); - } - - @Test - public void upsertEdgeChangeProperty() throws Exception { - // given - final EntityFoo foo = new EntityFoo().setUrn(makeFooUrn(1)); - final EntityBar bar = new EntityBar().setUrn(makeBarUrn(2)).setValue("bar"); - _dao.addEntity(foo); - _dao.addEntity(bar); - - final RelationshipFoo initialRelationship = - new RelationshipFoo().setSource(foo.getUrn()).setDestination(bar.getUrn()).setType("before"); - final RelationshipFoo updatedRelationship = - new RelationshipFoo().setSource(foo.getUrn()).setDestination(bar.getUrn()).setType("after"); - _dao.addRelationship(initialRelationship, false); - - // when - _dao.addRelationship(updatedRelationship, false); - - // then - assertEquals(_queryDao.findRelationships(EntityFoo.class, - new Filter().setCriteria(new CriterionArray(new Criterion().setField("urn").setValue(foo.getUrn().toString()))), - EntityBar.class, - new Filter().setCriteria(new CriterionArray(new Criterion().setField("urn").setValue(bar.getUrn().toString()))), - RelationshipFoo.class, new Filter().setCriteria(new CriterionArray()), 0, 10), - Collections.singletonList(updatedRelationship)); - } - - @Test - public void upsertNodeRemovedProperty() throws Exception { - // given - final FooUrn urn = makeFooUrn(1); - final EntityFoo initialEntity = new EntityFoo().setUrn(urn).setValue("before"); - final EntityFoo updatedEntity = new EntityFoo().setUrn(urn); - _dao.addEntity(initialEntity); - - // when - _dao.addEntity(updatedEntity); - - // then - // Upsert won't ever delete properties. - assertEntityFoo(_helper.getNode(urn).get(), initialEntity); - } - - @Test - public void upsertEdgeRemoveProperty() throws Exception { - // given - final EntityFoo foo = new EntityFoo().setUrn(makeFooUrn(1)); - final EntityBar bar = new EntityBar().setUrn(makeBarUrn(2)).setValue("bar"); - _dao.addEntity(foo); - _dao.addEntity(bar); - - final RelationshipFoo initialRelationship = - new RelationshipFoo().setSource(foo.getUrn()).setDestination(bar.getUrn()).setType("before"); - final RelationshipFoo updatedRelationship = - new RelationshipFoo().setSource(foo.getUrn()).setDestination(bar.getUrn()); - _dao.addRelationship(initialRelationship, false); - - // when - _dao.addRelationship(updatedRelationship, false); - - // then - assertEquals(_queryDao.findRelationships(EntityFoo.class, - new Filter().setCriteria(new CriterionArray(new Criterion().setField("urn").setValue(foo.getUrn().toString()))), - EntityBar.class, - new Filter().setCriteria(new CriterionArray(new Criterion().setField("urn").setValue(bar.getUrn().toString()))), - RelationshipFoo.class, new Filter().setCriteria(new CriterionArray()), 0, 10), - // Upsert won't ever delete properties. - Collections.singletonList(initialRelationship)); - } - - private void assertEntityFoo(@Nonnull Map node, @Nonnull EntityFoo entity) { - assertEquals(node.get("urn"), entity.getUrn().toString()); - assertEquals(node.get("value"), entity.getValue()); - } - - private void assertEntityBar(@Nonnull Map node, @Nonnull EntityBar entity) { - assertEquals(node.get("urn"), entity.getUrn().toString()); - assertEquals(node.get("value"), entity.getValue()); - } - - private void assertRelationshipFoo(@Nonnull List> edges, int count) { - assertEquals(edges.size(), count); - edges.forEach(edge -> assertTrue(edge.isEmpty())); - } -} diff --git a/dao-impl/neo4j-dao/src/test/java/com/linkedin/metadata/dao/internal/Neo4jQueriesTransformerTest.java b/dao-impl/neo4j-dao/src/test/java/com/linkedin/metadata/dao/internal/Neo4jQueriesTransformerTest.java deleted file mode 100644 index d31268866..000000000 --- a/dao-impl/neo4j-dao/src/test/java/com/linkedin/metadata/dao/internal/Neo4jQueriesTransformerTest.java +++ /dev/null @@ -1,27 +0,0 @@ -package com.linkedin.metadata.dao.internal; - -import com.linkedin.testing.EntityBar; -import com.linkedin.testing.EntityFoo; -import org.testng.annotations.Test; - -import static com.linkedin.metadata.dao.Neo4jUtil.*; -import static com.linkedin.testing.TestUtils.*; -import static org.testng.Assert.*; - - -public class Neo4jQueriesTransformerTest { - - @Test - public void testGetNodeTypeFromUrn() { - final Neo4jQueriesTransformer transformer = new Neo4jQueriesTransformer(getAllTestEntities()); - - assertEquals(transformer.getNodeType(makeBarUrn(1)), ":`com.linkedin.testing.EntityBar`"); - assertEquals(transformer.getNodeType(makeFooUrn(1)), ":`com.linkedin.testing.EntityFoo`"); - assertEquals(transformer.getNodeType(makeUrn(1, "foo")), ":`com.linkedin.testing.EntityFoo`"); - assertEquals(transformer.getNodeType(makeUrn("1")), ":UNKNOWN"); - - // test consistency !! - assertEquals(transformer.getNodeType(makeBarUrn(1)), getTypeOrEmptyString(EntityBar.class)); - assertEquals(transformer.getNodeType(makeFooUrn(1)), getTypeOrEmptyString(EntityFoo.class)); - } -} \ No newline at end of file diff --git a/dao-impl/neo4j-dao/src/test/java/com/linkedin/metadata/dao/internal/Neo4jTestHelper.java b/dao-impl/neo4j-dao/src/test/java/com/linkedin/metadata/dao/internal/Neo4jTestHelper.java deleted file mode 100644 index 809b1b076..000000000 --- a/dao-impl/neo4j-dao/src/test/java/com/linkedin/metadata/dao/internal/Neo4jTestHelper.java +++ /dev/null @@ -1,106 +0,0 @@ -package com.linkedin.metadata.dao.internal; - -import com.linkedin.common.urn.Urn; -import com.linkedin.data.template.RecordTemplate; -import java.util.HashMap; -import java.util.List; -import java.util.Map; -import java.util.Optional; -import java.util.Set; -import java.util.stream.Collectors; -import javax.annotation.Nonnull; -import org.neo4j.driver.Driver; -import org.neo4j.driver.Query; -import org.neo4j.driver.Session; - -import static com.linkedin.metadata.dao.Neo4jUtil.*; -import static com.linkedin.metadata.dao.utils.ModelUtils.*; - - -/** - * Helper for making queries in unit tests. - */ -public final class Neo4jTestHelper { - private final Driver _driver; - private final Neo4jQueriesTransformer _neo4jQueriesTransformer; - - private Neo4jTestHelper(@Nonnull Driver driver, @Nonnull Neo4jQueriesTransformer neo4jQueriesTransformer) { - _driver = driver; - _neo4jQueriesTransformer = neo4jQueriesTransformer; - } - - public Neo4jTestHelper(@Nonnull Driver driver) { - this(driver, new Neo4jQueriesTransformer()); - } - - public Neo4jTestHelper(@Nonnull Driver driver, @Nonnull Set> entitiesSet) { - this(driver, new Neo4jQueriesTransformer(entitiesSet)); - } - - private List> execute(@Nonnull Query query) { - try (Session session = _driver.session()) { - return session.run(query) - .list() - .stream() - .map(record -> record.values().get(0).asMap()) - .collect(Collectors.toList()); - } - } - - @Nonnull - public Optional> getNode(@Nonnull Urn urn) { - List> nodes = getAllNodes(urn); - if (nodes.isEmpty()) { - return Optional.empty(); - } - return Optional.of(nodes.get(0)); - } - - @Nonnull - public List> getAllNodes(@Nonnull Urn urn) { - final String matchTemplate = "MATCH (node%s {urn: $urn}) RETURN node"; - - final String sourceType = _neo4jQueriesTransformer.getNodeType(urn); - final String statement = String.format(matchTemplate, sourceType); - - final Map params = new HashMap<>(); - params.put("urn", urn.toString()); - - return execute(new Query(statement, params)); - } - - @Nonnull - public List> getEdges(@Nonnull RecordTemplate relationship) { - final Urn sourceUrn = getSourceUrnFromRelationship(relationship); - final Urn destinationUrn = getDestinationUrnFromRelationship(relationship); - final String relationshipType = getType(relationship); - - final String sourceType = _neo4jQueriesTransformer.getNodeType(sourceUrn); - final String destinationType = _neo4jQueriesTransformer.getNodeType(destinationUrn); - - final String matchTemplate = - "MATCH (source%s {urn: $sourceUrn})-[r:%s]->(destination%s {urn: $destinationUrn}) RETURN r"; - final String statement = String.format(matchTemplate, sourceType, relationshipType, destinationType); - - final Map params = new HashMap<>(); - params.put("sourceUrn", sourceUrn.toString()); - params.put("destinationUrn", destinationUrn.toString()); - - return execute(new Query(statement, params)); - } - - @Nonnull - public List> getEdgesFromSource(@Nonnull Urn sourceUrn, - @Nonnull Class relationshipClass) { - final String relationshipType = getType(relationshipClass); - final String sourceType = _neo4jQueriesTransformer.getNodeType(sourceUrn); - - final String matchTemplate = "MATCH (source%s {urn: $sourceUrn})-[r:%s]->() RETURN r"; - final String statement = String.format(matchTemplate, sourceType, relationshipType); - - final Map params = new HashMap<>(); - params.put("sourceUrn", sourceUrn.toString()); - - return execute(new Query(statement, params)); - } -} 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 39f8ac214..d1da709c8 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 @@ -49,6 +49,7 @@ import com.linkedin.testing.InternalEntityAspectUnion; import com.linkedin.testing.InternalEntitySnapshot; import com.linkedin.testing.localrelationship.AspectFooBar; +import com.linkedin.testing.localrelationship.AspectFooBaz; import com.linkedin.testing.localrelationship.AspectFooBarBaz; import com.linkedin.testing.localrelationship.BelongsTo; import com.linkedin.testing.urn.BarUrn; @@ -252,10 +253,11 @@ public void testGet() { AspectKey aspect1Key = new AspectKey<>(AspectFoo.class, urn, LATEST_VERSION); AspectKey aspect2Key = new AspectKey<>(AspectBar.class, urn, LATEST_VERSION); AspectKey aspect3Key = new AspectKey<>(AspectFooBar.class, urn, LATEST_VERSION); - AspectKey aspect4Key = new AspectKey<>(AspectFooBarBaz.class, urn, LATEST_VERSION); - AspectKey aspect5Key = new AspectKey<>(AspectAttributes.class, urn, LATEST_VERSION); + AspectKey aspect4Key = new AspectKey<>(AspectFooBaz.class, urn, LATEST_VERSION); + AspectKey aspect5Key = new AspectKey<>(AspectFooBarBaz.class, urn, LATEST_VERSION); + AspectKey aspect6Key = new AspectKey<>(AspectAttributes.class, urn, LATEST_VERSION); when(_mockLocalDAO.exists(urn)).thenReturn(true); - when(_mockLocalDAO.get(new HashSet<>(Arrays.asList(aspect1Key, aspect2Key, aspect3Key, aspect4Key, aspect5Key)))).thenReturn( + when(_mockLocalDAO.get(new HashSet<>(Arrays.asList(aspect1Key, aspect2Key, aspect3Key, aspect4Key, aspect5Key, aspect6Key)))).thenReturn( Collections.singletonMap(aspect1Key, Optional.of(foo))); EntityValue value = runAndWait(_resource.get(makeResourceKey(urn), null)); @@ -359,16 +361,18 @@ public void testBatchGet() { AspectKey aspectFooKey1 = new AspectKey<>(AspectFoo.class, urn1, LATEST_VERSION); AspectKey aspectBarKey1 = new AspectKey<>(AspectBar.class, urn1, LATEST_VERSION); AspectKey aspectFooBarKey1 = new AspectKey<>(AspectFooBar.class, urn1, LATEST_VERSION); + AspectKey aspectFooBazKey1 = new AspectKey<>(AspectFooBaz.class, urn1, LATEST_VERSION); AspectKey aspectFooBarBazKey1 = new AspectKey<>(AspectFooBarBaz.class, urn1, LATEST_VERSION); AspectKey aspectAttKey1 = new AspectKey<>(AspectAttributes.class, urn1, LATEST_VERSION); AspectKey aspectFooKey2 = new AspectKey<>(AspectFoo.class, urn2, LATEST_VERSION); AspectKey aspectBarKey2 = new AspectKey<>(AspectBar.class, urn2, LATEST_VERSION); AspectKey aspectFooBarKey2 = new AspectKey<>(AspectFooBar.class, urn2, LATEST_VERSION); + AspectKey aspectFooBazKey2 = new AspectKey<>(AspectFooBaz.class, urn2, LATEST_VERSION); AspectKey aspectFooBarBazKey2 = new AspectKey<>(AspectFooBarBaz.class, urn2, LATEST_VERSION); AspectKey aspectAttKey2 = new AspectKey<>(AspectAttributes.class, urn2, LATEST_VERSION); - when(_mockLocalDAO.get(ImmutableSet.of(aspectFooBarKey1, aspectFooBarKey2, aspectFooBarBazKey1, aspectFooBarBazKey2, - aspectFooKey1, aspectBarKey1, aspectFooKey2, aspectBarKey2, aspectAttKey1, aspectAttKey2))) + when(_mockLocalDAO.get(ImmutableSet.of(aspectFooBarKey1, aspectFooBarKey2, aspectFooBazKey1, aspectFooBazKey2, + aspectFooBarBazKey1, aspectFooBarBazKey2, aspectFooKey1, aspectBarKey1, aspectFooKey2, aspectBarKey2, aspectAttKey1, aspectAttKey2))) .thenReturn(ImmutableMap.of(aspectFooKey1, Optional.of(foo), aspectFooKey2, Optional.of(bar))); Map keyValueMap = @@ -715,18 +719,20 @@ public void testGetSnapshotWithAllAspects() { AspectFoo foo = new AspectFoo().setValue("foo"); AspectFoo bar = new AspectFoo().setValue("bar"); AspectFooBar fooBar = new AspectFooBar().setBars(new BarUrnArray(new BarUrn(1))); + AspectFooBaz fooBaz = new AspectFooBaz().setBars(new BarUrnArray(new BarUrn(1))); AspectFooBarBaz fooBarBaz = new AspectFooBarBaz().setBars(new BarUrnArray(new BarUrn(1))); AspectAttributes attributes = new AspectAttributes().setAttributes(new StringArray("a")); AspectKey fooKey = new AspectKey<>(AspectFoo.class, urn, LATEST_VERSION); AspectKey barKey = new AspectKey<>(AspectBar.class, urn, LATEST_VERSION); AspectKey fooBarKey = new AspectKey<>(AspectFooBar.class, urn, LATEST_VERSION); + AspectKey fooBazKey = new AspectKey<>(AspectFooBaz.class, urn, LATEST_VERSION); AspectKey fooBarBazKey = new AspectKey<>(AspectFooBarBaz.class, urn, LATEST_VERSION); AspectKey attKey = new AspectKey<>(AspectAttributes.class, urn, LATEST_VERSION); - Set> aspectKeys = ImmutableSet.of(fooKey, barKey, fooBarKey, fooBarBazKey, attKey); + Set> aspectKeys = ImmutableSet.of(fooKey, barKey, fooBarKey, fooBazKey, fooBarBazKey, attKey); when(_mockLocalDAO.get(aspectKeys)).thenReturn(ImmutableMap.of(fooKey, Optional.of(foo), barKey, Optional.of(bar), - fooBarKey, Optional.of(fooBar), fooBarBazKey, Optional.of(fooBarBaz), attKey, Optional.of(attributes))); + fooBarKey, Optional.of(fooBar), fooBazKey, Optional.of(fooBaz), fooBarBazKey, Optional.of(fooBarBaz), attKey, Optional.of(attributes))); EntitySnapshot snapshot = runAndWait(_resource.getSnapshot(urn.toString(), null)); @@ -734,7 +740,7 @@ public void testGetSnapshotWithAllAspects() { Set aspects = snapshot.getAspects().stream().map(RecordUtils::getSelectedRecordTemplateFromUnion).collect(Collectors.toSet()); - assertEquals(aspects, ImmutableSet.of(foo, bar, fooBar, attributes)); + assertEquals(aspects, ImmutableSet.of(foo, bar, fooBar, fooBaz, attributes)); } @Test @@ -1158,8 +1164,8 @@ public void testFilterFromIndexWithAspects() { // case 2: null aspects is provided i.e. all aspects in the aspect union will be returned, non-null last urn List> listResult2 = Collections.singletonList(entry2); - when(_mockLocalDAO.getAspects(ImmutableSet.of(AspectFoo.class, AspectBar.class, AspectFooBar.class, AspectFooBarBaz.class, - AspectAttributes.class), indexFilter, null, urn1, 2)) + when(_mockLocalDAO.getAspects(ImmutableSet.of(AspectFoo.class, AspectBar.class, AspectFooBar.class, AspectFooBaz.class, + AspectFooBarBaz.class, AspectAttributes.class), indexFilter, null, urn1, 2)) .thenReturn(listResult2); List actual2 = @@ -1307,10 +1313,11 @@ public void testParseAspectsParam() { // All aspects aspectClasses = _resource.parseAspectsParam(null, false); - assertEquals(aspectClasses.size(), 5); + assertEquals(aspectClasses.size(), 6); assertTrue(aspectClasses.contains(AspectFoo.class)); assertTrue(aspectClasses.contains(AspectBar.class)); assertTrue(aspectClasses.contains(AspectFooBar.class)); + assertTrue(aspectClasses.contains(AspectFooBaz.class)); assertTrue(aspectClasses.contains(AspectFooBarBaz.class)); assertTrue(aspectClasses.contains(AspectAttributes.class)); } diff --git a/restli-resources/src/test/java/com/linkedin/metadata/restli/BaseEntitySimpleKeyResourceTest.java b/restli-resources/src/test/java/com/linkedin/metadata/restli/BaseEntitySimpleKeyResourceTest.java index 823ec0d90..0e590a576 100644 --- a/restli-resources/src/test/java/com/linkedin/metadata/restli/BaseEntitySimpleKeyResourceTest.java +++ b/restli-resources/src/test/java/com/linkedin/metadata/restli/BaseEntitySimpleKeyResourceTest.java @@ -28,6 +28,7 @@ import com.linkedin.testing.InternalEntityAspectUnion; import com.linkedin.testing.InternalEntitySnapshot; import com.linkedin.testing.localrelationship.AspectFooBar; +import com.linkedin.testing.localrelationship.AspectFooBaz; import com.linkedin.testing.localrelationship.AspectFooBarBaz; import com.linkedin.testing.urn.BarUrn; import java.net.URISyntaxException; @@ -69,11 +70,12 @@ public void testGet() { AspectKey aspect1Key = new AspectKey<>(AspectFoo.class, urn, LATEST_VERSION); AspectKey aspect2Key = new AspectKey<>(AspectBar.class, urn, LATEST_VERSION); AspectKey aspect3Key = new AspectKey<>(AspectFooBar.class, urn, LATEST_VERSION); - AspectKey aspect4Key = new AspectKey<>(AspectFooBarBaz.class, urn, LATEST_VERSION); - AspectKey aspect5Key = new AspectKey<>(AspectAttributes.class, urn, LATEST_VERSION); + AspectKey aspect4Key = new AspectKey<>(AspectFooBaz.class, urn, LATEST_VERSION); + AspectKey aspect5Key = new AspectKey<>(AspectFooBarBaz.class, urn, LATEST_VERSION); + AspectKey aspect6Key = new AspectKey<>(AspectAttributes.class, urn, LATEST_VERSION); when(_mockLocalDAO.exists(urn)).thenReturn(true); - when(_mockLocalDAO.get(new HashSet<>(Arrays.asList(aspect1Key, aspect2Key, aspect3Key, aspect4Key, aspect5Key)))) + when(_mockLocalDAO.get(new HashSet<>(Arrays.asList(aspect1Key, aspect2Key, aspect3Key, aspect4Key, aspect5Key, aspect6Key)))) .thenReturn(Collections.singletonMap(aspect1Key, Optional.of(foo))); EntityValue value = runAndWait(_resource.get(id, null)); @@ -189,13 +191,15 @@ public void testBatchGet() { AspectKey aspectBarKey2 = new AspectKey<>(AspectBar.class, urn2, LATEST_VERSION); AspectKey aspectFooBarKey1 = new AspectKey<>(AspectFooBar.class, urn1, LATEST_VERSION); AspectKey aspectFooBarKey2 = new AspectKey<>(AspectFooBar.class, urn2, LATEST_VERSION); + AspectKey aspectFooBazKey1 = new AspectKey<>(AspectFooBaz.class, urn1, LATEST_VERSION); + AspectKey aspectFooBazKey2 = new AspectKey<>(AspectFooBaz.class, urn2, LATEST_VERSION); AspectKey aspectFooBarBazKey1 = new AspectKey<>(AspectFooBarBaz.class, urn1, LATEST_VERSION); AspectKey aspectFooBarBazKey2 = new AspectKey<>(AspectFooBarBaz.class, urn2, LATEST_VERSION); AspectKey aspectAttKey1 = new AspectKey<>(AspectAttributes.class, urn1, LATEST_VERSION); AspectKey aspectAttKey2 = new AspectKey<>(AspectAttributes.class, urn2, LATEST_VERSION); when(_mockLocalDAO.get(ImmutableSet.of(aspectFooKey1, aspectBarKey1, aspectAttKey1, aspectFooKey2, aspectBarKey2, - aspectAttKey2, aspectFooBarKey1, aspectFooBarKey2, aspectFooBarBazKey1, aspectFooBarBazKey2))).thenReturn( + aspectAttKey2, aspectFooBarKey1, aspectFooBarKey2, aspectFooBazKey1, aspectFooBazKey2, aspectFooBarBazKey1, aspectFooBarBazKey2))).thenReturn( ImmutableMap.of(aspectFooKey1, Optional.of(foo), aspectFooKey2, Optional.of(bar))); Map keyValueMap = runAndWait(_resource.batchGet(ImmutableSet.of(id1, id2), null)) @@ -301,23 +305,25 @@ public void testGetSnapshotWithAllAspects() { AspectFoo foo = new AspectFoo().setValue("foo"); AspectBar bar = new AspectBar().setValue("bar"); AspectFooBar fooBar = new AspectFooBar().setBars(new BarUrnArray(new BarUrn(1))); + AspectFooBaz fooBaz = new AspectFooBaz().setBars(new BarUrnArray(new BarUrn(1))); AspectFooBarBaz fooBarBaz = new AspectFooBarBaz().setBars(new BarUrnArray(new BarUrn(1))); AspectAttributes att = new AspectAttributes().setAttributes(new StringArray("a")); AspectKey fooKey = new AspectKey<>(AspectFoo.class, urn, LATEST_VERSION); AspectKey barKey = new AspectKey<>(AspectBar.class, urn, LATEST_VERSION); AspectKey fooBarKey = new AspectKey<>(AspectFooBar.class, urn, LATEST_VERSION); + AspectKey fooBazKey = new AspectKey<>(AspectFooBaz.class, urn, LATEST_VERSION); AspectKey fooBarBazKey = new AspectKey<>(AspectFooBarBaz.class, urn, LATEST_VERSION); AspectKey attKey = new AspectKey<>(AspectAttributes.class, urn, LATEST_VERSION); - Set> aspectKeys = ImmutableSet.of(fooKey, barKey, fooBarKey, fooBarBazKey, attKey); + Set> aspectKeys = ImmutableSet.of(fooKey, barKey, fooBarKey, fooBazKey, fooBarBazKey, attKey); when(_mockLocalDAO.get(aspectKeys)).thenReturn(ImmutableMap.of(fooKey, Optional.of(foo), barKey, Optional.of(bar), - fooBarKey, Optional.of(fooBar), fooBarBazKey, Optional.of(fooBarBaz), attKey, Optional.of(att))); + fooBarKey, Optional.of(fooBar), fooBazKey, Optional.of(fooBaz), fooBarBazKey, Optional.of(fooBarBaz), attKey, Optional.of(att))); EntitySnapshot snapshot = runAndWait(_resource.getSnapshot(urn.toString(), null)); assertEquals(snapshot.getUrn(), urn); Set aspects = snapshot.getAspects().stream().map(RecordUtils::getSelectedRecordTemplateFromUnion).collect(Collectors.toSet()); - assertEquals(aspects, ImmutableSet.of(foo, bar, fooBar, att)); + assertEquals(aspects, ImmutableSet.of(foo, bar, fooBar, fooBaz, att)); } @Test diff --git a/settings.gradle b/settings.gradle index 5770faa8b..0d71bf903 100644 --- a/settings.gradle +++ b/settings.gradle @@ -4,7 +4,6 @@ include 'dao-api' include 'dao-impl:ebean-dao' include 'dao-impl:elasticsearch-dao' include 'dao-impl:elasticsearch-dao-7' -include 'dao-impl:neo4j-dao' include 'restli-resources' include 'testing:core-models-testing' include 'testing:elasticsearch-dao-integ-testing' diff --git a/testing/test-models/src/main/pegasus/com/linkedin/testing/EntityAspectUnion.pdl b/testing/test-models/src/main/pegasus/com/linkedin/testing/EntityAspectUnion.pdl index 1f3530e46..3c23b2777 100644 --- a/testing/test-models/src/main/pegasus/com/linkedin/testing/EntityAspectUnion.pdl +++ b/testing/test-models/src/main/pegasus/com/linkedin/testing/EntityAspectUnion.pdl @@ -1,8 +1,9 @@ namespace com.linkedin.testing import com.linkedin.testing.localrelationship.AspectFooBar +import com.linkedin.testing.localrelationship.AspectFooBaz import com.linkedin.testing.localrelationship.AspectFooBarBaz /** * For unit tests */ -typeref EntityAspectUnion = union[AspectFoo, AspectBar, AspectFooBar, AspectFooBarBaz, AspectAttributes] \ No newline at end of file +typeref EntityAspectUnion = union[AspectFoo, AspectBar, AspectFooBar, AspectFooBaz, AspectFooBarBaz, AspectAttributes] \ No newline at end of file diff --git a/testing/test-models/src/main/pegasus/com/linkedin/testing/localrelationship/AspectFooBar.pdl b/testing/test-models/src/main/pegasus/com/linkedin/testing/localrelationship/AspectFooBar.pdl index 9dda84b79..281a62c0b 100644 --- a/testing/test-models/src/main/pegasus/com/linkedin/testing/localrelationship/AspectFooBar.pdl +++ b/testing/test-models/src/main/pegasus/com/linkedin/testing/localrelationship/AspectFooBar.pdl @@ -4,7 +4,7 @@ import com.linkedin.testing.BarUrn @gma.aspect.column.name = "aspectfoobar" record AspectFooBar { - bars: array[BarUrn] + bars: optional array[BarUrn] belongsTos: optional array[BelongsToV2] reportsTos: optional array[ReportsTo] } \ No newline at end of file diff --git a/testing/test-models/src/main/pegasus/com/linkedin/testing/localrelationship/AspectFooBaz.pdl b/testing/test-models/src/main/pegasus/com/linkedin/testing/localrelationship/AspectFooBaz.pdl new file mode 100644 index 000000000..9d3e472dc --- /dev/null +++ b/testing/test-models/src/main/pegasus/com/linkedin/testing/localrelationship/AspectFooBaz.pdl @@ -0,0 +1,9 @@ +namespace com.linkedin.testing.localrelationship + +import com.linkedin.testing.BarUrn + +@gma.aspect.column.name = "aspectfoobaz" +record AspectFooBaz { + bars: optional array[BarUrn] + belongsTos: optional array[BelongsToV2] +} \ No newline at end of file diff --git a/validators/src/main/java/com/linkedin/metadata/validator/RelationshipValidator.java b/validators/src/main/java/com/linkedin/metadata/validator/RelationshipValidator.java index 47c5471b3..66da3fdbc 100644 --- a/validators/src/main/java/com/linkedin/metadata/validator/RelationshipValidator.java +++ b/validators/src/main/java/com/linkedin/metadata/validator/RelationshipValidator.java @@ -80,18 +80,11 @@ public static void validateRelationshipSchema(@Nonnull RecordDataSchema schema, } } - /** - * Validates a specific relationship model defined in com.linkedin.metadata.relationship. + * Validates a specific relationship model defined in com.linkedin.metadata.relationship. and cache the results. * - * @param schema schema for the model - */ - public static void validateRelationshipSchema(@Nonnull RecordDataSchema schema) { - validateRelationshipSchema(schema, false); - } - - /** - * Similar to {@link #validateRelationshipSchema(RecordDataSchema)} but take a {@link Class} instead and caches results. + * @param clazz class of the relationship model + * @param isRelationshipInV2 flag indicating whether or not the relationship model is v1 or v2 */ public static void validateRelationshipSchema(@Nonnull Class clazz, boolean isRelationshipInV2) { if (VALIDATED.contains(clazz)) { @@ -102,13 +95,6 @@ public static void validateRelationshipSchema(@Nonnull Class clazz) { - validateRelationshipSchema(clazz, false); - } - /** * Similar to {@link #validateRelationshipUnionSchema(UnionDataSchema, String)} but take a {@link Class} instead and caches results. */