feat(sink): refator es sink to support opensearch

java/connector-node/risingwave-connector-service/src/main/java/com/risingwave/connector/SinkUtils.java

@@ -40,6 +40,7 @@ public static SinkFactory getSinkFactory(String sinkName) {
  case "jdbc":
  return new JDBCSinkFactory();
  case "elasticsearch":
+ case "opensearch":
  return new EsSinkFactory();
  case "cassandra":
  return new CassandraFactory();

java/connector-node/risingwave-connector-test/pom.xml

@@ -128,13 +128,13 @@
  <dependency>
  <groupId>com.fasterxml.jackson.core</groupId>
  <artifactId>jackson-databind</artifactId>
- <version>${jackson.version}</version>
+ <version>2.13.5</version>
  <scope>test</scope>
  </dependency>
  <dependency>
  <groupId>com.fasterxml.jackson.core</groupId>
  <artifactId>jackson-core</artifactId>
- <version>${jackson.version}</version>
+ <version>2.13.5</version>
  <scope>test</scope>
  </dependency>
 

java/connector-node/risingwave-connector-test/src/test/java/com/risingwave/connector/sink/elasticsearch/EsSinkTest.java

@@ -19,6 +19,7 @@
 
 import com.google.common.collect.Iterators;
 import com.google.common.collect.Lists;
+import com.risingwave.connector.ElasticRestHighLevelClientAdapter;
 import com.risingwave.connector.EsSink;
 import com.risingwave.connector.EsSinkConfig;
 import com.risingwave.connector.api.TableSchema;
@@ -31,7 +32,6 @@
 import org.elasticsearch.action.search.SearchRequest;
 import org.elasticsearch.action.search.SearchResponse;
 import org.elasticsearch.client.RequestOptions;
-import org.elasticsearch.client.RestHighLevelClient;
 import org.elasticsearch.index.query.QueryBuilders;
 import org.elasticsearch.search.SearchHit;
 import org.elasticsearch.search.SearchHits;
@@ -55,6 +55,7 @@ public void testEsSink(ElasticsearchContainer container, String username, String
  EsSink sink =
  new EsSink(
  new EsSinkConfig(container.getHttpHostAddress())
+ .setConnector("elasticsearch")
  .withIndex("test")
  .withDelimiter("$")
  .withUsername(username)
@@ -74,7 +75,8 @@ public void testEsSink(ElasticsearchContainer container, String username, String
  fail(e.getMessage());
  }
 
- RestHighLevelClient client = sink.getClient();
+ ElasticRestHighLevelClientAdapter client =
+ (ElasticRestHighLevelClientAdapter) sink.getClient();
  SearchRequest searchRequest = new SearchRequest("test");
  SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
  searchSourceBuilder.query(QueryBuilders.matchAllQuery());

java/connector-node/risingwave-sink-es-7/pom.xml

@@ -51,6 +51,14 @@
  <groupId>org.elasticsearch.client</groupId>
  <artifactId>elasticsearch-rest-high-level-client</artifactId>
  </dependency>
+ <dependency>
+ <groupId>org.opensearch</groupId>
+ <artifactId>opensearch</artifactId>
+ </dependency>
+ <dependency>
+ <groupId>org.opensearch.client</groupId>
+ <artifactId>opensearch-rest-high-level-client</artifactId>
+ </dependency>
  <dependency>
  <groupId>org.apache.httpcomponents</groupId>
  <artifactId>httpclient</artifactId>

java/connector-node/risingwave-sink-es-7/src/main/java/com/risingwave/connector/BulkRequestConsumerFactory.java → java/connector-node/risingwave-sink-es-7/src/main/java/com/risingwave/connector/BulkProcessorAdapter.java

@@ -16,14 +16,14 @@
 
 package com.risingwave.connector;
 
-import java.util.function.BiConsumer;
-import org.elasticsearch.action.ActionListener;
-import org.elasticsearch.action.bulk.BulkRequest;
-import org.elasticsearch.action.bulk.BulkResponse;
+import java.util.concurrent.TimeUnit;
 
-/**
- * {@link BulkRequestConsumerFactory} is used to bridge incompatible Elasticsearch Java API calls
- * across different Elasticsearch versions.
- */
-interface BulkRequestConsumerFactory
- extends BiConsumer<BulkRequest, ActionListener<BulkResponse>> {}
+public interface BulkProcessorAdapter {
+ public void addRow(String index, String key, String doc);
+
+ public void deleteRow(String index, String key);
+
+ public void flush();
+
+ public void awaitClose(long timeout, TimeUnit unit) throws InterruptedException;
+}

java/connector-node/risingwave-sink-es-7/src/main/java/com/risingwave/connector/ElasticBulkProcessorAdapter.java

@@ -0,0 +1,167 @@
+/*
+ * Copyright 2024 RisingWave Labs
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.risingwave.connector;
+
+import com.risingwave.connector.EsSink.RequestTracker;
+import java.util.concurrent.TimeUnit;
+import org.elasticsearch.action.bulk.BackoffPolicy;
+import org.elasticsearch.action.bulk.BulkProcessor;
+import org.elasticsearch.action.bulk.BulkRequest;
+import org.elasticsearch.action.delete.DeleteRequest;
+import org.elasticsearch.action.update.UpdateRequest;
+import org.elasticsearch.client.RequestOptions;
+import org.elasticsearch.common.unit.ByteSizeUnit;
+import org.elasticsearch.common.unit.ByteSizeValue;
+import org.elasticsearch.core.TimeValue;
+import org.elasticsearch.xcontent.XContentType;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+public class ElasticBulkProcessorAdapter implements BulkProcessorAdapter {
+ private static final Logger LOG = LoggerFactory.getLogger(EsSink.class);
+ BulkProcessor esBulkProcessor;
+ private final RequestTracker requestTracker;
+
+ public ElasticBulkProcessorAdapter(
+ RequestTracker requestTracker, ElasticRestHighLevelClientAdapter client) {
+ BulkProcessor.Builder builder =
+ BulkProcessor.builder(
+ (bulkRequest, bulkResponseActionListener) ->
+ client.bulkAsync(
+ bulkRequest,
+ RequestOptions.DEFAULT,
+ bulkResponseActionListener),
+ new BulkListener(requestTracker));
+ // Possible feature: move these to config
+ // execute the bulk every 10 000 requests
+ builder.setBulkActions(1000);
+ // flush the bulk every 5mb
+ builder.setBulkSize(new ByteSizeValue(5, ByteSizeUnit.MB));
+ // flush the bulk every 5 seconds whatever the number of requests
+ builder.setFlushInterval(TimeValue.timeValueSeconds(5));
+ // Set the number of concurrent requests
+ builder.setConcurrentRequests(1);
+ // Set a custom backoff policy which will initially wait for 100ms, increase exponentially
+ // and retries up to three times.
+ builder.setBackoffPolicy(
+ BackoffPolicy.exponentialBackoff(TimeValue.timeValueMillis(100), 3));
+ this.esBulkProcessor = builder.build();
+ this.requestTracker = requestTracker;
+ }
+
+ @Override
+ public void flush() {
+ esBulkProcessor.flush();
+ }
+
+ @Override
+ public void awaitClose(long timeout, TimeUnit unit) throws InterruptedException {
+ esBulkProcessor.awaitClose(timeout, unit);
+ }
+
+ @Override
+ public void addRow(String index, String key, String doc) {
+ UpdateRequest updateRequest;
+ updateRequest = new UpdateRequest(index, "_doc", key).doc(doc, XContentType.JSON);
+ updateRequest.docAsUpsert(true);
+ this.requestTracker.addWriteTask();
+ this.esBulkProcessor.add(updateRequest);
+ }
+
+ @Override
+ public void deleteRow(String index, String key) {
+ DeleteRequest deleteRequest;
+ deleteRequest = new DeleteRequest(index, "_doc", key);
+ this.requestTracker.addWriteTask();
+ this.esBulkProcessor.add(deleteRequest);
+ }
+}
+
+class BulkListener
+ implements org.elasticsearch.action.bulk.BulkProcessor.Listener,
+ org.opensearch.action.bulk.BulkProcessor.Listener {
+ private static final Logger LOG = LoggerFactory.getLogger(EsSink.class);
+ private final RequestTracker requestTracker;
+
+ public BulkListener(RequestTracker requestTracker) {
+ this.requestTracker = requestTracker;
+ }
+
+ @Override
+ public void beforeBulk(long executionId, org.elasticsearch.action.bulk.BulkRequest request) {
+ LOG.debug("Sending bulk of {} actions to Elasticsearch.", request.numberOfActions());
+ }
+
+ @Override
+ public void afterBulk(
+ long executionId,
+ org.elasticsearch.action.bulk.BulkRequest request,
+ org.elasticsearch.action.bulk.BulkResponse response) {
+ if (response.hasFailures()) {
+ String errMessage =
+ String.format(
+ "Bulk of %d actions failed. Failure: %s",
+ request.numberOfActions(), response.buildFailureMessage());
+ this.requestTracker.addErrResult(errMessage);
+ } else {
+ this.requestTracker.addOkResult(request.numberOfActions());
+ LOG.debug("Sent bulk of {} actions to Elasticsearch.", request.numberOfActions());
+ }
+ }
+
+ /** This method is called when the bulk failed and raised a Throwable */
+ @Override
+ public void afterBulk(long executionId, BulkRequest request, Throwable failure) {
+ String errMessage =
+ String.format(
+ "Bulk of %d actions failed. Failure: %s",
+ request.numberOfActions(), failure.getMessage());
+ this.requestTracker.addErrResult(errMessage);
+ }
+
+ @Override
+ public void beforeBulk(long executionId, org.opensearch.action.bulk.BulkRequest request) {
+ LOG.debug("Sending bulk of {} actions to Opensearch.", request.numberOfActions());
+ }
+
+ @Override
+ public void afterBulk(
+ long executionId,
+ org.opensearch.action.bulk.BulkRequest request,
+ org.opensearch.action.bulk.BulkResponse response) {
+ if (response.hasFailures()) {
+ String errMessage =
+ String.format(
+ "Bulk of %d actions failed. Failure: %s",
+ request.numberOfActions(), response.buildFailureMessage());
+ this.requestTracker.addErrResult(errMessage);
+ } else {
+ this.requestTracker.addOkResult(request.numberOfActions());
+ LOG.debug("Sent bulk of {} actions to Opensearch.", request.numberOfActions());
+ }
+ }
+
+ @Override
+ public void afterBulk(
+ long executionId, org.opensearch.action.bulk.BulkRequest request, Throwable failure) {
+ String errMessage =
+ String.format(
+ "Bulk of %d actions failed. Failure: %s",
+ request.numberOfActions(), failure.getMessage());
+ this.requestTracker.addErrResult(errMessage);
+ }
+}

java/connector-node/risingwave-sink-es-7/src/main/java/com/risingwave/connector/ElasticRestHighLevelClientAdapter.java

@@ -0,0 +1,89 @@
+/*
+ * Copyright 2024 RisingWave Labs
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.risingwave.connector;
+
+import java.io.IOException;
+import org.apache.http.HttpHost;
+import org.apache.http.auth.AuthScope;
+import org.apache.http.auth.UsernamePasswordCredentials;
+import org.apache.http.client.CredentialsProvider;
+import org.apache.http.impl.client.BasicCredentialsProvider;
+import org.elasticsearch.action.ActionListener;
+import org.elasticsearch.action.bulk.BulkRequest;
+import org.elasticsearch.action.bulk.BulkResponse;
+import org.elasticsearch.action.search.SearchRequest;
+import org.elasticsearch.action.search.SearchResponse;
+import org.elasticsearch.client.Cancellable;
+import org.elasticsearch.client.RequestOptions;
+import org.elasticsearch.client.RestClient;
+import org.elasticsearch.client.RestClientBuilder;
+import org.elasticsearch.client.RestHighLevelClient;
+import org.elasticsearch.client.RestHighLevelClientBuilder;
+
+public class ElasticRestHighLevelClientAdapter implements RestHighLevelClientAdapter {
+ RestHighLevelClient esClient;
+
+ private static RestClientBuilder configureRestClientBuilder(
+ RestClientBuilder builder, EsSinkConfig config) {
+ // Possible config:
+ // 1. Connection path prefix
+ // 2. Username and password
+ if (config.getPassword() != null && config.getUsername() != null) {
+ final CredentialsProvider credentialsProvider = new BasicCredentialsProvider();
+ credentialsProvider.setCredentials(
+ AuthScope.ANY,
+ new UsernamePasswordCredentials(config.getUsername(), config.getPassword()));
+ builder.setHttpClientConfigCallback(
+ httpClientBuilder ->
+ httpClientBuilder.setDefaultCredentialsProvider(credentialsProvider));
+ }
+ // 3. Timeout
+ return builder;
+ }
+
+ public ElasticRestHighLevelClientAdapter(HttpHost host, EsSinkConfig config) {
+ this.esClient =
+ new RestHighLevelClientBuilder(
+ configureRestClientBuilder(RestClient.builder(host), config)
+ .build())
+ .setApiCompatibilityMode(true)
+ .build();
+ }
+
+ @Override
+ public void close() throws IOException {
+ esClient.close();
+ }
+
+ public boolean ping(RequestOptions options) throws IOException {
+ boolean flag = esClient.ping(options);
+ return flag;
+ }
+
+ public Cancellable bulkAsync(
+ BulkRequest bulkRequest,
+ RequestOptions options,
+ ActionListener<BulkResponse> listener) {
+ Cancellable cancellable = esClient.bulkAsync(bulkRequest, options, listener);
+ return cancellable;
+ }
+
+ public SearchResponse search(SearchRequest searchRequest, RequestOptions options)
+ throws IOException {
+ return this.esClient.search(searchRequest, options);
+ }
+}