From 82a50cbedfcc44c6dba1eafc465123ed72e530b1 Mon Sep 17 00:00:00 2001 From: owenhalpert Date: Sun, 2 Mar 2025 18:05:35 -0800 Subject: [PATCH] Move common constants to Remote constants class, simplify RemoteBuildRequest using serialization logic (XContent) Signed-off-by: owenhalpert --- .../opensearch/knn/common/KNNConstants.java | 21 ---- .../DefaultVectorRepositoryAccessor.java | 6 +- .../remote/RemoteIndexBuildStrategy.java | 16 +-- .../knn/index/engine/faiss/Faiss.java | 4 +- .../index/remote/HTTPRemoteBuildRequest.java | 59 ---------- .../knn/index/remote/KNNRemoteConstants.java | 27 +++++ .../knn/index/remote/RemoteBuildRequest.java | 58 ++++++++-- .../remote/RemoteBuildRequestBuilder.java | 65 ----------- .../knn/index/remote/RemoteBuildResponse.java | 42 ++++++- .../remote/RemoteIndexClientFactory.java | 10 +- .../index/remote/RemoteIndexHTTPClient.java | 70 +++++------ .../DefaultVectorRepositoryAccessorTests.java | 4 +- .../remote/RemoteIndexHTTPClientTests.java | 109 ++++++++++++------ 13 files changed, 224 insertions(+), 267 deletions(-) delete mode 100644 src/main/java/org/opensearch/knn/index/remote/HTTPRemoteBuildRequest.java create mode 100644 src/main/java/org/opensearch/knn/index/remote/KNNRemoteConstants.java delete mode 100644 src/main/java/org/opensearch/knn/index/remote/RemoteBuildRequestBuilder.java diff --git a/src/main/java/org/opensearch/knn/common/KNNConstants.java b/src/main/java/org/opensearch/knn/common/KNNConstants.java index c2d4289685..14e95887c7 100644 --- a/src/main/java/org/opensearch/knn/common/KNNConstants.java +++ b/src/main/java/org/opensearch/knn/common/KNNConstants.java @@ -165,25 +165,4 @@ public class KNNConstants { public static final String DERIVED_VECTOR_FIELD_ATTRIBUTE_KEY = "knn-derived-source-enabled"; public static final String DERIVED_VECTOR_FIELD_ATTRIBUTE_TRUE_VALUE = "true"; public static final String DERIVED_VECTOR_FIELD_ATTRIBUTE_FALSE_VALUE = "false"; - - // Remote build constants - public static final String BUILD_ENDPOINT = "/_build"; - public static final String STATUS_ENDPOINT = "/_status"; - public static final String S3 = "s3"; - public static final String BUCKET = "bucket"; - // Build request keys - public static final String ALGORITHM = "algorithm"; - public static final String ALGORITHM_PARAMETERS = "algorithm_parameters"; - public static final String INDEX_PARAMETERS = "index_parameters"; - public static final String DOC_COUNT = "doc_count"; - public static final String TENANT_ID = "tenant_id"; - public static final String DOC_ID_PATH = "doc_id_path"; - public static final String VECTOR_PATH = "vector_path"; - public static final String CONTAINER_NAME = "container_name"; - public static final String REPOSITORY_TYPE = "repository_type"; - // Server responses - public static final String JOB_ID = "job_id"; - public static final String TASK_STATUS = "task_status"; - public static final String INDEX_PATH = "index_path"; - public static final String ERROR_MESSAGE = "error_message"; } diff --git a/src/main/java/org/opensearch/knn/index/codec/nativeindex/remote/DefaultVectorRepositoryAccessor.java b/src/main/java/org/opensearch/knn/index/codec/nativeindex/remote/DefaultVectorRepositoryAccessor.java index 1ad077b5cb..ae41e07cf0 100644 --- a/src/main/java/org/opensearch/knn/index/codec/nativeindex/remote/DefaultVectorRepositoryAccessor.java +++ b/src/main/java/org/opensearch/knn/index/codec/nativeindex/remote/DefaultVectorRepositoryAccessor.java @@ -33,10 +33,10 @@ import java.util.concurrent.atomic.AtomicReference; import java.util.function.Supplier; -import static org.opensearch.knn.index.codec.nativeindex.remote.RemoteIndexBuildStrategy.DOC_ID_FILE_EXTENSION; -import static org.opensearch.knn.index.codec.nativeindex.remote.RemoteIndexBuildStrategy.VECTORS_PATH; -import static org.opensearch.knn.index.codec.nativeindex.remote.RemoteIndexBuildStrategy.VECTOR_BLOB_FILE_EXTENSION; import static org.opensearch.knn.index.codec.util.KNNCodecUtil.initializeVectorValues; +import static org.opensearch.knn.index.remote.KNNRemoteConstants.DOC_ID_FILE_EXTENSION; +import static org.opensearch.knn.index.remote.KNNRemoteConstants.VECTORS_PATH; +import static org.opensearch.knn.index.remote.KNNRemoteConstants.VECTOR_BLOB_FILE_EXTENSION; @Log4j2 @AllArgsConstructor diff --git a/src/main/java/org/opensearch/knn/index/codec/nativeindex/remote/RemoteIndexBuildStrategy.java b/src/main/java/org/opensearch/knn/index/codec/nativeindex/remote/RemoteIndexBuildStrategy.java index 3b1244f3b2..e001fbd074 100644 --- a/src/main/java/org/opensearch/knn/index/codec/nativeindex/remote/RemoteIndexBuildStrategy.java +++ b/src/main/java/org/opensearch/knn/index/codec/nativeindex/remote/RemoteIndexBuildStrategy.java @@ -13,9 +13,7 @@ import org.opensearch.knn.index.KNNSettings; import org.opensearch.knn.index.codec.nativeindex.NativeIndexBuildStrategy; import org.opensearch.knn.index.codec.nativeindex.model.BuildIndexParams; -import org.opensearch.knn.index.remote.HTTPRemoteBuildRequest; import org.opensearch.knn.index.remote.RemoteBuildRequest; -import org.opensearch.knn.index.remote.RemoteBuildRequestBuilder; import org.opensearch.knn.index.remote.RemoteBuildResponse; import org.opensearch.knn.index.remote.RemoteIndexClient; import org.opensearch.knn.index.remote.RemoteIndexClientFactory; @@ -43,10 +41,6 @@ public class RemoteIndexBuildStrategy implements NativeIndexBuildStrategy { private final NativeIndexBuildStrategy fallbackStrategy; private final IndexSettings indexSettings; - public static final String VECTOR_BLOB_FILE_EXTENSION = ".knnvec"; - public static final String DOC_ID_FILE_EXTENSION = ".knndid"; - static final String VECTORS_PATH = "_vectors"; - /** * Public constructor, intended to be called by {@link org.opensearch.knn.index.codec.nativeindex.NativeIndexBuildStrategyFactory} based in * part on the return value from {@link RemoteIndexBuildStrategy#shouldBuildIndexRemotely} @@ -130,14 +124,8 @@ public void buildAndWriteIndex(BuildIndexParams indexInfo) throws IOException { time_in_millis = stopWatch.stop().totalTime().millis(); log.debug("Repository write took {} ms for vector field [{}]", time_in_millis, indexInfo.getFieldName()); - // TODO future implementations will set the following two params depending on some setting to denote the protocol - RemoteIndexClient client = RemoteIndexClientFactory.getRemoteIndexClient(RemoteIndexClientFactory.TYPE_HTTP); - RemoteBuildRequest request = RemoteBuildRequestBuilder.builder(HTTPRemoteBuildRequest.class) - .indexSettings(indexSettings) - .indexInfo(indexInfo) - .repositoryMetadata(getRepository().getMetadata()) - .blobName(blobName) - .build(); + RemoteIndexClient client = RemoteIndexClientFactory.getRemoteIndexClient(); + RemoteBuildRequest request = new RemoteBuildRequest(indexSettings, indexInfo, getRepository().getMetadata(), blobName); stopWatch = new StopWatch().start(); RemoteBuildResponse remoteBuildResponse = client.submitVectorBuild(request); time_in_millis = stopWatch.stop().totalTime().millis(); diff --git a/src/main/java/org/opensearch/knn/index/engine/faiss/Faiss.java b/src/main/java/org/opensearch/knn/index/engine/faiss/Faiss.java index 43977ec208..bdbf46a0b8 100644 --- a/src/main/java/org/opensearch/knn/index/engine/faiss/Faiss.java +++ b/src/main/java/org/opensearch/knn/index/engine/faiss/Faiss.java @@ -19,8 +19,6 @@ import java.util.Map; import java.util.function.Function; -import static org.opensearch.knn.common.KNNConstants.ALGORITHM; -import static org.opensearch.knn.common.KNNConstants.ALGORITHM_PARAMETERS; import static org.opensearch.knn.common.KNNConstants.INDEX_DESCRIPTION_PARAMETER; import static org.opensearch.knn.common.KNNConstants.METHOD_HNSW; import static org.opensearch.knn.common.KNNConstants.METHOD_IVF; @@ -35,6 +33,8 @@ import static org.opensearch.knn.common.KNNConstants.NAME; import static org.opensearch.knn.common.KNNConstants.PARAMETERS; import static org.opensearch.knn.common.KNNConstants.SPACE_TYPE; +import static org.opensearch.knn.index.remote.KNNRemoteConstants.ALGORITHM; +import static org.opensearch.knn.index.remote.KNNRemoteConstants.ALGORITHM_PARAMETERS; import static org.opensearch.knn.index.KNNSettings.INDEX_KNN_DEFAULT_ALGO_PARAM_EF_CONSTRUCTION; import static org.opensearch.knn.index.KNNSettings.INDEX_KNN_DEFAULT_ALGO_PARAM_EF_SEARCH; import static org.opensearch.knn.index.KNNSettings.INDEX_KNN_DEFAULT_SPACE_TYPE; diff --git a/src/main/java/org/opensearch/knn/index/remote/HTTPRemoteBuildRequest.java b/src/main/java/org/opensearch/knn/index/remote/HTTPRemoteBuildRequest.java deleted file mode 100644 index 380e6ff434..0000000000 --- a/src/main/java/org/opensearch/knn/index/remote/HTTPRemoteBuildRequest.java +++ /dev/null @@ -1,59 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.knn.index.remote; - -import lombok.Getter; -import org.opensearch.cluster.metadata.RepositoryMetadata; -import org.opensearch.common.xcontent.json.JsonXContent; -import org.opensearch.core.xcontent.XContentBuilder; -import org.opensearch.index.IndexSettings; -import org.opensearch.knn.index.codec.nativeindex.model.BuildIndexParams; - -import java.io.IOException; - -import static org.opensearch.knn.common.KNNConstants.CONTAINER_NAME; -import static org.opensearch.knn.common.KNNConstants.DIMENSION; -import static org.opensearch.knn.common.KNNConstants.DOC_COUNT; -import static org.opensearch.knn.common.KNNConstants.DOC_ID_PATH; -import static org.opensearch.knn.common.KNNConstants.INDEX_PARAMETERS; -import static org.opensearch.knn.common.KNNConstants.KNN_ENGINE; -import static org.opensearch.knn.common.KNNConstants.REPOSITORY_TYPE; -import static org.opensearch.knn.common.KNNConstants.TENANT_ID; -import static org.opensearch.knn.common.KNNConstants.VECTOR_DATA_TYPE_FIELD; -import static org.opensearch.knn.common.KNNConstants.VECTOR_PATH; - -/** - * HTTP-specific implementation of RemoteBuildRequest. - */ -@Getter -public class HTTPRemoteBuildRequest extends RemoteBuildRequest { - public HTTPRemoteBuildRequest( - IndexSettings indexSettings, - BuildIndexParams indexInfo, - RepositoryMetadata repositoryMetadata, - String blobName - ) throws IOException { - super(indexSettings, indexInfo, repositoryMetadata, blobName); - } - - public String toJson() throws IOException { - try (XContentBuilder builder = JsonXContent.contentBuilder()) { - builder.startObject(); - builder.field(REPOSITORY_TYPE, repositoryType); - builder.field(CONTAINER_NAME, containerName); - builder.field(VECTOR_PATH, vectorPath); - builder.field(DOC_ID_PATH, docIdPath); - builder.field(TENANT_ID, tenantId); - builder.field(DIMENSION, dimension); - builder.field(DOC_COUNT, docCount); - builder.field(VECTOR_DATA_TYPE_FIELD, vectorDataType); - builder.field(KNN_ENGINE, engine); - builder.field(INDEX_PARAMETERS, indexParameters); - builder.endObject(); - return builder.toString(); - } - } -} diff --git a/src/main/java/org/opensearch/knn/index/remote/KNNRemoteConstants.java b/src/main/java/org/opensearch/knn/index/remote/KNNRemoteConstants.java new file mode 100644 index 0000000000..3a00b6e6ff --- /dev/null +++ b/src/main/java/org/opensearch/knn/index/remote/KNNRemoteConstants.java @@ -0,0 +1,27 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.knn.index.remote; + +// Public class to define the constants used by Remote Index Build in 2 or more classes. +public class KNNRemoteConstants { + // Repository filepath constants + public static final String VECTOR_BLOB_FILE_EXTENSION = ".knnvec"; + public static final String DOC_ID_FILE_EXTENSION = ".knndid"; + public static final String VECTORS_PATH = "_vectors"; + + // Repository-S3 + public static final String S3 = "s3"; + public static final String BUCKET = "bucket"; + + // Build request keys + public static final String ALGORITHM = "algorithm"; + public static final String ALGORITHM_PARAMETERS = "algorithm_parameters"; + public static final String INDEX_PARAMETERS = "index_parameters"; + + // HTTP implementation + public static final String BUILD_ENDPOINT = "/_build"; + public static final String STATUS_ENDPOINT = "/_status"; +} diff --git a/src/main/java/org/opensearch/knn/index/remote/RemoteBuildRequest.java b/src/main/java/org/opensearch/knn/index/remote/RemoteBuildRequest.java index bcd4151a61..42896fcc66 100644 --- a/src/main/java/org/opensearch/knn/index/remote/RemoteBuildRequest.java +++ b/src/main/java/org/opensearch/knn/index/remote/RemoteBuildRequest.java @@ -8,23 +8,38 @@ import lombok.Getter; import org.opensearch.cluster.ClusterName; import org.opensearch.cluster.metadata.RepositoryMetadata; +import org.opensearch.core.xcontent.ToXContentObject; +import org.opensearch.core.xcontent.XContentBuilder; import org.opensearch.index.IndexSettings; import org.opensearch.knn.index.codec.nativeindex.model.BuildIndexParams; -import org.opensearch.knn.index.codec.nativeindex.remote.RemoteIndexBuildStrategy; import org.opensearch.knn.index.codec.util.KNNCodecUtil; import org.opensearch.knn.index.vectorvalues.KNNVectorValues; import java.io.IOException; import java.util.Map; -import static org.opensearch.knn.common.KNNConstants.BUCKET; -import static org.opensearch.knn.common.KNNConstants.S3; +import static org.opensearch.knn.common.KNNConstants.DIMENSION; +import static org.opensearch.knn.common.KNNConstants.KNN_ENGINE; +import static org.opensearch.knn.common.KNNConstants.VECTOR_DATA_TYPE_FIELD; +import static org.opensearch.knn.index.remote.KNNRemoteConstants.BUCKET; +import static org.opensearch.knn.index.remote.KNNRemoteConstants.DOC_ID_FILE_EXTENSION; +import static org.opensearch.knn.index.remote.KNNRemoteConstants.INDEX_PARAMETERS; +import static org.opensearch.knn.index.remote.KNNRemoteConstants.S3; +import static org.opensearch.knn.index.remote.KNNRemoteConstants.VECTOR_BLOB_FILE_EXTENSION; /** - * Abstract base class for Remote Build Requests. + * Request object for sending build requests to the remote build service, encapsulating all the required parameters + * in a generic XContent format. */ @Getter -public abstract class RemoteBuildRequest { +public class RemoteBuildRequest implements ToXContentObject { + private final String DOC_COUNT = "doc_count"; + private final String TENANT_ID = "tenant_id"; + private final String DOC_ID_PATH = "doc_id_path"; + public final String VECTOR_PATH = "vector_path"; + public final String CONTAINER_NAME = "container_name"; + public final String REPOSITORY_TYPE = "repository_type"; + protected String repositoryType; protected String containerName; protected String vectorPath; @@ -36,6 +51,15 @@ public abstract class RemoteBuildRequest { protected String engine; protected Map indexParameters; + /** + * Constructor for RemoteBuildRequest. + * + * @param indexSettings IndexSettings object + * @param indexInfo BuildIndexParams object + * @param repositoryMetadata RepositoryMetadata object + * @param blobName Name of the blob + * @throws IOException if an I/O error occurs + */ public RemoteBuildRequest( IndexSettings indexSettings, BuildIndexParams indexInfo, @@ -60,8 +84,8 @@ public RemoteBuildRequest( this.repositoryType = repositoryType; this.containerName = containerName; - this.vectorPath = blobName + RemoteIndexBuildStrategy.VECTOR_BLOB_FILE_EXTENSION; - this.docIdPath = blobName + RemoteIndexBuildStrategy.DOC_ID_FILE_EXTENSION; + this.vectorPath = blobName + VECTOR_BLOB_FILE_EXTENSION; + this.docIdPath = blobName + DOC_ID_FILE_EXTENSION; this.tenantId = indexSettings.getSettings().get(ClusterName.CLUSTER_NAME_SETTING.getKey()); this.dimension = vectorValues.dimension(); this.docCount = indexInfo.getTotalLiveDocs(); @@ -70,8 +94,20 @@ public RemoteBuildRequest( this.indexParameters = indexParameters; } - /** - * Convert the request to JSON format. - */ - public abstract String toJson() throws IOException; + @Override + public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException { + builder.startObject(); + builder.field(REPOSITORY_TYPE, repositoryType); + builder.field(CONTAINER_NAME, containerName); + builder.field(VECTOR_PATH, vectorPath); + builder.field(DOC_ID_PATH, docIdPath); + builder.field(TENANT_ID, tenantId); + builder.field(DIMENSION, dimension); + builder.field(DOC_COUNT, docCount); + builder.field(VECTOR_DATA_TYPE_FIELD, vectorDataType); + builder.field(KNN_ENGINE, engine); + builder.field(INDEX_PARAMETERS, indexParameters); + builder.endObject(); + return builder; + } } diff --git a/src/main/java/org/opensearch/knn/index/remote/RemoteBuildRequestBuilder.java b/src/main/java/org/opensearch/knn/index/remote/RemoteBuildRequestBuilder.java deleted file mode 100644 index 484b032f27..0000000000 --- a/src/main/java/org/opensearch/knn/index/remote/RemoteBuildRequestBuilder.java +++ /dev/null @@ -1,65 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ -package org.opensearch.knn.index.remote; - -import org.opensearch.cluster.metadata.RepositoryMetadata; -import org.opensearch.index.IndexSettings; -import org.opensearch.knn.index.codec.nativeindex.model.BuildIndexParams; - -import java.io.IOException; - -/** - * Generic Builder for constructing RemoteBuildRequest objects per requestType (HTTP, etc). - */ -public class RemoteBuildRequestBuilder { - private IndexSettings indexSettings; - private BuildIndexParams indexInfo; - private RepositoryMetadata repositoryMetadata; - private String blobName; - private final Class requestType; - - // Private constructor - private RemoteBuildRequestBuilder(Class requestType) { - this.requestType = requestType; - } - - /** - * Static factory method to create a builder instance. - * @param requestType Class type of the RemoteBuildRequest - * @return RemoteBuildRequestBuilder instance - */ - public static RemoteBuildRequestBuilder builder(Class requestType) { - return new RemoteBuildRequestBuilder<>(requestType); - } - - public RemoteBuildRequestBuilder indexSettings(IndexSettings indexSettings) { - this.indexSettings = indexSettings; - return this; - } - - public RemoteBuildRequestBuilder indexInfo(BuildIndexParams indexInfo) { - this.indexInfo = indexInfo; - return this; - } - - public RemoteBuildRequestBuilder repositoryMetadata(RepositoryMetadata repositoryMetadata) { - this.repositoryMetadata = repositoryMetadata; - return this; - } - - public RemoteBuildRequestBuilder blobName(String blobName) { - this.blobName = blobName; - return this; - } - - public T build() throws IOException { - try { - return requestType.getConstructor(IndexSettings.class, BuildIndexParams.class, RepositoryMetadata.class, String.class) - .newInstance(indexSettings, indexInfo, repositoryMetadata, blobName); - } catch (Exception e) { - throw new IOException("Failed to instantiate RemoteBuildRequest", e); - } - } -} diff --git a/src/main/java/org/opensearch/knn/index/remote/RemoteBuildResponse.java b/src/main/java/org/opensearch/knn/index/remote/RemoteBuildResponse.java index f2ac9715ae..e28889d5d4 100644 --- a/src/main/java/org/opensearch/knn/index/remote/RemoteBuildResponse.java +++ b/src/main/java/org/opensearch/knn/index/remote/RemoteBuildResponse.java @@ -5,9 +5,45 @@ package org.opensearch.knn.index.remote; +import lombok.Builder; +import lombok.Value; +import org.apache.commons.lang.StringUtils; +import org.opensearch.core.ParseField; +import org.opensearch.core.xcontent.XContentParser; + +import java.io.IOException; + /** - * Remote build response. Currently, this just contains the jobId from the server. - * In the future, this may be an interface if different clients expect different responses. + * Response from the remote index build service. This class is used to parse the response from the remote index build service. */ -public record RemoteBuildResponse(String jobId) { +@Value +@Builder +public class RemoteBuildResponse { + private static final String JOB_ID_FIELD = "job_id"; + private static final ParseField JOB_ID = new ParseField(JOB_ID_FIELD); + String jobId; + + public static RemoteBuildResponse fromXContent(XContentParser parser) throws IOException { + final RemoteBuildResponseBuilder builder = new RemoteBuildResponseBuilder(); + 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 (JOB_ID.match(currentFieldName, parser.getDeprecationHandler())) { + builder.jobId(parser.text()); + } else { + throw new IOException("Invalid response format, unknown field: " + currentFieldName); + } + } + } + if (StringUtils.isBlank(builder.jobId)) { + throw new IOException("Invalid response format, missing + " + JOB_ID_FIELD); + } + return builder.build(); + } } diff --git a/src/main/java/org/opensearch/knn/index/remote/RemoteIndexClientFactory.java b/src/main/java/org/opensearch/knn/index/remote/RemoteIndexClientFactory.java index 9efa167722..ddf4a23984 100644 --- a/src/main/java/org/opensearch/knn/index/remote/RemoteIndexClientFactory.java +++ b/src/main/java/org/opensearch/knn/index/remote/RemoteIndexClientFactory.java @@ -7,12 +7,8 @@ public class RemoteIndexClientFactory { - public static final String TYPE_HTTP = "HTTP"; - - public static RemoteIndexClient getRemoteIndexClient(String type) { - if (TYPE_HTTP.equalsIgnoreCase(type)) { - return new RemoteIndexHTTPClient(); - } - throw new IllegalArgumentException("Unsupported RemoteIndexClient type: " + type); + // Default to HTTP client + public static RemoteIndexClient getRemoteIndexClient() { + return new RemoteIndexHTTPClient(); } } diff --git a/src/main/java/org/opensearch/knn/index/remote/RemoteIndexHTTPClient.java b/src/main/java/org/opensearch/knn/index/remote/RemoteIndexHTTPClient.java index 80bf8e78d8..cda7909d15 100644 --- a/src/main/java/org/opensearch/knn/index/remote/RemoteIndexHTTPClient.java +++ b/src/main/java/org/opensearch/knn/index/remote/RemoteIndexHTTPClient.java @@ -17,12 +17,13 @@ import org.apache.hc.core5.http.io.entity.EntityUtils; import org.apache.hc.core5.http.io.entity.StringEntity; import org.opensearch.common.settings.Settings; +import org.opensearch.common.xcontent.LoggingDeprecationHandler; import org.opensearch.common.xcontent.json.JsonXContent; import org.opensearch.core.common.settings.SecureString; -import org.opensearch.core.xcontent.DeprecationHandler; import org.opensearch.core.xcontent.NamedXContentRegistry; +import org.opensearch.core.xcontent.ToXContentObject; +import org.opensearch.core.xcontent.XContentBuilder; import org.opensearch.core.xcontent.XContentParser; -import org.opensearch.knn.common.KNNConstants; import org.opensearch.knn.index.KNNSettings; import org.opensearch.knn.index.codec.nativeindex.remote.RemoteStatusResponse; @@ -32,13 +33,12 @@ import java.nio.charset.StandardCharsets; import java.security.AccessController; import java.security.PrivilegedExceptionAction; -import java.util.Map; import static org.apache.hc.core5.http.HttpStatus.SC_OK; -import static org.opensearch.knn.common.KNNConstants.JOB_ID; import static org.opensearch.knn.index.KNNSettings.KNN_REMOTE_BUILD_CLIENT_PASSWORD_SETTING; import static org.opensearch.knn.index.KNNSettings.KNN_REMOTE_BUILD_CLIENT_USERNAME_SETTING; import static org.opensearch.knn.index.KNNSettings.KNN_REMOTE_BUILD_SERVICE_ENDPOINT_SETTING; +import static org.opensearch.knn.index.remote.KNNRemoteConstants.BUILD_ENDPOINT; /** * Class to handle all interactions with the remote vector build service. @@ -46,7 +46,8 @@ */ @Log4j2 public class RemoteIndexHTTPClient implements RemoteIndexClient, Closeable { - public static final String BASIC_PREFIX = "Basic "; + private static final String BASIC_PREFIX = "Basic "; + private static volatile String authHeader = null; private final String endpoint; @@ -81,10 +82,7 @@ public RemoteIndexHTTPClient() { */ @Override public RemoteBuildResponse submitVectorBuild(RemoteBuildRequest remoteBuildRequest) throws IOException { - assert (remoteBuildRequest instanceof HTTPRemoteBuildRequest); - HTTPRemoteBuildRequest request = (HTTPRemoteBuildRequest) remoteBuildRequest; - HttpPost buildRequest = getHttpPost(request); - + HttpPost buildRequest = getHttpPost(toJson(remoteBuildRequest)); try { String response = AccessController.doPrivileged( (PrivilegedExceptionAction) () -> getHttpClient().execute(buildRequest, body -> { @@ -94,15 +92,12 @@ public RemoteBuildResponse submitVectorBuild(RemoteBuildRequest remoteBuildReque return EntityUtils.toString(body.getEntity()); }) ); - - if (response == null || response.isEmpty()) { - throw new IOException("Received success status code but response is null or empty."); - } - String jobId = getValueFromResponse(response, JOB_ID); - if (jobId == null || jobId.isEmpty()) { - throw new IOException("Received success status code but " + JOB_ID + " is null or empty."); - } - return new RemoteBuildResponse(jobId); + XContentParser parser = JsonXContent.jsonXContent.createParser( + NamedXContentRegistry.EMPTY, + LoggingDeprecationHandler.INSTANCE, + response + ); + return RemoteBuildResponse.fromXContent(parser); } catch (Exception e) { throw new IOException("Failed to execute HTTP request", e); } @@ -110,14 +105,14 @@ public RemoteBuildResponse submitVectorBuild(RemoteBuildRequest remoteBuildReque /** * Helper method to form the HttpPost request from the HTTPRemoteBuildRequest - * @param request HTTPRemoteBuildRequest to be submitted + * @param jsonRequest JSON converted request body to be submitted * @return HttpPost request to be submitted * @throws IOException if the request cannot be formed */ - private HttpPost getHttpPost(HTTPRemoteBuildRequest request) throws IOException { - HttpPost buildRequest = new HttpPost(URI.create(endpoint) + KNNConstants.BUILD_ENDPOINT); + private HttpPost getHttpPost(String jsonRequest) { + HttpPost buildRequest = new HttpPost(URI.create(endpoint) + BUILD_ENDPOINT); buildRequest.setHeader(HttpHeaders.CONTENT_TYPE, ContentType.APPLICATION_JSON.toString()); - buildRequest.setEntity(new StringEntity(request.toJson())); + buildRequest.setEntity(new StringEntity(jsonRequest, ContentType.APPLICATION_JSON)); if (authHeader != null) { buildRequest.setHeader(HttpHeaders.AUTHORIZATION, authHeader); } @@ -135,28 +130,15 @@ public RemoteStatusResponse awaitVectorBuild(RemoteBuildResponse remoteBuildResp } /** - * Given a JSON response string, get a value for a specific key. Converts json {@literal } to Java null. - * @param responseBody The response to read - * @param key The key to lookup - * @return The value for the key - */ - static String getValueFromResponse(String responseBody, String key) throws IOException { - try ( - XContentParser parser = JsonXContent.jsonXContent.createParser( - NamedXContentRegistry.EMPTY, - DeprecationHandler.THROW_UNSUPPORTED_OPERATION, - responseBody - ) - ) { - Map responseMap = parser.map(); - if (responseMap.containsKey(key)) { - Object value = responseMap.get(key); - if (value == null) { - return null; - } - return value.toString(); - } - throw new IllegalArgumentException("Key " + key + " not found in response"); + * Convert the RemoteBuildRequest object to a JSON object for this specific HTTP implementation. + * @param object RemoteBuildRequest with parameters + * @return JSON String representation of the request body + * @throws IOException if the request cannot be converted to JSON + */ + private String toJson(ToXContentObject object) throws IOException { + try (XContentBuilder builder = JsonXContent.contentBuilder()) { + object.toXContent(builder, ToXContentObject.EMPTY_PARAMS); + return builder.toString(); } } diff --git a/src/test/java/org/opensearch/knn/index/codec/nativeindex/remote/DefaultVectorRepositoryAccessorTests.java b/src/test/java/org/opensearch/knn/index/codec/nativeindex/remote/DefaultVectorRepositoryAccessorTests.java index 2eb3510e8a..e401603eab 100644 --- a/src/test/java/org/opensearch/knn/index/codec/nativeindex/remote/DefaultVectorRepositoryAccessorTests.java +++ b/src/test/java/org/opensearch/knn/index/codec/nativeindex/remote/DefaultVectorRepositoryAccessorTests.java @@ -35,8 +35,8 @@ import static org.mockito.Mockito.mock; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; -import static org.opensearch.knn.index.codec.nativeindex.remote.RemoteIndexBuildStrategy.DOC_ID_FILE_EXTENSION; -import static org.opensearch.knn.index.codec.nativeindex.remote.RemoteIndexBuildStrategy.VECTOR_BLOB_FILE_EXTENSION; +import static org.opensearch.knn.index.remote.KNNRemoteConstants.DOC_ID_FILE_EXTENSION; +import static org.opensearch.knn.index.remote.KNNRemoteConstants.VECTOR_BLOB_FILE_EXTENSION; import static org.opensearch.knn.index.codec.util.KNNCodecUtil.initializeVectorValues; public class DefaultVectorRepositoryAccessorTests extends RemoteIndexBuildTests { diff --git a/src/test/java/org/opensearch/knn/index/remote/RemoteIndexHTTPClientTests.java b/src/test/java/org/opensearch/knn/index/remote/RemoteIndexHTTPClientTests.java index e86a299ff7..d4d46b3435 100644 --- a/src/test/java/org/opensearch/knn/index/remote/RemoteIndexHTTPClientTests.java +++ b/src/test/java/org/opensearch/knn/index/remote/RemoteIndexHTTPClientTests.java @@ -27,6 +27,8 @@ import org.opensearch.common.xcontent.json.JsonXContent; import org.opensearch.core.xcontent.DeprecationHandler; import org.opensearch.core.xcontent.NamedXContentRegistry; +import org.opensearch.core.xcontent.ToXContentObject; +import org.opensearch.core.xcontent.XContentBuilder; import org.opensearch.core.xcontent.XContentParser; import org.opensearch.index.IndexSettings; import org.opensearch.knn.index.KNNSettings; @@ -50,20 +52,29 @@ import static org.mockito.Mockito.mock; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; -import static org.opensearch.knn.common.KNNConstants.*; +import static org.opensearch.knn.common.KNNConstants.ENCODER_FLAT; +import static org.opensearch.knn.common.KNNConstants.FAISS_NAME; +import static org.opensearch.knn.common.KNNConstants.INDEX_DESCRIPTION_PARAMETER; +import static org.opensearch.knn.common.KNNConstants.METHOD_HNSW; +import static org.opensearch.knn.common.KNNConstants.METHOD_PARAMETER_EF_CONSTRUCTION; +import static org.opensearch.knn.common.KNNConstants.METHOD_PARAMETER_EF_SEARCH; +import static org.opensearch.knn.common.KNNConstants.NAME; +import static org.opensearch.knn.common.KNNConstants.PARAMETERS; import static org.opensearch.knn.common.KNNConstants.SPACE_TYPE; +import static org.opensearch.knn.common.KNNConstants.VECTOR_DATA_TYPE_FIELD; import static org.opensearch.knn.index.KNNSettings.KNN_REMOTE_BUILD_CLIENT_PASSWORD_SETTING; import static org.opensearch.knn.index.KNNSettings.KNN_REMOTE_BUILD_CLIENT_USERNAME_SETTING; import static org.opensearch.knn.index.KNNSettings.KNN_REMOTE_BUILD_SERVICE_ENDPOINT_SETTING; import static org.opensearch.knn.index.SpaceType.L2; import static org.opensearch.knn.index.VectorDataType.FLOAT; -import static org.opensearch.knn.index.codec.nativeindex.remote.RemoteIndexBuildStrategy.DOC_ID_FILE_EXTENSION; -import static org.opensearch.knn.index.codec.nativeindex.remote.RemoteIndexBuildStrategy.VECTOR_BLOB_FILE_EXTENSION; import static org.opensearch.knn.index.engine.faiss.Faiss.getMFromIndexDescription; +import static org.opensearch.knn.index.remote.KNNRemoteConstants.BUCKET; +import static org.opensearch.knn.index.remote.KNNRemoteConstants.BUILD_ENDPOINT; +import static org.opensearch.knn.index.remote.KNNRemoteConstants.DOC_ID_FILE_EXTENSION; +import static org.opensearch.knn.index.remote.KNNRemoteConstants.S3; +import static org.opensearch.knn.index.remote.KNNRemoteConstants.VECTOR_BLOB_FILE_EXTENSION; public class RemoteIndexHTTPClientTests extends OpenSearchSingleNodeTestCase { - - public static final String S3 = "s3"; public static final String TEST_BUCKET = "test-bucket"; public static final String TEST_CLUSTER = "test-cluster"; public static final String MOCK_JOB_ID_RESPONSE = "{\"job_id\": \"job-1739930402\"}"; @@ -72,6 +83,7 @@ public class RemoteIndexHTTPClientTests extends OpenSearchSingleNodeTestCase { public static final String MOCK_ENDPOINT = "https://mock-build-service.com"; public static final String USERNAME = "username"; public static final String PASSWORD = "password"; + @Mock protected ClusterService clusterService; @@ -93,25 +105,15 @@ public void testGetAndCloseHttpclient_success() throws IOException { client.close(); } - public void testGetValueFromResponse() throws IOException { - String jobID = "{\"job_id\": \"job-1739930402\"}"; - assertEquals("job-1739930402", RemoteIndexHTTPClient.getValueFromResponse(jobID, JOB_ID)); - String failedIndexBuild = "{" - + "\"task_status\":\"FAILED_INDEX_BUILD\"," - + "\"error_message\":\"Index build process interrupted.\"," - + "\"index_path\": null" - + "}"; - String error = RemoteIndexHTTPClient.getValueFromResponse(failedIndexBuild, ERROR_MESSAGE); - assertEquals("Index build process interrupted.", error); - assertNull(RemoteIndexHTTPClient.getValueFromResponse(failedIndexBuild, INDEX_PATH)); - } - public void testGetMFromIndexDescription() { assertEquals(16, getMFromIndexDescription("HNSW16,Flat")); assertEquals(8, getMFromIndexDescription("HNSW8,SQ")); assertThrows(IllegalArgumentException.class, () -> getMFromIndexDescription("Invalid description")); } + /** + * Test the construction of the build request by comparing it to an explicitly created JSON object. + */ public void testBuildRequest() { RepositoryMetadata metadata = createTestRepositoryMetadata(); KNNSettings knnSettingsMock = mock(KNNSettings.class); @@ -125,12 +127,7 @@ public void testBuildRequest() { BuildIndexParams indexInfo = createTestBuildIndexParams(); - HTTPRemoteBuildRequest request = RemoteBuildRequestBuilder.builder(HTTPRemoteBuildRequest.class) - .indexSettings(mockIndexSettings) - .indexInfo(indexInfo) - .repositoryMetadata(metadata) - .blobName(MOCK_BLOB_NAME) - .build(); + RemoteBuildRequest request = new RemoteBuildRequest(mockIndexSettings, indexInfo, metadata, MOCK_BLOB_NAME); assertEquals(S3, request.getRepositoryType()); assertEquals(TEST_BUCKET, request.getContainerName()); @@ -162,24 +159,61 @@ public void testBuildRequest() { + "}" + "}" + "}"; - XContentParser parser1 = JsonXContent.jsonXContent.createParser( + XContentParser expectedParser = JsonXContent.jsonXContent.createParser( NamedXContentRegistry.EMPTY, DeprecationHandler.THROW_UNSUPPORTED_OPERATION, expectedJson ); + Map expectedMap = expectedParser.map(); + + String jsonRequest; + try (XContentBuilder builder = JsonXContent.contentBuilder()) { + request.toXContent(builder, ToXContentObject.EMPTY_PARAMS); + jsonRequest = builder.toString(); + } - XContentParser parser2 = JsonXContent.jsonXContent.createParser( + XContentParser generatedParser = JsonXContent.jsonXContent.createParser( NamedXContentRegistry.EMPTY, DeprecationHandler.THROW_UNSUPPORTED_OPERATION, - request.toJson() + jsonRequest ); + Map generatedMap = generatedParser.map(); - assertEquals(parser1.map(), parser2.map()); + assertEquals(expectedMap, generatedMap); } catch (IOException e) { throw new RuntimeException(e); } } + public void testRemoteBuildResponseParsing() throws IOException { + String jsonResponse = "{\"job_id\":\"test-job-123\"}"; + + try ( + XContentParser parser = JsonXContent.jsonXContent.createParser( + NamedXContentRegistry.EMPTY, + DeprecationHandler.THROW_UNSUPPORTED_OPERATION, + jsonResponse + ) + ) { + RemoteBuildResponse response = RemoteBuildResponse.fromXContent(parser); + assertNotNull(response); + assertEquals("test-job-123", response.getJobId()); + } + } + + public void testRemoteBuildResponseParsingError() throws IOException { + String jsonResponse = "{\"error\":\"test-error\"}"; + try ( + XContentParser parser = JsonXContent.jsonXContent.createParser( + NamedXContentRegistry.EMPTY, + DeprecationHandler.THROW_UNSUPPORTED_OPERATION, + jsonResponse + ) + ) { + assertThrows(IOException.class, () -> RemoteBuildResponse.fromXContent(parser)); + } + } + public void testSubmitVectorBuild() throws IOException, URISyntaxException { RepositoryMetadata metadata = createTestRepositoryMetadata(); KNNSettings knnSettingsMock = mock(KNNSettings.class); @@ -202,19 +236,17 @@ public void testSubmitVectorBuild() throws IOException, URISyntaxException { ); RemoteIndexHTTPClient client = new RemoteIndexHTTPClient(); - clientStaticMock.when(() -> RemoteIndexHTTPClient.getValueFromResponse(any(String.class), any(String.class))) - .thenCallRealMethod(); RemoteBuildResponse remoteBuildResponse = client.submitVectorBuild( - new HTTPRemoteBuildRequest(mockIndexSettings, buildIndexParams, metadata, MOCK_BLOB_NAME) + new RemoteBuildRequest(mockIndexSettings, buildIndexParams, metadata, MOCK_BLOB_NAME) ); - assertEquals(MOCK_JOB_ID, remoteBuildResponse.jobId()); + assertEquals(MOCK_JOB_ID, remoteBuildResponse.getJobId()); ArgumentCaptor requestCaptor = ArgumentCaptor.forClass(HttpPost.class); verify(mockHttpClient).execute(requestCaptor.capture(), any(HttpClientResponseHandler.class)); HttpPost capturedRequest = requestCaptor.getValue(); assertEquals(MOCK_ENDPOINT + BUILD_ENDPOINT, capturedRequest.getUri().toString()); - assert (!capturedRequest.containsHeader(HttpHeaders.AUTHORIZATION)); + assertFalse(capturedRequest.containsHeader(HttpHeaders.AUTHORIZATION)); } } } @@ -245,14 +277,12 @@ public void testSecureSettingsReloadAndException() throws IOException { BuildIndexParams buildIndexParams = createTestBuildIndexParams(); clientStaticMock.when(() -> RemoteIndexHTTPClient.reloadAuthHeader(any(Settings.class))).thenCallRealMethod(); - clientStaticMock.when(() -> RemoteIndexHTTPClient.getValueFromResponse(any(String.class), any(String.class))) - .thenCallRealMethod(); RemoteIndexHTTPClient client = new RemoteIndexHTTPClient(); RemoteIndexHTTPClient.reloadAuthHeader(settings); ArgumentCaptor requestCaptor = ArgumentCaptor.forClass(HttpPost.class); - client.submitVectorBuild(new HTTPRemoteBuildRequest(mockIndexSettings, buildIndexParams, metadata, MOCK_BLOB_NAME)); + client.submitVectorBuild(new RemoteBuildRequest(mockIndexSettings, buildIndexParams, metadata, MOCK_BLOB_NAME)); verify(mockHttpClient).execute(requestCaptor.capture(), any(HttpClientResponseHandler.class)); HttpPost capturedRequest = requestCaptor.getValue(); @@ -266,6 +296,7 @@ public void testSecureSettingsReloadAndException() throws IOException { throw new RuntimeException(e); } } + clearAuthHeader(); } // Utility methods to populate settings for build requests @@ -323,4 +354,10 @@ private void setupTestClusterSettings() { when(clusterService.getClusterSettings()).thenReturn(clusterSettings); KNNSettings.state().setClusterService(clusterService); } + + private void clearAuthHeader() { + final MockSecureSettings secureSettings = new MockSecureSettings(); + final Settings settings = Settings.builder().setSecureSettings(secureSettings).build(); + RemoteIndexHTTPClient.reloadAuthHeader(settings); + } }