Allow multiple @EntityScan annotations to be used

Update EntityScanRegistrar so that multiple @EntityScan annotations can
be used with a single application. Previously, when an application used
multiple annotations only the first one found would get applied. This
changes alters that to augment the packages that will be scanned.

Fixes gh-2757

spring-boot/src/main/java/org/springframework/boot/orm/jpa/EntityScanRegistrar.java

@@ -16,8 +16,8 @@
 
 package org.springframework.boot.orm.jpa;
 
-import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Collections;
 import java.util.LinkedHashSet;
 import java.util.Set;
 
@@ -25,6 +25,7 @@ import org.springframework.beans.BeansException;
 import org.springframework.beans.factory.SmartInitializingSingleton;
 import org.springframework.beans.factory.config.BeanDefinition;
 import org.springframework.beans.factory.config.BeanPostProcessor;
+import org.springframework.beans.factory.config.ConstructorArgumentValues.ValueHolder;
 import org.springframework.beans.factory.support.BeanDefinitionRegistry;
 import org.springframework.beans.factory.support.GenericBeanDefinition;
 import org.springframework.context.annotation.ImportBeanDefinitionRegistrar;
@@ -40,6 +41,7 @@ import org.springframework.util.ObjectUtils;
  * {@link ImportBeanDefinitionRegistrar} used by {@link EntityScan}.
  *
  * @author Phillip Webb
+ * @author Oliver Gierke
  */
 class EntityScanRegistrar implements ImportBeanDefinitionRegistrar {
 
@@ -48,31 +50,25 @@ class EntityScanRegistrar implements ImportBeanDefinitionRegistrar {
 	@Override
 	public void registerBeanDefinitions(AnnotationMetadata importingClassMetadata,
 			BeanDefinitionRegistry registry) {
+		Set<String> packagesToScan = getPackagesToScan(importingClassMetadata);
 		if (!registry.containsBeanDefinition(BEAN_NAME)) {
-			GenericBeanDefinition beanDefinition = new GenericBeanDefinition();
-			beanDefinition.setBeanClass(EntityScanBeanPostProcessor.class);
-			beanDefinition.getConstructorArgumentValues().addGenericArgumentValue(
-					getPackagesToScan(importingClassMetadata));
-			beanDefinition.setRole(BeanDefinition.ROLE_INFRASTRUCTURE);
-			// We don't need this one to be post processed otherwise it can cause a
-			// cascade of bean instantiation that we would rather avoid.
-			beanDefinition.setSynthetic(true);
-			registry.registerBeanDefinition(BEAN_NAME, beanDefinition);
+			addEntityScanBeanPostProcessor(registry, packagesToScan);
+		}
+		else {
+			updateEntityScanBeanPostProcessor(registry, packagesToScan);
 		}
 	}
 
-	private String[] getPackagesToScan(AnnotationMetadata metadata) {
+	private Set<String> getPackagesToScan(AnnotationMetadata metadata) {
 		AnnotationAttributes attributes = AnnotationAttributes.fromMap(metadata
 				.getAnnotationAttributes(EntityScan.class.getName()));
 		String[] value = attributes.getStringArray("value");
 		String[] basePackages = attributes.getStringArray("basePackages");
 		Class<?>[] basePackageClasses = attributes.getClassArray("basePackageClasses");
-
 		if (!ObjectUtils.isEmpty(value)) {
 			Assert.state(ObjectUtils.isEmpty(basePackages),
 					"@EntityScan basePackages and value attributes are mutually exclusive");
 		}
-
 		Set<String> packagesToScan = new LinkedHashSet<String>();
 		packagesToScan.addAll(Arrays.asList(value));
 		packagesToScan.addAll(Arrays.asList(basePackages));
@@ -80,10 +76,38 @@ class EntityScanRegistrar implements ImportBeanDefinitionRegistrar {
 			packagesToScan.add(ClassUtils.getPackageName(basePackageClass));
 		}
 		if (packagesToScan.isEmpty()) {
-			return new String[] { ClassUtils.getPackageName(metadata.getClassName()) };
+			return Collections.singleton(ClassUtils.getPackageName(metadata
+					.getClassName()));
 		}
-		return new ArrayList<String>(packagesToScan).toArray(new String[packagesToScan
-				.size()]);
+		return packagesToScan;
+	}
+
+	private void addEntityScanBeanPostProcessor(BeanDefinitionRegistry registry,
+			Set<String> packagesToScan) {
+		GenericBeanDefinition beanDefinition = new GenericBeanDefinition();
+		beanDefinition.setBeanClass(EntityScanBeanPostProcessor.class);
+		beanDefinition.getConstructorArgumentValues().addGenericArgumentValue(
+				toArray(packagesToScan));
+		beanDefinition.setRole(BeanDefinition.ROLE_INFRASTRUCTURE);
+		// We don't need this one to be post processed otherwise it can cause a
+		// cascade of bean instantiation that we would rather avoid.
+		beanDefinition.setSynthetic(true);
+		registry.registerBeanDefinition(BEAN_NAME, beanDefinition);
+	}
+
+	private void updateEntityScanBeanPostProcessor(BeanDefinitionRegistry registry,
+			Set<String> packagesToScan) {
+		BeanDefinition definition = registry.getBeanDefinition(BEAN_NAME);
+		ValueHolder constructorArguments = definition.getConstructorArgumentValues()
+				.getGenericArgumentValue(String[].class);
+		Set<String> mergedPackages = new LinkedHashSet<String>();
+		mergedPackages.addAll(Arrays.asList((String[]) constructorArguments.getValue()));
+		mergedPackages.addAll(packagesToScan);
+		constructorArguments.setValue(toArray(mergedPackages));
+	}
+
+	private String[] toArray(Set<String> set) {
+		return set.toArray(new String[set.size()]);
 	}
 
 	/**

spring-boot/src/test/java/org/springframework/boot/orm/jpa/EntityScanTests.java

@@ -109,6 +109,13 @@ public class EntityScanTests {
 		assertSetPackagesToScan("com.mycorp.entity");
 	}
 
+	@Test
+	public void considersMultipleEntityScanAnnotations() {
+		this.context = new AnnotationConfigApplicationContext(MultiScanFirst.class,
+				MultiScanSecond.class);
+		assertSetPackagesToScan("foo", "bar");
+	}
+
 	private void assertSetPackagesToScan(String... expected) {
 		String[] actual = this.context.getBean(
 				TestLocalContainerEntityManagerFactoryBean.class).getPackagesToScan();
@@ -185,6 +192,16 @@ public class EntityScanTests {
 		}
 	}
 
+	@EntityScan(basePackages = "foo")
+	static class MultiScanFirst extends BaseConfig {
+
+	}
+
+	@EntityScan(basePackages = "bar")
+	static class MultiScanSecond extends BaseConfig {
+
+	}
+
 	private static class TestLocalContainerEntityManagerFactoryBean extends
 			LocalContainerEntityManagerFactoryBean {