forked from opensearch-project/k-NN
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Added functionalities like IndexBuildServiceClient, parsing of reques…
…t, response for create index API Signed-off-by: Navneet Verma <[email protected]>
- Loading branch information
Showing
8 changed files
with
337 additions
and
14 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
108 changes: 108 additions & 0 deletions
108
src/main/java/org/opensearch/knn/remote/index/client/IndexBuildServiceClient.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,108 @@ | ||
| /* | ||
| * Copyright OpenSearch Contributors | ||
| * SPDX-License-Identifier: Apache-2.0 | ||
| */ | ||
|
|
||
| package org.opensearch.knn.remote.index.client; | ||
|
|
||
| import org.apache.http.HttpEntity; | ||
| import org.apache.http.HttpHost; | ||
| import org.apache.http.HttpRequest; | ||
| import org.apache.http.HttpResponse; | ||
| import org.apache.http.client.HttpClient; | ||
| import org.apache.http.client.methods.HttpPost; | ||
| import org.apache.http.entity.StringEntity; | ||
| import org.apache.http.impl.client.HttpClientBuilder; | ||
| import org.apache.http.util.EntityUtils; | ||
| import org.opensearch.common.xcontent.XContentFactory; | ||
| import org.opensearch.core.xcontent.DeprecationHandler; | ||
| import org.opensearch.core.xcontent.MediaTypeRegistry; | ||
| import org.opensearch.core.xcontent.NamedXContentRegistry; | ||
| import org.opensearch.core.xcontent.XContent; | ||
| import org.opensearch.core.xcontent.XContentBuilder; | ||
| import org.opensearch.core.xcontent.XContentParser; | ||
| import org.opensearch.knn.index.KNNSettings; | ||
| import org.opensearch.knn.remote.index.model.CreateIndexRequest; | ||
| import org.opensearch.knn.remote.index.model.CreateIndexResponse; | ||
| import org.opensearch.knn.remote.index.s3.S3Client; | ||
|
|
||
| import java.io.IOException; | ||
|
|
||
| /** | ||
| * Main class to class the IndexBuildServiceAPIs | ||
| */ | ||
| public class IndexBuildServiceClient { | ||
| private static volatile IndexBuildServiceClient INSTANCE; | ||
| private static final String CONTENT_TYPE = "Content-Type"; | ||
| private static final String APPLICATION_JSON = "application/json"; | ||
| private static final String ACCEPT = "Accept"; | ||
| private final HttpClient httpClient; | ||
| private final HttpHost httpHost; | ||
|
|
||
| public static IndexBuildServiceClient getInstance() throws IOException { | ||
| IndexBuildServiceClient result = INSTANCE; | ||
| if (result == null) { | ||
| synchronized (S3Client.class) { | ||
| result = INSTANCE; | ||
| if (result == null) { | ||
| INSTANCE = result = new IndexBuildServiceClient(); | ||
| } | ||
| } | ||
| } | ||
| return result; | ||
| } | ||
|
|
||
| private IndexBuildServiceClient() { | ||
| this.httpClient = HttpClientBuilder.create().build(); | ||
| this.httpHost = new HttpHost(KNNSettings.getRemoteServiceEndpoint(), KNNSettings.getRemoteServicePort(), "http"); | ||
| } | ||
|
|
||
| /** | ||
| * API to be called to create the Vector Index using remote endpoint | ||
| * @param createIndexRequest {@link CreateIndexRequest} | ||
| * @throws IOException Exception called if createIndex request is not successful | ||
| */ | ||
| public CreateIndexResponse createIndex(final CreateIndexRequest createIndexRequest) throws IOException { | ||
| HttpPost request = new HttpPost(); | ||
| request.setHeader(CONTENT_TYPE, APPLICATION_JSON); | ||
| request.setHeader(ACCEPT, APPLICATION_JSON); | ||
| XContentBuilder builder = XContentFactory.jsonBuilder(); | ||
| builder = createIndexRequest.toXContent(builder, null); | ||
| request.setEntity(new StringEntity(builder.toString())); | ||
|
|
||
| HttpResponse response = makeHTTPRequest(request); | ||
| HttpEntity httpEntity = response.getEntity(); | ||
| String responseString = EntityUtils.toString(httpEntity); | ||
| return parseCreateIndexResponse(responseString); | ||
| } | ||
|
|
||
| // TODO: To be implemented | ||
| public void checkIndexBuildStatus() { | ||
|
|
||
| } | ||
|
|
||
| private HttpResponse makeHTTPRequest(final HttpRequest request) throws IOException { | ||
| HttpResponse response = httpClient.execute(httpHost, request); | ||
| HttpEntity entity = response.getEntity(); | ||
| int statusCode = response.getStatusLine().getStatusCode(); | ||
|
|
||
| if (statusCode >= 400) { | ||
| String errorBody = entity != null ? EntityUtils.toString(entity) : "No response body"; | ||
| throw new IOException("Request failed with status code: " + statusCode + ", body: " + errorBody); | ||
| } | ||
|
|
||
| return response; | ||
| } | ||
|
|
||
| // Keeping it package private for doing the unit testing for now. | ||
| static CreateIndexResponse parseCreateIndexResponse(final String responseString) throws IOException { | ||
| final XContent xContent = MediaTypeRegistry.getDefaultMediaType().xContent(); | ||
| final XContentParser parser = xContent.createParser( | ||
| NamedXContentRegistry.EMPTY, | ||
| DeprecationHandler.THROW_UNSUPPORTED_OPERATION, | ||
| responseString | ||
| ); | ||
| return CreateIndexResponse.fromXContent(parser); | ||
| } | ||
|
|
||
| } |
32 changes: 32 additions & 0 deletions
32
src/main/java/org/opensearch/knn/remote/index/model/CreateIndexRequest.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,32 @@ | ||
| /* | ||
| * Copyright OpenSearch Contributors | ||
| * SPDX-License-Identifier: Apache-2.0 | ||
| */ | ||
|
|
||
| package org.opensearch.knn.remote.index.model; | ||
|
|
||
| import lombok.Builder; | ||
| import lombok.Value; | ||
| import org.opensearch.core.xcontent.ToXContentObject; | ||
| import org.opensearch.core.xcontent.XContentBuilder; | ||
|
|
||
| import java.io.IOException; | ||
|
|
||
| @Value | ||
| @Builder | ||
| public class CreateIndexRequest implements ToXContentObject { | ||
| String bucketName; | ||
| String objectLocation; | ||
| long numberOfVectors; | ||
| int dimensions; | ||
|
|
||
| @Override | ||
| public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException { | ||
| return builder.startObject() | ||
| .field("bucket_name", bucketName) | ||
| .field("object_location", objectLocation) | ||
| .field("number_of_vectors", numberOfVectors) | ||
| .field("dimensions", dimensions) | ||
| .endObject(); | ||
| } | ||
| } |
45 changes: 45 additions & 0 deletions
45
src/main/java/org/opensearch/knn/remote/index/model/CreateIndexResponse.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,45 @@ | ||
| /* | ||
| * Copyright OpenSearch Contributors | ||
| * SPDX-License-Identifier: Apache-2.0 | ||
| */ | ||
|
|
||
| package org.opensearch.knn.remote.index.model; | ||
|
|
||
| import lombok.Builder; | ||
| import lombok.Value; | ||
| import org.opensearch.core.ParseField; | ||
| import org.opensearch.core.xcontent.XContentParser; | ||
|
|
||
| import java.io.IOException; | ||
|
|
||
| @Value | ||
| @Builder | ||
| public class CreateIndexResponse { | ||
| private static final ParseField INDEX_CREATION_REQUEST_ID = new ParseField("indexCreationRequestId"); | ||
| private static final ParseField STATUS = new ParseField("status"); | ||
| String indexCreationRequestId; | ||
| String status; | ||
|
|
||
| public static CreateIndexResponse fromXContent(XContentParser parser) throws IOException { | ||
| final CreateIndexResponseBuilder builder = new CreateIndexResponseBuilder(); | ||
| XContentParser.Token token = parser.nextToken(); | ||
| if (token != XContentParser.Token.START_OBJECT) { | ||
| throw new IOException("Invalid response format, was expecting a " + XContentParser.Token.START_OBJECT); | ||
| } | ||
| String currentFieldName = null; | ||
| while ((token = parser.nextToken()) != XContentParser.Token.END_OBJECT) { | ||
| if (token == XContentParser.Token.FIELD_NAME) { | ||
| currentFieldName = parser.currentName(); | ||
| } else if (token.isValue()) { | ||
| if (INDEX_CREATION_REQUEST_ID.match(currentFieldName, parser.getDeprecationHandler())) { | ||
| builder.indexCreationRequestId(parser.text()); | ||
| } else if (STATUS.match(currentFieldName, parser.getDeprecationHandler())) { | ||
| builder.status(parser.text()); | ||
| } else { | ||
| throw new IOException("Invalid response format, unknown field: " + currentFieldName); | ||
| } | ||
| } | ||
| } | ||
| return builder.build(); | ||
| } | ||
| } |
Oops, something went wrong.