From c04d640452efdf641e2b55349acd535ef987529b Mon Sep 17 00:00:00 2001
From: Stephane Nicoll <snicoll@vmware.com>
Date: Wed, 14 Apr 2021 14:42:32 +0200
Subject: [PATCH] Scan for ElasticSearch's @Document type

Closes gh-26028
---
 .../ElasticsearchDataConfiguration.java          |  8 ++++++--
 .../ElasticsearchDataAutoConfigurationTests.java | 16 ++++++++++++++++
 2 files changed, 22 insertions(+), 2 deletions(-)

diff --git a/spring-boot-project/spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/data/elasticsearch/ElasticsearchDataConfiguration.java b/spring-boot-project/spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/data/elasticsearch/ElasticsearchDataConfiguration.java
index d6d4adc3da..a7e4178242 100644
--- a/spring-boot-project/spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/data/elasticsearch/ElasticsearchDataConfiguration.java
+++ b/spring-boot-project/spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/data/elasticsearch/ElasticsearchDataConfiguration.java
@@ -24,8 +24,11 @@ import org.elasticsearch.client.RestHighLevelClient;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnBean;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
+import org.springframework.boot.autoconfigure.domain.EntityScanner;
+import org.springframework.context.ApplicationContext;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
+import org.springframework.data.elasticsearch.annotations.Document;
 import org.springframework.data.elasticsearch.client.reactive.ReactiveElasticsearchClient;
 import org.springframework.data.elasticsearch.core.ElasticsearchOperations;
 import org.springframework.data.elasticsearch.core.ElasticsearchRestTemplate;
@@ -61,9 +64,10 @@ abstract class ElasticsearchDataConfiguration {
 
 		@Bean
 		@ConditionalOnMissingBean
-		SimpleElasticsearchMappingContext mappingContext(
-				ElasticsearchCustomConversions elasticsearchCustomConversions) {
+		SimpleElasticsearchMappingContext mappingContext(ApplicationContext applicationContext,
+				ElasticsearchCustomConversions elasticsearchCustomConversions) throws ClassNotFoundException {
 			SimpleElasticsearchMappingContext mappingContext = new SimpleElasticsearchMappingContext();
+			mappingContext.setInitialEntitySet(new EntityScanner(applicationContext).scan(Document.class));
 			mappingContext.setSimpleTypeHolder(elasticsearchCustomConversions.getSimpleTypeHolder());
 			return mappingContext;
 		}
diff --git a/spring-boot-project/spring-boot-autoconfigure/src/test/java/org/springframework/boot/autoconfigure/data/elasticsearch/ElasticsearchDataAutoConfigurationTests.java b/spring-boot-project/spring-boot-autoconfigure/src/test/java/org/springframework/boot/autoconfigure/data/elasticsearch/ElasticsearchDataAutoConfigurationTests.java
index 66d0cc0d9d..794f856ea0 100644
--- a/spring-boot-project/spring-boot-autoconfigure/src/test/java/org/springframework/boot/autoconfigure/data/elasticsearch/ElasticsearchDataAutoConfigurationTests.java
+++ b/spring-boot-project/spring-boot-autoconfigure/src/test/java/org/springframework/boot/autoconfigure/data/elasticsearch/ElasticsearchDataAutoConfigurationTests.java
@@ -25,6 +25,8 @@ import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 
 import org.springframework.boot.autoconfigure.AutoConfigurations;
+import org.springframework.boot.autoconfigure.TestAutoConfigurationPackage;
+import org.springframework.boot.autoconfigure.data.elasticsearch.city.City;
 import org.springframework.boot.autoconfigure.elasticsearch.ElasticsearchRestClientAutoConfiguration;
 import org.springframework.boot.test.context.runner.ApplicationContextRunner;
 import org.springframework.context.annotation.Bean;
@@ -105,6 +107,14 @@ class ElasticsearchDataAutoConfigurationTests {
 						.contains("reactiveElasticsearchTemplate"));
 	}
 
+	@Test
+	void shouldFilterInitialEntityScanWithDocumentAnnotation() {
+		this.contextRunner.withUserConfiguration(EntityScanConfig.class).run((context) -> {
+			SimpleElasticsearchMappingContext mappingContext = context.getBean(SimpleElasticsearchMappingContext.class);
+			assertThat(mappingContext.hasPersistentEntityFor(City.class)).isTrue();
+		});
+	}
+
 	@Configuration(proxyBeanMethods = false)
 	static class CustomElasticsearchCustomConversions {
 
@@ -135,6 +145,12 @@ class ElasticsearchDataAutoConfigurationTests {
 
 	}
 
+	@Configuration(proxyBeanMethods = false)
+	@TestAutoConfigurationPackage(City.class)
+	static class EntityScanConfig {
+
+	}
+
 	static class MyConverter implements Converter<ElasticsearchRestTemplate, Boolean> {
 
 		@Override