Skip to content

Commit 849fecf

Browse files
dbwiddisdbwiddis
authored
Add IndexMapping Tool (opensearch-project#1891)
* Add IndexMapping Tool Signed-off-by: Daniel Widdis <widdis@gmail.com> * Improve description and make index a required parameter Signed-off-by: Daniel Widdis <widdis@gmail.com> * Add some more test coverage Signed-off-by: Daniel Widdis <widdis@gmail.com> --------- Signed-off-by: Daniel Widdis <widdis@gmail.com>
1 parent 521b880 commit 849fecf

File tree

4 files changed

+404
-8
lines changed

4 files changed

+404
-8
lines changed

‎ml-algorithms/src/main/java/org/opensearch/ml/engine/tools/CatIndexTool.java

+3-8
Original file line numberDiff line numberDiff line change
@@ -100,9 +100,7 @@ public <T> void run(Map<String, String> parameters, ActionListener<T> listener)
100100
final IndicesOptions indicesOptions = IndicesOptions.strictExpand();
101101
final boolean local = parameters.containsKey("local") ? Boolean.parseBoolean("local") : false;
102102
final TimeValue clusterManagerNodeTimeout = DEFAULT_CLUSTER_MANAGER_NODE_TIMEOUT;
103-
final boolean includeUnloadedSegments = parameters.containsKey("include_unloaded_segments")
104-
? Boolean.parseBoolean(parameters.get("include_unloaded_segments"))
105-
: false;
103+
final boolean includeUnloadedSegments = Boolean.parseBoolean(parameters.get("include_unloaded_segments"));
106104

107105
final ActionListener<Table> internalListener = ActionListener.notifyOnce(ActionListener.wrap(table -> {
108106
// Handle empty table
@@ -297,10 +295,7 @@ public void onFailure(final Exception e) {
297295

298296
@Override
299297
public boolean validate(Map<String, String> parameters) {
300-
if (parameters == null || parameters.size() == 0) {
301-
return false;
302-
}
303-
return true;
298+
return parameters != null && !parameters.isEmpty();
304299
}
305300

306301
/**
@@ -388,7 +383,7 @@ private Table buildTable(
388383
final Table table = getTableWithHeader();
389384

390385
indicesSettings.forEach((indexName, settings) -> {
391-
if (indicesMetadatas.containsKey(indexName) == false) {
386+
if (!indicesMetadatas.containsKey(indexName)) {
392387
// the index exists in the Get Indices response but is not present in the cluster state:
393388
// it is likely that the index was deleted in the meanwhile, so we ignore it.
394389
return;

‎ml-algorithms/src/main/java/org/opensearch/ml/engine/tools/IndexMappingTool.java

+208
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,208 @@
1+
/*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package org.opensearch.ml.engine.tools;
7+
8+
import static org.opensearch.action.support.clustermanager.ClusterManagerNodeRequest.DEFAULT_CLUSTER_MANAGER_NODE_TIMEOUT;
9+
import static org.opensearch.ml.common.utils.StringUtils.gson;
10+
11+
import java.util.Collections;
12+
import java.util.List;
13+
import java.util.Map;
14+
import java.util.Map.Entry;
15+
16+
import org.apache.logging.log4j.util.Strings;
17+
import org.opensearch.action.admin.indices.get.GetIndexRequest;
18+
import org.opensearch.action.admin.indices.get.GetIndexResponse;
19+
import org.opensearch.action.support.IndicesOptions;
20+
import org.opensearch.client.Client;
21+
import org.opensearch.cluster.metadata.MappingMetadata;
22+
import org.opensearch.common.settings.Settings;
23+
import org.opensearch.common.unit.TimeValue;
24+
import org.opensearch.core.action.ActionListener;
25+
import org.opensearch.ml.common.output.model.ModelTensors;
26+
import org.opensearch.ml.common.spi.tools.Parser;
27+
import org.opensearch.ml.common.spi.tools.Tool;
28+
import org.opensearch.ml.common.spi.tools.ToolAnnotation;
29+
30+
import lombok.Getter;
31+
import lombok.Setter;
32+
33+
@ToolAnnotation(IndexMappingTool.TYPE)
34+
public class IndexMappingTool implements Tool {
35+
public static final String TYPE = "IndexMappingTool";
36+
private static final String DEFAULT_DESCRIPTION = String
37+
.join(
38+
" ",
39+
"This tool gets index mapping information from a certain index.",
40+
"It takes 1 required argument named `index` which is a comma-delimited list of one or more indices to get mapping information from, which expands wildcards.",
41+
"It takes 1 optional argument named `local` which means whether to return information from the local node only instead of the cluster manager node (Default is false).",
42+
"The tool returns the index mapping information, which is about how documents and their fields are stored and indexed, and also returns the index settings."
43+
);
44+
45+
@Setter
46+
@Getter
47+
private String name = IndexMappingTool.TYPE;
48+
@Getter
49+
@Setter
50+
private String description = DEFAULT_DESCRIPTION;
51+
@Getter
52+
private String version;
53+
54+
private Client client;
55+
@Setter
56+
private Parser<?, ?> inputParser;
57+
@Setter
58+
private Parser<?, ?> outputParser;
59+
60+
public IndexMappingTool(Client client) {
61+
this.client = client;
62+
63+
outputParser = new Parser<>() {
64+
@Override
65+
public Object parse(Object o) {
66+
@SuppressWarnings("unchecked")
67+
List<ModelTensors> mlModelOutputs = (List<ModelTensors>) o;
68+
return mlModelOutputs.get(0).getMlModelTensors().get(0).getDataAsMap().get("response");
69+
}
70+
};
71+
}
72+
73+
@Override
74+
public <T> void run(Map<String, String> parameters, ActionListener<T> listener) {
75+
@SuppressWarnings("unchecked")
76+
List<String> indexList = parameters.containsKey("index")
77+
? gson.fromJson(parameters.get("index"), List.class)
78+
: Collections.emptyList();
79+
if (indexList.isEmpty()) {
80+
@SuppressWarnings("unchecked")
81+
T empty = (T) ("There were no results searching the index parameter [" + parameters.get("index") + "].");
82+
listener.onResponse(empty);
83+
return;
84+
}
85+
86+
final String[] indices = indexList.toArray(Strings.EMPTY_ARRAY);
87+
88+
final IndicesOptions indicesOptions = IndicesOptions.strictExpand();
89+
final boolean local = Boolean.parseBoolean(parameters.get("local"));
90+
final TimeValue clusterManagerNodeTimeout = DEFAULT_CLUSTER_MANAGER_NODE_TIMEOUT;
91+
92+
ActionListener<GetIndexResponse> internalListener = new ActionListener<GetIndexResponse>() {
93+
94+
@Override
95+
public void onResponse(GetIndexResponse getIndexResponse) {
96+
try {
97+
// Handle empty response
98+
if (getIndexResponse.indices().length == 0) {
99+
@SuppressWarnings("unchecked")
100+
T empty = (T) ("There were no results searching the index parameter [" + parameters.get("index") + "].");
101+
listener.onResponse(empty);
102+
return;
103+
}
104+
StringBuilder sb = new StringBuilder();
105+
for (String index : getIndexResponse.indices()) {
106+
sb.append("index: ").append(index).append("\n\n");
107+
108+
MappingMetadata mapping = getIndexResponse.mappings().get(index);
109+
if (mapping != null) {
110+
sb.append("mappings:\n");
111+
for (Entry<String, Object> entry : mapping.sourceAsMap().entrySet()) {
112+
sb.append(entry.getKey()).append("=").append(entry.getValue()).append('\n');
113+
}
114+
sb.append("\n\n");
115+
}
116+
117+
Settings settings = getIndexResponse.settings().get(index);
118+
if (settings != null) {
119+
sb.append("settings:\n").append(settings.toDelimitedString('\n')).append("\n\n");
120+
}
121+
}
122+
123+
@SuppressWarnings("unchecked")
124+
T response = (T) sb.toString();
125+
listener.onResponse(response);
126+
} catch (Exception e) {
127+
onFailure(e);
128+
}
129+
}
130+
131+
@Override
132+
public void onFailure(final Exception e) {
133+
listener.onFailure(e);
134+
}
135+
136+
};
137+
final GetIndexRequest getIndexRequest = new GetIndexRequest()
138+
.indices(indices)
139+
.indicesOptions(indicesOptions)
140+
.local(local)
141+
.clusterManagerNodeTimeout(clusterManagerNodeTimeout);
142+
143+
client.admin().indices().getIndex(getIndexRequest, internalListener);
144+
}
145+
146+
@Override
147+
public String getType() {
148+
return TYPE;
149+
}
150+
151+
@Override
152+
public boolean validate(Map<String, String> parameters) {
153+
return parameters != null && parameters.containsKey("index");
154+
}
155+
156+
/**
157+
* Factory for the {@link IndexMappingTool}
158+
*/
159+
public static class Factory implements Tool.Factory<IndexMappingTool> {
160+
private Client client;
161+
162+
private static Factory INSTANCE;
163+
164+
/**
165+
* Create or return the singleton factory instance
166+
*/
167+
public static Factory getInstance() {
168+
if (INSTANCE != null) {
169+
return INSTANCE;
170+
}
171+
synchronized (IndexMappingTool.class) {
172+
if (INSTANCE != null) {
173+
return INSTANCE;
174+
}
175+
INSTANCE = new Factory();
176+
return INSTANCE;
177+
}
178+
}
179+
180+
/**
181+
* Initialize this factory
182+
* @param client The OpenSearch client
183+
*/
184+
public void init(Client client) {
185+
this.client = client;
186+
}
187+
188+
@Override
189+
public IndexMappingTool create(Map<String, Object> map) {
190+
return new IndexMappingTool(client);
191+
}
192+
193+
@Override
194+
public String getDefaultDescription() {
195+
return DEFAULT_DESCRIPTION;
196+
}
197+
198+
@Override
199+
public String getDefaultType() {
200+
return TYPE;
201+
}
202+
203+
@Override
204+
public String getDefaultVersion() {
205+
return null;
206+
}
207+
}
208+
}

0 commit comments

Comments
 (0)