[Fix][Connector-V2] Fix partitioning column selection logic (#10319)

Author: dybyte

seatunnel-connectors-v2/connector-cdc/connector-cdc-base/src/main/java/org/apache/seatunnel/connectors/cdc/base/source/enumerator/splitter/AbstractJdbcSourceChunkSplitter.java

@@ -67,6 +67,10 @@ public Collection<SnapshotSplit> generateSplits(TableId tableId) {
             long start = System.currentTimeMillis();
 
             Column splitColumn = getSplitColumn(jdbc, dialect, tableId);
+            log.info(
+                    "Chosen split column {} for table {}",
+                    splitColumn != null ? splitColumn.name() : "null",
+                    tableId);
             List<SnapshotSplit> splits = new ArrayList<>();
             if (splitColumn == null) {
                 if (sourceConfig.isExactlyOnce()) {
@@ -429,15 +433,11 @@ protected Column getSplitColumn(
 
         Optional<PrimaryKey> primaryKey = dialect.getPrimaryKey(jdbc, tableId);
         if (primaryKey.isPresent()) {
-            List<String> pkColumns = primaryKey.get().getColumnNames();
-
-            for (String pkColumn : pkColumns) {
-                Column column = table.columnWithName(pkColumn);
-                if (isEvenlySplitColumn(column)) {
-                    splitColumn = columnComparable(splitColumn, column);
-                    if (sqlTypePriority(splitColumn) == 1) {
-                        return splitColumn;
-                    }
+            Column firstColumn = table.columnWithName(primaryKey.get().getColumnNames().get(0));
+            if (isEvenlySplitColumn(firstColumn)) {
+                splitColumn = columnComparable(splitColumn, firstColumn);
+                if (sqlTypePriority(splitColumn) == 1) {
+                    return splitColumn;
                 }
             }
         } else {
@@ -447,15 +447,12 @@ protected Column getSplitColumn(
         List<ConstraintKey> uniqueKeys = dialect.getUniqueKeys(jdbc, tableId);
         if (!uniqueKeys.isEmpty()) {
             for (ConstraintKey uniqueKey : uniqueKeys) {
-                List<ConstraintKey.ConstraintKeyColumn> uniqueKeyColumns =
-                        uniqueKey.getColumnNames();
-                for (ConstraintKey.ConstraintKeyColumn uniqueKeyColumn : uniqueKeyColumns) {
-                    Column column = table.columnWithName(uniqueKeyColumn.getColumnName());
-                    if (isEvenlySplitColumn(column)) {
-                        splitColumn = columnComparable(splitColumn, column);
-                        if (sqlTypePriority(splitColumn) == 1) {
-                            return splitColumn;
-                        }
+                Column firstColumn =
+                        table.columnWithName(uniqueKey.getColumnNames().get(0).getColumnName());
+                if (isEvenlySplitColumn(firstColumn)) {
+                    splitColumn = columnComparable(splitColumn, firstColumn);
+                    if (sqlTypePriority(splitColumn) == 1) {
+                        return splitColumn;
                     }
                 }
             }

seatunnel-connectors-v2/connector-cdc/connector-cdc-base/src/test/java/jdbc/source/JdbcSourceChunkSplitterTest.java

@@ -48,16 +48,36 @@
 import java.util.List;
 import java.util.Optional;
 
-public class JdbcSourceChunkSplitterTest {
+class JdbcSourceChunkSplitterTest {
 
     @Test
-    public void splitColumnTest() throws SQLException {
+    void splitColumnTest() throws SQLException {
         TestJdbcSourceChunkSplitter testJdbcSourceChunkSplitter =
                 new TestJdbcSourceChunkSplitter(null, new TestSourceDialect());
         Column splitColumn =
                 testJdbcSourceChunkSplitter.getSplitColumn(
                         null, new TestSourceDialect(), new TableId("", "", ""));
-        Assertions.assertEquals(splitColumn.typeName(), "tinyint");
+        Assertions.assertEquals("varchar", splitColumn.typeName());
+    }
+
+    @Test
+    void splitColumnTestWithUniqueKey() throws SQLException {
+        TestJdbcSourceChunkSplitter testJdbcSourceChunkSplitter =
+                new TestJdbcSourceChunkSplitter(null, new TestSourceDialectWithUniqueKey());
+        Column splitColumn =
+                testJdbcSourceChunkSplitter.getSplitColumn(
+                        null, new TestSourceDialectWithUniqueKey(), new TableId("", "", ""));
+        Assertions.assertEquals("bigint", splitColumn.typeName());
+    }
+
+    @Test
+    void splitColumnTestWithUniqueKey_2() throws SQLException {
+        TestJdbcSourceChunkSplitter testJdbcSourceChunkSplitter =
+                new TestJdbcSourceChunkSplitter(null, new TestSourceDialectWithUniqueKey_2());
+        Column splitColumn =
+                testJdbcSourceChunkSplitter.getSplitColumn(
+                        null, new TestSourceDialectWithUniqueKey_2(), new TableId("", "", ""));
+        Assertions.assertEquals("int", splitColumn.typeName());
     }
 
     private class TestJdbcSourceChunkSplitter extends AbstractJdbcSourceChunkSplitter {
@@ -247,4 +267,66 @@ public List<ConstraintKey> getUniqueKeys(JdbcConnection jdbcConnection, TableId
             return new ArrayList<ConstraintKey>();
         }
     }
+
+    private class TestSourceDialectWithUniqueKey extends TestSourceDialect {
+
+        @Override
+        public Optional<PrimaryKey> getPrimaryKey(JdbcConnection jdbcConnection, TableId tableId)
+                throws SQLException {
+            return Optional.of(PrimaryKey.of("pkName", Arrays.asList("bigint_col")));
+        }
+
+        @Override
+        public List<ConstraintKey> getUniqueKeys(JdbcConnection jdbcConnection, TableId tableId)
+                throws SQLException {
+            List<ConstraintKey> keys = new ArrayList<>();
+
+            keys.add(
+                    ConstraintKey.of(
+                            ConstraintKey.ConstraintType.UNIQUE_KEY,
+                            "uk_1",
+                            Arrays.asList(
+                                    ConstraintKey.ConstraintKeyColumn.of(
+                                            "string_col", ConstraintKey.ColumnSortType.ASC),
+                                    ConstraintKey.ConstraintKeyColumn.of(
+                                            "int", ConstraintKey.ColumnSortType.ASC))));
+
+            return keys;
+        }
+    }
+
+    private class TestSourceDialectWithUniqueKey_2 extends TestSourceDialect {
+
+        @Override
+        public Optional<PrimaryKey> getPrimaryKey(JdbcConnection jdbcConnection, TableId tableId)
+                throws SQLException {
+            return Optional.of(PrimaryKey.of("pkName", Arrays.asList("bigint_col")));
+        }
+
+        @Override
+        public List<ConstraintKey> getUniqueKeys(JdbcConnection jdbcConnection, TableId tableId)
+                throws SQLException {
+            List<ConstraintKey> keys = new ArrayList<>();
+
+            keys.add(
+                    ConstraintKey.of(
+                            ConstraintKey.ConstraintType.UNIQUE_KEY,
+                            "uk_1",
+                            Arrays.asList(
+                                    ConstraintKey.ConstraintKeyColumn.of(
+                                            "string_col", ConstraintKey.ColumnSortType.ASC))));
+
+            keys.add(
+                    ConstraintKey.of(
+                            ConstraintKey.ConstraintType.UNIQUE_KEY,
+                            "uk_2",
+                            Arrays.asList(
+                                    ConstraintKey.ConstraintKeyColumn.of(
+                                            "int", ConstraintKey.ColumnSortType.ASC),
+                                    ConstraintKey.ConstraintKeyColumn.of(
+                                            "smallint", ConstraintKey.ColumnSortType.ASC))));
+
+            return keys;
+        }
+    }
 }