[Improve][JDBC Source] Fix Split can not be cancel (apache#6825)

chaorongzhi · Aug 21, 2024 · 9091545 · 9091545
1 parent f2864b1
commit 9091545
Show file tree

Hide file tree

Showing 19 changed files with 43 additions and 26 deletions.
diff --git a/...unnel/connectors/cdc/base/source/enumerator/splitter/AbstractJdbcSourceChunkSplitter.java b/...unnel/connectors/cdc/base/source/enumerator/splitter/AbstractJdbcSourceChunkSplitter.java
@@ -114,7 +114,7 @@ public Collection<SnapshotSplit> generateSplits(TableId tableId) {
     }
 
     private List<ChunkRange> splitTableIntoChunks(
-            JdbcConnection jdbc, TableId tableId, Column splitColumn) throws SQLException {
+            JdbcConnection jdbc, TableId tableId, Column splitColumn) throws Exception {
         final String splitColumnName = splitColumn.name();
         final Object[] minMax = queryMinMax(jdbc, tableId, splitColumn);
         final Object min = minMax[0];

diff --git a/...che/seatunnel/connectors/cdc/base/source/enumerator/splitter/JdbcSourceChunkSplitter.java b/...che/seatunnel/connectors/cdc/base/source/enumerator/splitter/JdbcSourceChunkSplitter.java
@@ -81,7 +81,7 @@ default Object queryMin(
     @Deprecated
     Object[] sampleDataFromColumn(
             JdbcConnection jdbc, TableId tableId, String columnName, int samplingRate)
-            throws SQLException;
+            throws Exception;
 
     /**
      * Performs a sampling operation on the specified column of a table in a JDBC-connected
@@ -97,7 +97,7 @@ Object[] sampleDataFromColumn(
      */
     default Object[] sampleDataFromColumn(
             JdbcConnection jdbc, TableId tableId, Column column, int samplingRate)
-            throws SQLException {
+            throws Exception {
         return sampleDataFromColumn(jdbc, tableId, column.name(), samplingRate);
     }
 

diff --git a/...nnector-cdc/connector-cdc-base/src/test/java/jdbc/source/JdbcSourceChunkSplitterTest.java b/...nnector-cdc/connector-cdc-base/src/test/java/jdbc/source/JdbcSourceChunkSplitterTest.java
@@ -83,7 +83,7 @@ public Object queryMin(
         @Override
         public Object[] sampleDataFromColumn(
                 JdbcConnection jdbc, TableId tableId, String columnName, int samplingRate)
-                throws SQLException {
+                throws Exception {
             return new Object[0];
         }
 

diff --git a/...l/connectors/cdc/base/source/enumerator/splitter/AbstractJdbcSourceChunkSplitterTest.java b/...l/connectors/cdc/base/source/enumerator/splitter/AbstractJdbcSourceChunkSplitterTest.java
@@ -195,7 +195,7 @@ public Object queryMin(
         @Override
         public Object[] sampleDataFromColumn(
                 JdbcConnection jdbc, TableId tableId, String columnName, int samplingRate)
-                throws SQLException {
+                throws Exception {
             return new Object[0];
         }
 

diff --git a/.../apache/seatunnel/connectors/seatunnel/cdc/mysql/source/eumerator/MySqlChunkSplitter.java b/.../apache/seatunnel/connectors/seatunnel/cdc/mysql/source/eumerator/MySqlChunkSplitter.java
@@ -57,7 +57,7 @@ public Object queryMin(
     @Override
     public Object[] sampleDataFromColumn(
             JdbcConnection jdbc, TableId tableId, String columnName, int inverseSamplingRate)
-            throws SQLException {
+            throws Exception {
         return MySqlUtils.skipReadAndSortSampleData(jdbc, tableId, columnName, inverseSamplingRate);
     }
 

diff --git a/...l/src/main/java/org/apache/seatunnel/connectors/seatunnel/cdc/mysql/utils/MySqlUtils.java b/...l/src/main/java/org/apache/seatunnel/connectors/seatunnel/cdc/mysql/utils/MySqlUtils.java
@@ -146,7 +146,7 @@ public static Object[] sampleDataFromColumn(
 
     public static Object[] skipReadAndSortSampleData(
             JdbcConnection jdbc, TableId tableId, String columnName, int inverseSamplingRate)
-            throws SQLException {
+            throws Exception {
         final String sampleQuery =
                 String.format("SELECT %s FROM %s", quote(columnName), quote(tableId));
 
@@ -172,6 +172,9 @@ public static Object[] skipReadAndSortSampleData(
                 if (count % inverseSamplingRate == 0) {
                     results.add(rs.getObject(1));
                 }
+                if (Thread.currentThread().isInterrupted()) {
+                    throw new InterruptedException("Thread interrupted");
+                }
             }
         } finally {
             if (rs != null) {

diff --git a/...pache/seatunnel/connectors/seatunnel/cdc/oracle/source/eumerator/OracleChunkSplitter.java b/...pache/seatunnel/connectors/seatunnel/cdc/oracle/source/eumerator/OracleChunkSplitter.java
@@ -61,7 +61,7 @@ public Object queryMin(
     @Override
     public Object[] sampleDataFromColumn(
             JdbcConnection jdbc, TableId tableId, String columnName, int inverseSamplingRate)
-            throws SQLException {
+            throws Exception {
         return OracleUtils.skipReadAndSortSampleData(
                 jdbc, tableId, columnName, inverseSamplingRate);
     }

diff --git a/...src/main/java/org/apache/seatunnel/connectors/seatunnel/cdc/oracle/utils/OracleUtils.java b/...src/main/java/org/apache/seatunnel/connectors/seatunnel/cdc/oracle/utils/OracleUtils.java
@@ -150,7 +150,7 @@ public static Object[] sampleDataFromColumn(
 
     public static Object[] skipReadAndSortSampleData(
             JdbcConnection jdbc, TableId tableId, String columnName, int inverseSamplingRate)
-            throws SQLException {
+            throws Exception {
         final String sampleQuery =
                 String.format("SELECT %s FROM %s", quote(columnName), quoteSchemaAndTable(tableId));
 
@@ -176,6 +176,9 @@ public static Object[] skipReadAndSortSampleData(
                 if (count % inverseSamplingRate == 0) {
                     results.add(rs.getObject(1));
                 }
+                if (Thread.currentThread().isInterrupted()) {
+                    throw new InterruptedException("Thread interrupted");
+                }
             }
         } finally {
             if (rs != null) {

diff --git a/.../seatunnel/connectors/seatunnel/cdc/postgres/source/enumerator/PostgresChunkSplitter.java b/.../seatunnel/connectors/seatunnel/cdc/postgres/source/enumerator/PostgresChunkSplitter.java
@@ -70,15 +70,15 @@ public Object queryMin(
     @Override
     public Object[] sampleDataFromColumn(
             JdbcConnection jdbc, TableId tableId, String columnName, int inverseSamplingRate)
-            throws SQLException {
+            throws Exception {
         return PostgresUtils.skipReadAndSortSampleData(
                 jdbc, tableId, columnName, null, inverseSamplingRate);
     }
 
     @Override
     public Object[] sampleDataFromColumn(
             JdbcConnection jdbc, TableId tableId, Column column, int inverseSamplingRate)
-            throws SQLException {
+            throws Exception {
         return PostgresUtils.skipReadAndSortSampleData(
                 jdbc, tableId, column.name(), column, inverseSamplingRate);
     }

diff --git a/...main/java/org/apache/seatunnel/connectors/seatunnel/cdc/postgres/utils/PostgresUtils.java b/...main/java/org/apache/seatunnel/connectors/seatunnel/cdc/postgres/utils/PostgresUtils.java
@@ -161,7 +161,7 @@ public static Object[] skipReadAndSortSampleData(
             String columnName,
             Column column,
             int inverseSamplingRate)
-            throws SQLException {
+            throws Exception {
         columnName = quote(columnName);
         if (column != null) {
             columnName = JDBC_DIALECT.convertType(columnName, column.typeName());
@@ -187,6 +187,9 @@ public static Object[] skipReadAndSortSampleData(
                 if (count % 100000 == 0) {
                     log.info("Processing row index: {}", count);
                 }
+                if (Thread.currentThread().isInterrupted()) {
+                    throw new InterruptedException("Thread interrupted");
+                }
                 if (count % inverseSamplingRate == 0) {
                     results.add(rs.getObject(1));
                 }

diff --git a/...el/connectors/seatunnel/cdc/sqlserver/source/source/eumerator/SqlServerChunkSplitter.java b/...el/connectors/seatunnel/cdc/sqlserver/source/source/eumerator/SqlServerChunkSplitter.java
@@ -57,7 +57,7 @@ public Object queryMin(
     @Override
     public Object[] sampleDataFromColumn(
             JdbcConnection jdbc, TableId tableId, String columnName, int inverseSamplingRate)
-            throws SQLException {
+            throws Exception {
         return SqlServerUtils.skipReadAndSortSampleData(
                 jdbc, tableId, columnName, inverseSamplingRate);
     }

diff --git a/.../org/apache/seatunnel/connectors/seatunnel/cdc/sqlserver/source/utils/SqlServerUtils.java b/.../org/apache/seatunnel/connectors/seatunnel/cdc/sqlserver/source/utils/SqlServerUtils.java
@@ -151,7 +151,7 @@ public static Object[] sampleDataFromColumn(
 
     public static Object[] skipReadAndSortSampleData(
             JdbcConnection jdbc, TableId tableId, String columnName, int inverseSamplingRate)
-            throws SQLException {
+            throws Exception {
         final String sampleQuery =
                 String.format("SELECT %s FROM %s", quote(columnName), quote(tableId));
 
@@ -177,6 +177,9 @@ public static Object[] skipReadAndSortSampleData(
                 if (count % inverseSamplingRate == 0) {
                     results.add(rs.getObject(1));
                 }
+                if (Thread.currentThread().isInterrupted()) {
+                    throw new InterruptedException("Thread interrupted");
+                }
             }
         } finally {
             if (rs != null) {

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
@@ -312,7 +312,7 @@ default Object[] sampleDataFromColumn(
             String columnName,
             int samplingRate,
             int fetchSize)
-            throws SQLException {
+            throws Exception {
         String sampleQuery;
         if (StringUtils.isNotBlank(table.getQuery())) {
             sampleQuery =
@@ -337,6 +337,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, TableSchema schema)
             throws SQLException {

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
@@ -62,7 +62,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);
     }
 
@@ -73,7 +73,7 @@ protected PreparedStatement createSplitStatement(JdbcSourceSplit split, TableSch
     }
 
     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);
@@ -105,7 +105,7 @@ private PreparedStatement createDynamicSplitStatement(JdbcSourceSplit split, Tab
 
     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();
@@ -136,7 +136,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/...c-e2e-part-7/src/test/java/org/apache/seatunnel/connectors/seatunnel/jdbc/JdbcIrisIT.java b/...c-e2e-part-7/src/test/java/org/apache/seatunnel/connectors/seatunnel/jdbc/JdbcIrisIT.java
@@ -288,7 +288,7 @@ public class JdbcIrisIT extends AbstractJdbcIT {
             };
 
     @Test
-    public void testSampleDataFromColumnSuccess() throws SQLException {
+    public void testSampleDataFromColumnSuccess() throws Exception {
         JdbcDialect dialect = new IrisDialect();
         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);