WIP: [BitSail][Connector] Support ElasticSearch Source connector

bitsail-connectors/connector-elasticsearch/pom.xml

@@ -63,6 +63,13 @@
             <version>${revision}</version>
             <scope>test</scope>
         </dependency>
+
+        <dependency>
+            <groupId>com.bytedance.bitsail</groupId>
+            <artifactId>bitsail-connector-print</artifactId>
+            <version>${revision}</version>
+            <scope>test</scope>
+        </dependency>
     </dependencies>
 
 </project>

bitsail-connectors/connector-elasticsearch/src/main/java/com/bytedance/bitsail/connector/elasticsearch/base/EsConstants.java

@@ -53,4 +53,6 @@ public class EsConstants {
   public static final String DEFAULT_OPERATION_TYPE = OPERATION_TYPE_INDEX;
 
   public static final String ES_CONNECTOR_NAME = "elasticsearch";
+
+  public static final String SPLIT_COMMA = ",\\s*";
 }

bitsail-connectors/connector-elasticsearch/src/main/java/com/bytedance/bitsail/connector/elasticsearch/error/ElasticsearchErrorCode.java

@@ -0,0 +1,53 @@
+/*
+ * Copyright 2022-2023 Bytedance Ltd. and/or its affiliates.
+ *
+ * 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.bytedance.bitsail.connector.elasticsearch.error;
+
+import com.bytedance.bitsail.common.exception.ErrorCode;
+
+public enum ElasticsearchErrorCode implements ErrorCode {
+
+  REQUIRED_VALUE("Elasticsearch-00", "The configuration file is lack of necessary options"),
+  VALID_INDEX_FAILED("Elasticsearch-01", "Try to connect index failed."),
+  NOT_SUPPORT_SPLIT_STRATEGY("Elasticsearch-02", "Split strategy not support yet."),
+  FETCH_DATA_FAILED("Elasticsearch-03", "Fetch data from elasticsearch cluster failed."),
+  DESERIALIZE_FAILED("Elasticsearch-04", "Deserialize data from elasticsearch cluster failed.");
+
+  private final String code;
+
+  private final String describe;
+
+  ElasticsearchErrorCode(String code, String describe) {
+    this.code = code;
+    this.describe = describe;
+  }
+
+  @Override
+  public String getCode() {
+    return code;
+  }
+
+  @Override
+  public String getDescription() {
+    return describe;
+  }
+
+  @Override
+  public String toString() {
+    return String.format("Code:[%s], Describe:[%s]", this.code,
+        this.describe);
+  }
+}

bitsail-connectors/connector-elasticsearch/src/main/java/com/bytedance/bitsail/connector/elasticsearch/option/ElasticsearchReaderOptions.java

@@ -0,0 +1,62 @@
+/*
+ * Copyright 2022-2023 Bytedance Ltd. and/or its affiliates.
+ *
+ * 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.bytedance.bitsail.connector.elasticsearch.option;
+
+import com.bytedance.bitsail.common.option.ConfigOption;
+import com.bytedance.bitsail.common.option.ReaderOptions;
+
+import com.alibaba.fastjson.TypeReference;
+
+import java.util.List;
+
+import static com.bytedance.bitsail.common.option.ConfigOptions.key;
+
+public interface ElasticsearchReaderOptions extends ReaderOptions.BaseReaderOptions {
+
+  ConfigOption<List<String>> ES_HOSTS =
+      key(ReaderOptions.READER_PREFIX + "es_hosts")
+          .onlyReference(new TypeReference<List<String>>() {
+          });
+
+  ConfigOption<String> ES_INDEX =
+      key(ReaderOptions.READER_PREFIX + "es_index")
+          .noDefaultValue(String.class);
+
+  ConfigOption<Integer> CONNECTION_REQUEST_TIMEOUT_MS =
+      key(ReaderOptions.READER_PREFIX + "connection_request_timeout_ms")
+          .defaultValue(10000);
+
+  ConfigOption<Integer> CONNECTION_TIMEOUT_MS =
+      key(ReaderOptions.READER_PREFIX + "connection_timeout_ms")
+          .defaultValue(10000);
+
+  ConfigOption<Integer> SOCKET_TIMEOUT_MS =
+      key(ReaderOptions.READER_PREFIX + "socket_timeout_ms")
+          .defaultValue(60000);
+
+  ConfigOption<String> SCROLL_TIME =
+      key(ReaderOptions.READER_PREFIX + "scroll_time")
+          .defaultValue("1m");
+
+  ConfigOption<Integer> SCROLL_SIZE =
+      key(ReaderOptions.READER_PREFIX + "scroll_size")
+          .defaultValue(100);
+
+  ConfigOption<String> SPLIT_STRATEGY =
+      key(ReaderOptions.READER_PREFIX + "split_strategy")
+          .defaultValue("round_robin");
+}

bitsail-connectors/connector-elasticsearch/src/main/java/com/bytedance/bitsail/connector/elasticsearch/rest/EsRestClientBuilder.java

@@ -19,8 +19,11 @@
 import com.bytedance.bitsail.common.BitSailException;
 import com.bytedance.bitsail.common.configuration.BitSailConfiguration;
 import com.bytedance.bitsail.common.exception.CommonErrorCode;
+import com.bytedance.bitsail.common.option.ReaderOptions;
+import com.bytedance.bitsail.common.option.WriterOptions;
 import com.bytedance.bitsail.common.util.Preconditions;
 import com.bytedance.bitsail.connector.elasticsearch.base.NetUtil;
+import com.bytedance.bitsail.connector.elasticsearch.option.ElasticsearchReaderOptions;
 import com.bytedance.bitsail.connector.elasticsearch.option.ElasticsearchWriterOptions;
 
 import org.apache.commons.lang.StringUtils;
@@ -46,7 +49,13 @@ public class EsRestClientBuilder {
   private final RestClientBuilder builder;
 
   public EsRestClientBuilder(BitSailConfiguration jobConf) {
-    List<String> hostAddressList = jobConf.get(ElasticsearchWriterOptions.ES_HOSTS);
+    List<String> hostAddressList = null;
+    if (jobConf.fieldExists(ReaderOptions.READER_PREFIX)) {
+      hostAddressList = jobConf.get(ElasticsearchReaderOptions.ES_HOSTS);
+    } else if (jobConf.fieldExists(WriterOptions.WRITER_PREFIX)) {
+      hostAddressList = jobConf.get(ElasticsearchWriterOptions.ES_HOSTS);
+    }
+
     List<HttpHost> hosts = parseHostsAddress(hostAddressList);
     Preconditions.checkState(!hosts.isEmpty(), "cannot find any valid host from configurations.");
     LOG.info("Elasticsearch http client hosts: {}", hosts);

bitsail-connectors/connector-elasticsearch/src/main/java/com/bytedance/bitsail/connector/elasticsearch/rest/source/EsSourceRequest.java

@@ -0,0 +1,138 @@
+/*
+ * Copyright 2022-2023 Bytedance Ltd. and/or its affiliates.
+ *
+ * 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.bytedance.bitsail.connector.elasticsearch.rest.source;
+
+import com.bytedance.bitsail.common.BitSailException;
+import com.bytedance.bitsail.connector.elasticsearch.error.ElasticsearchErrorCode;
+
+import org.apache.commons.compress.utils.Lists;
+import org.elasticsearch.action.search.ClearScrollRequest;
+import org.elasticsearch.action.search.ClearScrollResponse;
+import org.elasticsearch.action.search.SearchRequest;
+import org.elasticsearch.action.search.SearchResponse;
+import org.elasticsearch.action.search.SearchScrollRequest;
+import org.elasticsearch.client.RequestOptions;
+import org.elasticsearch.client.RestHighLevelClient;
+import org.elasticsearch.client.core.CountRequest;
+import org.elasticsearch.client.core.CountResponse;
+import org.elasticsearch.index.query.QueryBuilders;
+import org.elasticsearch.rest.RestStatus;
+import org.elasticsearch.search.SearchHit;
+import org.elasticsearch.search.builder.SearchSourceBuilder;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.io.IOException;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+
+public class EsSourceRequest {
+
+  private static final Logger LOG = LoggerFactory.getLogger(EsSourceRequest.class);
+  private final RestHighLevelClient restHighLevelClient;
+
+  public EsSourceRequest(RestHighLevelClient restHighLevelClient) {
+    this.restHighLevelClient = restHighLevelClient;
+  }
+
+  public Long validateIndex(String index) {
+    CountRequest countRequest = new CountRequest(index);
+    countRequest.query(QueryBuilders.matchAllQuery());
+    CountResponse response = null;
+    try {
+      response = restHighLevelClient.count(countRequest, RequestOptions.DEFAULT);
+      if (response == null) {
+        throw new BitSailException(ElasticsearchErrorCode.VALID_INDEX_FAILED,
+            "GET " + index + " metadata failed");
+      }
+      if (response.status() != RestStatus.OK) {
+        throw new BitSailException(ElasticsearchErrorCode.VALID_INDEX_FAILED,
+            String.format("Get %s response status code = %d", index, response.status().getStatus()));
+      }
+    } catch (IOException e) {
+      throw new BitSailException(ElasticsearchErrorCode.VALID_INDEX_FAILED, e.getMessage());
+    }
+    return response.getCount();
+  }
+
+  /**
+   * Returns a list of all records, Uses scroll API for pagination.
+   * According to: <a href="https://www.elastic.co/guide/en/elasticsearch/client/java-rest/current/java-rest-high-search-scroll.html">...</a>
+   *
+   * @param index index name in Elasticsearch cluster
+   * @param columnNames source field names
+   * @param scrollSize scroll size
+   * @param scrollTime scroll time
+   * @return list of documents
+   * @throws IOException throws IOException if Elasticsearch request fails
+   */
+  public List<Map<String, Object>> getAllDocuments(String index, List<String> columnNames, int scrollSize, String scrollTime) throws IOException {
+    LOG.info("Start get all records from index: {}, scroll size: {}, scroll time: {}", index, scrollSize, scrollTime);
+    SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
+    searchSourceBuilder.size(scrollSize);
+    searchSourceBuilder.query(QueryBuilders.matchAllQuery());
+    searchSourceBuilder.fetchSource(columnNames.toArray(new String[0]), null);
+    searchSourceBuilder.sort("_doc");
+
+    SearchRequest searchRequest = new SearchRequest(index);
+    searchRequest.scroll(scrollTime);
+    searchRequest.source(searchSourceBuilder);
+
+    List<Map<String, Object>> allData = Lists.newArrayList();
+
+    SearchResponse searchResponse = restHighLevelClient.search(searchRequest, RequestOptions.DEFAULT);
+    String scrollId = searchResponse.getScrollId();
+    SearchHit[] hits = addSearchHits(searchResponse, allData);
+    LOG.info("Finish first time search, get {} hits.", Objects.isNull(hits) ? 0 : hits.length);
+
+    while (Objects.nonNull(hits) && hits.length > 0) {
+      LOG.info("Continue scroll query with scrollId: {}", scrollId);
+      SearchScrollRequest searchScrollRequest = new SearchScrollRequest(scrollId);
+      searchScrollRequest.scroll(scrollTime);
+      searchResponse = restHighLevelClient.scroll(searchScrollRequest, RequestOptions.DEFAULT);
+      scrollId = searchResponse.getScrollId();
+      hits = addSearchHits(searchResponse, allData);
+    }
+
+    clearScroll(scrollId);
+    return allData;
+  }
+
+  private void clearScroll(String scrollId) throws IOException {
+    ClearScrollRequest clearScrollRequest = new ClearScrollRequest();
+    clearScrollRequest.addScrollId(scrollId);
+    ClearScrollResponse clearScrollResponse = restHighLevelClient.clearScroll(clearScrollRequest, RequestOptions.DEFAULT);
+    boolean succeeded = clearScrollResponse.isSucceeded();
+    if (succeeded) {
+      LOG.info("Scroll response cleared successfully");
+    } else {
+      LOG.error("Fail to clear scroll response");
+    }
+  }
+
+  private SearchHit[] addSearchHits(SearchResponse searchResponse, List<Map<String, Object>> allData) {
+    SearchHit[] hits = searchResponse.getHits().getHits();
+    if (Objects.isNull(hits)) {
+      return null;
+    }
+    for (SearchHit hit : hits) {
+      allData.add(hit.getSourceAsMap());
+    }
+    return hits;
+  }
+}

bitsail-connectors/connector-elasticsearch/src/main/java/com/bytedance/bitsail/connector/elasticsearch/source/ElasticsearchSource.java

@@ -0,0 +1,94 @@
+/*
+ * Copyright 2022-2023 Bytedance Ltd. and/or its affiliates.
+ *
+ * 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.bytedance.bitsail.connector.elasticsearch.source;
+
+import com.bytedance.bitsail.base.connector.reader.v1.Boundedness;
+import com.bytedance.bitsail.base.connector.reader.v1.Source;
+import com.bytedance.bitsail.base.connector.reader.v1.SourceReader;
+import com.bytedance.bitsail.base.connector.reader.v1.SourceSplitCoordinator;
+import com.bytedance.bitsail.base.connector.writer.v1.state.EmptyState;
+import com.bytedance.bitsail.base.execution.ExecutionEnviron;
+import com.bytedance.bitsail.base.extension.ParallelismComputable;
+import com.bytedance.bitsail.base.parallelism.ParallelismAdvice;
+import com.bytedance.bitsail.common.configuration.BitSailConfiguration;
+import com.bytedance.bitsail.common.row.Row;
+import com.bytedance.bitsail.common.type.TypeInfoConverter;
+import com.bytedance.bitsail.common.type.filemapping.FileMappingTypeInfoConverter;
+import com.bytedance.bitsail.connector.elasticsearch.base.EsConstants;
+import com.bytedance.bitsail.connector.elasticsearch.option.ElasticsearchReaderOptions;
+import com.bytedance.bitsail.connector.elasticsearch.source.coordinator.ElasticsearchSourceSplitCoordinator;
+import com.bytedance.bitsail.connector.elasticsearch.source.reader.ElasticsearchReader;
+import com.bytedance.bitsail.connector.elasticsearch.source.split.ElasticsearchSourceSplit;
+import com.bytedance.bitsail.connector.elasticsearch.source.split.ElasticsearchSplitByIndexStrategy;
+import com.bytedance.bitsail.connector.elasticsearch.source.split.ElasticsearchSplitStrategy;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.io.IOException;
+
+public class ElasticsearchSource implements Source<Row, ElasticsearchSourceSplit, EmptyState>, ParallelismComputable {
+
+  private static final Logger LOG = LoggerFactory.getLogger(ElasticsearchSource.class);
+
+  private BitSailConfiguration jobConf;
+
+  @Override
+  public void configure(ExecutionEnviron execution, BitSailConfiguration readerConfiguration) throws IOException {
+    this.jobConf = readerConfiguration;
+  }
+
+  @Override
+  public Boundedness getSourceBoundedness() {
+    return Boundedness.BOUNDEDNESS;
+  }
+
+  @Override
+  public SourceReader<Row, ElasticsearchSourceSplit> createReader(SourceReader.Context readerContext) {
+    return new ElasticsearchReader(jobConf, readerContext);
+  }
+
+  @Override
+  public SourceSplitCoordinator<ElasticsearchSourceSplit, EmptyState> createSplitCoordinator(
+      SourceSplitCoordinator.Context<ElasticsearchSourceSplit, EmptyState> coordinatorContext) {
+    return new ElasticsearchSourceSplitCoordinator(coordinatorContext, jobConf);
+  }
+
+  @Override
+  public String getReaderName() {
+    return EsConstants.ES_CONNECTOR_NAME;
+  }
+
+  @Override
+  public TypeInfoConverter createTypeInfoConverter() {
+    return new FileMappingTypeInfoConverter(getReaderName());
+  }
+
+  @Override
+  public ParallelismAdvice getParallelismAdvice(BitSailConfiguration commonConf, BitSailConfiguration selfConf, ParallelismAdvice upstreamAdvice) throws Exception {
+    int parallelism;
+    if (selfConf.fieldExists(ElasticsearchReaderOptions.READER_PARALLELISM_NUM)) {
+      parallelism = selfConf.get(ElasticsearchReaderOptions.READER_PARALLELISM_NUM);
+      LOG.info("Use user-defined reader parallelism: {}", parallelism);
+    } else {
+      ElasticsearchSplitStrategy splitStrategy = new ElasticsearchSplitByIndexStrategy();
+      parallelism = splitStrategy.estimateSplitNum(jobConf);
+      LOG.info("Use index number as parallelism: {}", parallelism);
+    }
+    return new ParallelismAdvice(false, parallelism);
+  }
+}