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/src/test/java/com/risingwave/connector/sink/elasticsearch/EsSinkTest.java

@@ -21,6 +21,7 @@
 import com.google.common.collect.Lists;
 import com.risingwave.connector.EsSink;
 import com.risingwave.connector.EsSinkConfig;
+import com.risingwave.connector.RestHighLevelClientAdapter;
 import com.risingwave.connector.api.TableSchema;
 import com.risingwave.connector.api.sink.ArraySinkRow;
 import com.risingwave.proto.Data;
@@ -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;
@@ -74,12 +74,13 @@ public void testEsSink(ElasticsearchContainer container, String username, String
             fail(e.getMessage());
         }
 
-        RestHighLevelClient client = sink.getClient();
+        RestHighLevelClientAdapter client = sink.getClient();
         SearchRequest searchRequest = new SearchRequest("test");
         SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
         searchSourceBuilder.query(QueryBuilders.matchAllQuery());
         searchRequest.source(searchSourceBuilder);
-        SearchResponse searchResponse = client.search(searchRequest, RequestOptions.DEFAULT);
+        SearchResponse searchResponse =
+                (SearchResponse) client.search(searchRequest, RequestOptions.DEFAULT);
 
         SearchHits hits = searchResponse.getHits();
         assertEquals(2, hits.getHits().length);

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/BulkProcessorAdapter.java

@@ -0,0 +1,31 @@
+/*
+ * 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 com.risingwave.connector.api.sink.SinkRow;
+import java.util.concurrent.TimeUnit;
+
+public interface BulkProcessorAdapter {
+    public void addRow(SinkRow row, String indexName, RequestTracker requestTracker);
+
+    public void deleteRow(SinkRow row, String indexName, RequestTracker requestTracker);
+
+    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/BulkRequestConsumerFactory.java

@@ -22,8 +22,18 @@
 import org.elasticsearch.action.bulk.BulkResponse;
 
 /**
- * {@link BulkRequestConsumerFactory} is used to bridge incompatible Elasticsearch Java API calls
- * across different Elasticsearch versions.
+ * {@link ElasticBulkRequestConsumerFactory} is used to bridge incompatible Elasticsearch Java API
+ * calls across different Elasticsearch versions.
  */
-interface BulkRequestConsumerFactory
+interface ElasticBulkRequestConsumerFactory
         extends BiConsumer<BulkRequest, ActionListener<BulkResponse>> {}
+
+/**
+ * {@link ElasticBulkRequestConsumerFactory} is used to bridge incompatible Elasticsearch Java API
+ * calls across different Elasticsearch versions.
+ */
+interface OpensearchBulkRequestConsumerFactory
+        extends BiConsumer<
+                org.opensearch.action.bulk.BulkRequest,
+                org.opensearch.core.action.ActionListener<
+                        org.opensearch.action.bulk.BulkResponse>> {}

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

@@ -0,0 +1,147 @@
+/*
+ * 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 com.risingwave.connector.api.sink.SinkRow;
+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.bulk.BulkResponse;
+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 class BulkListener implements BulkProcessor.Listener {
+        private final RequestTracker requestTracker;
+
+        public BulkListener(RequestTracker requestTracker) {
+            this.requestTracker = requestTracker;
+        }
+
+        /** This method is called just before bulk is executed. */
+        @Override
+        public void beforeBulk(long executionId, BulkRequest request) {
+            LOG.debug("Sending bulk of {} actions to Elasticsearch.", request.numberOfActions());
+        }
+
+        /** This method is called after bulk execution. */
+        @Override
+        public void afterBulk(long executionId, BulkRequest request, 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);
+        }
+    }
+
+    public ElasticBulkProcessorAdapter(
+            RequestTracker requestTracker, ElasticRestHighLevelClientAdapter client) {
+        BulkProcessor.Builder builder =
+                BulkProcessor.builder(
+                        (ElasticBulkRequestConsumerFactory)
+                                (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();
+    }
+
+    @Override
+    public void flush() {
+        esBulkProcessor.flush();
+    }
+
+    @Override
+    public void awaitClose(long timeout, TimeUnit unit) throws InterruptedException {
+        esBulkProcessor.awaitClose(timeout, unit);
+    }
+
+    @Override
+    public void addRow(SinkRow row, String indexName, RequestTracker requestTracker) {
+        final String index = (String) row.get(0);
+        final String key = (String) row.get(1);
+        String doc = (String) row.get(2);
+
+        UpdateRequest updateRequest;
+        if (indexName != null) {
+            updateRequest = new UpdateRequest(indexName, "_doc", key).doc(doc, XContentType.JSON);
+        } else {
+            updateRequest = new UpdateRequest(indexName, "_doc", key).doc(doc, XContentType.JSON);
+        }
+        updateRequest.docAsUpsert(true);
+        requestTracker.addWriteTask();
+        this.esBulkProcessor.add(updateRequest);
+    }
+
+    @Override
+    public void deleteRow(SinkRow row, String indexName, RequestTracker requestTracker) {
+        final String index = (String) row.get(0);
+        final String key = (String) row.get(1);
+
+        DeleteRequest deleteRequest;
+        if (indexName != null) {
+            deleteRequest = new DeleteRequest(indexName, "_doc", key);
+        } else {
+            deleteRequest = new DeleteRequest(index, "_doc", key);
+        }
+        requestTracker.addWriteTask();
+        this.esBulkProcessor.add(deleteRequest);
+    }
+}