DATAGRAPH-1412 - Deeper relationship mapping for custom queries.

meistermeier · meistermeier · commit a62d92b088de · 2020-10-23T11:53:48.000+02:00
diff --git a/src/main/java/org/springframework/data/neo4j/core/mapping/DefaultNeo4jEntityConverter.java b/src/main/java/org/springframework/data/neo4j/core/mapping/DefaultNeo4jEntityConverter.java
@@ -118,7 +118,7 @@ public <R> R read(Class<R> targetType, MapAccessor mapAccessor) {
 			if (queryRoot == null) {
 				throw new MappingException(String.format("Could not find mappable nodes or relationships inside %s for %s", mapAccessor, rootNodeDescription));
 			} else {
-				return map(queryRoot, rootNodeDescription, new KnownObjects(), new HashSet<>());
+				return map(queryRoot, queryRoot, rootNodeDescription, new KnownObjects(), new HashSet<>());
 			}
 		} catch (Exception e) {
 			throw new MappingException("Error mapping " + mapAccessor.toString(), e);
@@ -195,33 +195,23 @@ private static MapAccessor mergeRootNodeWithRecord(Node node, MapAccessor record
 	}
 
 	/**
-	 * @param queryResult The original query result
+	 * @param queryResult The original query result or a reduced form like a node or similar
+	 * @param allValues The original query result
 	 * @param nodeDescription The node description of the current entity to be mapped from the result
 	 * @param knownObjects The current list of known objects
+	 * @param processedSegments Path segments already processed in the mapping process. Only applies to path-based queries
 	 * @param <ET> As in entity type
-	 * @return
+	 * @return The mapped entity
 	 */
-	private <ET> ET map(MapAccessor queryResult, Neo4jPersistentEntity<ET> nodeDescription, KnownObjects knownObjects, Set<Path.Segment> processedSegments) {
-		return map(queryResult, nodeDescription, knownObjects, null, processedSegments);
+	private <ET> ET map(MapAccessor queryResult, MapAccessor allValues, Neo4jPersistentEntity<ET> nodeDescription, KnownObjects knownObjects, Set<Path.Segment> processedSegments) {
+		return map(queryResult, allValues, nodeDescription, knownObjects, null, processedSegments);
 	}
 
-	/**
-	 * @param queryResult The original query result
-	 * @param nodeDescription The node description of the current entity to be mapped from the result
-	 * @param knownObjects The current list of known objects
-	 * @param lastMappedEntity Previous created entity for relationships, can be null.
-	 * @param <ET> As in entity type
-	 * @return
-	 */
-	private <ET> ET map(MapAccessor queryResult, Neo4jPersistentEntity<ET> nodeDescription, KnownObjects knownObjects,
+	private <ET> ET map(MapAccessor queryResult, MapAccessor allValues, Neo4jPersistentEntity<ET> nodeDescription, KnownObjects knownObjects,
 						@Nullable Object lastMappedEntity, Set<Path.Segment> processedSegments) {
 
 		// if the given result does not contain an identifier to the mapped object cannot get temporarily saved
-		Long internalId = queryResult.get(Constants.NAME_OF_INTERNAL_ID).isNull()
-			? null
-			:	queryResult instanceof Node
-				? ((Node) queryResult).id()
-				: queryResult.get(Constants.NAME_OF_INTERNAL_ID).asLong();
+		Long internalId = getInternalId(queryResult);
 
 		Supplier<Object> mappedObjectSupplier = () -> {
 
@@ -233,7 +223,7 @@ private <ET> ET map(MapAccessor queryResult, Neo4jPersistentEntity<ET> nodeDescr
 
 			Collection<RelationshipDescription> relationships = concreteNodeDescription.getRelationships();
 
-			ET instance = instantiate(concreteNodeDescription, queryResult, knownObjects, relationships,
+			ET instance = instantiate(concreteNodeDescription, queryResult, allValues, knownObjects, relationships,
 					nodeDescriptionAndLabels.getDynamicLabels(), lastMappedEntity, processedSegments);
 
 			PersistentPropertyAccessor<ET> propertyAccessor = concreteNodeDescription.getPropertyAccessor(instance);
@@ -253,7 +243,7 @@ private <ET> ET map(MapAccessor queryResult, Neo4jPersistentEntity<ET> nodeDescr
 				knownObjects.storeObject(internalId, instance);
 				// Fill associations
 				concreteNodeDescription.doWithAssociations(
-						populateFrom(queryResult, propertyAccessor, isConstructorParameter, relationships, knownObjects, processedSegments));
+						populateFrom(queryResult, allValues, propertyAccessor, isConstructorParameter, relationships, knownObjects, processedSegments));
 			}
 			ET bean = propertyAccessor.getBean();
 
@@ -270,6 +260,15 @@ private <ET> ET map(MapAccessor queryResult, Neo4jPersistentEntity<ET> nodeDescr
 		return (ET) mappedObject;
 	}
 
+	@Nullable
+	private Long getInternalId(@NonNull MapAccessor queryResult) {
+		return queryResult instanceof Node
+				? (Long) ((Node) queryResult).id()
+				: queryResult.get(Constants.NAME_OF_INTERNAL_ID) == null || queryResult.get(Constants.NAME_OF_INTERNAL_ID).isNull()
+				? null
+				: queryResult.get(Constants.NAME_OF_INTERNAL_ID).asLong();
+	}
+
 	/**
 	 * Returns the list of labels for the entity to be created from the "main" node returned.
 	 *
@@ -289,7 +288,7 @@ private List<String> getLabels(MapAccessor queryResult) {
 		return labels;
 	}
 
-	private <ET> ET instantiate(Neo4jPersistentEntity<ET> nodeDescription, MapAccessor values, KnownObjects knownObjects,
+	private <ET> ET instantiate(Neo4jPersistentEntity<ET> nodeDescription, MapAccessor values, MapAccessor allValues, KnownObjects knownObjects,
 			Collection<RelationshipDescription> relationships, Collection<String> surplusLabels, Object lastMappedEntity,
 								Set<Path.Segment> processedSegments) {
 
@@ -300,7 +299,7 @@ public Object getParameterValue(PreferredConstructor.Parameter parameter) {
 				Neo4jPersistentProperty matchingProperty = nodeDescription.getRequiredPersistentProperty(parameter.getName());
 
 				if (matchingProperty.isRelationship()) {
-					return createInstanceOfRelationships(matchingProperty, values, knownObjects, relationships, processedSegments).orElse(null);
+					return createInstanceOfRelationships(matchingProperty, values, allValues, knownObjects, relationships, processedSegments).orElse(null);
 				} else if (matchingProperty.isDynamicLabels()) {
 					return createDynamicLabelsProperty(matchingProperty.getTypeInformation(), surplusLabels);
 				} else if (matchingProperty.isEntityInRelationshipWithProperties()) {
@@ -335,7 +334,7 @@ private PropertyHandler<Neo4jPersistentProperty> populateFrom(MapAccessor queryR
 		};
 	}
 
-	private AssociationHandler<Neo4jPersistentProperty> populateFrom(MapAccessor queryResult,
+	private AssociationHandler<Neo4jPersistentProperty> populateFrom(MapAccessor queryResult, MapAccessor allValues,
 			PersistentPropertyAccessor<?> propertyAccessor, Predicate<Neo4jPersistentProperty> isConstructorParameter,
 			Collection<RelationshipDescription> relationshipDescriptions, KnownObjects knownObjects, Set<Path.Segment> processedSegments) {
 		return association -> {
@@ -345,13 +344,13 @@ private AssociationHandler<Neo4jPersistentProperty> populateFrom(MapAccessor que
 				return;
 			}
 
-			createInstanceOfRelationships(persistentProperty, queryResult, knownObjects, relationshipDescriptions, processedSegments)
+			createInstanceOfRelationships(persistentProperty, queryResult, allValues, knownObjects, relationshipDescriptions, processedSegments)
 					.ifPresent(value -> propertyAccessor.setProperty(persistentProperty, value));
 		};
 	}
 
 	private Optional<Object> createInstanceOfRelationships(Neo4jPersistentProperty persistentProperty, MapAccessor values,
-			KnownObjects knownObjects, Collection<RelationshipDescription> relationshipDescriptions, Set<Path.Segment> processedSegments) {
+			MapAccessor allValues, KnownObjects knownObjects, Collection<RelationshipDescription> relationshipDescriptions, Set<Path.Segment> processedSegments) {
 
 		RelationshipDescription relationshipDescription = relationshipDescriptions.stream()
 				.filter(r -> r.getFieldName().equals(persistentProperty.getName())).findFirst().get();
@@ -392,11 +391,11 @@ private Optional<Object> createInstanceOfRelationships(Neo4jPersistentProperty p
 			mappedObjectHandler = (type, mappedObject) -> value.add(mappedObject);
 		}
 
-		Value list = values.get(relationshipDescription.generateRelatedNodesCollectionName());
+		Value list = allValues.get(relationshipDescription.generateRelatedNodesCollectionName());
 
 		List<Object> relationshipsAndProperties = new ArrayList<>();
 
-		boolean isGeneratedPathBased = values.containsKey(Constants.NAME_OF_PATHS);
+		boolean isGeneratedPathBased = allValues.containsKey(Constants.NAME_OF_PATHS);
 
 		Predicate<Value> isList = entry -> typeSystem.LIST().isTypeOf(entry);
 
@@ -426,11 +425,11 @@ private Optional<Object> createInstanceOfRelationships(Neo4jPersistentProperty p
 					continue;
 				}
 				processedSegments.add(segment);
-				Object mappedObject = map(extractNextNodeAndAppendPath(segment.end(), allPaths),
+				Object mappedObject = map(extractNextNodeAndAppendPath(segment.end(), allPaths), allValues,
 						concreteTargetNodeDescription, knownObjects, processedSegments);
 				if (relationshipDescription.hasRelationshipProperties()) {
 
-					Object relationshipProperties = map(segment.relationship(),
+					Object relationshipProperties = map(segment.relationship(), allValues,
 							(Neo4jPersistentEntity) relationshipDescription.getRelationshipPropertiesEntity(),
 							knownObjects, mappedObject, processedSegments);
 					relationshipsAndProperties.add(relationshipProperties);
@@ -449,11 +448,11 @@ private Optional<Object> createInstanceOfRelationships(Neo4jPersistentProperty p
 
 			// find relationships in the result
 			List<Relationship> allMatchingTypeRelationshipsInResult = StreamSupport
-					.stream(values.values().spliterator(), false).filter(isList.and(containsOnlyRelationships))
+					.stream(allValues.values().spliterator(), false).filter(isList.and(containsOnlyRelationships))
 					.flatMap(entry -> entry.asList(Value::asRelationship).stream()).filter(r -> r.type().equals(relationshipType))
 					.collect(Collectors.toList());
 
-			List<Node> allNodesWithMatchingLabelInResult = StreamSupport.stream(values.values().spliterator(), false)
+			List<Node> allNodesWithMatchingLabelInResult = StreamSupport.stream(allValues.values().spliterator(), false)
 					.filter(isList.and(containsOnlyNodes)).flatMap(entry -> entry.asList(Value::asNode).stream())
 					.filter(n -> n.hasLabel(targetLabel)).collect(Collectors.toList());
 
@@ -462,16 +461,17 @@ private Optional<Object> createInstanceOfRelationships(Neo4jPersistentProperty p
 			}
 
 			Function<Relationship, Long> targetIdSelector = relationshipDescription.isIncoming() ? Relationship::startNodeId : Relationship::endNodeId;
-
+			Function<Relationship, Long> sourceIdSelector = relationshipDescription.isIncoming() ? Relationship::endNodeId : Relationship::startNodeId;
+			Long sourceNodeId = getInternalId(values);
 			for (Node possibleValueNode : allNodesWithMatchingLabelInResult) {
-				long nodeId = possibleValueNode.id();
+				long targetNodeId = possibleValueNode.id();
 
 				for (Relationship possibleRelationship : allMatchingTypeRelationshipsInResult) {
-					if (targetIdSelector.apply(possibleRelationship) == nodeId) {
-						Object mappedObject = map(possibleValueNode, concreteTargetNodeDescription, knownObjects, processedSegments);
+					if (targetIdSelector.apply(possibleRelationship) == targetNodeId && sourceIdSelector.apply(possibleRelationship).equals(sourceNodeId)) {
+						Object mappedObject = map(possibleValueNode, values, concreteTargetNodeDescription, knownObjects, processedSegments);
 						if (relationshipDescription.hasRelationshipProperties()) {
 
-							Object relationshipProperties = map(possibleRelationship,
+							Object relationshipProperties = map(possibleRelationship, allValues,
 									(Neo4jPersistentEntity) relationshipDescription.getRelationshipPropertiesEntity(),
 									knownObjects, mappedObject, processedSegments);
 							relationshipsAndProperties.add(relationshipProperties);
@@ -487,13 +487,13 @@ private Optional<Object> createInstanceOfRelationships(Neo4jPersistentProperty p
 		} else {
 			for (Value relatedEntity : list.asList(Function.identity())) {
 
-				Object valueEntry = map(relatedEntity, concreteTargetNodeDescription, knownObjects, processedSegments);
+				Object valueEntry = map(relatedEntity, allValues, concreteTargetNodeDescription, knownObjects, processedSegments);
 
 				if (relationshipDescription.hasRelationshipProperties()) {
 					Relationship relatedEntityRelationship = relatedEntity.get(RelationshipDescription.NAME_OF_RELATIONSHIP)
 							.asRelationship();
 
-					Object relationshipProperties = map(relatedEntityRelationship,
+					Object relationshipProperties = map(relatedEntityRelationship, allValues,
 							(Neo4jPersistentEntity) relationshipDescription.getRelationshipPropertiesEntity(),
 							knownObjects, valueEntry, processedSegments);
 					relationshipsAndProperties.add(relationshipProperties);
diff --git a/src/test/java/org/springframework/data/neo4j/integration/imperative/RepositoryIT.java b/src/test/java/org/springframework/data/neo4j/integration/imperative/RepositoryIT.java
@@ -563,6 +563,31 @@ void findSliceShouldWork(@Autowired PersonRepository repository) {
 			assertThat(slice.get()).hasSize(1).extracting("name").containsExactly(TEST_PERSON1_NAME);
 			assertThat(slice.hasNext()).isFalse();
 		}
+
+		@Test // DATAGRAPH-1412
+		void customFindMapsDeepRelationships(@Autowired PetRepository repository) {
+			long petNode1Id;
+			long petNode2Id;
+			long petNode3Id;
+
+			try (Session session = createSession()) {
+				Record record = session.run("CREATE " + "(p1:Pet{name: 'Pet1'})-[:Has]->(p2:Pet{name: 'Pet2'}), "
+						+ "(p2)-[:Has]->(p3:Pet{name: 'Pet3'}) " + "RETURN p1, p2, p3").single();
+
+				petNode1Id = record.get("p1").asNode().id();
+				petNode2Id = record.get("p2").asNode().id();
+				petNode3Id = record.get("p3").asNode().id();
+			}
+
+			Pet loadedPet = repository.customQueryWithDeepRelationshipMapping(petNode1Id);
+
+			Pet comparisonPet2 = new Pet(petNode2Id, "Pet2");
+			Pet comparisonPet3 = new Pet(petNode3Id, "Pet3");
+			assertThat(loadedPet.getFriends()).containsExactlyInAnyOrder(comparisonPet2);
+
+			Pet pet2 = loadedPet.getFriends().get(loadedPet.getFriends().indexOf(comparisonPet2));
+			assertThat(pet2.getFriends()).containsExactly(comparisonPet3);
+		}
 	}
 
 	@Nested
@@ -645,7 +670,6 @@ void findDeepSameLabelsAndTypeRelationships(@Autowired PetRepository repository)
 			}
 
 			Pet loadedPet = repository.findById(petNode1Id).get();
-
 			Pet comparisonPet2 = new Pet(petNode2Id, "Pet2");
 			Pet comparisonPet3 = new Pet(petNode3Id, "Pet3");
 			assertThat(loadedPet.getFriends()).containsExactlyInAnyOrder(comparisonPet2);
@@ -3147,7 +3171,12 @@ interface PersonWithRelationshipWithPropertiesRepository
 		PersonWithRelationshipWithProperties findByHobbiesHobbyName(String hobbyName);
 	}
 
-	interface PetRepository extends Neo4jRepository<Pet, Long> {}
+	interface PetRepository extends Neo4jRepository<Pet, Long> {
+
+		@Query("MATCH (p:Pet)-[r1:Has]->(p2:Pet)-[r2:Has]->(p3:Pet) " +
+				"where id(p) = $petNode1Id return p, collect(r1), collect(p2), collect(r2), collect(p3)")
+		Pet customQueryWithDeepRelationshipMapping(@Param("petNode1Id") long petNode1Id);
+	}
 
 	interface RelationshipRepository extends Neo4jRepository<PersonWithRelationship, Long> {
 
