This is an automated email from the ASF dual-hosted git repository.
davidzollo pushed a commit to branch dev
in repository https://gitbox.apache.org/repos/asf/seatunnel.git
The following commit(s) were added to refs/heads/dev by this push:
new b1c53562b0 [Improve][connector-elasticsearchv2]Optimize Elasticsearch
source che… (#10446)
b1c53562b0 is described below
commit b1c53562b06be6dd71131dffd73b0439aed9640b
Author: cosmosni <[email protected]>
AuthorDate: Sun Jul 19 09:55:03 2026 +0800
[Improve][connector-elasticsearchv2]Optimize Elasticsearch source che…
(#10446)
---
.../elasticsearch/client/EsRestClient.java | 44 ++++++++-
.../source/ElasticsearchSourceReader.java | 30 ++++--
.../source/ElasticsearchSourceReaderTest.java | 109 +++++++++++++++++++++
3 files changed, 173 insertions(+), 10 deletions(-)
diff --git
a/seatunnel-connectors-v2/connector-elasticsearch/src/main/java/org/apache/seatunnel/connectors/seatunnel/elasticsearch/client/EsRestClient.java
b/seatunnel-connectors-v2/connector-elasticsearch/src/main/java/org/apache/seatunnel/connectors/seatunnel/elasticsearch/client/EsRestClient.java
index 6a43db1791..96432b315c 100644
---
a/seatunnel-connectors-v2/connector-elasticsearch/src/main/java/org/apache/seatunnel/connectors/seatunnel/elasticsearch/client/EsRestClient.java
+++
b/seatunnel-connectors-v2/connector-elasticsearch/src/main/java/org/apache/seatunnel/connectors/seatunnel/elasticsearch/client/EsRestClient.java
@@ -345,6 +345,48 @@ public class EsRestClient implements Closeable {
}
}
+ /**
+ * Close SQL cursor to release server-side resources.
+ *
+ * @param cursor The SQL cursor to close
+ * @return True if the cursor was successfully closed
+ */
+ public boolean closeSqlCursor(String cursor) {
+ if (StringUtils.isEmpty(cursor)) {
+ log.debug("SQL cursor is empty; skip closing.");
+ return false;
+ }
+
+ String endpoint = "/_sql/close";
+ Request request = new Request("POST", endpoint);
+ Map<String, String> requestBody = new HashMap<>();
+ requestBody.put("cursor", cursor);
+ request.setJsonEntity(JsonUtils.toJsonString(requestBody));
+
+ try {
+ Response response = restClient.performRequest(request);
+ if (response == null) {
+ log.warn("POST {} response null for cursor: {}", endpoint,
cursor);
+ return false;
+ }
+ if (response.getStatusLine().getStatusCode() == HttpStatus.SC_OK) {
+ String entity = EntityUtils.toString(response.getEntity());
+ JsonNode jsonNode = JsonUtils.parseObject(entity);
+ return jsonNode.get("succeeded").asBoolean();
+ } else {
+ log.warn(
+ "POST {} response status code={} for cursor: {}",
+ endpoint,
+ response.getStatusLine().getStatusCode(),
+ cursor);
+ return false;
+ }
+ } catch (Exception ex) {
+ log.warn("Failed to close SQL cursor: {}", cursor, ex);
+ return false;
+ }
+ }
+
private ScrollResult getDocsFromSqlResult(
String endpoint, String requestBody, JsonNode columnNodes) {
Request request = new Request("POST", endpoint);
@@ -1110,7 +1152,7 @@ public class EsRestClient implements Closeable {
String updatedPitId = rootNode.has("pit_id") ?
rootNode.get("pit_id").asText() : pitId;
// Determine if there are more results
- boolean hasMore = docs.size() > 0 && totalHits > 0 && docs.size() <
totalHits;
+ boolean hasMore = !docs.isEmpty();
return new PointInTimeResult(updatedPitId, docs, totalHits,
searchAfter, hasMore);
}
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
index 1faca89f15..020f1767d3 100644
---
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
@@ -102,16 +102,28 @@ public class ElasticsearchSourceReader
// SQL client
if (SearchTypeEnum.SQL.equals(sourceIndexInfo.getSearchType())) {
log.info("Using SQL query for index: {}",
sourceIndexInfo.getIndex());
- ScrollResult scrollResult =
- esRestClient.searchBySql(
- sourceIndexInfo.getSqlQuery(),
sourceIndexInfo.getScrollSize());
-
- outputFromScrollResult(scrollResult, sourceIndexInfo, output,
deserializer);
- while (StringUtils.isNotEmpty(scrollResult.getScrollId())) {
- scrollResult =
- esRestClient.searchWithSql(
- scrollResult.getScrollId(),
scrollResult.getColumnNodes());
+ String cursor = null;
+ try {
+ ScrollResult scrollResult =
+ esRestClient.searchBySql(
+ sourceIndexInfo.getSqlQuery(),
sourceIndexInfo.getScrollSize());
+
outputFromScrollResult(scrollResult, sourceIndexInfo, output,
deserializer);
+ cursor = scrollResult.getScrollId();
+ while (StringUtils.isNotEmpty(cursor)) {
+ scrollResult =
+ esRestClient.searchWithSql(cursor,
scrollResult.getColumnNodes());
+ outputFromScrollResult(scrollResult, sourceIndexInfo,
output, deserializer);
+ cursor = scrollResult.getScrollId();
+ }
+ } finally {
+ if (StringUtils.isNotEmpty(cursor)) {
+ try {
+ esRestClient.closeSqlCursor(cursor);
+ } catch (Exception e) {
+ log.warn("Failed to close SQL cursor: {}", cursor, e);
+ }
+ }
}
} else {
// Check if we should use PIT API
diff --git
a/seatunnel-connectors-v2/connector-elasticsearch/src/test/java/org/apache/seatunnel/connectors/seatunnel/elasticsearch/source/ElasticsearchSourceReaderTest.java
b/seatunnel-connectors-v2/connector-elasticsearch/src/test/java/org/apache/seatunnel/connectors/seatunnel/elasticsearch/source/ElasticsearchSourceReaderTest.java
new file mode 100644
index 0000000000..7521af1e7d
--- /dev/null
+++
b/seatunnel-connectors-v2/connector-elasticsearch/src/test/java/org/apache/seatunnel/connectors/seatunnel/elasticsearch/source/ElasticsearchSourceReaderTest.java
@@ -0,0 +1,109 @@
+/*
+ * 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.source;
+
+import org.apache.seatunnel.api.configuration.ReadonlyConfig;
+import org.apache.seatunnel.api.source.Collector;
+import org.apache.seatunnel.api.source.SourceReader;
+import org.apache.seatunnel.api.table.catalog.CatalogTable;
+import org.apache.seatunnel.api.table.catalog.PhysicalColumn;
+import org.apache.seatunnel.api.table.catalog.TableIdentifier;
+import org.apache.seatunnel.api.table.catalog.TableSchema;
+import org.apache.seatunnel.api.table.type.BasicType;
+import org.apache.seatunnel.api.table.type.SeaTunnelRow;
+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.SearchTypeEnum;
+import
org.apache.seatunnel.connectors.seatunnel.elasticsearch.dto.source.ScrollResult;
+
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+
+import java.lang.reflect.Field;
+import java.util.Collections;
+
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+public class ElasticsearchSourceReaderTest {
+
+ @Test
+ public void testCloseSqlCursorWhenSearchWithSqlThrowsException() throws
Exception {
+ SourceReader.Context context = mock(SourceReader.Context.class);
+ ElasticsearchSourceReader reader =
+ new ElasticsearchSourceReader(
+ context,
ReadonlyConfig.fromMap(Collections.emptyMap()));
+
+ EsRestClient esRestClient = mock(EsRestClient.class);
+ injectEsRestClient(reader, esRestClient);
+
+ ScrollResult firstPage = new ScrollResult();
+ firstPage.setScrollId("cursor-1");
+ firstPage.setDocs(Collections.emptyList());
+
+ when(esRestClient.searchBySql(anyString(),
anyInt())).thenReturn(firstPage);
+ when(esRestClient.searchWithSql(eq("cursor-1"), any()))
+ .thenThrow(new RuntimeException("mock sql scroll exception"));
+
+ reader.addSplits(
+ Collections.singletonList(new ElasticsearchSourceSplit("0",
createSqlSplit())));
+
+ Collector<SeaTunnelRow> collector = mock(Collector.class);
+ when(collector.getCheckpointLock()).thenReturn(new Object());
+
+ Assertions.assertThrows(RuntimeException.class, () ->
reader.pollNext(collector));
+
+ verify(esRestClient).closeSqlCursor("cursor-1");
+ }
+
+ private static ElasticsearchConfig createSqlSplit() {
+ ElasticsearchConfig config = new ElasticsearchConfig();
+ config.setIndex("test_index");
+ config.setSqlQuery("SELECT id FROM test_index");
+ config.setScrollSize(100);
+ config.setSearchType(SearchTypeEnum.SQL);
+ config.setSource(Collections.singletonList("id"));
+ config.setCatalogTable(createCatalogTable());
+ return config;
+ }
+
+ private static CatalogTable createCatalogTable() {
+ return CatalogTable.of(
+ TableIdentifier.of("test_catalog", "test_database",
"test_table"),
+ TableSchema.builder()
+ .column(
+ PhysicalColumn.of(
+ "id", BasicType.INT_TYPE, (Long) null,
true, null, ""))
+ .build(),
+ Collections.emptyMap(),
+ Collections.emptyList(),
+ "");
+ }
+
+ private static void injectEsRestClient(ElasticsearchSourceReader reader,
EsRestClient client)
+ throws Exception {
+ Field field =
ElasticsearchSourceReader.class.getDeclaredField("esRestClient");
+ field.setAccessible(true);
+ field.set(reader, client);
+ }
+}
