forbid write to table with generated column (pingcap#907)

core/src/main/scala/com/pingcap/tispark/TiBatchWrite.scala

@@ -437,14 +437,22 @@ class TiBatchWrite(@transient val df: DataFrame,
   }
 
   @throws(classOf[TiBatchWriteException])
-  private def checkUnsupported(): Unit =
+  private def checkUnsupported(): Unit = {
     // write to partition table
     if (tiTableInfo.isPartitionEnabled) {
       throw new TiBatchWriteException(
         "tispark currently does not support write data to partition table!"
       )
     }
 
+    // write to table with generated column
+    if (tiTableInfo.hasGeneratedColumn) {
+      throw new TiBatchWriteException(
+        "tispark currently does not support write data to table with generated column!"
+      )
+    }
+  }
+
   private def checkColumnNumbers(): Unit = {
     if (!tiTableInfo.hasAutoIncrementColumn && colsInDf.length != tableColSize) {
       throw new TiBatchWriteException(

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

@@ -5,16 +5,6 @@ import org.apache.spark.sql.Row
 import org.apache.spark.sql.types.{IntegerType, StringType, StructField, StructType}
 
 class CheckUnsupportedSuite extends BaseDataSourceTest("test_datasource_check_unsupported") {
-  private val row1 = Row(null, "Hello")
-  private val row2 = Row(2, "TiDB")
-  private val row3 = Row(3, "Spark")
-
-  private val schema = StructType(
-    List(
-      StructField("i", IntegerType),
-      StructField("s", StringType)
-    )
-  )
 
   override def beforeAll(): Unit =
     super.beforeAll()
@@ -31,6 +21,17 @@ class CheckUnsupportedSuite extends BaseDataSourceTest("test_datasource_check_un
       s"insert into $dbtable values(null, 'Hello')"
     )
 
+    val row1 = Row(null, "Hello")
+    val row2 = Row(2, "TiDB")
+    val row3 = Row(3, "Spark")
+
+    val schema = StructType(
+      List(
+        StructField("i", IntegerType),
+        StructField("s", StringType)
+      )
+    )
+
     {
       val caught = intercept[TiBatchWriteException] {
         tidbWrite(List(row2, row3), schema)
@@ -44,6 +45,55 @@ class CheckUnsupportedSuite extends BaseDataSourceTest("test_datasource_check_un
     testTiDBSelect(Seq(row1))
   }
 
+  test("Check Virtual Generated Column") {
+    dropTable()
+    jdbcUpdate(
+      s"create table $dbtable(i INT, c1 INT, c2 INT,  c3 INT AS (c1 + c2))"
+    )
+
+    val row1 = Row(1, 2, 3)
+    val schema = StructType(
+      List(
+        StructField("i", IntegerType),
+        StructField("c1", IntegerType),
+        StructField("c2", IntegerType)
+      )
+    )
+
+    val caught = intercept[TiBatchWriteException] {
+      tidbWrite(List(row1), schema)
+    }
+    assert(
+      caught.getMessage
+        .equals("tispark currently does not support write data to table with generated column!")
+    )
+
+  }
+
+  test("Check Stored Generated Column") {
+    dropTable()
+    jdbcUpdate(
+      s"create table $dbtable(i INT, c1 INT, c2 INT,  c3 INT AS (c1 + c2) STORED)"
+    )
+
+    val row1 = Row(1, 2, 3)
+    val schema = StructType(
+      List(
+        StructField("i", IntegerType),
+        StructField("c1", IntegerType),
+        StructField("c2", IntegerType)
+      )
+    )
+    val caught = intercept[TiBatchWriteException] {
+      tidbWrite(List(row1), schema)
+    }
+    assert(
+      caught.getMessage
+        .equals("tispark currently does not support write data to table with generated column!")
+    )
+
+  }
+
   override def afterAll(): Unit =
     try {
       dropTable()

tikv-client/src/main/java/com/pingcap/tikv/meta/TiColumnInfo.java

@@ -47,6 +47,7 @@ public class TiColumnInfo implements Serializable {
   // If version is 0 then timestamp's default value will be read and decoded as local timezone.
   // if version is 1 then timestamp's default value will be read and decoded as utc.
   private final long version;
+  private final String generatedExprString;
 
   static TiColumnInfo getRowIdColumn(int offset) {
     return new TiColumnInfo(-1, "_tidb_rowid", offset, IntegerType.ROW_ID_TYPE, true);
@@ -64,7 +65,8 @@ public TiColumnInfo(
       @JsonProperty("origin_default") String originalDefaultValue,
       @JsonProperty("default") String defaultValue,
       @JsonProperty("comment") String comment,
-      @JsonProperty("version") long version) {
+      @JsonProperty("version") long version,
+      @JsonProperty("generated_expr_string") String generatedExprString) {
     this.id = id;
     this.name = requireNonNull(name, "column name is null").getL();
     this.offset = offset;
@@ -77,6 +79,7 @@ public TiColumnInfo(
     // Refactor against original tidb code
     this.isPrimaryKey = (type.getFlag() & PK_MASK) > 0;
     this.version = version;
+    this.generatedExprString = generatedExprString;
   }
 
   public TiColumnInfo(
@@ -88,7 +91,8 @@ public TiColumnInfo(
       String originalDefaultValue,
       String defaultValue,
       String comment,
-      long version) {
+      long version,
+      String generatedExprString) {
     this.id = id;
     this.name = requireNonNull(name, "column name is null").toLowerCase();
     this.offset = offset;
@@ -99,6 +103,7 @@ public TiColumnInfo(
     this.originDefaultValue = originalDefaultValue;
     this.isPrimaryKey = (type.getFlag() & PK_MASK) > 0;
     this.version = version;
+    this.generatedExprString = generatedExprString;
   }
 
   TiColumnInfo copyWithoutPrimaryKey() {
@@ -114,7 +119,8 @@ TiColumnInfo copyWithoutPrimaryKey() {
         this.originDefaultValue,
         this.defaultValue,
         this.comment,
-        this.version);
+        this.version,
+        this.generatedExprString);
   }
 
   @VisibleForTesting
@@ -129,6 +135,7 @@ public TiColumnInfo(long id, String name, int offset, DataType type, boolean isP
     this.originDefaultValue = "1";
     this.defaultValue = "";
     this.version = DataType.COLUMN_VERSION_FLAG;
+    this.generatedExprString = "";
   }
 
   public long getId() {
@@ -341,4 +348,12 @@ public boolean canSkip(boolean isPkHandle) {
     //      2. if generated or generatedStored is false
     return isPrimaryKey & isPkHandle;
   }
+
+  public String getGeneratedExprString() {
+    return generatedExprString;
+  }
+
+  public boolean isGeneratedColumn() {
+    return generatedExprString != null && !generatedExprString.isEmpty();
+  }
 }

tikv-client/src/main/java/com/pingcap/tikv/meta/TiTableInfo.java

@@ -214,7 +214,8 @@ private TiColumnInfo copyColumn(TiColumnInfo col) {
             col.getOriginDefaultValue(),
             col.getDefaultValue(),
             col.getComment(),
-            col.getVersion())
+            col.getVersion(),
+            col.getGeneratedExprString())
         .copyWithoutPrimaryKey();
   }
 
@@ -253,4 +254,13 @@ public boolean isPartitionEnabled() {
     if (partitionInfo == null) return false;
     return partitionInfo.isEnable();
   }
+
+  public boolean hasGeneratedColumn() {
+    for (TiColumnInfo col : getColumns()) {
+      if (col.isGeneratedColumn()) {
+        return true;
+      }
+    }
+    return false;
+  }
 }