-
Notifications
You must be signed in to change notification settings - Fork 146
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Signed-off-by: Dhrubo Saha <[email protected]>
- Loading branch information
Showing
29 changed files
with
1,079 additions
and
157 deletions.
There are no files selected for viewing
70 changes: 70 additions & 0 deletions
70
common/src/main/java/org/opensearch/ml/common/transport/search/MLSearchActionRequest.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,70 @@ | ||
| package org.opensearch.ml.common.transport.search; | ||
|
|
||
| import static org.opensearch.ml.common.CommonValue.VERSION_2_19_0; | ||
|
|
||
| import java.io.ByteArrayInputStream; | ||
| import java.io.ByteArrayOutputStream; | ||
| import java.io.IOException; | ||
| import java.io.UncheckedIOException; | ||
|
|
||
| import org.opensearch.Version; | ||
| import org.opensearch.action.ActionRequest; | ||
| import org.opensearch.action.search.SearchRequest; | ||
| 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 lombok.Builder; | ||
| import lombok.Getter; | ||
|
|
||
| @Getter | ||
| public class MLSearchActionRequest extends SearchRequest { | ||
| SearchRequest searchRequest; | ||
| String tenantId; | ||
|
|
||
| @Builder | ||
| public MLSearchActionRequest(SearchRequest searchRequest, String tenantId) { | ||
| this.searchRequest = searchRequest; | ||
| this.tenantId = tenantId; | ||
| } | ||
|
|
||
| public MLSearchActionRequest(StreamInput input) throws IOException { | ||
| super(input); | ||
| Version streamInputVersion = input.getVersion(); | ||
| if (input.readBoolean()) { | ||
| searchRequest = new SearchRequest(input); | ||
| } | ||
| this.tenantId = streamInputVersion.onOrAfter(VERSION_2_19_0) ? input.readOptionalString() : null; | ||
| } | ||
|
|
||
| @Override | ||
| public void writeTo(StreamOutput output) throws IOException { | ||
| super.writeTo(output); | ||
| Version streamOutputVersion = output.getVersion(); | ||
| if (searchRequest != null) { | ||
| output.writeBoolean(true); // user exists | ||
| searchRequest.writeTo(output); | ||
| } else { | ||
| output.writeBoolean(false); // user does not exist | ||
| } | ||
| if (streamOutputVersion.onOrAfter(VERSION_2_19_0)) { | ||
| output.writeOptionalString(tenantId); | ||
| } | ||
| } | ||
|
|
||
| public static MLSearchActionRequest fromActionRequest(ActionRequest actionRequest) { | ||
| if (actionRequest instanceof MLSearchActionRequest) { | ||
| return (MLSearchActionRequest) 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 MLSearchActionRequest(input); | ||
| } | ||
| } catch (IOException e) { | ||
| throw new UncheckedIOException("failed to parse ActionRequest into MLSearchActionRequest", e); | ||
| } | ||
| } | ||
| } |
147 changes: 147 additions & 0 deletions
147
...on/src/test/java/org/opensearch/ml/common/transport/search/MLSearchActionRequestTest.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,147 @@ | ||
| package org.opensearch.ml.common.transport.search; | ||
|
|
||
| import static org.junit.Assert.assertEquals; | ||
| import static org.junit.Assert.assertNotSame; | ||
| import static org.junit.Assert.assertNull; | ||
| import static org.junit.Assert.assertSame; | ||
|
|
||
| import java.io.IOException; | ||
| import java.io.UncheckedIOException; | ||
|
|
||
| import org.junit.Before; | ||
| import org.junit.Test; | ||
| import org.opensearch.Version; | ||
| import org.opensearch.action.ActionRequest; | ||
| import org.opensearch.action.ActionRequestValidationException; | ||
| import org.opensearch.action.search.SearchRequest; | ||
| import org.opensearch.common.io.stream.BytesStreamOutput; | ||
| import org.opensearch.core.common.io.stream.StreamInput; | ||
| import org.opensearch.core.common.io.stream.StreamOutput; | ||
|
|
||
| public class MLSearchActionRequestTest { | ||
|
|
||
| private SearchRequest searchRequest; | ||
|
|
||
| @Before | ||
| public void setUp() { | ||
| searchRequest = new SearchRequest("test-index"); | ||
| } | ||
|
|
||
| @Test | ||
| public void testConstructorAndGetters() { | ||
| MLSearchActionRequest request = MLSearchActionRequest.builder().searchRequest(searchRequest).tenantId("test-tenant").build(); | ||
| assertEquals("test-index", request.getSearchRequest().indices()[0]); | ||
| assertEquals("test-tenant", request.getTenantId()); | ||
| } | ||
|
|
||
| @Test | ||
| public void testStreamConstructorAndWriteTo() throws IOException { | ||
| MLSearchActionRequest request = MLSearchActionRequest.builder().searchRequest(searchRequest).tenantId("test-tenant").build(); | ||
| BytesStreamOutput out = new BytesStreamOutput(); | ||
| request.writeTo(out); | ||
|
|
||
| MLSearchActionRequest deserializedRequest = new MLSearchActionRequest(out.bytes().streamInput()); | ||
| assertEquals("test-index", deserializedRequest.getSearchRequest().indices()[0]); | ||
| assertEquals("test-tenant", deserializedRequest.getTenantId()); | ||
| } | ||
|
|
||
| @Test | ||
| public void testWriteToWithNullSearchRequest() throws IOException { | ||
| MLSearchActionRequest request = MLSearchActionRequest.builder().tenantId("test-tenant").build(); | ||
| BytesStreamOutput out = new BytesStreamOutput(); | ||
| request.writeTo(out); | ||
|
|
||
| MLSearchActionRequest deserializedRequest = new MLSearchActionRequest(out.bytes().streamInput()); | ||
| assertNull(deserializedRequest.getSearchRequest()); | ||
| assertEquals("test-tenant", deserializedRequest.getTenantId()); | ||
| } | ||
|
|
||
| @Test | ||
| public void testFromActionRequestWithMLSearchActionRequest() { | ||
| MLSearchActionRequest request = MLSearchActionRequest.builder().searchRequest(searchRequest).tenantId("test-tenant").build(); | ||
| MLSearchActionRequest result = MLSearchActionRequest.fromActionRequest(request); | ||
| assertSame(result, request); | ||
| } | ||
|
|
||
| @Test | ||
| public void testFromActionRequestWithNonMLSearchActionRequest() throws IOException { | ||
| MLSearchActionRequest request = MLSearchActionRequest.builder().searchRequest(searchRequest).tenantId("test-tenant").build(); | ||
| ActionRequest actionRequest = new ActionRequest() { | ||
| @Override | ||
| public ActionRequestValidationException validate() { | ||
| return null; | ||
| } | ||
|
|
||
| @Override | ||
| public void writeTo(StreamOutput out) throws IOException { | ||
| request.writeTo(out); | ||
| } | ||
| }; | ||
|
|
||
| MLSearchActionRequest result = MLSearchActionRequest.fromActionRequest(actionRequest); | ||
| assertNotSame(result, request); | ||
| assertEquals(request.getSearchRequest().indices()[0], result.getSearchRequest().indices()[0]); | ||
| assertEquals(request.getTenantId(), result.getTenantId()); | ||
| } | ||
|
|
||
| @Test(expected = UncheckedIOException.class) | ||
| public void testFromActionRequestIOException() { | ||
| ActionRequest actionRequest = new ActionRequest() { | ||
| @Override | ||
| public ActionRequestValidationException validate() { | ||
| return null; | ||
| } | ||
|
|
||
| @Override | ||
| public void writeTo(StreamOutput out) throws IOException { | ||
| throw new IOException("test"); | ||
| } | ||
| }; | ||
| MLSearchActionRequest.fromActionRequest(actionRequest); | ||
| } | ||
|
|
||
| @Test | ||
| public void testBackwardCompatibility() throws IOException { | ||
| MLSearchActionRequest request = MLSearchActionRequest.builder().searchRequest(searchRequest).tenantId("test-tenant").build(); | ||
|
|
||
| BytesStreamOutput out = new BytesStreamOutput(); | ||
| out.setVersion(Version.V_2_18_0); // Older version | ||
| request.writeTo(out); | ||
|
|
||
| StreamInput in = out.bytes().streamInput(); | ||
| in.setVersion(Version.V_2_18_0); | ||
|
|
||
| MLSearchActionRequest deserializedRequest = new MLSearchActionRequest(in); | ||
| assertNull(deserializedRequest.getTenantId()); // Ensure tenantId is ignored | ||
| } | ||
|
|
||
| @Test | ||
| public void testFromActionRequestWithValidRequest() { | ||
| MLSearchActionRequest request = MLSearchActionRequest.builder().searchRequest(searchRequest).tenantId("test-tenant").build(); | ||
|
|
||
| MLSearchActionRequest result = MLSearchActionRequest.fromActionRequest(request); | ||
| assertSame(request, result); | ||
| } | ||
|
|
||
| @Test | ||
| public void testMixedVersionCompatibility() throws IOException { | ||
| MLSearchActionRequest originalRequest = MLSearchActionRequest | ||
| .builder() | ||
| .searchRequest(searchRequest) | ||
| .tenantId("test-tenant") | ||
| .build(); | ||
|
|
||
| // Serialize with a newer version | ||
| BytesStreamOutput out = new BytesStreamOutput(); | ||
| out.setVersion(Version.V_2_19_0); | ||
| originalRequest.writeTo(out); | ||
|
|
||
| // Deserialize with an older version | ||
| StreamInput in = out.bytes().streamInput(); | ||
| in.setVersion(Version.V_2_18_0); | ||
|
|
||
| MLSearchActionRequest deserializedRequest = new MLSearchActionRequest(in); | ||
| assertNull(deserializedRequest.getTenantId()); // tenantId should not exist in older versions | ||
| } | ||
|
|
||
| } |
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
Oops, something went wrong.