Fix Split can not be cancel

apache · May 9, 2024 · 959232d · 959232d
1 parent d826cf9
commit 959232d
Show file tree

Hide file tree

Showing 6 changed files with 16 additions and 11 deletions.
diff --git a/...ain/java/org/apache/seatunnel/connectors/seatunnel/jdbc/internal/dialect/JdbcDialect.java b/...ain/java/org/apache/seatunnel/connectors/seatunnel/jdbc/internal/dialect/JdbcDialect.java
@@ -308,7 +308,7 @@ default Object[] sampleDataFromColumn(
             String columnName,
             int samplingRate,
             int fetchSize)
-            throws SQLException {
+            throws Exception {
         String sampleQuery;
         if (StringUtils.isNotBlank(table.getQuery())) {
             sampleQuery =
@@ -333,6 +333,9 @@ default Object[] sampleDataFromColumn(
                     if (count % samplingRate == 0) {
                         results.add(rs.getObject(1));
                     }
+                    if (Thread.currentThread().isInterrupted()) {
+                        throw new InterruptedException("Thread interrupted");
+                    }
                 }
                 Object[] resultsArray = results.toArray();
                 Arrays.sort(resultsArray);

diff --git a/...a/org/apache/seatunnel/connectors/seatunnel/jdbc/internal/dialect/mysql/MysqlDialect.java b/...a/org/apache/seatunnel/connectors/seatunnel/jdbc/internal/dialect/mysql/MysqlDialect.java
@@ -131,7 +131,7 @@ public Object[] sampleDataFromColumn(
             String columnName,
             int samplingRate,
             int fetchSize)
-            throws SQLException {
+            throws Exception {
         String sampleQuery;
         if (StringUtils.isNotBlank(table.getQuery())) {
             sampleQuery =
@@ -158,6 +158,9 @@ public Object[] sampleDataFromColumn(
                     if (count % samplingRate == 0) {
                         results.add(rs.getObject(1));
                     }
+                    if (Thread.currentThread().isInterrupted()) {
+                        throw new InterruptedException("Thread interrupted");
+                    }
                 }
                 Object[] resultsArray = results.toArray();
                 Arrays.sort(resultsArray);

diff --git a/...bc/src/main/java/org/apache/seatunnel/connectors/seatunnel/jdbc/source/ChunkSplitter.java b/...bc/src/main/java/org/apache/seatunnel/connectors/seatunnel/jdbc/source/ChunkSplitter.java
@@ -85,7 +85,7 @@ public synchronized void close() {
         }
     }
 
-    public Collection<JdbcSourceSplit> generateSplits(JdbcSourceTable table) throws SQLException {
+    public Collection<JdbcSourceSplit> generateSplits(JdbcSourceTable table) throws Exception {
         log.info("Start splitting table {} into chunks...", table.getTablePath());
         long start = System.currentTimeMillis();
 
@@ -111,7 +111,7 @@ public Collection<JdbcSourceSplit> generateSplits(JdbcSourceTable table) throws
     }
 
     protected abstract Collection<JdbcSourceSplit> createSplits(
-            JdbcSourceTable table, SeaTunnelRowType splitKeyType) throws SQLException;
+            JdbcSourceTable table, SeaTunnelRowType splitKeyType) throws SQLException, Exception;
 
     public PreparedStatement generateSplitStatement(JdbcSourceSplit split) throws SQLException {
         if (split.getSplitKeyName() == null) {

diff --git a/...main/java/org/apache/seatunnel/connectors/seatunnel/jdbc/source/DynamicChunkSplitter.java b/...main/java/org/apache/seatunnel/connectors/seatunnel/jdbc/source/DynamicChunkSplitter.java
@@ -60,7 +60,7 @@ public DynamicChunkSplitter(JdbcSourceConfig config) {
 
     @Override
     protected Collection<JdbcSourceSplit> createSplits(
-            JdbcSourceTable table, SeaTunnelRowType splitKey) throws SQLException {
+            JdbcSourceTable table, SeaTunnelRowType splitKey) throws Exception {
         return createDynamicSplits(table, splitKey);
     }
 
@@ -70,7 +70,7 @@ protected PreparedStatement createSplitStatement(JdbcSourceSplit split) throws S
     }
 
     private Collection<JdbcSourceSplit> createDynamicSplits(
-            JdbcSourceTable table, SeaTunnelRowType splitKey) throws SQLException {
+            JdbcSourceTable table, SeaTunnelRowType splitKey) throws Exception {
         String splitKeyName = splitKey.getFieldNames()[0];
         SeaTunnelDataType splitKeyType = splitKey.getFieldType(0);
         List<ChunkRange> chunks = splitTableIntoChunks(table, splitKeyName, splitKeyType);
@@ -102,7 +102,7 @@ private PreparedStatement createDynamicSplitStatement(JdbcSourceSplit split)
 
     private List<ChunkRange> splitTableIntoChunks(
             JdbcSourceTable table, String splitColumnName, SeaTunnelDataType splitColumnType)
-            throws SQLException {
+            throws Exception {
         Pair<Object, Object> minMax = queryMinMax(table, splitColumnName);
         Object min = minMax.getLeft();
         Object max = minMax.getRight();
@@ -133,7 +133,7 @@ private List<ChunkRange> splitTableIntoChunks(
 
     private List<ChunkRange> evenlyColumnSplitChunks(
             JdbcSourceTable table, String splitColumnName, Object min, Object max, int chunkSize)
-            throws SQLException {
+            throws Exception {
         TablePath tablePath = table.getTablePath();
         double distributionFactorUpper = config.getSplitEvenDistributionFactorUpperBound();
         double distributionFactorLower = config.getSplitEvenDistributionFactorLowerBound();

diff --git a/...e2e-part-1/src/test/java/org/apache/seatunnel/connectors/seatunnel/jdbc/JdbcOracleIT.java b/...e2e-part-1/src/test/java/org/apache/seatunnel/connectors/seatunnel/jdbc/JdbcOracleIT.java
@@ -40,7 +40,6 @@
 
 import java.math.BigDecimal;
 import java.sql.Date;
-import java.sql.SQLException;
 import java.sql.Timestamp;
 import java.time.LocalDate;
 import java.time.LocalDateTime;
@@ -104,7 +103,7 @@ public class JdbcOracleIT extends AbstractJdbcIT {
             };
 
     @Test
-    public void testSampleDataFromColumnSuccess() throws SQLException {
+    public void testSampleDataFromColumnSuccess() throws Exception {
         JdbcDialect dialect = new OracleDialect();
         JdbcSourceTable table =
                 JdbcSourceTable.builder()

diff --git a/...part-7/src/test/java/org/apache/seatunnel/connectors/seatunnel/jdbc/JdbcMysqlSplitIT.java b/...part-7/src/test/java/org/apache/seatunnel/connectors/seatunnel/jdbc/JdbcMysqlSplitIT.java
@@ -468,7 +468,7 @@ public void testSplit() throws Exception {
 
     private JdbcSourceSplit[] getCheckedSplitArray(
             Map<String, Object> configMap, CatalogTable table, String splitKey, int splitNum)
-            throws SQLException {
+            throws Exception {
         configMap.put("partition_column", splitKey);
         DynamicChunkSplitter splitter = getDynamicChunkSplitter(configMap);