diff --git a/core/src/main/scala/com/pingcap/tispark/TiBatchWrite.scala b/core/src/main/scala/com/pingcap/tispark/TiBatchWrite.scala index 51806a88eb..3e6b6b48d0 100644 --- a/core/src/main/scala/com/pingcap/tispark/TiBatchWrite.scala +++ b/core/src/main/scala/com/pingcap/tispark/TiBatchWrite.scala @@ -437,7 +437,7 @@ 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( @@ -445,6 +445,14 @@ class TiBatchWrite(@transient val df: DataFrame, ) } + // 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( diff --git a/core/src/test/scala/com/pingcap/tispark/datasource/CheckUnsupportedSuite.scala b/core/src/test/scala/com/pingcap/tispark/datasource/CheckUnsupportedSuite.scala index 50df3897be..8687080a59 100644 --- a/core/src/test/scala/com/pingcap/tispark/datasource/CheckUnsupportedSuite.scala +++ b/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() diff --git a/tikv-client/src/main/java/com/pingcap/tikv/meta/TiColumnInfo.java b/tikv-client/src/main/java/com/pingcap/tikv/meta/TiColumnInfo.java index 34914cb0b3..d1c9e39848 100644 --- a/tikv-client/src/main/java/com/pingcap/tikv/meta/TiColumnInfo.java +++ b/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(); + } } diff --git a/tikv-client/src/main/java/com/pingcap/tikv/meta/TiTableInfo.java b/tikv-client/src/main/java/com/pingcap/tikv/meta/TiTableInfo.java index 39e0827bf6..792ae976a7 100644 --- a/tikv-client/src/main/java/com/pingcap/tikv/meta/TiTableInfo.java +++ b/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; + } }