Skip to content

Commit fbad6a4

Browse files
committed
Run the Get Connector request in a stashed threadcontext
Signed-off-by: Craig Perkins <cwperx@amazon.com>
1 parent 9e014fa commit fbad6a4

File tree

1 file changed

+33
-30
lines changed

1 file changed

+33
-30
lines changed

plugin/src/main/java/org/opensearch/ml/model/MLModelManager.java

+33-30
Original file line numberDiff line numberDiff line change
@@ -2077,40 +2077,43 @@ public void getConnector(String connectorId, String tenantId, ActionListener<Con
20772077
.tenantId(tenantId)
20782078
.build();
20792079

2080-
sdkClient.getDataObjectAsync(getDataObjectRequest).whenComplete((r, throwable) -> {
2081-
log.debug("Completed Get Connector Request, id:{}", connectorId);
2082-
if (throwable != null) {
2083-
Exception cause = SdkClientUtils.unwrapAndConvertToException(throwable);
2084-
if (ExceptionsHelper.unwrap(cause, IndexNotFoundException.class) != null) {
2085-
log.error("Failed to get connector index", cause);
2086-
listener.onFailure(new OpenSearchStatusException("Failed to find connector", RestStatus.NOT_FOUND));
2080+
try (ThreadContext.StoredContext ctx = client.threadPool().getThreadContext().stashContext()) {
2081+
sdkClient.getDataObjectAsync(getDataObjectRequest).whenComplete((r, throwable) -> {
2082+
log.debug("Completed Get Connector Request, id:{}", connectorId);
2083+
if (throwable != null) {
2084+
Exception cause = SdkClientUtils.unwrapAndConvertToException(throwable);
2085+
if (ExceptionsHelper.unwrap(cause, IndexNotFoundException.class) != null) {
2086+
log.error("Failed to get connector index", cause);
2087+
listener.onFailure(new OpenSearchStatusException("Failed to find connector", RestStatus.NOT_FOUND));
2088+
} else {
2089+
log.error("Failed to get ML connector {}", connectorId, cause);
2090+
listener.onFailure(cause);
2091+
}
20872092
} else {
2088-
log.error("Failed to get ML connector {}", connectorId, cause);
2089-
listener.onFailure(cause);
2090-
}
2091-
} else {
2092-
try {
2093-
GetResponse gr = r.parser() == null ? null : GetResponse.fromXContent(r.parser());
2094-
if (gr != null && gr.isExists()) {
2095-
try (
2096-
XContentParser parser = MLNodeUtils
2097-
.createXContentParserFromRegistry(NamedXContentRegistry.EMPTY, gr.getSourceAsBytesRef())
2098-
) {
2099-
ensureExpectedToken(XContentParser.Token.START_OBJECT, parser.nextToken(), parser);
2100-
Connector connector = Connector.createConnector(parser);
2101-
listener.onResponse(connector);
2102-
} catch (Exception e) {
2103-
log.error("Failed to parse connector:{}", connectorId);
2104-
listener.onFailure(e);
2093+
try {
2094+
GetResponse gr = r.parser() == null ? null : GetResponse.fromXContent(r.parser());
2095+
if (gr != null && gr.isExists()) {
2096+
try (
2097+
XContentParser parser = MLNodeUtils
2098+
.createXContentParserFromRegistry(NamedXContentRegistry.EMPTY, gr.getSourceAsBytesRef())
2099+
) {
2100+
ensureExpectedToken(XContentParser.Token.START_OBJECT, parser.nextToken(), parser);
2101+
Connector connector = Connector.createConnector(parser);
2102+
listener.onResponse(connector);
2103+
} catch (Exception e) {
2104+
log.error("Failed to parse connector:{}", connectorId);
2105+
listener.onFailure(e);
2106+
}
2107+
} else {
2108+
listener
2109+
.onFailure(new OpenSearchStatusException("Failed to find connector:" + connectorId, RestStatus.NOT_FOUND));
21052110
}
2106-
} else {
2107-
listener.onFailure(new OpenSearchStatusException("Failed to find connector:" + connectorId, RestStatus.NOT_FOUND));
2111+
} catch (Exception e) {
2112+
listener.onFailure(e);
21082113
}
2109-
} catch (Exception e) {
2110-
listener.onFailure(e);
21112114
}
2112-
}
2113-
});
2115+
});
2116+
}
21142117
}
21152118

21162119
/**

0 commit comments

Comments
 (0)