expression: support tidb encode key function

pkg/expression/BUILD.bazel

@@ -90,6 +90,7 @@ go_library(
         "//pkg/parser/types",
         "//pkg/sessionctx/stmtctx",
         "//pkg/sessionctx/variable",
+        "//pkg/store/helper",
         "//pkg/types",
         "//pkg/types/parser_driver",
         "//pkg/util",

pkg/expression/builtin.go

@@ -910,7 +910,10 @@ var funcs = map[string]functionClass{
 	ast.JSONLength:        &jsonLengthFunctionClass{baseFunctionClass{ast.JSONLength, 1, 2}},
 
 	// TiDB internal function.
-	ast.TiDBDecodeKey: &tidbDecodeKeyFunctionClass{baseFunctionClass{ast.TiDBDecodeKey, 1, 1}},
+	ast.TiDBDecodeKey:       &tidbDecodeKeyFunctionClass{baseFunctionClass{ast.TiDBDecodeKey, 1, 1}},
+	ast.TiDBMVCCInfo:        &tidbMVCCInfoFunctionClass{baseFunctionClass: baseFunctionClass{ast.TiDBMVCCInfo, 1, 1}},
+	ast.TiDBEncodeRecordKey: &tidbEncodeRecordKeyClass{baseFunctionClass{ast.TiDBEncodeRecordKey, 3, -1}},
+	ast.TiDBEncodeIndexKey:  &tidbEncodeIndexKeyClass{baseFunctionClass{ast.TiDBEncodeIndexKey, 4, -1}},
 	// This function is used to show tidb-server version info.
 	ast.TiDBVersion:          &tidbVersionFunctionClass{baseFunctionClass{ast.TiDBVersion, 0, 0}},
 	ast.TiDBIsDDLOwner:       &tidbIsDDLOwnerFunctionClass{baseFunctionClass{ast.TiDBIsDDLOwner, 0, 0}},

pkg/expression/builtin_info.go

@@ -20,6 +20,7 @@ package expression
 
 import (
 	"context"
+	"encoding/hex"
 	"encoding/json"
 	"slices"
 	"strings"
@@ -32,6 +33,7 @@ import (
 	"github.com/pingcap/tidb/pkg/parser/ast"
 	"github.com/pingcap/tidb/pkg/parser/mysql"
 	"github.com/pingcap/tidb/pkg/sessionctx/variable"
+	"github.com/pingcap/tidb/pkg/store/helper"
 	"github.com/pingcap/tidb/pkg/types"
 	"github.com/pingcap/tidb/pkg/util/chunk"
 	"github.com/pingcap/tidb/pkg/util/plancodec"
@@ -57,6 +59,9 @@ var (
 	_ functionClass = &tidbVersionFunctionClass{}
 	_ functionClass = &tidbIsDDLOwnerFunctionClass{}
 	_ functionClass = &tidbDecodePlanFunctionClass{}
+	_ functionClass = &tidbMVCCInfoFunctionClass{}
+	_ functionClass = &tidbEncodeRecordKeyClass{}
+	_ functionClass = &tidbEncodeIndexKeyClass{}
 	_ functionClass = &tidbDecodeKeyFunctionClass{}
 	_ functionClass = &tidbDecodeSQLDigestsFunctionClass{}
 	_ functionClass = &nextValFunctionClass{}
@@ -78,6 +83,9 @@ var (
 	_ builtinFunc = &builtinVersionSig{}
 	_ builtinFunc = &builtinTiDBVersionSig{}
 	_ builtinFunc = &builtinRowCountSig{}
+	_ builtinFunc = &builtinTiDBMVCCInfoSig{}
+	_ builtinFunc = &builtinTiDBEncodeRecordKeySig{}
+	_ builtinFunc = &builtinTiDBEncodeIndexKeySig{}
 	_ builtinFunc = &builtinTiDBDecodeKeySig{}
 	_ builtinFunc = &builtinTiDBDecodeSQLDigestsSig{}
 	_ builtinFunc = &builtinNextValSig{}
@@ -892,6 +900,173 @@ func (b *builtinRowCountSig) evalInt(ctx EvalContext, row chunk.Row) (res int64,
 	return res, false, nil
 }
 
+type tidbMVCCInfoFunctionClass struct {
+	baseFunctionClass
+	contextopt.KVStorePropReader
+}
+
+// RequiredOptionalEvalProps implements the RequireOptionalEvalProps interface.
+func (c *tidbMVCCInfoFunctionClass) RequiredOptionalEvalProps() OptionalEvalPropKeySet {
+	return c.KVStorePropReader.RequiredOptionalEvalProps()
+}
+
+func (c *tidbMVCCInfoFunctionClass) getFunction(ctx BuildContext, args []Expression) (builtinFunc, error) {
+	bf, err := newBaseBuiltinFuncWithTp(ctx, c.funcName, args, types.ETString, types.ETString)
+	if err != nil {
+		return nil, err
+	}
+	store, err := c.GetKVStore(ctx)
+	if err != nil {
+		return nil, err
+	}
+	hStore, ok := store.(helper.Storage)
+	if !ok {
+		return nil, errors.New("storage is not a helper.Storage")
+	}
+	sig := &builtinTiDBMVCCInfoSig{baseBuiltinFunc: bf, helper: helper.NewHelper(hStore)}
+	return sig, nil
+}
+
+type builtinTiDBMVCCInfoSig struct {
+	baseBuiltinFunc
+	helper *helper.Helper
+}
+
+func (b *builtinTiDBMVCCInfoSig) Clone() builtinFunc {
+	newSig := &builtinTiDBMVCCInfoSig{}
+	newSig.cloneFrom(&b.baseBuiltinFunc)
+	newSig.helper = helper.NewHelper(b.helper.Store)
+	return newSig
+}
+
+// evalString evals a builtinTiDBMVCCInfoSig.
+func (b *builtinTiDBMVCCInfoSig) evalString(ctx EvalContext, row chunk.Row) (string, bool, error) {
+	s, isNull, err := b.args[0].EvalString(ctx, row)
+	if isNull || err != nil {
+		return "", isNull, err
+	}
+
+	encodedKey, err := hex.DecodeString(s)
+	if err != nil {
+		return "", false, err
+	}
+	resp, err := b.helper.GetMvccByEncodedKey(encodedKey)
+	if err != nil {
+		return "", false, err
+	}
+	js, err := json.Marshal(resp)
+	if err != nil {
+		return "", false, err
+	}
+	return string(js), false, nil
+}
+
+type tidbEncodeRecordKeyClass struct {
+	baseFunctionClass
+}
+
+func (c *tidbEncodeRecordKeyClass) getFunction(ctx BuildContext, args []Expression) (builtinFunc, error) {
+	if err := c.verifyArgs(args); err != nil {
+		return nil, err
+	}
+	evalTps := make([]types.EvalType, 0, len(args))
+	evalTps = append(evalTps, types.ETString, types.ETString)
+	for _, arg := range args[2:] {
+		evalTps = append(evalTps, arg.GetType().EvalType())
+	}
+	bf, err := newBaseBuiltinFuncWithTp(ctx, c.funcName, args, types.ETString, evalTps...)
+	if err != nil {
+		return nil, err
+	}
+	sig := &builtinTiDBEncodeRecordKeySig{baseBuiltinFunc: bf}
+	return sig, nil
+}
+
+type builtinTiDBEncodeRecordKeySig struct {
+	baseBuiltinFunc
+	contextopt.InfoSchemaPropReader
+}
+
+// RequiredOptionalEvalProps implements the RequireOptionalEvalProps interface.
+func (b *builtinTiDBEncodeRecordKeySig) RequiredOptionalEvalProps() OptionalEvalPropKeySet {
+	return b.InfoSchemaPropReader.RequiredOptionalEvalProps()
+}
+
+func (b *builtinTiDBEncodeRecordKeySig) Clone() builtinFunc {
+	newSig := &builtinTiDBEncodeRecordKeySig{}
+	newSig.cloneFrom(&b.baseBuiltinFunc)
+	return newSig
+}
+
+// evalString evals a builtinTiDBEncodeRecordKeySig.
+func (b *builtinTiDBEncodeRecordKeySig) evalString(ctx EvalContext, row chunk.Row) (string, bool, error) {
+	is, err := b.GetDomainInfoSchema(ctx)
+	if err != nil {
+		return "", true, err
+	}
+	if EncodeRecordKeyFromRow == nil {
+		return "", false, errors.New("EncodeRecordKeyFromRow is not initialized")
+	}
+	recordKey, isNull, err := EncodeRecordKeyFromRow(ctx, is, b.args, row)
+	if isNull || err != nil {
+		return "", isNull, err
+	}
+	return hex.EncodeToString(recordKey), false, nil
+}
+
+type tidbEncodeIndexKeyClass struct {
+	baseFunctionClass
+}
+
+func (c *tidbEncodeIndexKeyClass) getFunction(ctx BuildContext, args []Expression) (builtinFunc, error) {
+	if err := c.verifyArgs(args); err != nil {
+		return nil, err
+	}
+	evalTps := make([]types.EvalType, 0, len(args))
+	evalTps = append(evalTps, types.ETString, types.ETString, types.ETString)
+	for _, arg := range args[3:] {
+		evalTps = append(evalTps, arg.GetType().EvalType())
+	}
+	bf, err := newBaseBuiltinFuncWithTp(ctx, c.funcName, args, types.ETString, evalTps...)
+	if err != nil {
+		return nil, err
+	}
+	sig := &builtinTiDBEncodeIndexKeySig{baseBuiltinFunc: bf}
+	return sig, nil
+}
+
+type builtinTiDBEncodeIndexKeySig struct {
+	baseBuiltinFunc
+	contextopt.InfoSchemaPropReader
+}
+
+// RequiredOptionalEvalProps implements the RequireOptionalEvalProps interface.
+func (b *builtinTiDBEncodeIndexKeySig) RequiredOptionalEvalProps() OptionalEvalPropKeySet {
+	return b.InfoSchemaPropReader.RequiredOptionalEvalProps()
+}
+
+func (b *builtinTiDBEncodeIndexKeySig) Clone() builtinFunc {
+	newSig := &builtinTiDBEncodeIndexKeySig{}
+	newSig.cloneFrom(&b.baseBuiltinFunc)
+	return newSig
+}
+
+// evalString evals a builtinTiDBEncodeIndexKeySig.
+func (b *builtinTiDBEncodeIndexKeySig) evalString(ctx EvalContext, row chunk.Row) (string, bool, error) {
+	is, err := b.GetDomainInfoSchema(ctx)
+	if err != nil {
+		return "", true, err
+	}
+	if EncodeIndexKeyFromRow == nil {
+		return "", false, errors.New("EncodeIndexKeyFromRow is not initialized")
+	}
+	idxKey, isNull, err := EncodeIndexKeyFromRow(ctx, is, b.args, row)
+	if isNull || err != nil {
+		return "", isNull, err
+	}
+	return hex.EncodeToString(idxKey), false, nil
+}
+
 type tidbDecodeKeyFunctionClass struct {
 	baseFunctionClass
 }
@@ -911,6 +1086,12 @@ func (c *tidbDecodeKeyFunctionClass) getFunction(ctx BuildContext, args []Expres
 // DecodeKeyFromString is used to decode key by expressions
 var DecodeKeyFromString func(types.Context, infoschema.InfoSchemaMetaVersion, string) string
 
+// EncodeRecordKeyFromRow is used to encode record key by expressions.
+var EncodeRecordKeyFromRow func(ctx EvalContext, isVer infoschema.InfoSchemaMetaVersion, args []Expression, row chunk.Row) ([]byte, bool, error)
+
+// EncodeIndexKeyFromRow is used to encode index key by expressions.
+var EncodeIndexKeyFromRow func(ctx EvalContext, isVer infoschema.InfoSchemaMetaVersion, args []Expression, row chunk.Row) ([]byte, bool, error)
+
 type builtinTiDBDecodeKeySig struct {
 	baseBuiltinFunc
 	contextopt.InfoSchemaPropReader

pkg/expression/builtin_info_vec.go

@@ -385,6 +385,30 @@ func (b *builtinVersionSig) vecEvalString(ctx EvalContext, input *chunk.Chunk, r
 	return nil
 }
 
+func (b *builtinTiDBMVCCInfoSig) vectorized() bool {
+	return false
+}
+
+func (b *builtinTiDBMVCCInfoSig) vecEvalString(ctx EvalContext, input *chunk.Chunk, result *chunk.Column) error {
+	return errors.Errorf("not implemented")
+}
+
+func (b *builtinTiDBEncodeRecordKeySig) vectorized() bool {
+	return false
+}
+
+func (b *builtinTiDBEncodeRecordKeySig) vecEvalString(ctx EvalContext, input *chunk.Chunk, result *chunk.Column) error {
+	return errors.Errorf("not implemented")
+}
+
+func (b *builtinTiDBEncodeIndexKeySig) vectorized() bool {
+	return false
+}
+
+func (b *builtinTiDBEncodeIndexKeySig) vecEvalString(ctx EvalContext, input *chunk.Chunk, result *chunk.Column) error {
+	return errors.Errorf("not implemented")
+}
+
 func (b *builtinTiDBDecodeKeySig) vectorized() bool {
 	return true
 }

pkg/expression/integration_test/BUILD.bazel

@@ -8,7 +8,7 @@ go_test(
         "main_test.go",
     ],
     flaky = True,
-    shard_count = 25,
+    shard_count = 26,
     deps = [
         "//pkg/config",
         "//pkg/domain",

pkg/expression/integration_test/integration_test.go

@@ -580,6 +580,31 @@ func TestTiDBDecodeKeyFunc(t *testing.T) {
 	tk.MustQuery(sql).Check(testkit.Rows(rs))
 }
 
+func TestTiDBEncodeKey(t *testing.T) {
+	store := testkit.CreateMockStore(t)
+	tk := testkit.NewTestKit(t, store)
+	tk.MustExec("use test")
+	tk.MustExec("create table t (a int primary key, b int);")
+	tk.MustExec("insert into t values (1, 1);")
+	err := tk.QueryToErr("select tidb_encode_record_key('test', 't1', 0);")
+	require.ErrorContains(t, err, "doesn't exist")
+	tk.MustQuery("select tidb_encode_record_key('test', 't', 1);").
+		Check(testkit.Rows("7480000000000000685f728000000000000001"))
+
+	tk.MustExec("alter table t add index i(b);")
+	err = tk.QueryToErr("select tidb_encode_index_key('test', 't', 'i1', 1);")
+	require.ErrorContains(t, err, "index not found")
+	tk.MustQuery("select tidb_encode_index_key('test', 't', 'i', 1, 1);").
+		Check(testkit.Rows("7480000000000000685f698000000000000001038000000000000001038000000000000001"))
+
+	tk.MustExec("create table t1 (a int primary key, b int) partition by hash(a) partitions 4;")
+	tk.MustExec("insert into t1 values (1, 1);")
+	tk.MustQuery("select tidb_encode_record_key('test', 't1(p1)', 1);").Check(testkit.Rows("74800000000000006d5f728000000000000001"))
+	rs := tk.MustQuery("select tidb_mvcc_info('74800000000000006d5f728000000000000001');")
+	mvccInfo := rs.Rows()[0][0].(string)
+	require.NotEqual(t, mvccInfo, `{"info":{}}`)
+}
+
 func TestIssue9710(t *testing.T) {
 	store := testkit.CreateMockStore(t)
 

pkg/parser/ast/functions.go

@@ -357,6 +357,9 @@ const (
 
 	// TiDB internal function.
 	TiDBDecodeKey       = "tidb_decode_key"
+	TiDBMVCCInfo        = "tidb_mvcc_info"
+	TiDBEncodeRecordKey = "tidb_encode_record_key"
+	TiDBEncodeIndexKey  = "tidb_encode_index_key"
 	TiDBDecodeBase64Key = "tidb_decode_base64_key"
 
 	// MVCC information fetching function.