more changes

Arpit-Bandejiya · Nov 18, 2022 · 2a2c038 · 2a2c038
1 parent 85eb684
commit 2a2c038
Show file tree

Hide file tree

Showing 2 changed files with 70 additions and 23 deletions.
diff --git a/client/rest-high-level/src/main/java/org/opensearch/client/RestHighLevelClient.java b/client/rest-high-level/src/main/java/org/opensearch/client/RestHighLevelClient.java
@@ -57,18 +57,7 @@
 import org.opensearch.action.get.MultiGetResponse;
 import org.opensearch.action.index.IndexRequest;
 import org.opensearch.action.index.IndexResponse;
-import org.opensearch.action.search.ClearScrollRequest;
-import org.opensearch.action.search.ClearScrollResponse;
-import org.opensearch.action.search.CreatePitRequest;
-import org.opensearch.action.search.CreatePitResponse;
-import org.opensearch.action.search.DeletePitRequest;
-import org.opensearch.action.search.DeletePitResponse;
-import org.opensearch.action.search.GetAllPitNodesResponse;
-import org.opensearch.action.search.MultiSearchRequest;
-import org.opensearch.action.search.MultiSearchResponse;
-import org.opensearch.action.search.SearchRequest;
-import org.opensearch.action.search.SearchResponse;
-import org.opensearch.action.search.SearchScrollRequest;
+import org.opensearch.action.search.*;
 import org.opensearch.action.support.master.AcknowledgedResponse;
 import org.opensearch.action.update.UpdateRequest;
 import org.opensearch.action.update.UpdateResponse;
@@ -1403,12 +1392,12 @@ public final Cancellable getAllPitsAsync(RequestOptions options, ActionListener<
         );
     }
 
-    public final CreatePitResponse updatePit(CreatePitRequest createPitRequest, RequestOptions options) throws IOException {
+    public final UpdatePitResponse updatePit(UpdatePitRequest updatePitRequest, RequestOptions options) throws IOException {
         return performRequestAndParseEntity(
-            createPitRequest,
-            RequestConverters::createPit,
+            updatePitRequest,
+            RequestConverters::updatePit,
             options,
-            CreatePitResponse::fromXContent,
+            UpdatePitResponse::fromXContent,
             emptySet()
         );
     }

diff --git a/server/src/main/java/org/opensearch/action/search/UpdatePitResponse.java b/server/src/main/java/org/opensearch/action/search/UpdatePitResponse.java
@@ -16,6 +16,7 @@
 import org.opensearch.common.xcontent.XContentBuilder;
 import org.opensearch.common.xcontent.XContentParser;
 import org.opensearch.rest.RestStatus;
+import org.opensearch.rest.action.RestActions;
 
 import java.io.IOException;
 
@@ -29,29 +30,52 @@ public class UpdatePitResponse extends ActionResponse implements StatusToXConten
     private static final ParseField SUCCESS_FULL = new ParseField("successfull");
 
     private final String id;
+    private final int totalShards;
+    private final int successfulShards;
+    private final int failedShards;
+    private final int skippedShards;
+    private final ShardSearchFailure[] shardFailures;
+
     private final long creationTime;
     private final long expiryTime;
     public UpdatePitResponse(StreamInput in) throws IOException {
         super(in);
         id = in.readString();
+        totalShards = in.readVInt();
+        successfulShards = in.readVInt();
+        failedShards = in.readVInt();
+        skippedShards = in.readVInt();
         creationTime = in.readLong();
         expiryTime = in.readLong();
-
+        int size = in.readVInt();
+        if(size==0){
+            shardFailures = ShardSearchFailure.EMPTY_ARRAY;
+        } else {
+            shardFailures = new ShardSearchFailure[size];
+            for (int i=0; i < shardFailures.length; i++){
+                shardFailures[i] = ShardSearchFailure.readShardSearchFailure(in);
+            }
+        }
     }
 
     public UpdatePitResponse(
         String id,
         long creationTime,
-        long expiryTime
+        long expiryTime,
+        int totalShards,
+        int successfulShards,
+        int skippedShards,
+        int failedShards,
+        ShardSearchFailure[] shardFailures
     ) {
         this.id = id;
         this.creationTime = creationTime;
         this.expiryTime = expiryTime;
-    }
-
-    @Override
-    public RestStatus status(){
-        return OK;
+        this.totalShards = totalShards;
+        this.successfulShards = successfulShards;
+        this.skippedShards = skippedShards;
+        this.failedShards = failedShards;
+        this.shardFailures = shardFailures;
     }
 
     @Override
@@ -65,8 +89,42 @@ public void writeTo(StreamOutput out) throws IOException{
     public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
         builder.startObject();
         builder.field(ID.getPreferredName(), id);
+        RestActions.buildBroadcastShardsHeader(
+            builder,
+            params,
+            getTotalShards(),
+            getSuccessfulShards(),
+            getSkippedShards(),
+            getFailedShards(),
+            getShardFailures()
+        );
         builder.endObject();
         return builder;
     }
+
+    public int getTotalShards(){
+        return totalShards;
+    }
+
+    public int getSuccessfulShards(){
+        return successfulShards;
+    }
+
+    public int getFailedShards(){
+        return shardFailures.length;
+    }
+
+    public ShardSearchFailure[] getShardFailures() {
+        return this.shardFailures;
+    }
+
+    public int getSkippedShards(){
+        return skippedShards;
+    }
+
+    @Override
+    public RestStatus status() {
+        return RestStatus.status(successfulShards, totalShards, shardFailures);
+    }
 }
 //