From 5afe8c07024b1c1db50484e5a25065216d131691 Mon Sep 17 00:00:00 2001
From: David Zane <38449481+dzane17@users.noreply.github.com>
Date: Thu, 19 Sep 2024 16:22:48 -0700
Subject: [PATCH] Add successfulSearchShardIndices in searchRequestContext
 before onRequestEnd (#15967)

Signed-off-by: David Zane <davizane@amazon.com>
Signed-off-by: Ankit Jain <akjain@amazon.com>
Co-authored-by: Ankit Jain <akjain@amazon.com>
---
 CHANGELOG.md                                       |  1 +
 .../action/search/AbstractSearchAsyncAction.java   |  3 +++
 .../action/search/SearchRequestContext.java        | 14 ++++++++++++++
 3 files changed, 18 insertions(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index f142978ddfa8a..4f8079c3c912b 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -10,6 +10,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - [Workload Management] Add QueryGroup Stats API Logic ([15777](https://github.com/opensearch-project/OpenSearch/pull/15777))
 - Fallback to Remote cluster-state on Term-Version check mismatch - ([#15424](https://github.com/opensearch-project/OpenSearch/pull/15424))
 - Implement WithFieldName interface in ValuesSourceAggregationBuilder & FieldSortBuilder ([#15916](https://github.com/opensearch-project/OpenSearch/pull/15916))
+- Add successfulSearchShardIndices in searchRequestContext ([#15967](https://github.com/opensearch-project/OpenSearch/pull/15967))
 - Remove identity-related feature flagged code from the RestController ([#15430](https://github.com/opensearch-project/OpenSearch/pull/15430))
 
 ### Dependencies
diff --git a/server/src/main/java/org/opensearch/action/search/AbstractSearchAsyncAction.java b/server/src/main/java/org/opensearch/action/search/AbstractSearchAsyncAction.java
index f0fc05c595d6f..836083f91b995 100644
--- a/server/src/main/java/org/opensearch/action/search/AbstractSearchAsyncAction.java
+++ b/server/src/main/java/org/opensearch/action/search/AbstractSearchAsyncAction.java
@@ -754,6 +754,9 @@ public void sendSearchResponse(InternalSearchResponse internalSearchResponse, At
             }
             searchRequestContext.setTotalHits(internalSearchResponse.hits().getTotalHits());
             searchRequestContext.setShardStats(results.getNumShards(), successfulOps.get(), skippedOps.get(), failures.length);
+            searchRequestContext.setSuccessfulSearchShardIndices(
+                results.getSuccessfulResults().map(result -> result.getSearchShardTarget().getIndex()).collect(Collectors.toSet())
+            );
             onPhaseEnd(searchRequestContext);
             onRequestEnd(searchRequestContext);
             listener.onResponse(buildSearchResponse(internalSearchResponse, failures, scrollId, searchContextId));
diff --git a/server/src/main/java/org/opensearch/action/search/SearchRequestContext.java b/server/src/main/java/org/opensearch/action/search/SearchRequestContext.java
index 111d9c64550b3..376cf71448d5c 100644
--- a/server/src/main/java/org/opensearch/action/search/SearchRequestContext.java
+++ b/server/src/main/java/org/opensearch/action/search/SearchRequestContext.java
@@ -20,6 +20,7 @@
 import java.util.List;
 import java.util.Locale;
 import java.util.Map;
+import java.util.Set;
 import java.util.concurrent.LinkedBlockingQueue;
 import java.util.function.Supplier;
 
@@ -36,6 +37,7 @@ public class SearchRequestContext {
     private final Map<String, Long> phaseTookMap;
     private TotalHits totalHits;
     private final EnumMap<ShardStatsFieldNames, Integer> shardStats;
+    private Set<String> successfulSearchShardIndices;
 
     private final SearchRequest searchRequest;
     private final LinkedBlockingQueue<TaskResourceInfo> phaseResourceUsage;
@@ -141,6 +143,18 @@ public List<TaskResourceInfo> getPhaseResourceUsage() {
     public SearchRequest getRequest() {
         return searchRequest;
     }
+
+    void setSuccessfulSearchShardIndices(Set<String> successfulSearchShardIndices) {
+        this.successfulSearchShardIndices = successfulSearchShardIndices;
+    }
+
+    /**
+     * @return A {@link List} of {@link String} representing the names of the indices that were
+     * successfully queried at the shard level.
+     */
+    public Set<String> getSuccessfulSearchShardIndices() {
+        return successfulSearchShardIndices;
+    }
 }
 
 enum ShardStatsFieldNames {