From 86cb76a0ce8ecb1532cde94a0164f8fa41dbc96b Mon Sep 17 00:00:00 2001
From: Boaz Leskes <b.leskes@gmail.com>
Date: Wed, 28 Aug 2013 20:21:22 +0200
Subject: [PATCH] Added a status method CountResponse that resolve shard
 failures into rest status code. That method is now used in RestCountAction to
 return proper status.

Closes #3585
---
 .../action/count/CountResponse.java           | 24 +++++++++++++++++--
 .../rest/action/count/RestCountAction.java    |  3 +--
 2 files changed, 23 insertions(+), 4 deletions(-)

diff --git a/src/main/java/org/elasticsearch/action/count/CountResponse.java b/src/main/java/org/elasticsearch/action/count/CountResponse.java
index 659c9625e2a..dfcde158231 100644
--- a/src/main/java/org/elasticsearch/action/count/CountResponse.java
+++ b/src/main/java/org/elasticsearch/action/count/CountResponse.java
@@ -23,14 +23,13 @@ import org.elasticsearch.action.ShardOperationFailedException;
 import org.elasticsearch.action.support.broadcast.BroadcastOperationResponse;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
+import org.elasticsearch.rest.RestStatus;
 
 import java.io.IOException;
 import java.util.List;
 
 /**
  * The response of the count action.
- *
- *
  */
 public class CountResponse extends BroadcastOperationResponse {
 
@@ -52,6 +51,27 @@ public class CountResponse extends BroadcastOperationResponse {
         return count;
     }
 
+    public RestStatus status() {
+        if (getFailedShards() == 0) {
+            if (getSuccessfulShards() == 0 && getTotalShards() > 0) {
+                return RestStatus.SERVICE_UNAVAILABLE;
+            }
+            return RestStatus.OK;
+        }
+        // if total failure, bubble up the status code to the response level
+        if (getSuccessfulShards() == 0 && getTotalShards() > 0) {
+            RestStatus status = RestStatus.OK;
+            for (ShardOperationFailedException shardFailure : getShardFailures()) {
+                RestStatus shardStatus = shardFailure.status();
+                if (shardStatus.getStatus() >= status.getStatus()) {
+                    status = shardStatus;
+                }
+            }
+            return status;
+        }
+        return RestStatus.OK;
+    }
+
     @Override
     public void readFrom(StreamInput in) throws IOException {
         super.readFrom(in);
diff --git a/src/main/java/org/elasticsearch/rest/action/count/RestCountAction.java b/src/main/java/org/elasticsearch/rest/action/count/RestCountAction.java
index cbb6d5e7156..ee101a2b21a 100644
--- a/src/main/java/org/elasticsearch/rest/action/count/RestCountAction.java
+++ b/src/main/java/org/elasticsearch/rest/action/count/RestCountAction.java
@@ -39,7 +39,6 @@ import static org.elasticsearch.action.count.CountRequest.DEFAULT_MIN_SCORE;
 import static org.elasticsearch.rest.RestRequest.Method.GET;
 import static org.elasticsearch.rest.RestRequest.Method.POST;
 import static org.elasticsearch.rest.RestStatus.BAD_REQUEST;
-import static org.elasticsearch.rest.RestStatus.OK;
 import static org.elasticsearch.rest.action.support.RestActions.buildBroadcastShardsHeader;
 import static org.elasticsearch.rest.action.support.RestActions.splitTypes;
 
@@ -111,7 +110,7 @@ public class RestCountAction extends BaseRestHandler {
                     buildBroadcastShardsHeader(builder, response);
 
                     builder.endObject();
-                    channel.sendResponse(new XContentRestResponse(request, OK, builder));
+                    channel.sendResponse(new XContentRestResponse(request, response.status(), builder));
                 } catch (Throwable e) {
                     onFailure(e);
                 }