kvencoder: add EncodeMetaAutoID method (pingcap#5272)

kv/mock.go

@@ -99,6 +99,14 @@ func (t *mockTxn) GetMemBuffer() MemBuffer {
 	return nil
 }
 
+// NewMockTxn new a mockTxn.
+func NewMockTxn() Transaction {
+	return &mockTxn{
+		opts:  make(map[Option]interface{}),
+		valid: true,
+	}
+}
+
 // mockStorage is used to start a must commit-failed txn.
 type mockStorage struct {
 }

meta/meta.go

@@ -140,6 +140,13 @@ func (m *Meta) parseTableID(key string) (int64, error) {
 	return n, errors.Trace(err)
 }
 
+// GenAutoTableIDIDKeyValue generate meta key by dbID, tableID and coresponding value by autoID.
+func (m *Meta) GenAutoTableIDIDKeyValue(dbID, tableID, autoID int64) (key, value []byte) {
+	dbKey := m.dbKey(dbID)
+	autoTableIDKey := m.autoTableIDKey(tableID)
+	return m.txn.EncodeHashAutoIDKeyValue(dbKey, autoTableIDKey, autoID)
+}
+
 // GenAutoTableID adds step to the auto ID of the table and returns the sum.
 func (m *Meta) GenAutoTableID(dbID, tableID, step int64) (int64, error) {
 	// Check if DB exists.

structure/hash.go

@@ -62,6 +62,15 @@ func (t *TxStructure) HGet(key []byte, field []byte) ([]byte, error) {
 	return value, errors.Trace(err)
 }
 
+func (t *TxStructure) hashFieldIntegerVal(val int64) []byte {
+	return []byte(strconv.FormatInt(val, 10))
+}
+
+// EncodeHashAutoIDKeyValue returns the hash key-value generated by the key and the field
+func (t *TxStructure) EncodeHashAutoIDKeyValue(key []byte, field []byte, val int64) (k, v []byte) {
+	return t.encodeHashDataKey(key, field), t.hashFieldIntegerVal(val)
+}
+
 // HInc increments the integer value of a hash field, by step, returns
 // the value after the increment.
 func (t *TxStructure) HInc(key []byte, field []byte, step int64) (int64, error) {
@@ -78,7 +87,7 @@ func (t *TxStructure) HInc(key []byte, field []byte, step int64) (int64, error)
 			}
 		}
 		base += step
-		return []byte(strconv.FormatInt(base, 10)), nil
+		return t.hashFieldIntegerVal(base), nil
 	})
 
 	return base, errors.Trace(err)

util/kvencoder/kv_encoder.go

@@ -23,6 +23,7 @@ import (
 	"github.com/pingcap/tidb"
 	"github.com/pingcap/tidb/domain"
 	"github.com/pingcap/tidb/kv"
+	"github.com/pingcap/tidb/meta"
 	"github.com/pingcap/tidb/meta/autoid"
 	"github.com/pingcap/tidb/mysql"
 	"github.com/pingcap/tidb/store/tikv"
@@ -44,7 +45,7 @@ type KvPair struct {
 // KvEncoder is an encoder that transfer sql to key-value pairs.
 type KvEncoder interface {
 	// Encode transfers sql to kv pairs.
-	// Before use Encode() method, please make sure you already call Schema() method.
+	// Before use Encode() method, please make sure you already created schame by calling ExecDDLSQL() method.
 	// NOTE: now we just support transfers insert statement to kv pairs.
 	// (if we wanna support other statement, we need to add a kv.Storage parameter,
 	// and pass tikv store in.)
@@ -54,6 +55,9 @@ type KvEncoder interface {
 	// ExecDDLSQL executes ddl sql, you must use it to create schema infos.
 	ExecDDLSQL(sql string) error
 
+	// EncodeMetaAutoID encode the table meta info, autoID to coresponding key-value pair.
+	EncodeMetaAutoID(dbID, tableID, autoID int64) (KvPair, error)
+
 	// Close cleanup the kvEncoder.
 	Close() error
 }
@@ -113,6 +117,13 @@ func (e *kvEncoder) Encode(sql string, tableID int64) (kvPairs []KvPair, affecte
 	return kvPairs, e.se.GetSessionVars().StmtCtx.AffectedRows(), nil
 }
 
+func (e *kvEncoder) EncodeMetaAutoID(dbID, tableID, autoID int64) (KvPair, error) {
+	mockTxn := kv.NewMockTxn()
+	m := meta.NewMeta(mockTxn)
+	k, v := m.GenAutoTableIDIDKeyValue(dbID, tableID, autoID)
+	return KvPair{Key: k, Val: v}, nil
+}
+
 func (e *kvEncoder) ExecDDLSQL(sql string) error {
 	_, err := e.se.Execute(goctx.Background(), sql)
 	if err != nil {

util/kvencoder/kv_encoder_test.go

@@ -16,6 +16,7 @@ package kvenc
 import (
 	"bytes"
 	"fmt"
+	"strconv"
 	"testing"
 
 	"github.com/juju/errors"
@@ -24,6 +25,7 @@ import (
 	"github.com/pingcap/tidb/domain"
 	"github.com/pingcap/tidb/kv"
 	"github.com/pingcap/tidb/store/tikv"
+	"github.com/pingcap/tidb/structure"
 	"github.com/pingcap/tidb/tablecodec"
 	"github.com/pingcap/tidb/types"
 	"github.com/pingcap/tidb/util/codec"
@@ -438,5 +440,60 @@ func (s *testKvEncoderSuite) TestSimpleKeyEncode(c *C) {
 			c.Assert(bytes.Compare(kv.Key, expectIdxKey), Equals, 0)
 		}
 	}
+}
+
+var (
+	mMetaPrefix    = []byte("m")
+	mDBPrefix      = "DB"
+	mTableIDPrefix = "TID"
+)
+
+func dbKey(dbID int64) []byte {
+	return []byte(fmt.Sprintf("%s:%d", mDBPrefix, dbID))
+}
+
+func autoTableIDKey(tableID int64) []byte {
+	return []byte(fmt.Sprintf("%s:%d", mTableIDPrefix, tableID))
+}
+
+func encodeHashDataKey(key []byte, field []byte) kv.Key {
+	ek := make([]byte, 0, len(mMetaPrefix)+len(key)+len(field)+30)
+	ek = append(ek, mMetaPrefix...)
+	ek = codec.EncodeBytes(ek, key)
+	ek = codec.EncodeUint(ek, uint64(structure.HashData))
+	return codec.EncodeBytes(ek, field)
+}
+
+func hashFieldIntegerVal(val int64) []byte {
+	return []byte(strconv.FormatInt(val, 10))
+}
+
+func (s *testKvEncoderSuite) TestEncodeMetaAutoID(c *C) {
+	encoder, err := New("test", nil)
+	c.Assert(err, IsNil)
+	defer encoder.Close()
+
+	dbID := int64(1)
+	tableID := int64(10)
+	autoID := int64(10000000111)
+	kvPair, err := encoder.EncodeMetaAutoID(dbID, tableID, autoID)
+	c.Assert(err, IsNil)
+
+	expectKey := encodeHashDataKey(dbKey(dbID), autoTableIDKey(tableID))
+	expectVal := hashFieldIntegerVal(autoID)
+
+	c.Assert(bytes.Compare(kvPair.Key, expectKey), Equals, 0)
+	c.Assert(bytes.Compare(kvPair.Val, expectVal), Equals, 0)
+
+	dbID = 10
+	tableID = 1
+	autoID = -1
+	kvPair, err = encoder.EncodeMetaAutoID(dbID, tableID, autoID)
+	c.Assert(err, IsNil)
+
+	expectKey = encodeHashDataKey(dbKey(dbID), autoTableIDKey(tableID))
+	expectVal = hashFieldIntegerVal(autoID)
 
+	c.Assert(bytes.Compare(kvPair.Key, expectKey), Equals, 0)
+	c.Assert(bytes.Compare(kvPair.Val, expectVal), Equals, 0)
 }