From 9b801c7b5027256d9886c8df83236f5510f0cadf Mon Sep 17 00:00:00 2001
From: Thomas Darimont <tdarimont@pivotal.io>
Date: Fri, 28 Nov 2014 12:21:36 +0100
Subject: [PATCH] DATAJPA-611 - Allow composite keys to be used in
 SimpleJpaRepository.findAll(Iterable<ID>).
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

We now support an Iterable of composite primary keys to be passed to findAll(…). Note that since there is no direct support in JPA to perform this query we have to execute a entityManager.find(…) query for every given id, which could lead to performance problems - use with care!

Original pull request: #127.
---
 .../support/SimpleJpaRepository.java          | 21 +++++++++-
 .../RepositoryWithCompositeKeyTests.java      | 38 +++++++++++++++++++
 2 files changed, 57 insertions(+), 2 deletions(-)

diff --git a/src/main/java/org/springframework/data/jpa/repository/support/SimpleJpaRepository.java b/src/main/java/org/springframework/data/jpa/repository/support/SimpleJpaRepository.java
index e78daa36b..eb9060e5b 100644
--- a/src/main/java/org/springframework/data/jpa/repository/support/SimpleJpaRepository.java
+++ b/src/main/java/org/springframework/data/jpa/repository/support/SimpleJpaRepository.java
@@ -302,7 +302,18 @@ public class SimpleJpaRepository<T, ID extends Serializable> implements JpaRepos
 			return Collections.emptyList();
 		}
 
-		ByIdsSpecification specification = new ByIdsSpecification();
+		if (entityInformation.hasCompositeId()) {
+
+			List<T> results = new ArrayList<T>();
+
+			for (ID id : ids) {
+				results.add(findOne(id));
+			}
+
+			return results;
+		}
+
+		ByIdsSpecification<T> specification = new ByIdsSpecification<T>(entityInformation);
 		TypedQuery<T> query = getQuery(specification, (Sort) null);
 
 		return query.setParameter(specification.parameter, ids).getResultList();
@@ -572,10 +583,16 @@ public class SimpleJpaRepository<T, ID extends Serializable> implements JpaRepos
 	 * @author Oliver Gierke
 	 */
 	@SuppressWarnings("rawtypes")
-	private final class ByIdsSpecification implements Specification<T> {
+	private static final class ByIdsSpecification<T> implements Specification<T> {
+
+		private final JpaEntityInformation<T, ?> entityInformation;
 
 		ParameterExpression<Iterable> parameter;
 
+		public ByIdsSpecification(JpaEntityInformation<T, ?> entityInformation) {
+			this.entityInformation = entityInformation;
+		}
+
 		/*
 		 * (non-Javadoc)
 		 * @see org.springframework.data.jpa.domain.Specification#toPredicate(javax.persistence.criteria.Root, javax.persistence.criteria.CriteriaQuery, javax.persistence.criteria.CriteriaBuilder)
diff --git a/src/test/java/org/springframework/data/jpa/repository/RepositoryWithCompositeKeyTests.java b/src/test/java/org/springframework/data/jpa/repository/RepositoryWithCompositeKeyTests.java
index 2871859cf..758ad751e 100644
--- a/src/test/java/org/springframework/data/jpa/repository/RepositoryWithCompositeKeyTests.java
+++ b/src/test/java/org/springframework/data/jpa/repository/RepositoryWithCompositeKeyTests.java
@@ -18,6 +18,7 @@ package org.springframework.data.jpa.repository;
 import static org.hamcrest.Matchers.*;
 import static org.junit.Assert.*;
 
+import java.util.Arrays;
 import java.util.List;
 
 import org.junit.Rule;
@@ -263,4 +264,41 @@ public class RepositoryWithCompositeKeyTests {
 
 		assertThat(employeeRepositoryWithEmbeddedId.exists(key), is(true));
 	}
+
+	/**
+	 * @see DATAJPA-611
+	 */
+	@Test
+	public void shouldAllowFindAllWithIdsForEntitiesWithCompoundIdClassKeys() {
+
+		IdClassExampleDepartment dep2 = new IdClassExampleDepartment();
+		dep2.setDepartmentId(2L);
+		dep2.setName("Dep2");
+
+		IdClassExampleEmployee emp1 = new IdClassExampleEmployee();
+		emp1.setEmpId(3L);
+		emp1.setDepartment(dep2);
+		emp1 = employeeRepositoryWithIdClass.save(emp1);
+
+		IdClassExampleDepartment dep1 = new IdClassExampleDepartment();
+		dep1.setDepartmentId(1L);
+		dep1.setName("Dep1");
+
+		IdClassExampleEmployee emp2 = new IdClassExampleEmployee();
+		emp2.setEmpId(2L);
+		emp2.setDepartment(dep1);
+		emp2 = employeeRepositoryWithIdClass.save(emp2);
+
+		IdClassExampleEmployeePK emp1PK = new IdClassExampleEmployeePK();
+		emp1PK.setDepartment(2L);
+		emp1PK.setEmpId(3L);
+
+		IdClassExampleEmployeePK emp2PK = new IdClassExampleEmployeePK();
+		emp1PK.setDepartment(1L);
+		emp1PK.setEmpId(2L);
+
+		List<IdClassExampleEmployee> result = employeeRepositoryWithIdClass.findAll(Arrays.asList(emp1PK, emp2PK));
+
+		assertThat(result, hasSize(2));
+	}
 }