DATAJPA-669 - Introduced JpaContext to abstract the current JPA setup.

DefaultJpaContext is set up of all EntityManager instances in the current ApplicationContext and exposed for injection. Its interface JpaContext then currently allows looking up EntityManagers by managed domain types.

If multiple EntityManagers of the current ApplicationContext manage the a single domain type the request is rejected.

src/main/java/org/springframework/data/jpa/repository/JpaContext.java

@@ -0,0 +1,38 @@
+/*
+ * Copyright 2015 the original author or 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
+ *
+ *      http://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.springframework.data.jpa.repository;
+
+import javax.persistence.EntityManager;
+
+/**
+ * Interface for components to provide useful information about the current JPA setup within the current
+ * {@link org.springframework.context.ApplicationContext}.
+ * 
+ * @author Oliver Gierke
+ * @soundtrack Marcus Miller - Water Dancer (Afrodeezia)
+ * @since 1.9
+ */
+public interface JpaContext {
+
+	/**
+	 * Returns the {@link EntityManager} managing the given domain type.
+	 * 
+	 * @param managedType must not be {@literal null}.
+	 * @return the {@link EntityManager} that manages the given type, will never be {@literal null}.
+	 * @throws IllegalArgumentException if the given type is not a JPA managed one no unique {@link EntityManager} managing this type can be resolved.
+	 */
+	EntityManager getEntityManagerByManagedType(Class<?> managedType);
+}

src/main/java/org/springframework/data/jpa/repository/config/BeanDefinitionNames.java

@@ -24,4 +24,5 @@ package org.springframework.data.jpa.repository.config;
 interface BeanDefinitionNames {
 
 	public static final String JPA_MAPPING_CONTEXT_BEAN_NAME = "jpaMappingContext";
+	public static final String JPA_CONTEXT_BEAN_NAME = "jpaContext";
 }

src/main/java/org/springframework/data/jpa/repository/config/JpaRepositoryConfigExtension.java

@@ -37,6 +37,7 @@ import org.springframework.core.annotation.AnnotationAttributes;
 import org.springframework.dao.DataAccessException;
 import org.springframework.dao.annotation.PersistenceExceptionTranslationPostProcessor;
 import org.springframework.data.jpa.repository.JpaRepository;
+import org.springframework.data.jpa.repository.support.DefaultJpaContext;
 import org.springframework.data.jpa.repository.support.EntityManagerBeanDefinitionRegistrarPostProcessor;
 import org.springframework.data.jpa.repository.support.JpaRepositoryFactoryBean;
 import org.springframework.data.repository.config.AnnotationRepositoryConfigurationSource;
@@ -170,6 +171,13 @@ public class JpaRepositoryConfigExtension extends RepositoryConfigurationExtensi
 
 		registerIfNotAlreadyRegistered(new RootBeanDefinition(PAB_POST_PROCESSOR), registry,
 				AnnotationConfigUtils.PERSISTENCE_ANNOTATION_PROCESSOR_BEAN_NAME, source);
+
+		// Register bean definition for DefaultJpaContext
+
+		RootBeanDefinition contextDefinition = new RootBeanDefinition(DefaultJpaContext.class);
+		contextDefinition.setAutowireMode(AbstractBeanDefinition.AUTOWIRE_CONSTRUCTOR);
+
+		registerIfNotAlreadyRegistered(contextDefinition, registry, JPA_CONTEXT_BEAN_NAME, source);
 	}
 
 	/**

src/main/java/org/springframework/data/jpa/repository/support/DefaultJpaContext.java

@@ -0,0 +1,81 @@
+/*
+ * Copyright 2015 the original author or 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
+ *
+ *      http://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.springframework.data.jpa.repository.support;
+
+import java.util.List;
+import java.util.Set;
+
+import javax.persistence.EntityManager;
+import javax.persistence.metamodel.ManagedType;
+
+import org.springframework.data.jpa.repository.JpaContext;
+import org.springframework.util.Assert;
+import org.springframework.util.LinkedMultiValueMap;
+import org.springframework.util.MultiValueMap;
+
+/**
+ * Default implementation of {@link JpaContext}.
+ * 
+ * @author Oliver Gierke
+ * @soundtrack Marcus Miller - Son Of Macbeth (Afrodeezia)
+ * @since 1.9
+ */
+public class DefaultJpaContext implements JpaContext {
+
+	private final MultiValueMap<Class<?>, EntityManager> entityManagers;
+
+	/**
+	 * Creates a new {@link DefaultJpaContext} for the given {@link Set} of {@link EntityManager}s.
+	 * 
+	 * @param entityManagers must not be {@literal null}.
+	 */
+	public DefaultJpaContext(Set<EntityManager> entityManagers) {
+
+		Assert.notNull(entityManagers, "EntityManagers must not be null!");
+		Assert.notEmpty(entityManagers, "EntityManagers must not be empty!");
+
+		this.entityManagers = new LinkedMultiValueMap<Class<?>, EntityManager>();
+
+		for (EntityManager em : entityManagers) {
+			for (ManagedType<?> managedType : em.getMetamodel().getManagedTypes()) {
+				this.entityManagers.add(managedType.getJavaType(), em);
+			}
+		}
+	}
+
+	/*
+	 * (non-Javadoc)
+	 * @see org.springframework.data.jpa.repository.JpaContext#getByManagedType(java.lang.Class)
+	 */
+	@Override
+	public EntityManager getEntityManagerByManagedType(Class<?> type) {
+
+		Assert.notNull(type, "Type must not be null!");
+
+		if (!entityManagers.containsKey(type)) {
+			throw new IllegalArgumentException(String.format("%s is not a managed type!", type));
+		}
+
+		List<EntityManager> candidates = this.entityManagers.get(type);
+
+		if (candidates.size() == 1) {
+			return candidates.get(0);
+		}
+
+		throw new IllegalArgumentException(
+				String.format("%s managed by more than one EntityManagers: %s!", type.getName(), candidates));
+	}
+}

src/test/java/org/springframework/data/jpa/repository/JavaConfigUserRepositoryTests.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2012-2014 the original author or authors.
+ * Copyright 2012-2015 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -16,6 +16,7 @@
 package org.springframework.data.jpa.repository;
 
 import java.io.IOException;
+import java.util.Collections;
 import java.util.List;
 
 import javax.persistence.EntityManager;
@@ -37,6 +38,7 @@ import org.springframework.data.jpa.repository.config.EnableJpaRepositories;
 import org.springframework.data.jpa.repository.sample.SampleEvaluationContextExtension;
 import org.springframework.data.jpa.repository.sample.UserRepository;
 import org.springframework.data.jpa.repository.sample.UserRepositoryImpl;
+import org.springframework.data.jpa.repository.support.DefaultJpaContext;
 import org.springframework.data.jpa.repository.support.JpaRepositoryFactoryBean;
 import org.springframework.data.repository.core.NamedQueries;
 import org.springframework.data.repository.core.support.PropertiesBasedNamedQueries;
@@ -78,7 +80,8 @@ public class JavaConfigUserRepositoryTests extends UserRepositoryTests {
 			factory.setEntityManager(entityManager);
 			factory.setBeanFactory(applicationContext);
 			factory.setRepositoryInterface(UserRepository.class);
-			factory.setCustomImplementation(new UserRepositoryImpl());
+			factory
+					.setCustomImplementation(new UserRepositoryImpl(new DefaultJpaContext(Collections.singleton(entityManager))));
 			factory.setNamedQueries(namedQueries());
 			factory.setEvaluationContextProvider(evaluationContextProvider);
 			factory.afterPropertiesSet();

src/test/java/org/springframework/data/jpa/repository/sample/UserRepositoryImpl.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2008-2011 the original author or authors.
+ * Copyright 2008-2015 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -17,7 +17,10 @@ package org.springframework.data.jpa.repository.sample;
 
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
+import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.data.jpa.domain.sample.User;
+import org.springframework.data.jpa.repository.JpaContext;
+import org.springframework.util.Assert;
 
 /**
  * Dummy implementation to allow check for invoking a custom implementation.
@@ -28,25 +31,24 @@ public class UserRepositoryImpl implements UserRepositoryCustom {
 
 	private static final Logger LOG = LoggerFactory.getLogger(UserRepositoryImpl.class);
 
+	@Autowired
+	public UserRepositoryImpl(JpaContext context) {
+		Assert.notNull(context, "JpaContext must not be null!");
+	}
+
 	/*
 	 * (non-Javadoc)
-	 * 
-	 * @see org.springframework.data.jpa.repository.sample.UserRepositoryCustom#
-	 * someCustomMethod(org.springframework.data.jpa.domain.sample.User)
+	 * @see org.springframework.data.jpa.repository.sample.UserRepositoryCustom#someCustomMethod(org.springframework.data.jpa.domain.sample.User)
 	 */
 	public void someCustomMethod(User u) {
-
 		LOG.debug("Some custom method was invoked!");
 	}
 
 	/*
 	 * (non-Javadoc)
-	 * 
-	 * @see org.springframework.data.jpa.repository.sample.UserRepositoryCustom#
-	 * findByOverrridingMethod()
+	 * @see org.springframework.data.jpa.repository.sample.UserRepositoryCustom#findByOverrridingMethod()
 	 */
 	public void findByOverrridingMethod() {
-
 		LOG.debug("A method overriding a finder was invoked!");
 	}
 }

src/test/java/org/springframework/data/jpa/repository/support/DefaultJpaContextIntegrationTests.java

@@ -0,0 +1,113 @@
+/*
+ * Copyright 2015 the original author or 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
+ *
+ *      http://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.springframework.data.jpa.repository.support;
+
+import static org.hamcrest.CoreMatchers.*;
+import static org.junit.Assert.*;
+
+import java.util.Arrays;
+import java.util.HashSet;
+
+import javax.persistence.EntityManager;
+import javax.persistence.EntityManagerFactory;
+
+import org.hibernate.ejb.HibernatePersistence;
+import org.junit.Before;
+import org.junit.BeforeClass;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.ExpectedException;
+import org.springframework.data.jpa.domain.sample.Category;
+import org.springframework.data.jpa.domain.sample.User;
+import org.springframework.data.jpa.repository.JpaContext;
+import org.springframework.jdbc.datasource.embedded.EmbeddedDatabaseBuilder;
+import org.springframework.jdbc.datasource.embedded.EmbeddedDatabaseType;
+import org.springframework.orm.jpa.LocalContainerEntityManagerFactoryBean;
+
+/**
+ * Integration tests for {@link DefaultJpaContext}.
+ * 
+ * @author Oliver Gierke
+ * @soundtrack Marcus Miller - Papa Was A Rolling Stone (Afrodeezia)
+ */
+public class DefaultJpaContextIntegrationTests {
+
+	public @Rule ExpectedException exception = ExpectedException.none();
+
+	static EntityManagerFactory firstEmf, secondEmf;
+
+	EntityManager firstEm, secondEm;
+	JpaContext jpaContext;
+
+	@BeforeClass
+	public static void bootstrapJpa() {
+
+		firstEmf = createEntityManagerFactory("spring-data-jpa");
+		secondEmf = createEntityManagerFactory("querydsl");
+	}
+
+	@Before
+	public void createEntityManagers() {
+
+		this.firstEm = firstEmf.createEntityManager();
+		this.secondEm = secondEmf.createEntityManager();
+
+		this.jpaContext = new DefaultJpaContext(new HashSet<EntityManager>(Arrays.asList(firstEm, secondEm)));
+	}
+
+	/**
+	 * @see DATAJPA-669
+	 */
+	@Test
+	public void rejectsUnmanagedType() {
+
+		exception.expect(IllegalArgumentException.class);
+		exception.expectMessage(Object.class.getSimpleName());
+
+		jpaContext.getEntityManagerByManagedType(Object.class);
+	}
+
+	/**
+	 * @see DATAJPA-669
+	 */
+	@Test
+	public void returnsEntitymanagerForUniqueType() {
+		assertThat(jpaContext.getEntityManagerByManagedType(Category.class), is(firstEm));
+	}
+
+	/**
+	 * @see DATAJPA-669
+	 */
+	@Test
+	public void rejectsRequestForTypeManagedByMultipleEntityManagers() {
+
+		exception.expect(IllegalArgumentException.class);
+		exception.expectMessage(User.class.getSimpleName());
+
+		jpaContext.getEntityManagerByManagedType(User.class);
+	}
+
+	private static final EntityManagerFactory createEntityManagerFactory(String persistenceUnitName) {
+
+		LocalContainerEntityManagerFactoryBean factoryBean = new LocalContainerEntityManagerFactoryBean();
+		factoryBean.setPersistenceProvider(new HibernatePersistence());
+		factoryBean.setDataSource(new EmbeddedDatabaseBuilder().setType(EmbeddedDatabaseType.HSQL).build());
+		factoryBean.setPersistenceUnitName(persistenceUnitName);
+		factoryBean.afterPropertiesSet();
+
+		return factoryBean.getObject();
+	}
+}

src/test/java/org/springframework/data/jpa/repository/support/DefaultJpaContextUnitTests.java

@@ -0,0 +1,48 @@
+/*
+ * Copyright 2015 the original author or 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
+ *
+ *      http://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.springframework.data.jpa.repository.support;
+
+import java.util.Collections;
+
+import javax.persistence.EntityManager;
+
+import org.junit.Test;
+
+/**
+ * Unit tests for {@link DefaultJpaContext}.
+ * 
+ * @author Oliver Gierke
+ * @soundtrack Marcus Miller - B's River (Afrodeezia)
+ * @since 1.9
+ */
+public class DefaultJpaContextUnitTests {
+
+	/**
+	 * @see DATAJPA-669
+	 */
+	@Test(expected = IllegalArgumentException.class)
+	public void rejectsNullEntityManagers() {
+		new DefaultJpaContext(null);
+	}
+
+	/**
+	 * @see DATAJPA-669
+	 */
+	@Test(expected = IllegalArgumentException.class)
+	public void rejectsEmptyEntityManagers() {
+		new DefaultJpaContext(Collections.<EntityManager> emptySet());
+	}
+}

src/test/resources/application-context.xml

@@ -12,7 +12,11 @@
 	<bean id="userDao" class="org.springframework.data.jpa.repository.support.JpaRepositoryFactoryBean">
 		<property name="repositoryInterface" value="org.springframework.data.jpa.repository.sample.UserRepository" />
 		<property name="customImplementation">
-			<bean class="org.springframework.data.jpa.repository.sample.UserRepositoryImpl" />
+			<bean class="org.springframework.data.jpa.repository.sample.UserRepositoryImpl">
+				<constructor-arg>
+					<bean class="org.springframework.data.jpa.repository.support.DefaultJpaContext" autowire="constructor" />
+				</constructor-arg>
+			</bean>
 		</property>
 		<property name="namedQueries">
 			<bean class="org.springframework.data.repository.core.support.PropertiesBasedNamedQueries">
@@ -40,5 +44,7 @@
 	<bean class="org.springframework.beans.factory.annotation.RequiredAnnotationBeanPostProcessor" />
 	
 	<bean id="expressionEvaluationContextProvider" class="org.springframework.data.repository.query.ExtensionAwareEvaluationContextProvider" />
+	
+	<bean class="org.springframework.data.jpa.repository.support.EntityManagerBeanDefinitionRegistrarPostProcessor" />
 
 </beans>