Fix alter primary key for write (#1644) (#1645)

.ci/integration_test.groovy

@@ -12,6 +12,7 @@ def call(ghprbActualCommit, ghprbCommentBody, ghprbPullId, ghprbPullTitle, ghprb
     def PARALLEL_NUMBER = 18
     def TEST_REGION_SIZE = "normal"
     def TEST_TIFLASH = "false"
+    def TEST_ALTER_PRIMARY_KEY = "true"
 
     // parse tidb branch
     def m1 = ghprbCommentBody =~ /tidb\s*=\s*([^\s\\]+)(\s|\\|$)/
@@ -59,12 +60,18 @@ def call(ghprbActualCommit, ghprbCommentBody, ghprbPullId, ghprbPullTitle, ghprb
         TEST_REGION_SIZE = "${m7[0][1]}"
     }
 
-    // parse test mode
+    // parse test tiflash
     def m8 = ghprbCommentBody =~ /test-flash\s*=\s*([^\s\\]+)(\s|\\|$)/
     if (m8) {
         TEST_TIFLASH = "${m8[0][1]}"
     }
 
+    // parse test alter primary key
+    def m9 = ghprbCommentBody =~ /test-alter-primary-key\s*=\s*([^\s\\]+)(\s|\\|$)/
+    if (m9) {
+        TEST_ALTER_PRIMARY_KEY = "${m9[0][1]}"
+    }
+
     groovy.lang.Closure readfile = { filename ->
         def file = readFile filename
         return file.split("\n") as List
@@ -127,6 +134,11 @@ def call(ghprbActualCommit, ghprbCommentBody, ghprbPullId, ghprbPullTitle, ghprb
                         """
                         stash includes: "tiflash/**", name: "tiflash_binary"
                     }
+                    // alter-primary-key
+                    if (TEST_ALTER_PRIMARY_KEY == "false") {
+                        sh "sed -i 's/alter-primary-key = true/alter-primary-key = false/' config/tidb.toml"
+                        sh "sed -i 's/alter-primary-key = true/alter-primary-key = false/' config/tidb-4.0.toml"
+                    }
 
                     stash includes: "bin/**", name: "binaries"
 

config/tidb-4.0.toml

@@ -8,4 +8,9 @@ split-table = true
 delay-clean-table-lock = 60000
 
 # enable-table-lock is used to control table lock feature. Default is false, indicate the table lock feature is disabled.
-enable-table-lock = true
+enable-table-lock = true
+
+# alter-primary-key is used to control alter primary key feature. Default is false, indicate the alter primary key feature is disabled.
+# If it is true, we can add the primary key by "alter table", but we may not be able to drop the primary key.
+# In order to support "drop primary key" operation , this flag must be true and the table does not have the pkIsHandle flag.
+alter-primary-key = true

config/tidb.toml

@@ -40,6 +40,11 @@ enable-chunk = true
 # enable-table-lock is used to control table lock feature. Default is false, indicate the table lock feature is disabled.
 enable-table-lock = true
 
+# alter-primary-key is used to control alter primary key feature. Default is false, indicate the alter primary key feature is disabled.
+# If it is true, we can add the primary key by "alter table", but we may not be able to drop the primary key.
+# In order to support "drop primary key" operation , this flag must be true and the table does not have the pkIsHandle flag.
+alter-primary-key = true
+
 [log]
 # Log level: info, debug, warn, error, fatal.
 level = "info"

core/src/main/scala/com/pingcap/tispark/write/TiBatchWriteTable.scala

@@ -153,7 +153,7 @@ class TiBatchWriteTable(
             if (colsMapInTiDB.contains(colsInDf(colOffset))) {
               if (colsMapInTiDB(colsInDf(colOffset)).isAutoIncrement) {
                 val index = row._2 + 1
-                rowIDAllocator.getShardRowId(index)
+                rowIDAllocator.getAutoIncId(index)
               } else {
                 data._1
               }

core/src/test/scala/com/pingcap/tispark/datasource/AutoIncrementSuite.scala

@@ -21,6 +21,37 @@ import org.apache.spark.sql.types._
 
 class AutoIncrementSuite extends BaseBatchWriteTest("test_datasource_auto_increment") {
 
+  test("alter primary key + auto increment + shard row bits") {
+    if (!isEnableAlterPrimaryKey) {
+      cancel("enable alter-primary-key by changing tidb.toml")
+    }
+
+    val schema = StructType(List(StructField("j", LongType)))
+
+    jdbcUpdate(
+      s"create table $dbtable(i int NOT NULL AUTO_INCREMENT, j int NOT NULL, primary key (i)) SHARD_ROW_ID_BITS=4")
+
+    val tiTableInfo = ti.tiSession.getCatalog.getTable(dbPrefix + database, table)
+    assert(!tiTableInfo.isPkHandle)
+
+    (1L until 10L).foreach { i =>
+      jdbcUpdate(s"insert into $dbtable (j) values(${i * 2 - 1})")
+
+      tidbWrite(List(Row(i * 2)), schema)
+    }
+
+    println(listToString(queryTiDBViaJDBC(s"select _tidb_rowid, i, j from $dbtable")))
+
+    spark.sql(s"select * from $table").show
+
+    val maxI = queryTiDBViaJDBC(s"select max(i) from $dbtable").head.head.toString.toLong
+    assert(maxI < 10000000)
+
+    val maxTiDBRowID =
+      queryTiDBViaJDBC(s"select max(_tidb_rowid) from $dbtable").head.head.toString.toLong
+    assert(maxTiDBRowID > 10000000)
+  }
+
   // Duplicate entry '2' for key 'PRIMARY'
   // currently user provided auto increment value is not supported!
   ignore("auto increment: user provide id") {
@@ -64,6 +95,10 @@ class AutoIncrementSuite extends BaseBatchWriteTest("test_datasource_auto_increm
   }
 
   test("bigint signed tidb overflow") {
+    if (isEnableAlterPrimaryKey) {
+      cancel()
+    }
+
     val schema = StructType(List(StructField("j", LongType)))
 
     jdbcUpdate(
@@ -91,6 +126,9 @@ class AutoIncrementSuite extends BaseBatchWriteTest("test_datasource_auto_increm
   }
 
   test("bigint signed tispark overflow") {
+    if (isEnableAlterPrimaryKey) {
+      cancel()
+    }
 
     val schema = StructType(List(StructField("j", LongType)))
 
@@ -121,6 +159,9 @@ class AutoIncrementSuite extends BaseBatchWriteTest("test_datasource_auto_increm
   }
 
   test("bigint unsigned tidb overflow") {
+    if (isEnableAlterPrimaryKey) {
+      cancel()
+    }
 
     val schema = StructType(List(StructField("j", LongType)))
 
@@ -148,6 +189,10 @@ class AutoIncrementSuite extends BaseBatchWriteTest("test_datasource_auto_increm
   }
 
   test("bigint unsigned tispark overflow") {
+    if (isEnableAlterPrimaryKey) {
+      cancel()
+    }
+
     val schema = StructType(List(StructField("j", LongType)))
 
     jdbcUpdate(
@@ -177,6 +222,10 @@ class AutoIncrementSuite extends BaseBatchWriteTest("test_datasource_auto_increm
   }
 
   test("tinyint signed tidb overflow") {
+    if (isEnableAlterPrimaryKey) {
+      cancel()
+    }
+
     val schema = StructType(List(StructField("j", LongType)))
 
     jdbcUpdate(
@@ -204,6 +253,10 @@ class AutoIncrementSuite extends BaseBatchWriteTest("test_datasource_auto_increm
   }
 
   test("tinyint signed tispark overflow") {
+    if (isEnableAlterPrimaryKey) {
+      cancel()
+    }
+
     val schema = StructType(List(StructField("j", LongType)))
 
     jdbcUpdate(
@@ -232,6 +285,10 @@ class AutoIncrementSuite extends BaseBatchWriteTest("test_datasource_auto_increm
   }
 
   test("tinyint unsigned tidb overflow") {
+    if (isEnableAlterPrimaryKey) {
+      cancel()
+    }
+
     val schema = StructType(List(StructField("j", LongType)))
 
     jdbcUpdate(
@@ -258,6 +315,10 @@ class AutoIncrementSuite extends BaseBatchWriteTest("test_datasource_auto_increm
   }
 
   test("tinyint unsigned tispark overflow") {
+    if (isEnableAlterPrimaryKey) {
+      cancel()
+    }
+
     val schema = StructType(List(StructField("j", LongType)))
 
     jdbcUpdate(

core/src/test/scala/com/pingcap/tispark/datasource/ColumnMappingSuite.scala

@@ -69,6 +69,10 @@ class ColumnMappingSuite
   }
 
   test("Test different column order without auto increment column") {
+    if (isEnableAlterPrimaryKey) {
+      cancel()
+    }
+
     jdbcUpdate(
       s"create table $dbtable(i int primary key auto_increment, s varchar(128), c varchar(128))")
 

core/src/test/scala/org/apache/spark/sql/BaseTiSparkTest.scala

@@ -115,6 +115,12 @@ class BaseTiSparkTest extends QueryTest with SharedSQLContext {
   def setLogLevel(level: String): Unit =
     spark.sparkContext.setLogLevel(level)
 
+  protected def isEnableAlterPrimaryKey: Boolean = {
+    val conn = TiDBUtils.createConnectionFactory(jdbcUrl)()
+    val tiDBJDBCClient = new TiDBJDBCClient(conn)
+    tiDBJDBCClient.isEnableAlterPrimaryKey
+  }
+
   protected def isEnableTableLock: Boolean = {
     val conn = TiDBUtils.createConnectionFactory(jdbcUrl)()
     val tiDBJDBCClient = new TiDBJDBCClient(conn)

core/src/test/scala/org/apache/spark/sql/catalyst/plans/statistics/StatisticsTestSuite.scala

@@ -158,6 +158,9 @@ class StatisticsTestSuite extends BasePlanTest {
     {
       val tableName = "full_data_type_table_idx"
       test(query) {
+        if (isEnableAlterPrimaryKey) {
+          cancel()
+        }
         val df = spark.sql(query)
         checkIsTableScan(df, tableName)
       }
@@ -178,6 +181,9 @@ class StatisticsTestSuite extends BasePlanTest {
     case (query, idxName) =>
       val tableName = "full_data_type_table_idx"
       test(query) {
+        if (isEnableAlterPrimaryKey) {
+          cancel()
+        }
         val df = spark.sql(query)
         checkIsCoveringIndexScan(df, tableName)
         checkIndex(df, idxName)

tikv-client/src/main/java/com/pingcap/tikv/TiDBJDBCClient.java

@@ -39,13 +39,20 @@ public class TiDBJDBCClient implements AutoCloseable {
   private static final int DELAY_CLEAN_TABLE_LOCK_DEFAULT = 0;
   private static final String TIDB_ROW_FORMAT_VERSION_SQL = "select @@tidb_row_format_version";
   private static final int TIDB_ROW_FORMAT_VERSION_DEFAULT = 1;
+  private static final String ALTER_PRIMARY_KEY_KEY = "alter-primary-key";
+  private static final Boolean ALTER_PRIMARY_KEY_DEFAULT = false;
   private final Logger logger = LoggerFactory.getLogger(getClass().getName());
   private final Connection connection;
 
   public TiDBJDBCClient(Connection connection) {
     this.connection = connection;
   }
 
+  public boolean isEnableAlterPrimaryKey() throws IOException, SQLException {
+    Map<String, Object> configMap = readConfMapFromTiDB();
+    return (Boolean) configMap.getOrDefault(ALTER_PRIMARY_KEY_KEY, ALTER_PRIMARY_KEY_DEFAULT);
+  }
+
   public boolean isEnableTableLock() throws IOException, SQLException {
     Map<String, Object> configMap = readConfMapFromTiDB();
     Object enableTableLock =

tikv-client/src/main/java/com/pingcap/tikv/allocator/RowIDAllocator.java

@@ -52,6 +52,10 @@ private RowIDAllocator(long maxShardRowIDBits, long dbId, long step, TiConfigura
     this.conf = conf;
   }
 
+  public long getAutoIncId(long index) {
+    return index + getStart();
+  }
+
   /**
    * @param index should >= 1
    * @return