Keep order for loadAll by objects or ids, fixes #196

Methods loadAll(Collection<T> objects, ..) and loadAll(Collection<ID>, ..) without SortOrder parameter now keep order of the input collection.
neo4j · Aug 15, 2017 · ccbe96c · ccbe96c
1 parent c38608c
commit ccbe96c
Show file tree

Hide file tree

Showing 4 changed files with 91 additions and 2 deletions.
diff --git a/core/src/main/java/org/neo4j/ogm/session/delegates/LoadByIdsDelegate.java b/core/src/main/java/org/neo4j/ogm/session/delegates/LoadByIdsDelegate.java
@@ -14,7 +14,9 @@
 
 import java.io.Serializable;
 import java.util.Collection;
+import java.util.HashMap;
 import java.util.LinkedHashSet;
+import java.util.Map;
 import java.util.Set;
 
 import org.neo4j.ogm.context.GraphEntityMapper;
@@ -55,6 +57,10 @@ public <T, ID extends Serializable> Collection<T> loadAll(Class<T> type, Collect
         GraphModelRequest request = new DefaultGraphModelRequest(qry.getStatement(), qry.getParameters());
         try (Response<GraphModel> response = session.requestHandler().execute(request)) {
             Iterable<T> mapped = new GraphEntityMapper(session.metaData(), session.context()).map(type, response);
+
+            if (sortOrder.sortClauses().isEmpty()) {
+                return sortResultsByIds(type, ids, mapped);
+            }
             Set<T> results = new LinkedHashSet<>();
             for (T entity : mapped) {
                 if (includeMappedEntity(ids, entity)) {
@@ -65,6 +71,32 @@ public <T, ID extends Serializable> Collection<T> loadAll(Class<T> type, Collect
         }
     }
 
+    private <T, ID extends Serializable> Set<T> sortResultsByIds(Class<T> type, Collection<ID> ids, Iterable<T> mapped) {
+        Map<ID, T> items = new HashMap<>();
+        ClassInfo classInfo = session.metaData().classInfo(type.getName());
+
+        FieldInfo idField = classInfo.primaryIndexField();
+        if (idField == null) {
+            idField = classInfo.identityField();
+        }
+
+        for (T t : mapped) {
+            Object id = idField.read(t);
+            if (id != null) {
+                items.put((ID) id, t);
+            }
+        }
+
+        Set<T> results = new LinkedHashSet<>();
+        for (ID id : ids) {
+            T item = items.get(id);
+            if (item != null) {
+                results.add(item);
+            }
+        }
+        return results;
+    }
+
     public <T, ID extends Serializable> Collection<T> loadAll(Class<T> type, Collection<ID> ids) {
         return loadAll(type, ids, new SortOrder(), null, 1);
     }

diff --git a/core/src/main/java/org/neo4j/ogm/session/delegates/LoadByInstancesDelegate.java b/core/src/main/java/org/neo4j/ogm/session/delegates/LoadByInstancesDelegate.java
@@ -14,7 +14,7 @@
 
 import java.lang.reflect.Field;
 import java.util.Collection;
-import java.util.HashSet;
+import java.util.LinkedHashSet;
 import java.util.Set;
 
 import org.neo4j.ogm.cypher.query.Pagination;
@@ -40,7 +40,7 @@ public <T> Collection<T> loadAll(Collection<T> objects, SortOrder sortOrder, Pag
             return objects;
         }
 
-        Set<Long> ids = new HashSet<>();
+        Set<Long> ids = new LinkedHashSet<>();
         Class type = objects.iterator().next().getClass();
         ClassInfo classInfo = session.metaData().classInfo(type.getName());
         Field identityField = classInfo.getField(classInfo.identityField());

diff --git a/test/src/test/java/org/neo4j/ogm/domain/music/Artist.java b/test/src/test/java/org/neo4j/ogm/domain/music/Artist.java
@@ -78,4 +78,12 @@ public void addAlbum(Album album) {
     public Long getId() {
         return id;
     }
+
+    @Override
+    public String toString() {
+        return "Artist{" +
+                "id=" + id +
+                ", name='" + name + '\'' +
+                '}';
+    }
 }
diff --git a/test/src/test/java/org/neo4j/ogm/persistence/session/capability/LoadCapabilityTest.java b/test/src/test/java/org/neo4j/ogm/persistence/session/capability/LoadCapabilityTest.java
@@ -13,6 +13,7 @@
 
 package org.neo4j.ogm.persistence.session.capability;
 
+import static com.google.common.collect.Lists.newArrayList;
 import static org.assertj.core.api.Assertions.*;
 
 import java.io.IOException;
@@ -678,4 +679,52 @@ public void shouldMaintainSortOrderWhenLoadingByIds() {
         assertThat(artistNames.get(1)).isEqualTo("Led Zeppelin");
         assertThat(artistNames.get(2)).isEqualTo("Bon Jovi");
     }
+
+    @Test
+    public void loadAllByIdsShouldSortByIdsIfSortOrderIsNotProvided() throws Exception {
+        Artist beatles = session.load(Artist.class, beatlesId);
+
+        Artist led = new Artist("Led Zeppelin");
+        session.save(led);
+        Artist bonJovi = new Artist("Bon Jovi");
+        session.save(bonJovi);
+
+        Long ledId = led.getId();
+        Long bonJoviId = bonJovi.getId();
+
+        Collection<Artist> artists;
+
+        artists = session.loadAll(Artist.class, newArrayList(beatlesId, ledId, bonJoviId));
+        assertThat(artists).containsExactly(beatles, led, bonJovi);
+
+        artists = session.loadAll(Artist.class, newArrayList(ledId, beatlesId, bonJoviId));
+        assertThat(artists).containsExactly(led, beatles, bonJovi);
+
+        artists = session.loadAll(Artist.class, newArrayList(ledId, bonJoviId, beatlesId));
+        assertThat(artists).containsExactly(led, bonJovi, beatles);
+    }
+
+    @Test
+    public void loadAllByInstancesShouldSortByIdsIfSortOrderIsNotProvided() throws Exception {
+        Artist beatles = session.load(Artist.class, beatlesId);
+
+        Artist led = new Artist("Led Zeppelin");
+        session.save(led);
+        Artist bonJovi = new Artist("Bon Jovi");
+        session.save(bonJovi);
+
+        Long ledId = led.getId();
+        Long bonJoviId = bonJovi.getId();
+
+        Collection<Artist> artists;
+
+        artists = session.loadAll(newArrayList(beatles, led, bonJovi));
+        assertThat(artists).containsExactly(beatles, led, bonJovi);
+
+        artists = session.loadAll(newArrayList(led, beatles, bonJovi));
+        assertThat(artists).containsExactly(led, beatles, bonJovi);
+
+        artists = session.loadAll(newArrayList(led, bonJovi, beatles));
+        assertThat(artists).containsExactly(led, bonJovi, beatles);
+    }
 }