From 4c89e3e2d03c56f74a15f5ddf2b716471e128e6e Mon Sep 17 00:00:00 2001
From: Peter-Josef Meisch <pj.meisch@sothawo.com>
Date: Thu, 28 Nov 2024 15:26:35 +0100
Subject: [PATCH] Add optional fetchSource flag to the SourceFilter.

Closes #3009
---
 .../client/elc/RequestConverter.java          | 14 +++++---
 .../MappingElasticsearchConverter.java        |  2 +-
 .../core/query/FetchSourceFilter.java         | 15 +++++++--
 .../core/query/FetchSourceFilterBuilder.java  |  8 ++++-
 .../core/query/SourceFilter.java              | 11 +++++++
 .../core/SourceFilterIntegrationTests.java    | 33 +++++++++++++++++++
 6 files changed, 74 insertions(+), 9 deletions(-)

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 ea42903a5..ec934452e 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
@@ -2014,9 +2014,12 @@ private VersionType retrieveVersionTypeFromPersistentEntity(@Nullable Class<?> c
 	private SourceConfig getSourceConfig(Query query) {
 
 		if (query.getSourceFilter() != null) {
-			return SourceConfig.of(s -> s //
-					.filter(sfb -> {
-						SourceFilter sourceFilter = query.getSourceFilter();
+			return SourceConfig.of(s -> {
+				SourceFilter sourceFilter = query.getSourceFilter();
+				if (sourceFilter.fetchSource() != null) {
+					s.fetch(sourceFilter.fetchSource());
+				} else {
+					s.filter(sfb -> {
 						String[] includes = sourceFilter.getIncludes();
 						String[] excludes = sourceFilter.getExcludes();
 
@@ -2029,7 +2032,10 @@ private SourceConfig getSourceConfig(Query query) {
 						}
 
 						return sfb;
-					}));
+					});
+				}
+				return s;
+			});
 		} else {
 			return null;
 		}
diff --git a/src/main/java/org/springframework/data/elasticsearch/core/convert/MappingElasticsearchConverter.java b/src/main/java/org/springframework/data/elasticsearch/core/convert/MappingElasticsearchConverter.java
index ccac971fd..0ae0e8928 100644
--- a/src/main/java/org/springframework/data/elasticsearch/core/convert/MappingElasticsearchConverter.java
+++ b/src/main/java/org/springframework/data/elasticsearch/core/convert/MappingElasticsearchConverter.java
@@ -1273,7 +1273,7 @@ private void updatePropertiesInFieldsAndSourceFilter(Query query, Class<?> domai
 							.toArray(new String[] {});
 				}
 
-				query.addSourceFilter(new FetchSourceFilter(includes, excludes));
+				query.addSourceFilter(new FetchSourceFilter(sourceFilter.fetchSource(), includes, excludes));
 			}
 		}
 	}
diff --git a/src/main/java/org/springframework/data/elasticsearch/core/query/FetchSourceFilter.java b/src/main/java/org/springframework/data/elasticsearch/core/query/FetchSourceFilter.java
index caca169c6..ced12f407 100644
--- a/src/main/java/org/springframework/data/elasticsearch/core/query/FetchSourceFilter.java
+++ b/src/main/java/org/springframework/data/elasticsearch/core/query/FetchSourceFilter.java
@@ -28,14 +28,16 @@
  */
 public class FetchSourceFilter implements SourceFilter {
 
+	@Nullable private final Boolean fetchSource;
 	@Nullable private final String[] includes;
 	@Nullable private final String[] excludes;
 
 	/**
 	 * @since 5.2
 	 */
-	public static SourceFilter of(@Nullable final String[] includes, @Nullable final String[] excludes) {
-		return new FetchSourceFilter(includes, excludes);
+	public static SourceFilter of(@Nullable Boolean fetchSource, @Nullable final String[] includes,
+			@Nullable final String[] excludes) {
+		return new FetchSourceFilter(fetchSource, includes, excludes);
 	}
 
 	/**
@@ -48,11 +50,18 @@ public static SourceFilter of(Function<FetchSourceFilterBuilder, FetchSourceFilt
 		return builderFunction.apply(new FetchSourceFilterBuilder()).build();
 	}
 
-	public FetchSourceFilter(@Nullable final String[] includes, @Nullable final String[] excludes) {
+	public FetchSourceFilter(@Nullable Boolean fetchSource, @Nullable final String[] includes,
+			@Nullable final String[] excludes) {
+		this.fetchSource = fetchSource;
 		this.includes = includes;
 		this.excludes = excludes;
 	}
 
+	@Override
+	public Boolean fetchSource() {
+		return fetchSource;
+	}
+
 	@Override
 	public String[] getIncludes() {
 		return includes;
diff --git a/src/main/java/org/springframework/data/elasticsearch/core/query/FetchSourceFilterBuilder.java b/src/main/java/org/springframework/data/elasticsearch/core/query/FetchSourceFilterBuilder.java
index 501f38004..d35d46a1f 100644
--- a/src/main/java/org/springframework/data/elasticsearch/core/query/FetchSourceFilterBuilder.java
+++ b/src/main/java/org/springframework/data/elasticsearch/core/query/FetchSourceFilterBuilder.java
@@ -25,6 +25,7 @@
  */
 public class FetchSourceFilterBuilder {
 
+	@Nullable private Boolean fetchSource;
 	@Nullable private String[] includes;
 	@Nullable private String[] excludes;
 
@@ -38,12 +39,17 @@ public FetchSourceFilterBuilder withExcludes(String... excludes) {
 		return this;
 	}
 
+	public FetchSourceFilterBuilder withFetchSource(Boolean fetchSource) {
+		this.fetchSource = fetchSource;
+		return this;
+	}
+
 	public SourceFilter build() {
 		if (includes == null)
 			includes = new String[0];
 		if (excludes == null)
 			excludes = new String[0];
 
-		return new FetchSourceFilter(includes, excludes);
+		return new FetchSourceFilter(fetchSource, includes, excludes);
 	}
 }
diff --git a/src/main/java/org/springframework/data/elasticsearch/core/query/SourceFilter.java b/src/main/java/org/springframework/data/elasticsearch/core/query/SourceFilter.java
index 286a8da80..ab01424b5 100644
--- a/src/main/java/org/springframework/data/elasticsearch/core/query/SourceFilter.java
+++ b/src/main/java/org/springframework/data/elasticsearch/core/query/SourceFilter.java
@@ -40,4 +40,15 @@ public interface SourceFilter {
 	 */
 	@Nullable
 	String[] getExcludes();
+
+	/**
+	 * Flag to set the _source parameter in a query to true or false. If this is not null, the values returned from
+	 * getIncludes() and getExcludes() are ignored
+	 *
+	 * @since 5.5
+	 */
+	@Nullable
+	default Boolean fetchSource() {
+		return null;
+	}
 }
diff --git a/src/test/java/org/springframework/data/elasticsearch/core/SourceFilterIntegrationTests.java b/src/test/java/org/springframework/data/elasticsearch/core/SourceFilterIntegrationTests.java
index bf5118f0b..321fcc419 100644
--- a/src/test/java/org/springframework/data/elasticsearch/core/SourceFilterIntegrationTests.java
+++ b/src/test/java/org/springframework/data/elasticsearch/core/SourceFilterIntegrationTests.java
@@ -30,6 +30,7 @@
 import org.springframework.data.elasticsearch.annotations.Field;
 import org.springframework.data.elasticsearch.annotations.FieldType;
 import org.springframework.data.elasticsearch.core.mapping.IndexCoordinates;
+import org.springframework.data.elasticsearch.core.query.FetchSourceFilter;
 import org.springframework.data.elasticsearch.core.query.FetchSourceFilterBuilder;
 import org.springframework.data.elasticsearch.core.query.Query;
 import org.springframework.data.elasticsearch.core.query.SourceFilter;
@@ -186,6 +187,38 @@ public String[] getExcludes() {
 		assertThat(entity.getField3()).isNull();
 	}
 
+	@Test // #3009
+	@DisplayName("should not return any fields when source is set to false")
+	void shouldNotReturnAnyFieldsWhenSourceIsSetToFalse() {
+
+		Query query = Query.findAll();
+		query.addSourceFilter(FetchSourceFilter.of(b -> b.withFetchSource(false)));
+
+		SearchHits<Entity> entities = operations.search(query, Entity.class);
+
+		assertThat(entities).hasSize(1);
+		Entity entity = entities.getSearchHit(0).getContent();
+		assertThat(entity.getField1()).isNull();
+		assertThat(entity.getField2()).isNull();
+		assertThat(entity.getField3()).isNull();
+	}
+
+	@Test // #3009
+	@DisplayName("should return all fields when source is set to true")
+	void shouldReturnAllFieldsWhenSourceIsSetToTrue() {
+
+		Query query = Query.findAll();
+        query.addSourceFilter(FetchSourceFilter.of(b -> b.withFetchSource(true)));
+
+		SearchHits<Entity> entities = operations.search(query, Entity.class);
+
+		assertThat(entities).hasSize(1);
+		Entity entity = entities.getSearchHit(0).getContent();
+		assertThat(entity.getField1()).isNotNull();
+		assertThat(entity.getField2()).isNotNull();
+		assertThat(entity.getField3()).isNotNull();
+	}
+
 	@Document(indexName = "#{@indexNameProvider.indexName()}")
 	public static class Entity {
 		@Nullable