From 1860d7353a1060115300740957917f1ac21cf190 Mon Sep 17 00:00:00 2001
From: Peter-Josef Meisch <pj.meisch@sothawo.com>
Date: Sun, 12 Feb 2023 17:50:55 +0100
Subject: [PATCH] Allow for Spring Data Elasticsearch queries to be added to
 NativeQuery.

Original Pull Request #2451
#Closes 2391
---
 .../elasticsearch/client/elc/NativeQuery.java |  21 +++
 .../client/elc/NativeQueryBuilder.java        |  20 +++
 .../client/elc/RequestConverter.java          |   2 +
 .../query/NativeQueryELCIntegrationTests.java |  39 +++++
 .../query/NativeQueryIntegrationTests.java    | 137 ++++++++++++++++++
 5 files changed, 219 insertions(+)
 create mode 100644 src/test/java/org/springframework/data/elasticsearch/core/query/NativeQueryELCIntegrationTests.java
 create mode 100644 src/test/java/org/springframework/data/elasticsearch/core/query/NativeQueryIntegrationTests.java

diff --git a/src/main/java/org/springframework/data/elasticsearch/client/elc/NativeQuery.java b/src/main/java/org/springframework/data/elasticsearch/client/elc/NativeQuery.java
index cd1c1d7d3..0dec41fb1 100644
--- a/src/main/java/org/springframework/data/elasticsearch/client/elc/NativeQuery.java
+++ b/src/main/java/org/springframework/data/elasticsearch/client/elc/NativeQuery.java
@@ -30,6 +30,7 @@ import java.util.Map;
 import org.springframework.data.elasticsearch.core.query.BaseQuery;
 import org.springframework.data.elasticsearch.core.query.ScriptedField;
 import org.springframework.lang.Nullable;
+import org.springframework.util.Assert;
 
 /**
  * A {@link org.springframework.data.elasticsearch.core.query.Query} implementation using query builders from the new
@@ -42,6 +43,7 @@ import org.springframework.lang.Nullable;
 public class NativeQuery extends BaseQuery {
 
 	@Nullable private final Query query;
+	@Nullable private org.springframework.data.elasticsearch.core.query.Query springDataQuery;
 	@Nullable private Query filter;
 	// note: the new client does not have pipeline aggs, these are just set up as normal aggs
 	private final Map<String, Aggregation> aggregations = new LinkedHashMap<>();
@@ -62,6 +64,12 @@ public class NativeQuery extends BaseQuery {
 		this.scriptedFields = builder.getScriptedFields();
 		this.sortOptions = builder.getSortOptions();
 		this.searchExtensions = builder.getSearchExtensions();
+
+		if (builder.getSpringDataQuery() != null) {
+			Assert.isTrue(!NativeQuery.class.isAssignableFrom(builder.getSpringDataQuery().getClass()),
+					"Cannot add an NativeQuery in a NativeQuery");
+		}
+		this.springDataQuery = builder.getSpringDataQuery();
 	}
 
 	public NativeQuery(@Nullable Query query) {
@@ -107,4 +115,17 @@ public class NativeQuery extends BaseQuery {
 	public Map<String, JsonData> getSearchExtensions() {
 		return searchExtensions;
 	}
+
+	/**
+	 * @see NativeQueryBuilder#withQuery(org.springframework.data.elasticsearch.core.query.Query).
+	 * @since 5.1
+	 */
+	public void setSpringDataQuery(@Nullable org.springframework.data.elasticsearch.core.query.Query springDataQuery) {
+		this.springDataQuery = springDataQuery;
+	}
+
+	@Nullable
+	public org.springframework.data.elasticsearch.core.query.Query getSpringDataQuery() {
+		return springDataQuery;
+	}
 }
diff --git a/src/main/java/org/springframework/data/elasticsearch/client/elc/NativeQueryBuilder.java b/src/main/java/org/springframework/data/elasticsearch/client/elc/NativeQueryBuilder.java
index 534d312bf..86de08fcd 100644
--- a/src/main/java/org/springframework/data/elasticsearch/client/elc/NativeQueryBuilder.java
+++ b/src/main/java/org/springframework/data/elasticsearch/client/elc/NativeQueryBuilder.java
@@ -51,6 +51,8 @@ public class NativeQueryBuilder extends BaseQueryBuilder<NativeQuery, NativeQuer
 	private List<SortOptions> sortOptions = new ArrayList<>();
 	private Map<String, JsonData> searchExtensions = new LinkedHashMap<>();
 
+	@Nullable private org.springframework.data.elasticsearch.core.query.Query springDataQuery;
+
 	public NativeQueryBuilder() {}
 
 	@Nullable
@@ -89,6 +91,11 @@ public class NativeQueryBuilder extends BaseQueryBuilder<NativeQuery, NativeQuer
 		return this.searchExtensions;
 	}
 
+	@Nullable
+	public org.springframework.data.elasticsearch.core.query.Query getSpringDataQuery() {
+		return springDataQuery;
+	}
+
 	public NativeQueryBuilder withQuery(Query query) {
 
 		Assert.notNull(query, "query must not be null");
@@ -188,7 +195,20 @@ public class NativeQueryBuilder extends BaseQueryBuilder<NativeQuery, NativeQuer
 		return this;
 	}
 
+	/**
+	 * Allows to use a {@link org.springframework.data.elasticsearch.core.query.Query} within a NativeQuery. Cannot be
+	 * used together with {@link #withQuery(Query)} that sets an Elasticsearch query. Passing in a {@link NativeQuery}
+	 * will result in an exception when {@link #build()} is called.
+	 *
+	 * @since 5.1
+	 */
+	public NativeQueryBuilder withQuery(org.springframework.data.elasticsearch.core.query.Query query) {
+		this.springDataQuery = query;
+		return this;
+	}
+
 	public NativeQuery build() {
+		Assert.isTrue(query == null || springDataQuery == null, "Cannot have both a native query and a Spring Data query");
 		return new NativeQuery(this);
 	}
 }
diff --git a/src/main/java/org/springframework/data/elasticsearch/client/elc/RequestConverter.java b/src/main/java/org/springframework/data/elasticsearch/client/elc/RequestConverter.java
index 5f6d6933d..93c29f1f3 100644
--- a/src/main/java/org/springframework/data/elasticsearch/client/elc/RequestConverter.java
+++ b/src/main/java/org/springframework/data/elasticsearch/client/elc/RequestConverter.java
@@ -1464,6 +1464,8 @@ class RequestConverter {
 
 			if (nativeQuery.getQuery() != null) {
 				esQuery = nativeQuery.getQuery();
+			} else if (nativeQuery.getSpringDataQuery() != null) {
+				esQuery = getQuery(nativeQuery.getSpringDataQuery(), clazz);
 			}
 		} else {
 			throw new IllegalArgumentException("unhandled Query implementation " + query.getClass().getName());
diff --git a/src/test/java/org/springframework/data/elasticsearch/core/query/NativeQueryELCIntegrationTests.java b/src/test/java/org/springframework/data/elasticsearch/core/query/NativeQueryELCIntegrationTests.java
new file mode 100644
index 000000000..eecab3f89
--- /dev/null
+++ b/src/test/java/org/springframework/data/elasticsearch/core/query/NativeQueryELCIntegrationTests.java
@@ -0,0 +1,39 @@
+/*
+ * Copyright 2023 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
+ *
+ *      https://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.elasticsearch.core.query;
+
+import org.springframework.context.annotation.Bean;
+import org.springframework.context.annotation.Configuration;
+import org.springframework.context.annotation.Import;
+import org.springframework.data.elasticsearch.junit.jupiter.ElasticsearchTemplateConfiguration;
+import org.springframework.data.elasticsearch.utils.IndexNameProvider;
+import org.springframework.test.context.ContextConfiguration;
+
+/**
+ * @author Peter-Josef Meisch
+ */
+@ContextConfiguration(classes = { NativeQueryELCIntegrationTests.Config.class })
+public class NativeQueryELCIntegrationTests extends NativeQueryIntegrationTests {
+	@Configuration
+	@Import({ ElasticsearchTemplateConfiguration.class })
+	static class Config {
+		@Bean
+		IndexNameProvider indexNameProvider() {
+			return new IndexNameProvider("criteria");
+		}
+	}
+
+}
diff --git a/src/test/java/org/springframework/data/elasticsearch/core/query/NativeQueryIntegrationTests.java b/src/test/java/org/springframework/data/elasticsearch/core/query/NativeQueryIntegrationTests.java
new file mode 100644
index 000000000..cfe2acbc2
--- /dev/null
+++ b/src/test/java/org/springframework/data/elasticsearch/core/query/NativeQueryIntegrationTests.java
@@ -0,0 +1,137 @@
+/*
+ * Copyright 2023 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
+ *
+ *      https://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.elasticsearch.core.query;
+
+import static org.assertj.core.api.Assertions.*;
+
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.DisplayName;
+import org.junit.jupiter.api.Order;
+import org.junit.jupiter.api.Test;
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.data.annotation.Id;
+import org.springframework.data.elasticsearch.annotations.Document;
+import org.springframework.data.elasticsearch.annotations.Field;
+import org.springframework.data.elasticsearch.annotations.FieldType;
+import org.springframework.data.elasticsearch.client.elc.NativeQuery;
+import org.springframework.data.elasticsearch.core.ElasticsearchOperations;
+import org.springframework.data.elasticsearch.core.mapping.IndexCoordinates;
+import org.springframework.data.elasticsearch.junit.jupiter.SpringIntegrationTest;
+import org.springframework.data.elasticsearch.utils.IndexNameProvider;
+import org.springframework.lang.Nullable;
+
+/**
+ * @author Peter-Josef Meisch
+ */
+@SpringIntegrationTest
+public abstract class NativeQueryIntegrationTests {
+	@Autowired private ElasticsearchOperations operations;
+	@Autowired private IndexNameProvider indexNameProvider;
+
+	@BeforeEach
+	public void before() {
+		indexNameProvider.increment();
+		operations.indexOps(SampleEntity.class).createWithMapping();
+	}
+
+	@Test
+	@Order(java.lang.Integer.MAX_VALUE)
+	void cleanup() {
+		operations.indexOps(IndexCoordinates.of(indexNameProvider.getPrefix() + "*")).delete();
+	}
+
+	@Test // #2391
+	@DisplayName("should be able to use CriteriaQuery in a NativeQuery")
+	void shouldBeAbleToUseCriteriaQueryInANativeQuery() {
+
+		var entity = new SampleEntity();
+		entity.setId("7");
+		entity.setText("seven");
+		operations.save(entity);
+		entity = new SampleEntity();
+		entity.setId("42");
+		entity.setText("criteria");
+		operations.save(entity);
+
+		var criteriaQuery = CriteriaQuery.builder(Criteria.where("text").is("criteria")).build();
+		var nativeQuery = NativeQuery.builder().withQuery(criteriaQuery).build();
+
+		var searchHits = operations.search(nativeQuery, SampleEntity.class);
+
+		assertThat(searchHits.getTotalHits()).isEqualTo(1);
+		assertThat(searchHits.getSearchHit(0).getId()).isEqualTo(entity.getId());
+	}
+
+	@Test // #2391
+	@DisplayName("should be able to use StringQuery in a NativeQuery")
+	void shouldBeAbleToUseStringQueryInANativeQuery() {
+
+		var entity = new SampleEntity();
+		entity.setId("7");
+		entity.setText("seven");
+		operations.save(entity);
+		entity = new SampleEntity();
+		entity.setId("42");
+		entity.setText("string");
+		operations.save(entity);
+
+		var stringQuery = StringQuery.builder("""
+				{
+				    "bool": {
+				      "must": [
+				        {
+				          "match": {
+				            "text": "string"
+				          }
+				        }
+				      ]
+				    }
+				  }
+								""").build();
+		var nativeQuery = NativeQuery.builder().withQuery(stringQuery).build();
+
+		var searchHits = operations.search(nativeQuery, SampleEntity.class);
+
+		assertThat(searchHits.getTotalHits()).isEqualTo(1);
+		assertThat(searchHits.getSearchHit(0).getId()).isEqualTo(entity.getId());
+	}
+
+	@Document(indexName = "#{@indexNameProvider.indexName()}")
+	static class SampleEntity {
+
+		@Nullable
+		public String getId() {
+			return id;
+		}
+
+		public void setId(@Nullable String id) {
+			this.id = id;
+		}
+
+		public String getText() {
+			return text;
+		}
+
+		public void setText(String text) {
+			this.text = text;
+		}
+
+		@Nullable
+		@Id private String id;
+
+		@Field(type = FieldType.Text) private String text;
+	}
+}