Raise a 404 exception when document source is not found (#33384)

server/src/main/java/org/elasticsearch/rest/action/document/RestGetSourceAction.java

@@ -19,6 +19,7 @@
 
 package org.elasticsearch.rest.action.document;
 
+import org.elasticsearch.ResourceNotFoundException;
 import org.elasticsearch.action.ActionRequestValidationException;
 import org.elasticsearch.action.get.GetRequest;
 import org.elasticsearch.action.get.GetResponse;
@@ -29,6 +30,7 @@
 import org.elasticsearch.common.xcontent.XContentHelper;
 import org.elasticsearch.rest.BaseRestHandler;
 import org.elasticsearch.rest.BytesRestResponse;
+import org.elasticsearch.rest.RestChannel;
 import org.elasticsearch.rest.RestController;
 import org.elasticsearch.rest.RestRequest;
 import org.elasticsearch.rest.RestResponse;
@@ -40,7 +42,6 @@
 
 import static org.elasticsearch.rest.RestRequest.Method.GET;
 import static org.elasticsearch.rest.RestRequest.Method.HEAD;
-import static org.elasticsearch.rest.RestStatus.NOT_FOUND;
 import static org.elasticsearch.rest.RestStatus.OK;
 
 /**
@@ -75,24 +76,45 @@ public RestChannelConsumer prepareRequest(final RestRequest request, final NodeC
                 validationError.addValidationError("fetching source can not be disabled");
                 channel.sendResponse(new BytesRestResponse(channel, validationError));
             } else {
-                client.get(getRequest, new RestResponseListener<GetResponse>(channel) {
-                    @Override
-                    public RestResponse buildResponse(final GetResponse response) throws Exception {
-                        final XContentBuilder builder = channel.newBuilder(request.getXContentType(), false);
-                        // check if doc source (or doc itself) is missing
-                        if (response.isSourceEmpty()) {
-                            return new BytesRestResponse(NOT_FOUND, builder);
-                        } else {
-                            final BytesReference source = response.getSourceInternal();
-                            try (InputStream stream = source.streamInput()) {
-                                builder.rawValue(stream, XContentHelper.xContentType(source));
-                            }
-                            return new BytesRestResponse(OK, builder);
-                        }
-                    }
-                });
+                client.get(getRequest, new RestGetSourceResponseListener(channel, request));
             }
         };
     }
 
+    static class RestGetSourceResponseListener extends RestResponseListener<GetResponse> {
+        private final RestRequest request;
+
+        RestGetSourceResponseListener(RestChannel channel, RestRequest request) {
+            super(channel);
+            this.request = request;
+        }
+
+        @Override
+        public RestResponse buildResponse(final GetResponse response) throws Exception {
+            checkIfSourceEmpty(response);
+
+            final XContentBuilder builder = channel.newBuilder(request.getXContentType(), false);
+            final BytesReference source = response.getSourceInternal();
+            try (InputStream stream = source.streamInput()) {
+                builder.rawValue(stream, XContentHelper.xContentType(source));
+            }
+            return new BytesRestResponse(OK, builder);
+        }
+
+        /**
+         * Checks if the requested source or document itself is missing.
+         *
+         * @param response a response
+         * @throws ResourceNotFoundException if source or doc itself is missing
+         */
+        private void checkIfSourceEmpty(final GetResponse response) {
+            if (response.isSourceEmpty()) {
+                final String index = response.getIndex();
+                final String type = response.getType();
+                final String id = response.getId();
+
+                throw new ResourceNotFoundException("Document or source not found [" + index + "]/[" + type + "]/[" + id + "]");
+            }
+        }
+    }
 }

server/src/test/java/org/elasticsearch/rest/action/document/RestGetSourceActionTests.java

@@ -0,0 +1,78 @@
+/*
+ * Licensed to Elasticsearch under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch licenses this file to you 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
+ *
+ *    http://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.elasticsearch.rest.action.document;
+
+import org.elasticsearch.ResourceNotFoundException;
+import org.elasticsearch.action.get.GetResponse;
+import org.elasticsearch.common.bytes.BytesArray;
+import org.elasticsearch.common.bytes.BytesReference;
+import org.elasticsearch.index.get.GetResult;
+import org.elasticsearch.rest.RestRequest;
+import org.elasticsearch.rest.RestResponse;
+import org.elasticsearch.test.ESTestCase;
+import org.elasticsearch.test.rest.FakeRestChannel;
+import org.elasticsearch.test.rest.FakeRestRequest;
+import org.junit.AfterClass;
+
+import static java.util.Collections.emptyMap;
+import static org.elasticsearch.rest.RestStatus.OK;
+import static org.elasticsearch.rest.action.document.RestGetSourceAction.RestGetSourceResponseListener;
+import static org.hamcrest.Matchers.equalTo;
+
+public class RestGetSourceActionTests extends ESTestCase {
+
+    private static RestRequest request = new FakeRestRequest();
+    private static FakeRestChannel channel = new FakeRestChannel(request, true, 0);
+    private static RestGetSourceResponseListener listener = new RestGetSourceResponseListener(channel, request);
+
+    @AfterClass
+    public static void cleanupReferences() {
+        request = null;
+        channel = null;
+        listener = null;
+    }
+
+    public void testRestGetSourceAction() throws Exception {
+        // GIVEN a REST Get Source action response with an existing result and a non-null source
+        final BytesReference source = new BytesArray("{\"foo\": \"bar\"}");
+        final GetResponse getResponse = new GetResponse(new GetResult("index1", "_doc", "1", -1, true, source, emptyMap()));
+
+        // WHEN building the REST response
+        final RestResponse restResponse = listener.buildResponse(getResponse);
+
+        // THEN expect to retrieve document source
+        assertThat(restResponse.status(), equalTo(OK));
+        assertThat(restResponse.contentType(), equalTo("application/json; charset=UTF-8"));
+        assertThat(restResponse.content(), equalTo(new BytesArray("{\"foo\": \"bar\"}")));
+    }
+
+    public void testRestGetSourceAction_withNullSource() {
+        // GIVEN a REST Get Source action response with a non-existing result and a null source
+        final GetResponse getResponse = new GetResponse(new GetResult("index1", "_doc", "1", -1, false, null, emptyMap()));
+
+        // WHEN building the REST response
+        // THEN expect a resource not found exception
+        final ResourceNotFoundException exception = expectThrows(ResourceNotFoundException.class,
+            () -> listener.buildResponse(getResponse));
+
+        // THEN expect a formatted error message
+        assertThat(exception.getMessage(), equalTo("Document or source not found [index1]/[_doc]/[1]"));
+    }
+}