From e737efde6bcbcf7f78ab583a4127c8f50d2c41fc Mon Sep 17 00:00:00 2001
From: Petr Strnad <pstrnad@idc.com>
Date: Sat, 18 Jun 2022 13:16:36 +0200
Subject: [PATCH] Fix EntityGraphFactory to not overwrite prior subgraphs.

Resolves #2527.
---
 .../support/EntityGraphFactory.java           | 17 +++++++---
 .../support/EntityGraphFactoryUnitTests.java  | 31 ++++++++++++++-----
 2 files changed, 36 insertions(+), 12 deletions(-)

diff --git a/spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/support/EntityGraphFactory.java b/spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/support/EntityGraphFactory.java
index 9e6c35a7e..f2c105fbc 100644
--- a/spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/support/EntityGraphFactory.java
+++ b/spring-data-jpa/src/main/java/org/springframework/data/jpa/repository/support/EntityGraphFactory.java
@@ -15,18 +15,21 @@
  */
 package org.springframework.data.jpa.repository.support;
 
-import java.util.Set;
-
 import jakarta.persistence.EntityGraph;
 import jakarta.persistence.EntityManager;
 import jakarta.persistence.Subgraph;
 
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Set;
+
 import org.springframework.data.mapping.PropertyPath;
 
 /**
  * Factory class to create an {@link EntityGraph} from a collection of property paths.
  *
  * @author Jens Schauder
+ * @author Petr Strnad
  * @since 2.6
  */
 abstract class EntityGraphFactory {
@@ -42,16 +45,22 @@ abstract class EntityGraphFactory {
 	public static <T> EntityGraph<T> create(EntityManager entityManager, Class<T> domainType, Set<String> properties) {
 
 		EntityGraph<T> entityGraph = entityManager.createEntityGraph(domainType);
+		Map<String, Subgraph<Object>> existingSubgraphs = new HashMap<>();
 
 		for (String property : properties) {
 
 			Subgraph<Object> current = null;
+			String currentFullPath = "";
 
 			for (PropertyPath path : PropertyPath.from(property, domainType)) {
 
+				currentFullPath += path.getSegment() + ".";
+
 				if (path.hasNext()) {
-					current = current == null ? entityGraph.addSubgraph(path.getSegment())
-							: current.addSubgraph(path.getSegment());
+					final Subgraph<Object> finalCurrent = current;
+					current = current == null
+							? existingSubgraphs.computeIfAbsent(currentFullPath, k -> entityGraph.addSubgraph(path.getSegment()))
+							: existingSubgraphs.computeIfAbsent(currentFullPath, k -> finalCurrent.addSubgraph(path.getSegment()));
 					continue;
 				}
 
diff --git a/spring-data-jpa/src/test/java/org/springframework/data/jpa/repository/support/EntityGraphFactoryUnitTests.java b/spring-data-jpa/src/test/java/org/springframework/data/jpa/repository/support/EntityGraphFactoryUnitTests.java
index 8dac763c7..c25caab81 100644
--- a/spring-data-jpa/src/test/java/org/springframework/data/jpa/repository/support/EntityGraphFactoryUnitTests.java
+++ b/spring-data-jpa/src/test/java/org/springframework/data/jpa/repository/support/EntityGraphFactoryUnitTests.java
@@ -15,15 +15,18 @@
  */
 package org.springframework.data.jpa.repository.support;
 
-import static java.util.Arrays.*;
-import static org.mockito.Mockito.*;
-
-import java.util.HashSet;
+import static java.util.Arrays.asList;
+import static org.mockito.Mockito.RETURNS_DEEP_STUBS;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
 
 import jakarta.persistence.EntityGraph;
 import jakarta.persistence.EntityManager;
 import jakarta.persistence.Subgraph;
 
+import java.util.HashSet;
+
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 
@@ -31,6 +34,7 @@ import org.junit.jupiter.api.Test;
  * Unit tests for {@link EntityGraphFactory}.
  *
  * @author Jens Schauder
+ * @author Petr Strnad
  */
 @SuppressWarnings("rawtypes")
 class EntityGraphFactoryUnitTests {
@@ -45,8 +49,7 @@ class EntityGraphFactoryUnitTests {
 		when(em.createEntityGraph(DummyEntity.class)).thenReturn(entityGraph);
 	}
 
-	// GH-2329
-	@Test
+	@Test // GH-2329
 	void simpleSetOfPropertiesGetRegistered() {
 
 		HashSet<String> properties = new HashSet<>(asList("one", "two"));
@@ -57,8 +60,7 @@ class EntityGraphFactoryUnitTests {
 		verify(entityGraph).addAttributeNodes("two");
 	}
 
-	// GH-2329
-	@Test
+	@Test // GH-2329
 	void setOfCompositePropertiesGetRegisteredPiecewise() {
 
 		HashSet<String> properties = new HashSet<>(asList("one.two", "eins.zwei.drei"));
@@ -76,6 +78,19 @@ class EntityGraphFactoryUnitTests {
 		verify(zwei).addAttributeNodes("drei");
 	}
 
+	@Test // GH-2571
+	void multipleSubNodesUnderSameParentNodeShouldWork() {
+
+		HashSet<String> properties = new HashSet<>(asList("one.one", "one.two"));
+
+		entityGraph = EntityGraphFactory.create(em, DummyEntity.class, properties);
+
+		verify(entityGraph).addSubgraph("one");
+		Subgraph<?> one = entityGraph.addSubgraph("one");
+		verify(one).addAttributeNodes("one");
+		verify(one).addAttributeNodes("two");
+	}
+
 	private static class DummyEntity {
 		DummyEntity one;
 		DummyEntity two;