Ensure fetch fields aren't dropped when rewriting search. (#61390)

Previously we didn't retain the requested fields when performing a shallow copy
of the search source. This meant that when a search was rewritten, we could drop
the requested fields and fail to return them in the response.

Author: jtibshirani

rest-api-spec/src/main/resources/rest-api-spec/test/search/330_fetch_fields.yml

@@ -92,6 +92,7 @@ setup:
           fields:
             - field: keyword
               format: "yyyy/MM/dd"
+
 ---
 "Test disable source":
   - do:
@@ -214,3 +215,43 @@ setup:
 
   - match: { hits.hits.0.fields.keyword.0: "a" }
   - match: { hits.hits.0.fields.integer.0: 42 }
+
+---
+"Test search rewrite":
+  - do:
+      indices.create:
+        index:  test
+        body:
+          settings:
+            index.number_of_shards: 1
+          mappings:
+            properties:
+              date:
+                type: date
+
+  - do:
+      index:
+        index:  test
+        id:     1
+        body:
+          date: "1990-12-29T22:30:00.000Z"
+
+  - do:
+      indices.refresh:
+        index: [ test ]
+
+  - do:
+      search:
+        index: test
+        body:
+          query:
+            range:
+              date:
+                from: "1990-12-29T22:30:00.000Z"
+          fields:
+            - field: date
+              format: "yyyy/MM/dd"
+
+  - is_true: hits.hits.0._id
+  - is_true: hits.hits.0._source
+  - match: { hits.hits.0.fields.date.0: "1990/12/29" }

server/src/main/java/org/elasticsearch/search/builder/SearchSourceBuilder.java

@@ -1045,6 +1045,7 @@ private SearchSourceBuilder shallowCopy(QueryBuilder queryBuilder, QueryBuilder
         rewrittenBuilder.explain = explain;
         rewrittenBuilder.extBuilders = extBuilders;
         rewrittenBuilder.fetchSourceContext = fetchSourceContext;
+        rewrittenBuilder.fetchFields = fetchFields;
         rewrittenBuilder.docValueFields = docValueFields;
         rewrittenBuilder.storedFieldsContext = storedFieldsContext;
         rewrittenBuilder.from = from;
@@ -1603,10 +1604,10 @@ public boolean equals(Object obj) {
 
     @Override
     public int hashCode() {
-        return Objects.hash(aggregations, explain, fetchSourceContext, docValueFields, storedFieldsContext, from, highlightBuilder,
-                indexBoosts, minScore, postQueryBuilder, queryBuilder, rescoreBuilders, scriptFields, size,
-                sorts, searchAfterBuilder, sliceBuilder, stats, suggestBuilder, terminateAfter, timeout, trackScores, version,
-                seqNoAndPrimaryTerm, profile, extBuilders, collapse, trackTotalHitsUpTo);
+        return Objects.hash(aggregations, explain, fetchSourceContext, fetchFields, docValueFields, storedFieldsContext, from,
+            highlightBuilder, indexBoosts, minScore, postQueryBuilder, queryBuilder, rescoreBuilders, scriptFields, size,
+            sorts, searchAfterBuilder, sliceBuilder, stats, suggestBuilder, terminateAfter, timeout, trackScores, version,
+            seqNoAndPrimaryTerm, profile, extBuilders, collapse, trackTotalHitsUpTo);
     }
 
     @Override
@@ -1621,6 +1622,7 @@ public boolean equals(Object obj) {
         return Objects.equals(aggregations, other.aggregations)
                 && Objects.equals(explain, other.explain)
                 && Objects.equals(fetchSourceContext, other.fetchSourceContext)
+                && Objects.equals(fetchFields, other.fetchFields)
                 && Objects.equals(docValueFields, other.docValueFields)
                 && Objects.equals(storedFieldsContext, other.storedFieldsContext)
                 && Objects.equals(from, other.from)

server/src/test/java/org/elasticsearch/search/builder/SearchSourceBuilderTests.java

@@ -105,6 +105,14 @@ public void testSerialization() throws IOException {
         }
     }
 
+    public void testShallowCopy() {
+        for (int i = 0; i < 10; i++) {
+            SearchSourceBuilder original = createSearchSourceBuilder();
+            SearchSourceBuilder copy = original.shallowCopy();
+            assertEquals(original, copy);
+        }
+    }
+
     public void testEqualsAndHashcode() throws IOException {
         // TODO add test checking that changing any member of this class produces an object that is not equal to the original
         EqualsHashCodeTestUtils.checkEqualsAndHashCode(createSearchSourceBuilder(), this::copyBuilder);

test/framework/src/main/java/org/elasticsearch/search/RandomSearchRequestGenerator.java

@@ -182,6 +182,13 @@ public static SearchSourceBuilder randomSearchSourceBuilder(
                 throw new IllegalStateException();
         }
 
+        if (randomBoolean()) {
+            int numFields = randomInt(5);
+            for (int i = 0; i < numFields; i++) {
+                builder.fetchField(randomAlphaOfLengthBetween(5, 10));
+            }
+        }
+
         if (randomBoolean()) {
             int scriptFieldsSize = randomInt(25);
             for (int i = 0; i < scriptFieldsSize; i++) {