From 056a4a76e7e0dc9f09b6c3657bb081a27e456f78 Mon Sep 17 00:00:00 2001
From: Onizuka <onizuka1229@gmail.com>
Date: Wed, 19 Jan 2022 03:17:18 +0800
Subject: [PATCH] Native query with ext section for SearchPlugin.

Original Pull Request #2064
Closes #2065
---
 .../elasticsearch/core/RequestFactory.java    |  5 +++++
 .../core/query/NativeSearchQuery.java         | 19 +++++++++++++++++++
 2 files changed, 24 insertions(+)

diff --git a/src/main/java/org/springframework/data/elasticsearch/core/RequestFactory.java b/src/main/java/org/springframework/data/elasticsearch/core/RequestFactory.java
index d7f340d41..0251bfeab 100644
--- a/src/main/java/org/springframework/data/elasticsearch/core/RequestFactory.java
+++ b/src/main/java/org/springframework/data/elasticsearch/core/RequestFactory.java
@@ -111,6 +111,7 @@ import org.springframework.util.StringUtils;
  * @author Subhobrata Dey
  * @author Farid Faoudi
  * @author Peer Mueller
+ * @author Sijia Liu
  * @since 4.0
  */
 class RequestFactory {
@@ -772,6 +773,10 @@ class RequestFactory {
 		if (query.getSuggestBuilder() != null) {
 			sourceBuilder.suggest(query.getSuggestBuilder());
 		}
+
+		if (!isEmpty(query.getSearchExtBuilders())) {
+			sourceBuilder.ext(query.getSearchExtBuilders());
+		}
 	}
 
 	@SuppressWarnings("rawtypes")
diff --git a/src/main/java/org/springframework/data/elasticsearch/core/query/NativeSearchQuery.java b/src/main/java/org/springframework/data/elasticsearch/core/query/NativeSearchQuery.java
index 593fb6c5e..5104f2bd8 100644
--- a/src/main/java/org/springframework/data/elasticsearch/core/query/NativeSearchQuery.java
+++ b/src/main/java/org/springframework/data/elasticsearch/core/query/NativeSearchQuery.java
@@ -21,6 +21,7 @@ import java.util.List;
 
 import org.elasticsearch.index.query.QueryBuilder;
 import org.elasticsearch.script.mustache.SearchTemplateRequestBuilder;
+import org.elasticsearch.search.SearchExtBuilder;
 import org.elasticsearch.search.aggregations.AbstractAggregationBuilder;
 import org.elasticsearch.search.aggregations.PipelineAggregationBuilder;
 import org.elasticsearch.search.collapse.CollapseBuilder;
@@ -41,6 +42,7 @@ import org.springframework.lang.Nullable;
  * @author Jean-Baptiste Nizet
  * @author Martin Choraine
  * @author Peter-Josef Meisch
+ * @author Sijia Liu
  */
 public class NativeSearchQuery extends BaseQuery {
 
@@ -56,6 +58,7 @@ public class NativeSearchQuery extends BaseQuery {
 	@Nullable private List<IndexBoost> indicesBoost;
 	@Nullable private SearchTemplateRequestBuilder searchTemplate;
 	@Nullable private SuggestBuilder suggestBuilder;
+	@Nullable private List<SearchExtBuilder> searchExtBuilders;
 
 	public NativeSearchQuery(@Nullable QueryBuilder query) {
 
@@ -201,4 +204,20 @@ public class NativeSearchQuery extends BaseQuery {
 	public SuggestBuilder getSuggestBuilder() {
 		return suggestBuilder;
 	}
+
+	public void setSearchExtBuilders(List<SearchExtBuilder> searchExtBuilders) {
+		this.searchExtBuilders = searchExtBuilders;
+	}
+
+	public void addSearchExtBuilder(SearchExtBuilder searchExtBuilder) {
+		if (searchExtBuilders == null) {
+			searchExtBuilders = new ArrayList<>();
+		}
+		searchExtBuilders.add(searchExtBuilder);
+	}
+
+	@Nullable
+	public List<SearchExtBuilder> getSearchExtBuilders() {
+		return searchExtBuilders;
+	}
 }