From bb05618fe06d42f25bd9d75d5bffd14f471e5b17 Mon Sep 17 00:00:00 2001
From: marsishandsome <marsishandsome@gmail.com>
Date: Fri, 19 Mar 2021 10:29:06 +0800
Subject: [PATCH] fix tidb_enable_clustered_index only support global variable

---
 .../pingcap/tispark/write/TiBatchWriteTable.scala   |  5 +++++
 .../scala/org/apache/spark/sql/IssueTestSuite.scala |  4 +++-
 .../spark/sql/clustered/ClusteredIndexTest.scala    |  6 +++---
 .../apache/spark/sql/test/SharedSQLContext.scala    | 13 +++++++++++++
 4 files changed, 24 insertions(+), 4 deletions(-)

diff --git a/core/src/main/scala/com/pingcap/tispark/write/TiBatchWriteTable.scala b/core/src/main/scala/com/pingcap/tispark/write/TiBatchWriteTable.scala
index 7fe060ac80..efa3ce73e9 100644
--- a/core/src/main/scala/com/pingcap/tispark/write/TiBatchWriteTable.scala
+++ b/core/src/main/scala/com/pingcap/tispark/write/TiBatchWriteTable.scala
@@ -317,6 +317,11 @@ class TiBatchWriteTable(
   }
 
   def checkUnsupported(): Unit = {
+    if (tiTableInfo.isCommonHandle) {
+      throw new TiBatchWriteException(
+        "tispark currently does not support write data to table with clustered index!")
+    }
+
     // write to table with auto random column
     if (tiTableInfo.hasAutoRandomColumn) {
       throw new TiBatchWriteException(
diff --git a/core/src/test/scala/org/apache/spark/sql/IssueTestSuite.scala b/core/src/test/scala/org/apache/spark/sql/IssueTestSuite.scala
index d58d8cbf1c..e58315f506 100644
--- a/core/src/test/scala/org/apache/spark/sql/IssueTestSuite.scala
+++ b/core/src/test/scala/org/apache/spark/sql/IssueTestSuite.scala
@@ -42,8 +42,8 @@ class IssueTestSuite extends BaseTiSparkTest {
       cancel("currently tidb instance does not support clustered index")
     }
     spark.sqlContext.setConf(TiConfigConst.USE_INDEX_SCAN_FIRST, "true")
+    enableClusteredIndex()
     tidbStmt.execute("""
-        |SET tidb_enable_clustered_index = 1;
         |drop table if exists `tispark_test`.`clustered0`;
         |CREATE TABLE `tispark_test`.`clustered0` (
         |  `col_bit0` bit(1) not null,
@@ -60,6 +60,8 @@ class IssueTestSuite extends BaseTiSparkTest {
     spark.sql(s"$sql").show(200, false)
     runTest(sql, skipJDBC = true)
     spark.sqlContext.setConf(TiConfigConst.USE_INDEX_SCAN_FIRST, "false")
+
+    disableClusteredIndex()
   }
 
   test("partition table with date partition column name") {
diff --git a/core/src/test/scala/org/apache/spark/sql/clustered/ClusteredIndexTest.scala b/core/src/test/scala/org/apache/spark/sql/clustered/ClusteredIndexTest.scala
index 1fe02ee9b4..421ba2aff0 100644
--- a/core/src/test/scala/org/apache/spark/sql/clustered/ClusteredIndexTest.scala
+++ b/core/src/test/scala/org/apache/spark/sql/clustered/ClusteredIndexTest.scala
@@ -41,8 +41,8 @@ trait ClusteredIndexTest extends BaseTiSparkTest with BaseEnumerateDataTypesTest
   override def test(): Unit = ???
 
   override def afterAll(): Unit = {
-    if(supportClusteredIndex) {
-      executeTiDBSQL("SET tidb_enable_clustered_index = 0;")
+    if (supportClusteredIndex) {
+      disableClusteredIndex()
     }
     super.afterAll()
   }
@@ -88,7 +88,7 @@ trait ClusteredIndexTest extends BaseTiSparkTest with BaseEnumerateDataTypesTest
   }
 
   protected def test(schema: Schema): Unit = {
-    executeTiDBSQL("SET tidb_enable_clustered_index = 1;")
+    enableClusteredIndex()
     executeTiDBSQL(s"drop table if exists `$dbName`.`${schema.tableName}`;")
     executeTiDBSQL(schema.toString)
 
diff --git a/core/src/test/scala/org/apache/spark/sql/test/SharedSQLContext.scala b/core/src/test/scala/org/apache/spark/sql/test/SharedSQLContext.scala
index d2f18d5622..c4d67ae0c7 100644
--- a/core/src/test/scala/org/apache/spark/sql/test/SharedSQLContext.scala
+++ b/core/src/test/scala/org/apache/spark/sql/test/SharedSQLContext.scala
@@ -175,6 +175,19 @@ trait SharedSQLContext
 
   protected def initializeStatement(): Unit = {
     _statement = _tidbConnection.createStatement()
+    disableClusteredIndex()
+  }
+
+  protected def enableClusteredIndex(): Unit = {
+    _statement.execute("SET GLOBAL tidb_enable_clustered_index = 1")
+    _statement.close()
+    _statement = _tidbConnection.createStatement()
+  }
+
+  protected def disableClusteredIndex(): Unit = {
+    _statement.execute("SET GLOBAL tidb_enable_clustered_index = 0")
+    _statement.close()
+    _statement = _tidbConnection.createStatement()
   }
 
   protected def timeZoneOffset: String = SharedSQLContext.timeZoneOffset