Fix self-referential association duplication issue (#1278)

Author: nakamura-to

doma-core/src/main/java/org/seasar/doma/jdbc/criteria/command/AssociateCommand.java

@@ -17,10 +17,13 @@
 
 import static java.util.stream.Collectors.toList;
 
+import java.util.HashMap;
 import java.util.LinkedHashMap;
+import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
+import java.util.Set;
 import java.util.function.BiFunction;
 import org.seasar.doma.internal.util.Combinations;
 import org.seasar.doma.internal.util.Pair;
@@ -47,23 +50,26 @@ public AssociateCommand(
   @Override
   @SuppressWarnings("unchecked")
   public List<ENTITY> execute() {
-    Map<EntityKey, Object> cache = new LinkedHashMap<>();
+    Set<CacheKey> rootEntityKeys = new LinkedHashSet<>();
+    Map<CacheKey, Object> cache = new HashMap<>();
     Combinations<EntityKey> combinations = new Combinations<>();
     SelectCommand<List<EntityPool>> command =
         new SelectCommand<>(
-            query, new EntityPoolIterationHandler(context.getProjectionEntityMetamodels()));
+            query,
+            new EntityPoolIterationHandler(
+                entityMetamodel, rootEntityKeys, context.getProjectionEntityMetamodels()));
     List<EntityPool> entityPools = command.execute();
     for (EntityPool entityPool : entityPools) {
       Map<EntityMetamodel<?>, Pair<EntityKey, Object>> associationCandidate = new LinkedHashMap<>();
       for (Map.Entry<EntityKey, EntityData> e : entityPool.entrySet()) {
         EntityKey key = e.getKey();
         EntityData data = e.getValue();
+        CacheKey cacheKey = CacheKey.of(key);
         Object entity =
             cache.computeIfAbsent(
-                key,
+                cacheKey,
                 k -> {
-                  EntityMetamodel<?> entityMetamodel = k.getEntityMetamodel();
-                  EntityType<Object> entityType = (EntityType<Object>) entityMetamodel.asType();
+                  EntityType<Object> entityType = (EntityType<Object>) k.entityType();
                   Object newEntity = entityType.newEntity(data.getStates());
                   if (!entityType.isImmutable()) {
                     entityType.saveCurrentStates(newEntity);
@@ -75,14 +81,11 @@ public List<ENTITY> execute() {
       associate(cache, combinations, associationCandidate);
     }
     return (List<ENTITY>)
-        cache.entrySet().stream()
-            .filter(e -> e.getKey().getEntityMetamodel() == entityMetamodel)
-            .map(Map.Entry::getValue)
-            .collect(toList());
+        rootEntityKeys.stream().map(cache::get).filter(Objects::nonNull).collect(toList());
   }
 
   private void associate(
-      Map<EntityKey, Object> cache,
+      Map<CacheKey, Object> cache,
       Combinations<EntityKey> combinations,
       Map<EntityMetamodel<?>, Pair<EntityKey, Object>> associationCandidate) {
     for (Map.Entry<Pair<EntityMetamodel<?>, EntityMetamodel<?>>, BiFunction<Object, Object, Object>>
@@ -100,7 +103,7 @@ private void associate(
       }
       Object newEntity = associator.apply(keyAndEntity1.snd, keyAndEntity2.snd);
       if (newEntity != null) {
-        cache.replace(keyAndEntity1.fst, newEntity);
+        cache.replace(CacheKey.of(keyAndEntity1.fst), newEntity);
         associationCandidate.replace(metamodelPair.fst, new Pair<>(keyAndEntity1.fst, newEntity));
       }
       combinations.add(keyPair);

doma-core/src/main/java/org/seasar/doma/jdbc/criteria/command/CacheKey.java

@@ -0,0 +1,25 @@
+/*
+ * Copyright Doma Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.seasar.doma.jdbc.criteria.command;
+
+import java.util.List;
+import org.seasar.doma.jdbc.entity.EntityType;
+
+public record CacheKey(EntityType<?> entityType, List<?> items) {
+  static CacheKey of(EntityKey key) {
+    return new CacheKey(key.getEntityMetamodel().asType(), key.getItems());
+  }
+}

doma-core/src/main/java/org/seasar/doma/jdbc/criteria/command/EntityKey.java

@@ -34,6 +34,10 @@ public EntityMetamodel<?> getEntityMetamodel() {
     return entityMetamodel;
   }
 
+  public List<?> getItems() {
+    return items;
+  }
+
   @Override
   public boolean equals(Object o) {
     if (this == o) return true;

doma-core/src/main/java/org/seasar/doma/jdbc/criteria/command/EntityPoolIterationHandler.java

@@ -18,6 +18,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
+import java.util.Set;
 import org.seasar.doma.internal.jdbc.command.AbstractIterationHandler;
 import org.seasar.doma.internal.jdbc.command.ResultListCallback;
 import org.seasar.doma.jdbc.ObjectProvider;
@@ -27,16 +28,23 @@
 
 public class EntityPoolIterationHandler
     extends AbstractIterationHandler<EntityPool, List<EntityPool>> {
+  private final EntityMetamodel<?> rootEntityMetamodel;
+  private final Set<CacheKey> rootEntityKeys;
   private final Map<EntityMetamodel<?>, List<PropertyMetamodel<?>>> projectionEntityMetamodels;
 
   public EntityPoolIterationHandler(
+      EntityMetamodel<?> rootEntityMetamodel,
+      Set<CacheKey> rootEntityKeys,
       Map<EntityMetamodel<?>, List<PropertyMetamodel<?>>> projectionEntityMetamodels) {
     super(new ResultListCallback<>());
+    this.rootEntityMetamodel = Objects.requireNonNull(rootEntityMetamodel);
+    this.rootEntityKeys = Objects.requireNonNull(rootEntityKeys);
     this.projectionEntityMetamodels = Objects.requireNonNull(projectionEntityMetamodels);
   }
 
   @Override
   protected ObjectProvider<EntityPool> createObjectProvider(SelectQuery query) {
-    return new EntityPoolProvider(projectionEntityMetamodels, query);
+    return new EntityPoolProvider(
+        rootEntityMetamodel, rootEntityKeys, projectionEntityMetamodels, query);
   }
 }

doma-core/src/main/java/org/seasar/doma/jdbc/criteria/command/EntityPoolProvider.java

@@ -24,6 +24,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
+import java.util.Set;
 import java.util.stream.Collectors;
 import org.seasar.doma.internal.jdbc.command.FetchSupport;
 import org.seasar.doma.jdbc.ObjectProvider;
@@ -35,11 +36,18 @@
 import org.seasar.doma.jdbc.query.Query;
 
 public class EntityPoolProvider implements ObjectProvider<EntityPool> {
+  private final EntityMetamodel<?> rootEntityMetamodel;
+  private final Set<CacheKey> rootEntityKeys;
   private final Map<EntityMetamodel<?>, List<PropertyMetamodel<?>>> projectionEntityMetamodels;
   private final FetchSupport fetchSupport;
 
   public EntityPoolProvider(
-      Map<EntityMetamodel<?>, List<PropertyMetamodel<?>>> projectionEntityMetamodels, Query query) {
+      EntityMetamodel<?> rootEntityMetamodel,
+      Set<CacheKey> rootEntityKeys,
+      Map<EntityMetamodel<?>, List<PropertyMetamodel<?>>> projectionEntityMetamodels,
+      Query query) {
+    this.rootEntityMetamodel = Objects.requireNonNull(rootEntityMetamodel);
+    this.rootEntityKeys = Objects.requireNonNull(rootEntityKeys);
     this.projectionEntityMetamodels = Objects.requireNonNull(projectionEntityMetamodels);
     Objects.requireNonNull(query);
     this.fetchSupport = new FetchSupport(query);
@@ -82,6 +90,9 @@ public EntityPool get(ResultSet resultSet) throws SQLException {
           props.stream().collect(Collectors.toMap(p -> p.propType.getName(), p -> p.prop));
       EntityData data = new EntityData(states);
       entityPool.put(key, data);
+      if (key.getEntityMetamodel() == rootEntityMetamodel) {
+        rootEntityKeys.add(CacheKey.of(key));
+      }
     }
     return entityPool;
   }

doma-core/src/main/java/org/seasar/doma/jdbc/criteria/metamodel/EntityTypeProxy.java

@@ -182,4 +182,17 @@ public void postUpdate(ENTITY entity, PostUpdateContext<ENTITY> context) {
   public void postDelete(ENTITY entity, PostDeleteContext<ENTITY> context) {
     entityType.postDelete(entity, context);
   }
+
+  @Override
+  public boolean equals(Object o) {
+    if (o == null || getClass() != o.getClass()) return false;
+    EntityTypeProxy<?> that = (EntityTypeProxy<?>) o;
+    return Objects.equals(entityType, that.entityType)
+        && Objects.equals(qualifiedTableName, that.qualifiedTableName);
+  }
+
+  @Override
+  public int hashCode() {
+    return Objects.hash(entityType, qualifiedTableName);
+  }
 }

integration-test-java/src/test/java/org/seasar/doma/it/criteria/QueryDslEntitySelectTest.java

@@ -20,6 +20,7 @@
 import static org.junit.jupiter.api.Assertions.assertIterableEquals;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.junit.jupiter.api.Assertions.assertNull;
+import static org.junit.jupiter.api.Assertions.assertSame;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.seasar.doma.it.criteria.CustomExpressions.addOne;
@@ -950,4 +951,24 @@ void iterableDomain() {
     assertEquals("TOKYO", iterator.next());
     assertEquals("KYOTO", iterator.next());
   }
+
+  @Test
+  void selfJoin_association() {
+    Employee_ e = new Employee_();
+    Employee_ m = new Employee_();
+
+    QueryDsl queryDsl = new QueryDsl(config);
+    List<Employee> employees =
+        queryDsl
+            .from(e)
+            .leftJoin(m, on -> on.eq(e.managerId, m.employeeId))
+            .where(c -> c.in(e.employeeId, List.of(6, 9)))
+            .orderBy(c -> c.asc(e.employeeId))
+            .associate(e, m, (Employee::setManager))
+            .fetch();
+    assertEquals(2, employees.size());
+    Employee blake = employees.get(0);
+    Employee king = employees.get(1);
+    assertSame(king, blake.getManager());
+  }
 }