-
Notifications
You must be signed in to change notification settings - Fork 190
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
* ENH: support es 6 for bulk API Signed-off-by: George Chen <[email protected]>
- Loading branch information
1 parent
808e239
commit 07f00c2
Showing
12 changed files
with
363 additions
and
4 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
33 changes: 33 additions & 0 deletions
33
...src/main/java/org/opensearch/dataprepper/plugins/sink/opensearch/DistributionVersion.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,33 @@ | ||
| package org.opensearch.dataprepper.plugins.sink.opensearch; | ||
|
|
||
| import java.util.Arrays; | ||
| import java.util.Map; | ||
| import java.util.stream.Collectors; | ||
|
|
||
| public enum DistributionVersion { | ||
| ES6("es6"), | ||
| DEFAULT("default"); | ||
|
|
||
| private static final Map<String, DistributionVersion> VERSION_MAP = Arrays.stream(DistributionVersion.values()) | ||
| .collect(Collectors.toMap( | ||
| value -> value.version, | ||
| value -> value | ||
| )); | ||
|
|
||
| private final String version; | ||
|
|
||
| DistributionVersion(final String version) { | ||
| this.version = version; | ||
| } | ||
|
|
||
| public static DistributionVersion fromTypeName(final String version) { | ||
| if (!VERSION_MAP.containsKey(version)) { | ||
| throw new IllegalArgumentException(String.format("Invalid distribution_version value: %s", version)); | ||
| } | ||
| return VERSION_MAP.get(version); | ||
| } | ||
|
|
||
| public String getVersion() { | ||
| return version; | ||
| } | ||
| } |
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
8 changes: 8 additions & 0 deletions
8
...src/main/java/org/opensearch/dataprepper/plugins/sink/opensearch/bulk/BulkApiWrapper.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,8 @@ | ||
| package org.opensearch.dataprepper.plugins.sink.opensearch.bulk; | ||
|
|
||
| import org.opensearch.client.opensearch.core.BulkRequest; | ||
| import org.opensearch.client.opensearch.core.BulkResponse; | ||
|
|
||
| public interface BulkApiWrapper { | ||
| BulkResponse bulk(BulkRequest request) throws Exception; | ||
| } |
16 changes: 16 additions & 0 deletions
16
...n/java/org/opensearch/dataprepper/plugins/sink/opensearch/bulk/BulkApiWrapperFactory.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,16 @@ | ||
| package org.opensearch.dataprepper.plugins.sink.opensearch.bulk; | ||
|
|
||
| import org.opensearch.client.opensearch.OpenSearchClient; | ||
| import org.opensearch.dataprepper.plugins.sink.opensearch.DistributionVersion; | ||
| import org.opensearch.dataprepper.plugins.sink.opensearch.index.IndexConfiguration; | ||
|
|
||
| public class BulkApiWrapperFactory { | ||
| public static BulkApiWrapper getWrapper(final IndexConfiguration indexConfiguration, | ||
| final OpenSearchClient openSearchClient) { | ||
| if (DistributionVersion.ES6.equals(indexConfiguration.getDistributionVersion())) { | ||
| return new Es6BulkApiWrapper(openSearchClient); | ||
| } else { | ||
| return new OpenSearchDefaultBulkApiWrapper(openSearchClient); | ||
| } | ||
| } | ||
| } |
86 changes: 86 additions & 0 deletions
86
.../main/java/org/opensearch/dataprepper/plugins/sink/opensearch/bulk/Es6BulkApiWrapper.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,86 @@ | ||
| package org.opensearch.dataprepper.plugins.sink.opensearch.bulk; | ||
|
|
||
| import org.opensearch.client.opensearch.OpenSearchClient; | ||
| import org.opensearch.client.opensearch._types.ErrorResponse; | ||
| import org.opensearch.client.opensearch._types.OpenSearchException; | ||
| import org.opensearch.client.opensearch.core.BulkRequest; | ||
| import org.opensearch.client.opensearch.core.BulkResponse; | ||
| import org.opensearch.client.transport.JsonEndpoint; | ||
| import org.opensearch.client.transport.endpoints.SimpleEndpoint; | ||
| import org.opensearch.client.util.ApiTypeHelper; | ||
|
|
||
| import javax.ws.rs.HttpMethod; | ||
| import java.io.IOException; | ||
| import java.util.HashMap; | ||
| import java.util.Map; | ||
| import java.util.stream.Collectors; | ||
|
|
||
| public class Es6BulkApiWrapper implements BulkApiWrapper { | ||
| private final OpenSearchClient openSearchClient; | ||
|
|
||
| public Es6BulkApiWrapper(final OpenSearchClient openSearchClient) { | ||
| this.openSearchClient = openSearchClient; | ||
| } | ||
|
|
||
| @Override | ||
| public BulkResponse bulk(BulkRequest request) throws IOException, OpenSearchException { | ||
| final JsonEndpoint<BulkRequest, BulkResponse, ErrorResponse> endpoint = es6BulkEndpoint(request); | ||
| return openSearchClient._transport().performRequest(request, endpoint, openSearchClient._transportOptions()); | ||
| } | ||
|
|
||
| private JsonEndpoint<BulkRequest, BulkResponse, ErrorResponse> es6BulkEndpoint(BulkRequest bulkRequest) { | ||
| return new SimpleEndpoint<>( | ||
| // Request method | ||
| request -> HttpMethod.POST, | ||
|
|
||
| // Request path | ||
| request -> { | ||
| final String index = request.index(); | ||
| if (index == null) { | ||
| throw new IllegalArgumentException("Bulk request index cannot be missing"); | ||
| } | ||
| StringBuilder buf = new StringBuilder(); | ||
| buf.append("/"); | ||
| SimpleEndpoint.pathEncode(index, buf); | ||
| buf.append("/_doc"); | ||
| buf.append("/_bulk"); | ||
| return buf.toString(); | ||
| }, | ||
|
|
||
| // Request parameters | ||
| request -> { | ||
| Map<String, String> params = new HashMap<>(); | ||
| if (request.pipeline() != null) { | ||
| params.put("pipeline", request.pipeline()); | ||
| } | ||
| if (request.routing() != null) { | ||
| params.put("routing", request.routing()); | ||
| } | ||
| if (request.requireAlias() != null) { | ||
| params.put("require_alias", String.valueOf(request.requireAlias())); | ||
| } | ||
| if (request.refresh() != null) { | ||
| params.put("refresh", request.refresh().jsonValue()); | ||
| } | ||
| if (request.waitForActiveShards() != null) { | ||
| params.put("wait_for_active_shards", request.waitForActiveShards()._toJsonString()); | ||
| } | ||
| if (request.source() != null) { | ||
| params.put("_source", request.source()._toJsonString()); | ||
| } | ||
| if (ApiTypeHelper.isDefined(request.sourceExcludes())) { | ||
| params.put("_source_excludes", | ||
| request.sourceExcludes().stream().map(v -> v).collect(Collectors.joining(","))); | ||
| } | ||
| if (ApiTypeHelper.isDefined(request.sourceIncludes())) { | ||
| params.put("_source_includes", | ||
| request.sourceIncludes().stream().map(v -> v).collect(Collectors.joining(","))); | ||
| } | ||
| if (request.timeout() != null) { | ||
| params.put("timeout", request.timeout()._toJsonString()); | ||
| } | ||
| return params; | ||
|
|
||
| }, SimpleEndpoint.emptyMap(), true, BulkResponse._DESERIALIZER); | ||
| } | ||
| } |
20 changes: 20 additions & 0 deletions
20
.../opensearch/dataprepper/plugins/sink/opensearch/bulk/OpenSearchDefaultBulkApiWrapper.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,20 @@ | ||
| package org.opensearch.dataprepper.plugins.sink.opensearch.bulk; | ||
|
|
||
| import org.opensearch.client.opensearch.OpenSearchClient; | ||
| import org.opensearch.client.opensearch.core.BulkRequest; | ||
| import org.opensearch.client.opensearch.core.BulkResponse; | ||
|
|
||
| import java.io.IOException; | ||
|
|
||
| public class OpenSearchDefaultBulkApiWrapper implements BulkApiWrapper { | ||
| private final OpenSearchClient openSearchClient; | ||
|
|
||
| public OpenSearchDefaultBulkApiWrapper(final OpenSearchClient openSearchClient) { | ||
| this.openSearchClient = openSearchClient; | ||
| } | ||
|
|
||
| @Override | ||
| public BulkResponse bulk(BulkRequest request) throws IOException { | ||
| return openSearchClient.bulk(request); | ||
| } | ||
| } |
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
36 changes: 36 additions & 0 deletions
36
...va/org/opensearch/dataprepper/plugins/sink/opensearch/bulk/BulkApiWrapperFactoryTest.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,36 @@ | ||
| package org.opensearch.dataprepper.plugins.sink.opensearch.bulk; | ||
|
|
||
| import org.junit.jupiter.api.Test; | ||
| import org.junit.jupiter.api.extension.ExtendWith; | ||
| import org.mockito.Mock; | ||
| import org.mockito.junit.jupiter.MockitoExtension; | ||
| import org.opensearch.client.opensearch.OpenSearchClient; | ||
| import org.opensearch.dataprepper.plugins.sink.opensearch.DistributionVersion; | ||
| import org.opensearch.dataprepper.plugins.sink.opensearch.index.IndexConfiguration; | ||
|
|
||
| import static org.hamcrest.CoreMatchers.instanceOf; | ||
| import static org.hamcrest.MatcherAssert.assertThat; | ||
| import static org.mockito.Mockito.when; | ||
|
|
||
| @ExtendWith(MockitoExtension.class) | ||
| class BulkApiWrapperFactoryTest { | ||
| @Mock | ||
| private IndexConfiguration indexConfiguration; | ||
|
|
||
| @Mock | ||
| private OpenSearchClient openSearchClient; | ||
|
|
||
| @Test | ||
| void testGetEs6BulkApiWrapper() { | ||
| when(indexConfiguration.getDistributionVersion()).thenReturn(DistributionVersion.ES6); | ||
| assertThat(BulkApiWrapperFactory.getWrapper(indexConfiguration, openSearchClient), | ||
| instanceOf(Es6BulkApiWrapper.class)); | ||
| } | ||
|
|
||
| @Test | ||
| void testGetOpenSearchDefaultBulkApiWrapper() { | ||
| when(indexConfiguration.getDistributionVersion()).thenReturn(DistributionVersion.DEFAULT); | ||
| assertThat(BulkApiWrapperFactory.getWrapper(indexConfiguration, openSearchClient), | ||
| instanceOf(OpenSearchDefaultBulkApiWrapper.class)); | ||
| } | ||
| } |
Oops, something went wrong.