[Feature][connector-elasticsearch] elasticsearch source support PIT

CosmosNi · CosmosNi · commit 1bb296f9df2d · 2025-04-14T11:45:31.000+08:00
diff --git a/docs/en/connector-v2/source/Elasticsearch.md b/docs/en/connector-v2/source/Elasticsearch.md
@@ -30,8 +30,9 @@ support version >= 2.x and <= 8.x.
 | index_list              | array   | no       | used to define a multiple table task                           |
 | source                  | array   | no       | -                                                              |
 | query                   | json    | no       | {"match_all": {}}                                              |
-| search_type             | json    | no       | Search method,SQL、PIT、SCROLL,default SCROLL                    |
-| sql_query               | json    | no       | sql query                                                      |
+| search_type             | enum    | no       | Query type, SQL or DSL, default DSL                            |
+| search_api_type         | enum    | no       | Pagination API type, SCROLL or PIT, default SCROLL             |
+| sql_query               | json    | no       | SQL query, required when search_type is SQL                    |
 | scroll_time             | string  | no       | 1m                                                             |
 | scroll_size             | int     | no       | 100                                                            |
 | tls_verify_certificate  | boolean | no       | true                                                           |
@@ -116,9 +117,14 @@ The path to PEM or JKS trust store. This file must be readable by the operating
 The key password for the trust store specified
 
 ### search_type
-PIT: Whether to use the Point in Time (PIT) API instead of the scroll API
-SQL: Use SQL query
-SCROLL: Use the scroll API by default
+Query type, available values:
+- DSL: Use Domain Specific Language query (default)
+- SQL: Use SQL query
+
+### search_api_type
+Pagination API type, available values:
+- SCROLL: Use Scroll API for pagination (default)
+- PIT: Use Point in Time (PIT) API for pagination
 
 ### pit_keep_alive [long]
 The amount of time (in milliseconds) for which the PIT should be keep alive
@@ -190,7 +196,7 @@ source {
            c_date2,
            c_null
            ]
-           
+
        }
 
     ]
@@ -227,7 +233,7 @@ source {
         hosts = ["https://localhost:9200"]
         username = "elastic"
         password = "elasticsearch"
-        
+
         tls_verify_certificate = false
     }
 }
@@ -241,7 +247,7 @@ source {
         hosts = ["https://localhost:9200"]
         username = "elastic"
         password = "elasticsearch"
-        
+
         tls_verify_hostname = false
     }
 }
@@ -255,7 +261,7 @@ source {
         hosts = ["https://localhost:9200"]
         username = "elastic"
         password = "elasticsearch"
-        
+
         tls_keystore_path = "${your elasticsearch home}/config/certs/http.p12"
         tls_keystore_password = "${your password}"
     }
@@ -291,9 +297,10 @@ source {
 
     index = "st_index"
     query = {"range": {"c_int": {"gte": 10, "lte": 20}}}
-   
-    # Enable PIT API
-    use_pit = true
+
+    # Use DSL query with PIT API
+    search_type = DSL
+    search_api_type = PIT
     pit_keep_alive = 60000  # 1 minute in milliseconds
     pit_batch_size = 100
 ```
diff --git a/docs/zh/connector-v2/source/Elasticsearch.md b/docs/zh/connector-v2/source/Elasticsearch.md
@@ -28,8 +28,9 @@ import ChangeLog from '../changelog/connector-elasticsearch.md';
 | index_list              | array   | no       | 用来定义多索引同步任务                         |
 | source                  | array   | no       | -                                   |
 | query                   | json    | no       | {"match_all": {}}                   |
-| search_type             | json    | no       | 查询方式，SQL、PIT、SCROLL,默认 SCROLL       |
-| sql_query               | json    | no       | sql 查询语句                            |
+| search_type             | enum    | no       | 查询类型，SQL 或 DSL，默认 DSL              |
+| search_api_type         | enum    | no       | 分页 API 类型，SCROLL 或 PIT，默认 SCROLL    |
+| sql_query               | json    | no       | SQL 查询语句，当 search_type 为 SQL 时必须    |
 | scroll_time             | string  | no       | 1m                                  |
 | scroll_size             | int     | no       | 100                                 |
 | tls_verify_certificate  | boolean | no       | true                                |
@@ -118,9 +119,14 @@ PEM 或 JKS 信任库的路径。该文件必须对运行 SeaTunnel 的操作系
 指定信任库的密钥密码。
 
 ### search_type
-PIT:是否使用时间点 (PIT) API 代替滚动 API
-SQL:使用sql 方式查询
-SCROLL:默认使用scroll API
+查询类型，可选值：
+- DSL: 使用 Domain Specific Language 查询（默认）
+- SQL: 使用 SQL 查询
+
+### search_api_type
+分页 API 类型，可选值：
+- SCROLL: 使用 Scroll API 进行分页（默认）
+- PIT: 使用 Point in Time (PIT) API 进行分页
 
 ### pit_keep_alive [long]
 PIT 应保持活动的时间量（以毫秒为单位）
@@ -193,7 +199,7 @@ source {
            c_date2,
            c_null
            ]
-           
+
        }
 
     ]
@@ -228,7 +234,7 @@ source {
         hosts = ["https://localhost:9200"]
         username = "elastic"
         password = "elasticsearch"
-        
+
         tls_verify_certificate = false
     }
 }
@@ -242,7 +248,7 @@ source {
         hosts = ["https://localhost:9200"]
         username = "elastic"
         password = "elasticsearch"
-        
+
         tls_verify_hostname = false
     }
 }
@@ -256,7 +262,7 @@ source {
         hosts = ["https://localhost:9200"]
         username = "elastic"
         password = "elasticsearch"
-        
+
         tls_keystore_path = "${your elasticsearch home}/config/certs/http.p12"
         tls_keystore_password = "${your password}"
     }
@@ -292,9 +298,10 @@ source {
 
     index = "st_index"
     query = {"range": {"c_int": {"gte": 10, "lte": 20}}}
-   
-    # Enable PIT API
-    search_type = PIT
+
+    # 使用 DSL 查询和 PIT API
+    search_type = DSL
+    search_api_type = PIT
     pit_keep_alive = 60000  # 1 minute in milliseconds
     pit_batch_size = 100
 ```
diff --git a/seatunnel-connectors-v2/connector-elasticsearch/src/main/java/org/apache/seatunnel/connectors/seatunnel/elasticsearch/config/ElasticsearchConfig.java b/seatunnel-connectors-v2/connector-elasticsearch/src/main/java/org/apache/seatunnel/connectors/seatunnel/elasticsearch/config/ElasticsearchConfig.java
@@ -38,6 +38,7 @@ public class ElasticsearchConfig implements Serializable {
     private String scrollTime;
     private int scrollSize;
     private SearchTypeEnum searchType;
+    private SearchApiTypeEnum searchApiType;
     private String sqlQuery;
 
     private long pitKeepAlive;
@@ -56,6 +57,7 @@ public ElasticsearchConfig clone() {
         elasticsearchConfig.setScrollSize(scrollSize);
         elasticsearchConfig.setCatalogTable(catalogTable);
         elasticsearchConfig.setSearchType(searchType);
+        elasticsearchConfig.setSearchApiType(searchApiType);
         elasticsearchConfig.setSqlQuery(sqlQuery);
         elasticsearchConfig.setPitKeepAlive(pitKeepAlive);
         elasticsearchConfig.setPitBatchSize(pitBatchSize);
diff --git a/seatunnel-connectors-v2/connector-elasticsearch/src/main/java/org/apache/seatunnel/connectors/seatunnel/elasticsearch/config/ElasticsearchSourceOptions.java b/seatunnel-connectors-v2/connector-elasticsearch/src/main/java/org/apache/seatunnel/connectors/seatunnel/elasticsearch/config/ElasticsearchSourceOptions.java
@@ -65,8 +65,14 @@ public class ElasticsearchSourceOptions extends ElasticsearchBaseOptions {
     public static final Option<SearchTypeEnum> SEARCH_TYPE =
             Options.key("search_type")
                     .enumType(SearchTypeEnum.class)
-                    .defaultValue(SearchTypeEnum.SCROLL)
-                    .withDescription("Choose dsl syntax or x-pack sql.");
+                    .defaultValue(SearchTypeEnum.DSL)
+                    .withDescription("Choose query type: DSL (Domain Specific Language) or SQL.");
+
+    public static final Option<SearchApiTypeEnum> SEARCH_API_TYPE =
+            Options.key("search_api_type")
+                    .enumType(SearchApiTypeEnum.class)
+                    .defaultValue(SearchApiTypeEnum.SCROLL)
+                    .withDescription("Choose API type for pagination: SCROLL or PIT (Point in Time).");
 
     public static final Option<String> SQL_QUERY =
             Options.key("sql_query")
diff --git a/seatunnel-connectors-v2/connector-elasticsearch/src/main/java/org/apache/seatunnel/connectors/seatunnel/elasticsearch/config/SearchApiTypeEnum.java b/seatunnel-connectors-v2/connector-elasticsearch/src/main/java/org/apache/seatunnel/connectors/seatunnel/elasticsearch/config/SearchApiTypeEnum.java
@@ -0,0 +1,26 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.seatunnel.connectors.seatunnel.elasticsearch.config;
+
+public enum SearchApiTypeEnum {
+    /** Use Scroll API for pagination */
+    SCROLL,
+
+    /** Use Point-in-Time (PIT) API for pagination */
+    PIT
+}
diff --git a/seatunnel-connectors-v2/connector-elasticsearch/src/main/java/org/apache/seatunnel/connectors/seatunnel/elasticsearch/config/SearchTypeEnum.java b/seatunnel-connectors-v2/connector-elasticsearch/src/main/java/org/apache/seatunnel/connectors/seatunnel/elasticsearch/config/SearchTypeEnum.java
@@ -18,12 +18,9 @@
 package org.apache.seatunnel.connectors.seatunnel.elasticsearch.config;
 
 public enum SearchTypeEnum {
-    /** Use Domain Specific Language (DSL) query with Scroll API */
-    SCROLL,
+    /** Use Domain Specific Language (DSL) query */
+    DSL,
 
     /** Use SQL query */
-    SQL,
-
-    /** Use Point-in-Time (PIT) API with DSL query */
-    PIT,
+    SQL
 }
diff --git a/seatunnel-connectors-v2/connector-elasticsearch/src/main/java/org/apache/seatunnel/connectors/seatunnel/elasticsearch/source/ElasticsearchSource.java b/seatunnel-connectors-v2/connector-elasticsearch/src/main/java/org/apache/seatunnel/connectors/seatunnel/elasticsearch/source/ElasticsearchSource.java
@@ -17,6 +17,7 @@
 
 package org.apache.seatunnel.connectors.seatunnel.elasticsearch.source;
 
+import org.apache.seatunnel.connectors.seatunnel.elasticsearch.config.SearchApiTypeEnum;
 import org.apache.seatunnel.shade.com.google.common.annotations.VisibleForTesting;
 
 import org.apache.seatunnel.api.configuration.ReadonlyConfig;
@@ -56,6 +57,7 @@
 import java.util.Map;
 import java.util.stream.Collectors;
 
+import static org.apache.seatunnel.connectors.seatunnel.elasticsearch.config.ElasticsearchSourceOptions.SEARCH_API_TYPE;
 import static org.apache.seatunnel.connectors.seatunnel.elasticsearch.config.ElasticsearchSourceOptions.SEARCH_TYPE;
 import static org.apache.seatunnel.connectors.seatunnel.elasticsearch.config.ElasticsearchSourceOptions.SQL_QUERY;
 
@@ -175,6 +177,7 @@ private ElasticsearchConfig parseOneIndexQueryConfig(ReadonlyConfig readonlyConf
                             "");
         }
         SearchTypeEnum searchType = readonlyConfig.get(SEARCH_TYPE);
+        SearchApiTypeEnum searchApiType = readonlyConfig.get(SEARCH_API_TYPE);
         String sqlQuery = readonlyConfig.get(ElasticsearchSourceOptions.SQL_QUERY);
         String scrollTime = readonlyConfig.get(ElasticsearchSourceOptions.SCROLL_TIME);
         int scrollSize = readonlyConfig.get(ElasticsearchSourceOptions.SCROLL_SIZE);
@@ -192,6 +195,7 @@ private ElasticsearchConfig parseOneIndexQueryConfig(ReadonlyConfig readonlyConf
         elasticsearchConfig.setCatalogTable(catalogTable);
         elasticsearchConfig.setSqlQuery(sqlQuery);
         elasticsearchConfig.setSearchType(searchType);
+        elasticsearchConfig.setSearchApiType(searchApiType);
 
         elasticsearchConfig.setPitKeepAlive(pitKeepAlive);
         elasticsearchConfig.setPitBatchSize(pitBatchSize);
diff --git a/seatunnel-connectors-v2/connector-elasticsearch/src/main/java/org/apache/seatunnel/connectors/seatunnel/elasticsearch/source/ElasticsearchSourceReader.java b/seatunnel-connectors-v2/connector-elasticsearch/src/main/java/org/apache/seatunnel/connectors/seatunnel/elasticsearch/source/ElasticsearchSourceReader.java
@@ -24,6 +24,7 @@
 import org.apache.seatunnel.api.table.type.SeaTunnelRowType;
 import org.apache.seatunnel.connectors.seatunnel.elasticsearch.client.EsRestClient;
 import org.apache.seatunnel.connectors.seatunnel.elasticsearch.config.ElasticsearchConfig;
+import org.apache.seatunnel.connectors.seatunnel.elasticsearch.config.SearchApiTypeEnum;
 import org.apache.seatunnel.connectors.seatunnel.elasticsearch.config.SearchTypeEnum;
 import org.apache.seatunnel.connectors.seatunnel.elasticsearch.dto.source.PointInTimeResult;
 import org.apache.seatunnel.connectors.seatunnel.elasticsearch.dto.source.ScrollResult;
@@ -98,13 +99,9 @@ private void scrollSearchResult(
         SeaTunnelRowDeserializer deserializer =
                 new DefaultSeaTunnelRowDeserializer(seaTunnelRowType);
 
-        // Check if we should use PIT API
-        if (SearchTypeEnum.PIT.equals(sourceIndexInfo.getSearchType())) {
-            log.info("Using Point-in-Time (PIT) API for index: {}", sourceIndexInfo.getIndex());
-            searchWithPointInTime(sourceIndexInfo, output, deserializer);
-        }
         // SQL client
-        else if (SearchTypeEnum.SQL.equals(sourceIndexInfo.getSearchType())) {
+        if (SearchTypeEnum.SQL.equals(sourceIndexInfo.getSearchType())) {
+            log.info("Using SQL query for index: {}", sourceIndexInfo.getIndex());
             ScrollResult scrollResult =
                     esRestClient.searchBySql(
                             sourceIndexInfo.getSqlQuery(), sourceIndexInfo.getScrollSize());
@@ -117,21 +114,30 @@ else if (SearchTypeEnum.SQL.equals(sourceIndexInfo.getSearchType())) {
                 outputFromScrollResult(scrollResult, sourceIndexInfo, output, deserializer);
             }
         }
-        // Default scroll API
+        // DSL query
         else {
-            ScrollResult scrollResult =
-                    esRestClient.searchByScroll(
-                            sourceIndexInfo.getIndex(),
-                            sourceIndexInfo.getSource(),
-                            sourceIndexInfo.getQuery(),
-                            sourceIndexInfo.getScrollTime(),
-                            sourceIndexInfo.getScrollSize());
-            outputFromScrollResult(scrollResult, sourceIndexInfo, output, deserializer);
-            while (scrollResult.getDocs() != null && !scrollResult.getDocs().isEmpty()) {
-                scrollResult =
-                        esRestClient.searchWithScrollId(
-                                scrollResult.getScrollId(), sourceIndexInfo.getScrollTime());
+            // Check if we should use PIT API
+            if (SearchApiTypeEnum.PIT.equals(sourceIndexInfo.getSearchApiType())) {
+                log.info("Using Point-in-Time (PIT) API for index: {}", sourceIndexInfo.getIndex());
+                searchWithPointInTime(sourceIndexInfo, output, deserializer);
+            }
+            // Default scroll API
+            else {
+                log.info("Using Scroll API for index: {}", sourceIndexInfo.getIndex());
+                ScrollResult scrollResult =
+                        esRestClient.searchByScroll(
+                                sourceIndexInfo.getIndex(),
+                                sourceIndexInfo.getSource(),
+                                sourceIndexInfo.getQuery(),
+                                sourceIndexInfo.getScrollTime(),
+                                sourceIndexInfo.getScrollSize());
                 outputFromScrollResult(scrollResult, sourceIndexInfo, output, deserializer);
+                while (scrollResult.getDocs() != null && !scrollResult.getDocs().isEmpty()) {
+                    scrollResult =
+                            esRestClient.searchWithScrollId(
+                                    scrollResult.getScrollId(), sourceIndexInfo.getScrollTime());
+                    outputFromScrollResult(scrollResult, sourceIndexInfo, output, deserializer);
+                }
             }
         }
     }
diff --git a/seatunnel-e2e/seatunnel-connector-v2-e2e/connector-elasticsearch-e2e/src/test/resources/elasticsearch/elasticsearch_source_with_pit.conf b/seatunnel-e2e/seatunnel-connector-v2-e2e/connector-elasticsearch-e2e/src/test/resources/elasticsearch/elasticsearch_source_with_pit.conf
@@ -34,12 +34,13 @@ source {
 
     index = "st_index"
     query = {"range": {"c_int": {"gte": 10, "lte": 20}}}
-    
-    # Enable PIT API
-    search_type = "PIT"
+
+    # Use DSL query with PIT API
+    search_type = "DSL"
+    search_api_type = "PIT"
     pit_keep_alive = 60000  # 1 minute in milliseconds
     pit_batch_size = 100
-    
+
     schema = {
       fields {
         c_map = "map<string, tinyint>"
