applying multi-tenancy in search [model, model group, agent, connector] #3433
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
| 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; | ||
|
There was a problem hiding this comment. A new MLSearchActionRequest class was added to wrap the original SearchRequest and include a tenantId field.So I want to understand does it impact other functionality using SearchRequest? There was a problem hiding this comment. Can you please give me an example which functionality are you referring to? There was a problem hiding this comment. I think the order matters, when did the SearchRequest converts to MLSearchActionRequest? I am worry about the tools and search pipelines that are using the search requests. if the transform is after them, it should be fine, but if that's before tools and search pipelines, does it impact them? |
||
|
|
||
| @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); | ||
|
There was a problem hiding this comment. In this exception handling, I think IllegalArgumentException is more proper in this case The method is expecting an ActionRequest that can be converted to MLSearchActionRequest. I think if using IllegalArgumentException, it provides clearer feedback to the caller about what went wrong. There was a problem hiding this comment. This is literally the same method we use for all other request classes. Example |
||
| } | ||
| } | ||
| } | ||
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.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Add java documentation
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
added