lightning: fix id allocator base for double auto_increment field (#1178)

pkg/kv/kv.go

@@ -15,6 +15,7 @@ package kv
 
 import (
 	"fmt"
+	"math"
 
 	"github.com/pingcap/errors"
 	"github.com/pingcap/log"
@@ -196,7 +197,7 @@ func (kvcodec *tableKVEncoder) AddRecord(
 		}
 		if isAutoIncCol {
 			// TODO use auto incremental type
-			_ = kvcodec.tbl.RebaseAutoID(kvcodec.se, value.GetInt64(), false, autoid.RowIDAllocType)
+			_ = kvcodec.tbl.RebaseAutoID(kvcodec.se, getAutoRecordID(value, &col.FieldType), false, autoid.RowIDAllocType)
 		}
 	}
 
@@ -228,6 +229,21 @@ func (kvcodec *tableKVEncoder) AddRecord(
 	return Pairs(pairs), size, nil
 }
 
+// get record value for auto-increment field
+//
+// See: https://github.com/pingcap/tidb/blob/47f0f15b14ed54fc2222f3e304e29df7b05e6805/executor/insert_common.go#L781-L852
+// TODO: merge this with pkg/lightning/backend/kv/sql2kv.go
+func getAutoRecordID(d types.Datum, target *types.FieldType) int64 {
+	switch target.Tp {
+	case mysql.TypeFloat, mysql.TypeDouble:
+		return int64(math.Round(d.GetFloat64()))
+	case mysql.TypeTiny, mysql.TypeShort, mysql.TypeInt24, mysql.TypeLong, mysql.TypeLonglong:
+		return d.GetInt64()
+	default:
+		panic(fmt.Sprintf("unsupported auto-increment field type '%d'", target.Tp))
+	}
+}
+
 // RemoveRecord encode a row of data into KV pairs.
 func (kvcodec *tableKVEncoder) RemoveRecord(
 	row []types.Datum,

pkg/lightning/backend/kv/sql2kv.go

@@ -16,6 +16,8 @@
 package kv
 
 import (
+	"fmt"
+	"math"
 	"math/rand"
 	"sort"
 
@@ -349,7 +351,7 @@ func (kvcodec *tableKVEncoder) Encode(
 			}
 		}
 		if isAutoIncCol {
-			if err := kvcodec.tbl.RebaseAutoID(kvcodec.se, value.GetInt64(), false, autoid.AutoIncrementType); err != nil {
+			if err := kvcodec.tbl.RebaseAutoID(kvcodec.se, getAutoRecordID(value, &col.FieldType), false, autoid.AutoIncrementType); err != nil {
 				return nil, errors.Trace(err)
 			}
 		}
@@ -402,6 +404,20 @@ func (kvcodec *tableKVEncoder) Encode(
 	return kvPairs(pairs), nil
 }
 
+// get record value for auto-increment field
+//
+// See: https://github.com/pingcap/tidb/blob/47f0f15b14ed54fc2222f3e304e29df7b05e6805/executor/insert_common.go#L781-L852
+func getAutoRecordID(d types.Datum, target *types.FieldType) int64 {
+	switch target.Tp {
+	case mysql.TypeFloat, mysql.TypeDouble:
+		return int64(math.Round(d.GetFloat64()))
+	case mysql.TypeTiny, mysql.TypeShort, mysql.TypeInt24, mysql.TypeLong, mysql.TypeLonglong:
+		return d.GetInt64()
+	default:
+		panic(fmt.Sprintf("unsupported auto-increment field type '%d'", target.Tp))
+	}
+}
+
 func (kvs kvPairs) Size() uint64 {
 	size := uint64(0)
 	for _, kv := range kvs {

pkg/lightning/backend/kv/sql2kv_test.go

@@ -214,6 +214,37 @@ func (s *kvSuite) TestEncodeTimestamp(c *C) {
 	}))
 }
 
+func (s *kvSuite) TestEncodeDoubleAutoIncrement(c *C) {
+	tblInfo := mockTableInfo(c, "create table t (id double not null auto_increment, unique key `u_id` (`id`));")
+	tbl, err := tables.TableFromMeta(NewPanickingAllocators(0), tblInfo)
+	c.Assert(err, IsNil)
+
+	logger := log.Logger{Logger: zap.NewNop()}
+
+	encoder, err := NewTableKVEncoder(tbl, &SessionOptions{
+		SQLMode: mysql.ModeStrictAllTables,
+		SysVars: map[string]string{
+			"tidb_row_format_version": "2",
+		},
+	})
+	c.Assert(err, IsNil)
+	pairs, err := encoder.Encode(logger, []types.Datum{
+		types.NewStringDatum("1"),
+	}, 70, []int{0, -1})
+	c.Assert(err, IsNil)
+	c.Assert(pairs, DeepEquals, kvPairs([]common.KvPair{
+		{
+			Key: []uint8{0x74, 0x80, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x1, 0x5f, 0x72, 0x80, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x46},
+			Val: []uint8{0x80, 0x0, 0x1, 0x0, 0x0, 0x0, 0x1, 0x8, 0x0, 0xbf, 0xf0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0},
+		},
+		{
+			Key: []uint8{0x74, 0x80, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x1, 0x5f, 0x69, 0x80, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x1, 0x5, 0xbf, 0xf0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0},
+			Val: []uint8{0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x46},
+		},
+	}))
+	c.Assert(tbl.Allocators(encoder.(*tableKVEncoder).se).Get(autoid.AutoIncrementType).Base(), Equals, int64(70))
+}
+
 func mockTableInfo(c *C, createSQL string) *model.TableInfo {
 	parser := parser.New()
 	node, err := parser.ParseOneStmt(createSQL, "", "")