-
Notifications
You must be signed in to change notification settings - Fork 140
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
* expose connector action parameter Signed-off-by: Yaliang Wu <[email protected]> * add connector tool Signed-off-by: Yaliang Wu <[email protected]> * fix ut Signed-off-by: Yaliang Wu <[email protected]> * fix it Signed-off-by: Yaliang Wu <[email protected]> * fix flaky test Signed-off-by: Yaliang Wu <[email protected]> --------- Signed-off-by: Yaliang Wu <[email protected]>
- Loading branch information
Showing
30 changed files
with
1,279 additions
and
204 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
18 changes: 18 additions & 0 deletions
18
.../src/main/java/org/opensearch/ml/common/transport/connector/MLExecuteConnectorAction.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,18 @@ | ||
| /* | ||
| * Copyright OpenSearch Contributors | ||
| * SPDX-License-Identifier: Apache-2.0 | ||
| */ | ||
|
|
||
| package org.opensearch.ml.common.transport.connector; | ||
|
|
||
| import org.opensearch.action.ActionType; | ||
| import org.opensearch.ml.common.transport.MLTaskResponse; | ||
|
|
||
| public class MLExecuteConnectorAction extends ActionType<MLTaskResponse> { | ||
| public static final MLExecuteConnectorAction INSTANCE = new MLExecuteConnectorAction(); | ||
| public static final String NAME = "cluster:admin/opensearch/ml/connectors/execute"; | ||
|
|
||
| private MLExecuteConnectorAction() { | ||
| super(NAME, MLTaskResponse::new); | ||
| } | ||
| } |
90 changes: 90 additions & 0 deletions
90
...src/main/java/org/opensearch/ml/common/transport/connector/MLExecuteConnectorRequest.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,90 @@ | ||
| /* | ||
| * Copyright OpenSearch Contributors | ||
| * SPDX-License-Identifier: Apache-2.0 | ||
| */ | ||
|
|
||
| package org.opensearch.ml.common.transport.connector; | ||
|
|
||
| import lombok.AccessLevel; | ||
| import lombok.Builder; | ||
| import lombok.Getter; | ||
| import lombok.ToString; | ||
| import lombok.experimental.FieldDefaults; | ||
| import org.opensearch.action.ActionRequest; | ||
| import org.opensearch.action.ActionRequestValidationException; | ||
| import org.opensearch.core.common.io.stream.InputStreamStreamInput; | ||
| import org.opensearch.core.common.io.stream.OutputStreamStreamOutput; | ||
| import org.opensearch.core.common.io.stream.StreamInput; | ||
| import org.opensearch.core.common.io.stream.StreamOutput; | ||
| import org.opensearch.ml.common.input.MLInput; | ||
| import org.opensearch.ml.common.transport.MLTaskRequest; | ||
|
|
||
| import java.io.ByteArrayInputStream; | ||
| import java.io.ByteArrayOutputStream; | ||
| import java.io.IOException; | ||
| import java.io.UncheckedIOException; | ||
|
|
||
| import static org.opensearch.action.ValidateActions.addValidationError; | ||
|
|
||
| @Getter | ||
| @FieldDefaults(level = AccessLevel.PRIVATE) | ||
| @ToString | ||
| public class MLExecuteConnectorRequest extends MLTaskRequest { | ||
|
|
||
| String connectorId; | ||
| MLInput mlInput; | ||
|
|
||
| @Builder | ||
| public MLExecuteConnectorRequest(String connectorId, MLInput mlInput, boolean dispatchTask) { | ||
| super(dispatchTask); | ||
| this.mlInput = mlInput; | ||
| this.connectorId = connectorId; | ||
| } | ||
|
|
||
| public MLExecuteConnectorRequest(String connectorId, MLInput mlInput) { | ||
| this(connectorId, mlInput, true); | ||
| } | ||
|
|
||
| public MLExecuteConnectorRequest(StreamInput in) throws IOException { | ||
| super(in); | ||
| this.connectorId = in.readString(); | ||
| this.mlInput = new MLInput(in); | ||
| } | ||
|
|
||
| @Override | ||
| public void writeTo(StreamOutput out) throws IOException { | ||
| super.writeTo(out); | ||
| out.writeString(this.connectorId); | ||
| this.mlInput.writeTo(out); | ||
| } | ||
|
|
||
| @Override | ||
| public ActionRequestValidationException validate() { | ||
| ActionRequestValidationException exception = null; | ||
| if (this.mlInput == null) { | ||
| exception = addValidationError("ML input can't be null", exception); | ||
| } else if (this.mlInput.getInputDataset() == null) { | ||
| exception = addValidationError("input data can't be null", exception); | ||
| } | ||
|
|
||
| return exception; | ||
| } | ||
|
|
||
|
|
||
| public static MLExecuteConnectorRequest fromActionRequest(ActionRequest actionRequest) { | ||
| if (actionRequest instanceof MLExecuteConnectorRequest) { | ||
| return (MLExecuteConnectorRequest) actionRequest; | ||
| } | ||
|
|
||
| try (ByteArrayOutputStream baos = new ByteArrayOutputStream(); | ||
| OutputStreamStreamOutput osso = new OutputStreamStreamOutput(baos)) { | ||
| actionRequest.writeTo(osso); | ||
| try (StreamInput input = new InputStreamStreamInput(new ByteArrayInputStream(baos.toByteArray()))) { | ||
| return new MLExecuteConnectorRequest(input); | ||
| } | ||
| } catch (IOException e) { | ||
| throw new UncheckedIOException("failed to parse ActionRequest into MLPredictionTaskRequest", e); | ||
| } | ||
|
|
||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.