From c063ae3ea5ad1a27dce94c41ac02d2893e942fb4 Mon Sep 17 00:00:00 2001
From: David Zane <davizane@amazon.com>
Date: Mon, 7 Oct 2024 16:27:23 -0700
Subject: [PATCH] Change successfulSearchShardIndices to Set<Index>

Signed-off-by: David Zane <davizane@amazon.com>
---
 CHANGELOG.md                                             | 1 +
 .../action/admin/cluster/stats/MappingVisitor.java       | 2 +-
 .../action/search/AbstractSearchAsyncAction.java         | 4 +++-
 .../opensearch/action/search/SearchRequestContext.java   | 9 +++++----
 4 files changed, 10 insertions(+), 6 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index c2d021273f271..150b35b744b4d 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -11,6 +11,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - 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))
+- Add successfulSearchShardIndices in searchRequestContext ([#15967](https://github.com/opensearch-project/OpenSearch/pull/15967)) ([#16110](https://github.com/opensearch-project/OpenSearch/pull/16110))
 - Add support for msearch API to pass search pipeline name - ([#15923](https://github.com/opensearch-project/OpenSearch/pull/15923))
 - Add _list/indices API as paginated alternate to _cat/indices ([#14718](https://github.com/opensearch-project/OpenSearch/pull/14718))
 - Add success and failure metrics for async shard fetch ([#15976](https://github.com/opensearch-project/OpenSearch/pull/15976))
diff --git a/server/src/main/java/org/opensearch/action/admin/cluster/stats/MappingVisitor.java b/server/src/main/java/org/opensearch/action/admin/cluster/stats/MappingVisitor.java
index b2a5acb235ace..9ca9e80f94f1d 100644
--- a/server/src/main/java/org/opensearch/action/admin/cluster/stats/MappingVisitor.java
+++ b/server/src/main/java/org/opensearch/action/admin/cluster/stats/MappingVisitor.java
@@ -40,7 +40,7 @@
  *
  * @opensearch.internal
  */
-final class MappingVisitor {
+public final class MappingVisitor {
 
     private MappingVisitor() {}
 
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 836083f91b995..85ea34e442c8f 100644
--- a/server/src/main/java/org/opensearch/action/search/AbstractSearchAsyncAction.java
+++ b/server/src/main/java/org/opensearch/action/search/AbstractSearchAsyncAction.java
@@ -755,7 +755,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())
+                results.getSuccessfulResults()
+                    .map(result -> result.getSearchShardTarget().getShardId().getIndex())
+                    .collect(Collectors.toSet())
             );
             onPhaseEnd(searchRequestContext);
             onRequestEnd(searchRequestContext);
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 376cf71448d5c..398896734280e 100644
--- a/server/src/main/java/org/opensearch/action/search/SearchRequestContext.java
+++ b/server/src/main/java/org/opensearch/action/search/SearchRequestContext.java
@@ -12,6 +12,7 @@
 import org.apache.logging.log4j.Logger;
 import org.apache.lucene.search.TotalHits;
 import org.opensearch.common.annotation.InternalApi;
+import org.opensearch.core.index.Index;
 import org.opensearch.core.tasks.resourcetracker.TaskResourceInfo;
 
 import java.util.ArrayList;
@@ -37,7 +38,7 @@ public class SearchRequestContext {
     private final Map<String, Long> phaseTookMap;
     private TotalHits totalHits;
     private final EnumMap<ShardStatsFieldNames, Integer> shardStats;
-    private Set<String> successfulSearchShardIndices;
+    private Set<Index> successfulSearchShardIndices;
 
     private final SearchRequest searchRequest;
     private final LinkedBlockingQueue<TaskResourceInfo> phaseResourceUsage;
@@ -144,15 +145,15 @@ public SearchRequest getRequest() {
         return searchRequest;
     }
 
-    void setSuccessfulSearchShardIndices(Set<String> successfulSearchShardIndices) {
+    void setSuccessfulSearchShardIndices(Set<Index> successfulSearchShardIndices) {
         this.successfulSearchShardIndices = successfulSearchShardIndices;
     }
 
     /**
-     * @return A {@link List} of {@link String} representing the names of the indices that were
+     * @return A {@link Set} of {@link Index} representing the names of the indices that were
      * successfully queried at the shard level.
      */
-    public Set<String> getSuccessfulSearchShardIndices() {
+    public Set<Index> getSuccessfulSearchShardIndices() {
         return successfulSearchShardIndices;
     }
 }