[Feature][Connector-V2]Jdbc chunk split add “snapshot.split.column” p…

…arams apache#7794
XenosK · Oct 15, 2024 · 1165447 · 1165447
1 parent 69bf03c
commit 1165447
Show file tree

Hide file tree

Showing 14 changed files with 70 additions and 4 deletions.
diff --git a/...-base/src/main/java/org/apache/seatunnel/connectors/cdc/base/config/BaseSourceConfig.java b/...-base/src/main/java/org/apache/seatunnel/connectors/cdc/base/config/BaseSourceConfig.java
@@ -34,6 +34,7 @@ public abstract class BaseSourceConfig implements SourceConfig {
     @Getter protected final StopConfig stopConfig;
 
     @Getter protected final int splitSize;
+    @Getter protected final String splitColumn;
 
     @Getter protected final double distributionFactorUpper;
     @Getter protected final double distributionFactorLower;
@@ -50,6 +51,7 @@ public BaseSourceConfig(
             StartupConfig startupConfig,
             StopConfig stopConfig,
             int splitSize,
+            String splitColumn,
             double distributionFactorUpper,
             double distributionFactorLower,
             int sampleShardingThreshold,
@@ -59,6 +61,7 @@ public BaseSourceConfig(
         this.startupConfig = startupConfig;
         this.stopConfig = stopConfig;
         this.splitSize = splitSize;
+        this.splitColumn = splitColumn;
         this.distributionFactorUpper = distributionFactorUpper;
         this.distributionFactorLower = distributionFactorLower;
         this.sampleShardingThreshold = sampleShardingThreshold;

diff --git a/...-base/src/main/java/org/apache/seatunnel/connectors/cdc/base/config/JdbcSourceConfig.java b/...-base/src/main/java/org/apache/seatunnel/connectors/cdc/base/config/JdbcSourceConfig.java
@@ -49,6 +49,7 @@ public JdbcSourceConfig(
             List<String> databaseList,
             List<String> tableList,
             int splitSize,
+            String splitColumn,
             double distributionFactorUpper,
             double distributionFactorLower,
             int sampleShardingThreshold,
@@ -70,6 +71,7 @@ public JdbcSourceConfig(
                 startupConfig,
                 stopConfig,
                 splitSize,
+                splitColumn,
                 distributionFactorUpper,
                 distributionFactorLower,
                 sampleShardingThreshold,

diff --git a/...rc/main/java/org/apache/seatunnel/connectors/cdc/base/config/JdbcSourceConfigFactory.java b/...rc/main/java/org/apache/seatunnel/connectors/cdc/base/config/JdbcSourceConfigFactory.java
@@ -51,6 +51,7 @@ public abstract class JdbcSourceConfigFactory implements SourceConfig.Factory<Jd
             JdbcSourceOptions.SAMPLE_SHARDING_THRESHOLD.defaultValue();
     protected int inverseSamplingRate = JdbcSourceOptions.INVERSE_SAMPLING_RATE.defaultValue();
     protected int splitSize = SourceOptions.SNAPSHOT_SPLIT_SIZE.defaultValue();
+    protected String splitColumn = SourceOptions.SNAPSHOT_SPLIT_COLUMN.defaultValue();;
     protected int fetchSize = SourceOptions.SNAPSHOT_FETCH_SIZE.defaultValue();
     protected String serverTimeZone = JdbcSourceOptions.SERVER_TIME_ZONE.defaultValue();
     protected long connectTimeoutMillis = JdbcSourceOptions.CONNECT_TIMEOUT_MS.defaultValue();
@@ -65,6 +66,11 @@ public JdbcSourceConfigFactory hostname(String hostname) {
         return this;
     }
 
+    public JdbcSourceConfigFactory splitColumn(String splitColumn) {
+        this.splitColumn = splitColumn;
+        return this;
+    }
+
     /** Integer port number of the database server. */
     public JdbcSourceConfigFactory port(int port) {
         this.port = port;
@@ -239,6 +245,7 @@ public JdbcSourceConfigFactory fromReadonlyConfig(ReadonlyConfig config) {
         this.sampleShardingThreshold = config.get(JdbcSourceOptions.SAMPLE_SHARDING_THRESHOLD);
         this.inverseSamplingRate = config.get(JdbcSourceOptions.INVERSE_SAMPLING_RATE);
         this.splitSize = config.get(SourceOptions.SNAPSHOT_SPLIT_SIZE);
+        this.splitColumn = config.get(SourceOptions.SNAPSHOT_SPLIT_COLUMN);
         this.fetchSize = config.get(SourceOptions.SNAPSHOT_FETCH_SIZE);
         this.serverTimeZone = config.get(JdbcSourceOptions.SERVER_TIME_ZONE);
         this.connectTimeoutMillis = config.get(JdbcSourceOptions.CONNECT_TIMEOUT_MS);

diff --git a/...src/main/java/org/apache/seatunnel/connectors/cdc/base/dialect/JdbcDataSourceDialect.java b/...src/main/java/org/apache/seatunnel/connectors/cdc/base/dialect/JdbcDataSourceDialect.java
@@ -119,6 +119,26 @@ default List<ConstraintKey> getUniqueKeys(JdbcConnection jdbcConnection, TableId
                 .collect(Collectors.toList());
     }
 
+    default Boolean isUniqueKey(JdbcConnection jdbcConnection, TableId tableId, String columnName)
+            throws SQLException {
+        boolean isUnique = false;
+        if (null != columnName) {
+            DatabaseMetaData metaData = jdbcConnection.connection().getMetaData();
+            ResultSet resultSet =
+                    metaData.getIndexInfo(
+                            tableId.catalog(), tableId.schema(), tableId.table(), false, false);
+
+            while (resultSet.next()) {
+                if (columnName.equalsIgnoreCase(resultSet.getString("COLUMN_NAME"))
+                        && resultSet.getBoolean("NON_UNIQUE") == false) {
+                    isUnique = true;
+                    break;
+                }
+            }
+        }
+        return isUnique;
+    }
+
     default List<ConstraintKey> getConstraintKeys(JdbcConnection jdbcConnection, TableId tableId)
             throws SQLException {
         DatabaseMetaData metaData = jdbcConnection.connection().getMetaData();

diff --git a/...cdc-base/src/main/java/org/apache/seatunnel/connectors/cdc/base/option/SourceOptions.java b/...cdc-base/src/main/java/org/apache/seatunnel/connectors/cdc/base/option/SourceOptions.java
@@ -37,6 +37,13 @@ public class SourceOptions {
                     .withDescription(
                             "The split size (number of rows) of table snapshot, captured tables are split into multiple splits when read the snapshot of table.");
 
+    public static final Option<String> SNAPSHOT_SPLIT_COLUMN =
+            Options.key("snapshot.split.column")
+                    .stringType()
+                    .noDefaultValue()
+                    .withDescription(
+                            "The split column of table snapshot, captured tables are split into multiple splits when read the snapshot of table.");
+
     public static final Option<Integer> SNAPSHOT_FETCH_SIZE =
             Options.key("snapshot.fetch.size")
                     .intType()
@@ -111,7 +118,7 @@ public class SourceOptions {
     public static OptionRule.Builder getBaseRule() {
         return OptionRule.builder()
                 .optional(FORMAT)
-                .optional(SNAPSHOT_SPLIT_SIZE, SNAPSHOT_FETCH_SIZE)
+                .optional(SNAPSHOT_SPLIT_SIZE, SNAPSHOT_FETCH_SIZE, SNAPSHOT_SPLIT_COLUMN)
                 .optional(INCREMENTAL_PARALLELISM)
                 .optional(DEBEZIUM_PROPERTIES);
     }

diff --git a/...unnel/connectors/cdc/base/source/enumerator/splitter/AbstractJdbcSourceChunkSplitter.java b/...unnel/connectors/cdc/base/source/enumerator/splitter/AbstractJdbcSourceChunkSplitter.java
@@ -379,12 +379,28 @@ protected SnapshotSplit createSnapshotSplit(
     protected Column getSplitColumn(
             JdbcConnection jdbc, JdbcDataSourceDialect dialect, TableId tableId)
             throws SQLException {
-        Optional<PrimaryKey> primaryKey = dialect.getPrimaryKey(jdbc, tableId);
         Column splitColumn = null;
+        Table table = dialect.queryTableSchema(jdbc, tableId).getTable();
+
+        // first , compare user defined split column is in the primary key or unique key
+        String sc = null;
+        try {
+            sc = sourceConfig.getSplitColumn();
+        } catch (Exception e) {
+            log.error("Config splitColumn get exception in {}:{}", tableId, e);
+        }
+        Boolean isUniqueKey = dialect.isUniqueKey(jdbc, tableId, sc);
+        if (isUniqueKey) {
+            Column column = table.columnWithName(sc);
+            return column;
+        } else {
+            log.warn("Config splitColumn not exists or nor unique key for table {}", tableId);
+        }
+
+        Optional<PrimaryKey> primaryKey = dialect.getPrimaryKey(jdbc, tableId);
         if (primaryKey.isPresent()) {
             List<String> pkColumns = primaryKey.get().getColumnNames();
 
-            Table table = dialect.queryTableSchema(jdbc, tableId).getTable();
             for (String pkColumn : pkColumns) {
                 Column column = table.columnWithName(pkColumn);
                 if (isEvenlySplitColumn(column)) {
@@ -400,7 +416,6 @@ protected Column getSplitColumn(
 
         List<ConstraintKey> uniqueKeys = dialect.getUniqueKeys(jdbc, tableId);
         if (!uniqueKeys.isEmpty()) {
-            Table table = dialect.queryTableSchema(jdbc, tableId).getTable();
             for (ConstraintKey uniqueKey : uniqueKeys) {
                 List<ConstraintKey.ConstraintKeyColumn> uniqueKeyColumns =
                         uniqueKey.getColumnNames();

diff --git a/...in/java/org/apache/seatunnel/connectors/seatunnel/cdc/mysql/config/MySqlSourceConfig.java b/...in/java/org/apache/seatunnel/connectors/seatunnel/cdc/mysql/config/MySqlSourceConfig.java
@@ -41,6 +41,7 @@ public MySqlSourceConfig(
             List<String> databaseList,
             List<String> tableList,
             int splitSize,
+            String splitColumn,
             double distributionFactorUpper,
             double distributionFactorLower,
             int sampleShardingThreshold,
@@ -64,6 +65,7 @@ public MySqlSourceConfig(
                 databaseList,
                 tableList,
                 splitSize,
+                splitColumn,
                 distributionFactorUpper,
                 distributionFactorLower,
                 sampleShardingThreshold,

diff --git a/.../org/apache/seatunnel/connectors/seatunnel/cdc/mysql/config/MySqlSourceConfigFactory.java b/.../org/apache/seatunnel/connectors/seatunnel/cdc/mysql/config/MySqlSourceConfigFactory.java
@@ -116,6 +116,7 @@ public MySqlSourceConfig create(int subtaskId) {
                 databaseList,
                 tableList,
                 splitSize,
+                splitColumn,
                 distributionFactorUpper,
                 distributionFactorLower,
                 sampleShardingThreshold,

diff --git a/.../java/org/apache/seatunnel/connectors/seatunnel/cdc/oracle/config/OracleSourceConfig.java b/.../java/org/apache/seatunnel/connectors/seatunnel/cdc/oracle/config/OracleSourceConfig.java
@@ -49,6 +49,7 @@ public OracleSourceConfig(
             List<String> databaseList,
             List<String> tableList,
             int splitSize,
+            String splitColumn,
             double distributionFactorUpper,
             double distributionFactorLower,
             int sampleShardingThreshold,
@@ -72,6 +73,7 @@ public OracleSourceConfig(
                 databaseList,
                 tableList,
                 splitSize,
+                splitColumn,
                 distributionFactorUpper,
                 distributionFactorLower,
                 sampleShardingThreshold,

diff --git a/...rg/apache/seatunnel/connectors/seatunnel/cdc/oracle/config/OracleSourceConfigFactory.java b/...rg/apache/seatunnel/connectors/seatunnel/cdc/oracle/config/OracleSourceConfigFactory.java
@@ -144,6 +144,7 @@ public OracleSourceConfig create(int subtask) {
                 databaseList,
                 tableList,
                 splitSize,
+                splitColumn,
                 distributionFactorUpper,
                 distributionFactorLower,
                 sampleShardingThreshold,

diff --git a/...a/org/apache/seatunnel/connectors/seatunnel/cdc/postgres/config/PostgresSourceConfig.java b/...a/org/apache/seatunnel/connectors/seatunnel/cdc/postgres/config/PostgresSourceConfig.java
@@ -36,6 +36,7 @@ public PostgresSourceConfig(
             List<String> databaseList,
             List<String> tableList,
             int splitSize,
+            String splitColumn,
             double distributionFactorUpper,
             double distributionFactorLower,
             int sampleShardingThreshold,
@@ -59,6 +60,7 @@ public PostgresSourceConfig(
                 databaseList,
                 tableList,
                 splitSize,
+                splitColumn,
                 distributionFactorUpper,
                 distributionFactorLower,
                 sampleShardingThreshold,

diff --git a/...pache/seatunnel/connectors/seatunnel/cdc/postgres/config/PostgresSourceConfigFactory.java b/...pache/seatunnel/connectors/seatunnel/cdc/postgres/config/PostgresSourceConfigFactory.java
@@ -115,6 +115,7 @@ public PostgresSourceConfig create(int subtask) {
                 databaseList,
                 tableList,
                 splitSize,
+                splitColumn,
                 distributionFactorUpper,
                 distributionFactorLower,
                 sampleShardingThreshold,

diff --git a/...org/apache/seatunnel/connectors/seatunnel/cdc/sqlserver/config/SqlServerSourceConfig.java b/...org/apache/seatunnel/connectors/seatunnel/cdc/sqlserver/config/SqlServerSourceConfig.java
@@ -41,6 +41,7 @@ public SqlServerSourceConfig(
             List<String> databaseList,
             List<String> tableList,
             int splitSize,
+            String splitColumn,
             double distributionFactorUpper,
             double distributionFactorLower,
             int sampleShardingThreshold,
@@ -64,6 +65,7 @@ public SqlServerSourceConfig(
                 databaseList,
                 tableList,
                 splitSize,
+                splitColumn,
                 distributionFactorUpper,
                 distributionFactorLower,
                 sampleShardingThreshold,

diff --git a/...che/seatunnel/connectors/seatunnel/cdc/sqlserver/config/SqlServerSourceConfigFactory.java b/...che/seatunnel/connectors/seatunnel/cdc/sqlserver/config/SqlServerSourceConfigFactory.java
@@ -83,6 +83,7 @@ public SqlServerSourceConfig create(int subtask) {
                 databaseList,
                 tableList,
                 splitSize,
+                splitColumn,
                 distributionFactorUpper,
                 distributionFactorLower,
                 sampleShardingThreshold,