Skip to content

Use GetMappingsRequest instead of GetIndexRequest in QueryPlanningTool #4785

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
owaiskazi19 wants to merge 2 commits into
base: main
Choose a base branch
from
Open

Use GetMappingsRequest instead of GetIndexRequest in QueryPlanningTool #4785

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
702c3c0
Use GetMappingsRequest instead of GetIndexRequest in QueryPlanningTool
owaiskazi19 Apr 7, 2026
00b97d2
Merge branch 'main' into fix/qpt-use-get-mappings
owaiskazi19 Apr 10, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 6 additions & 6 deletions ml-algorithms/src/main/java/org/opensearch/ml/engine/tools/QueryPlanningTool.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,8 +28,8 @@

import org.apache.commons.text.StringSubstitutor;
import org.opensearch.action.admin.cluster.storedscripts.GetStoredScriptRequest;
import org.opensearch.action.admin.indices.get.GetIndexRequest;
import org.opensearch.action.admin.indices.get.GetIndexResponse;
import org.opensearch.action.admin.indices.mapping.get.GetMappingsRequest;
import org.opensearch.action.admin.indices.mapping.get.GetMappingsResponse;
import org.opensearch.action.search.SearchRequest;
import org.opensearch.action.search.SearchResponse;
import org.opensearch.action.support.IndicesOptions;
Expand Down Expand Up @@ -363,20 +363,20 @@ public void onFailure(Exception e) {

private void getIndexMappingAsync(String indexName, ActionListener<String> listener) {
try {
GetIndexRequest getIndexRequest = new GetIndexRequest()
GetMappingsRequest getMappingsRequest = new GetMappingsRequest()
.indices(indexName)
.indicesOptions(IndicesOptions.strictExpand())
.local(false)
.clusterManagerNodeTimeout(DEFAULT_CLUSTER_MANAGER_NODE_TIMEOUT);

client.admin().indices().getIndex(getIndexRequest, new ActionListener<GetIndexResponse>() {
client.admin().indices().getMappings(getMappingsRequest, new ActionListener<GetMappingsResponse>() {
@Override
public void onResponse(GetIndexResponse getIndexResponse) {
public void onResponse(GetMappingsResponse getMappingsResponse) {
try {
// When indexName is a wildcard pattern or alias, the response keys are
// concrete index names, not the original pattern/alias. Pick the first
// mapping since indices matching a pattern/alias generally share the same mapping.
Map<String, MappingMetadata> mappings = getIndexResponse.mappings();
Map<String, MappingMetadata> mappings = getMappingsResponse.mappings();
if (mappings == null || mappings.isEmpty()) {
listener
.onFailure(
Expand Down
47 changes: 23 additions & 24 deletions ml-algorithms/src/test/java/org/opensearch/ml/engine/tools/QueryPlanningToolTests.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -57,7 +57,7 @@
import org.mockito.Mock;
import org.mockito.MockitoAnnotations;
import org.opensearch.action.admin.cluster.storedscripts.GetStoredScriptResponse;
import org.opensearch.action.admin.indices.get.GetIndexResponse;
import org.opensearch.action.admin.indices.mapping.get.GetMappingsResponse;
import org.opensearch.action.search.SearchResponse;
import org.opensearch.cluster.metadata.MappingMetadata;
import org.opensearch.common.xcontent.XContentHelper;
Expand Down Expand Up @@ -105,8 +105,8 @@ public class QueryPlanningToolTests {
private QueryPlanningTool.Factory factory;

// Common test objects
private ArgumentCaptor<ActionListener<GetIndexResponse>> actionListenerCaptor;
private GetIndexResponse getIndexResponse;
private ArgumentCaptor<ActionListener<GetMappingsResponse>> actionListenerCaptor;
private GetMappingsResponse getIndexResponse;
private MappingMetadata mapping;
private String mockedSearchResponseString;

Expand Down Expand Up @@ -157,13 +157,12 @@ public void mockSampleDoc() throws IOException {

public void mockGetIndexMapping() {
// Mock the getIndex operation for getIndexMappingAsync (following IndexMappingToolTests pattern)
ArgumentCaptor<ActionListener<GetIndexResponse>> captor = ArgumentCaptor.forClass(ActionListener.class);
ArgumentCaptor<ActionListener<GetMappingsResponse>> captor = ArgumentCaptor.forClass(ActionListener.class);
this.actionListenerCaptor = captor;
doNothing().when(indicesAdminClient).getIndex(any(), actionListenerCaptor.capture());
doNothing().when(indicesAdminClient).getMappings(any(), actionListenerCaptor.capture());

// Create a real GetIndexResponse with real MappingMetadata
this.getIndexResponse = mock(GetIndexResponse.class);
when(getIndexResponse.indices()).thenReturn(new String[] { "testIndex" });
// Create a real GetMappingsResponse with real MappingMetadata
this.getIndexResponse = mock(GetMappingsResponse.class);

// Create real MappingMetadata with actual source
String mappingSource = "{\"properties\":{\"title\":{\"type\":\"text\"}}}";
Expand Down Expand Up @@ -887,11 +886,11 @@ public void testCreateWithExtraFieldsInSearchTemplates() {
public void testRunWithNonExistentIndex() throws InterruptedException {

doAnswer(invocation -> {
org.opensearch.core.action.ActionListener<org.opensearch.action.admin.indices.get.GetIndexResponse> listener = invocation
.getArgument(1);
org.opensearch.core.action.ActionListener<org.opensearch.action.admin.indices.mapping.get.GetMappingsResponse> listener =
invocation.getArgument(1);
listener.onFailure(new org.opensearch.index.IndexNotFoundException("non_existent_index"));
return null;
}).when(indicesAdminClient).getIndex(any(), any());
}).when(indicesAdminClient).getMappings(any(), any());

QueryPlanningTool tool = new QueryPlanningTool("llmGenerated", queryGenerationTool, client, null, null);
final CompletableFuture<String> future = new CompletableFuture<>();
Expand Down Expand Up @@ -1083,19 +1082,19 @@ public void testGetIndexMapping_ErrorPaths() throws ExecutionException, Interrup
}).when(queryGenerationTool).run(any(), any());

// Case 1: onResponse mapping null -> NPE inside try caught
ArgumentCaptor<ActionListener<GetIndexResponse>> captor1 = ArgumentCaptor.forClass(ActionListener.class);
ArgumentCaptor<ActionListener<GetMappingsResponse>> captor1 = ArgumentCaptor.forClass(ActionListener.class);
doAnswer(invocation -> {
// capture and call onResponse with a response missing mapping
ActionListener<GetIndexResponse> al = invocation.getArgument(1);
ActionListener<GetMappingsResponse> al = invocation.getArgument(1);
if (actionListenerCaptor == null) {
actionListenerCaptor = captor1;
}
al.onResponse(getIndexResponse); // getIndexResponse will be mocked below
return null;
}).when(indicesAdminClient).getIndex(any(), any());
}).when(indicesAdminClient).getMappings(any(), any());

// Mock getIndexResponse to return no mapping for index
getIndexResponse = mock(GetIndexResponse.class);
getIndexResponse = mock(GetMappingsResponse.class);
when(getIndexResponse.mappings()).thenReturn(Collections.emptyMap());

CompletableFuture<String> f1 = new CompletableFuture<>();
Expand All @@ -1110,10 +1109,10 @@ public void testGetIndexMapping_ErrorPaths() throws ExecutionException, Interrup

// Case 2: onFailure path with generic exception (non IndexNotFoundException)
doAnswer(invocation -> {
ActionListener<GetIndexResponse> al = invocation.getArgument(1);
ActionListener<GetMappingsResponse> al = invocation.getArgument(1);
al.onFailure(new RuntimeException("boom"));
return null;
}).when(indicesAdminClient).getIndex(any(), any());
}).when(indicesAdminClient).getMappings(any(), any());

CompletableFuture<String> f2 = new CompletableFuture<>();
ActionListener<String> l2 = ActionListener.wrap(f2::complete, f2::completeExceptionally);
Expand All @@ -1126,7 +1125,7 @@ public void testGetIndexMapping_ErrorPaths() throws ExecutionException, Interrup
assertTrue(ix2.getCause().getMessage().contains("Failed to extract index mapping"));

// Case 3: outer try-catch (getIndex throws synchronously)
doAnswer(invocation -> { throw new RuntimeException("sync"); }).when(indicesAdminClient).getIndex(any(), any());
doAnswer(invocation -> { throw new RuntimeException("sync"); }).when(indicesAdminClient).getMappings(any(), any());

CompletableFuture<String> f3 = new CompletableFuture<>();
ActionListener<String> l3 = ActionListener.wrap(f3::complete, f3::completeExceptionally);
Expand Down Expand Up @@ -1511,11 +1510,11 @@ public void testRunWithAliasResolvingToMultipleIndices() throws ExecutionExcepti
mockSampleDoc();

// Mock getIndex to return mappings keyed by multiple concrete indices (simulating alias resolution)
ArgumentCaptor<ActionListener<GetIndexResponse>> captor = ArgumentCaptor.forClass(ActionListener.class);
doNothing().when(indicesAdminClient).getIndex(any(), captor.capture());
ArgumentCaptor<ActionListener<GetMappingsResponse>> captor = ArgumentCaptor.forClass(ActionListener.class);
doNothing().when(indicesAdminClient).getMappings(any(), captor.capture());
this.actionListenerCaptor = captor;

GetIndexResponse multiIndexResponse = mock(GetIndexResponse.class);
GetMappingsResponse multiIndexResponse = mock(GetMappingsResponse.class);
String mappingSource = "{\"properties\":{\"title\":{\"type\":\"text\"}}}";
MappingMetadata mapping1 = new MappingMetadata(
"logs_march",
Expand Down Expand Up @@ -1557,10 +1556,10 @@ public void testRunWithAliasResolvingToMultipleIndices() throws ExecutionExcepti

@Test
public void testGetIndexMapping_NullMappingsFails() throws ExecutionException, InterruptedException {
ArgumentCaptor<ActionListener<GetIndexResponse>> captor = ArgumentCaptor.forClass(ActionListener.class);
doNothing().when(indicesAdminClient).getIndex(any(), captor.capture());
ArgumentCaptor<ActionListener<GetMappingsResponse>> captor = ArgumentCaptor.forClass(ActionListener.class);
doNothing().when(indicesAdminClient).getMappings(any(), captor.capture());

GetIndexResponse response = mock(GetIndexResponse.class);
GetMappingsResponse response = mock(GetMappingsResponse.class);
when(response.mappings()).thenReturn(null);

QueryPlanningTool tool = new QueryPlanningTool(LLM_GENERATED_TYPE_FIELD, queryGenerationTool, client, null, null);
Expand Down
Loading
⚡