expression: support tidb encode key function

pkg/expression/BUILD.bazel

@@ -92,6 +92,7 @@ go_library(
         "//pkg/planner/cascades/base",
         "//pkg/sessionctx/stmtctx",
         "//pkg/sessionctx/variable",
+        "//pkg/store/helper",
         "//pkg/types",
         "//pkg/types/parser_driver",
         "//pkg/util",

pkg/expression/builtin.go

@@ -941,7 +941,10 @@ var funcs = map[string]functionClass{
 	ast.VecAsText:               &vecAsTextFunctionClass{baseFunctionClass{ast.VecAsText, 1, 1}},
 
 	// 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,8 +33,10 @@ 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/dbterror/plannererrors"
 	"github.com/pingcap/tidb/pkg/util/plancodec"
 	"github.com/pingcap/tidb/pkg/util/printer"
 	"github.com/pingcap/tipb/go-tipb"
@@ -57,6 +60,9 @@ var (
 	_ functionClass = &tidbVersionFunctionClass{}
 	_ functionClass = &tidbIsDDLOwnerFunctionClass{}
 	_ functionClass = &tidbDecodePlanFunctionClass{}
+	_ functionClass = &tidbMVCCInfoFunctionClass{}
+	_ functionClass = &tidbEncodeRecordKeyClass{}
+	_ functionClass = &tidbEncodeIndexKeyClass{}
 	_ functionClass = &tidbDecodeKeyFunctionClass{}
 	_ functionClass = &tidbDecodeSQLDigestsFunctionClass{}
 	_ functionClass = &nextValFunctionClass{}
@@ -78,6 +84,9 @@ var (
 	_ builtinFunc = &builtinVersionSig{}
 	_ builtinFunc = &builtinTiDBVersionSig{}
 	_ builtinFunc = &builtinRowCountSig{}
+	_ builtinFunc = &builtinTiDBMVCCInfoSig{}
+	_ builtinFunc = &builtinTiDBEncodeRecordKeySig{}
+	_ builtinFunc = &builtinTiDBEncodeIndexKeySig{}
 	_ builtinFunc = &builtinTiDBDecodeKeySig{}
 	_ builtinFunc = &builtinTiDBDecodeSQLDigestsSig{}
 	_ builtinFunc = &builtinNextValSig{}
@@ -892,6 +901,200 @@ func (b *builtinRowCountSig) evalInt(ctx EvalContext, row chunk.Row) (res int64,
 	return res, false, nil
 }
 
+type tidbMVCCInfoFunctionClass struct {
+	baseFunctionClass
+}
+
+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
+	}
+	sig := &builtinTiDBMVCCInfoSig{baseBuiltinFunc: bf}
+	return sig, nil
+}
+
+type builtinTiDBMVCCInfoSig struct {
+	baseBuiltinFunc
+	contextopt.KVStorePropReader
+}
+
+// RequiredOptionalEvalProps implements the RequireOptionalEvalProps interface.
+func (b *builtinTiDBMVCCInfoSig) RequiredOptionalEvalProps() OptionalEvalPropKeySet {
+	return b.KVStorePropReader.RequiredOptionalEvalProps()
+}
+
+func (b *builtinTiDBMVCCInfoSig) Clone() builtinFunc {
+	newSig := &builtinTiDBMVCCInfoSig{}
+	newSig.cloneFrom(&b.baseBuiltinFunc)
+	return newSig
+}
+
+// evalString evals a builtinTiDBMVCCInfoSig.
+func (b *builtinTiDBMVCCInfoSig) evalString(ctx EvalContext, row chunk.Row) (string, bool, error) {
+	if !ctx.RequestVerification("", "", "", mysql.SuperPriv) {
+		return "", false, plannererrors.ErrSpecificAccessDenied.FastGenByArgs("SUPER")
+	}
+	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
+	}
+	store, err := b.GetKVStore(ctx)
+	if err != nil {
+		return "", isNull, err
+	}
+	hStore, ok := store.(helper.Storage)
+	if !ok {
+		return "", isNull, errors.New("storage is not a helper.Storage")
+	}
+	h := helper.NewHelper(hStore)
+	resp, err := h.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(ctx.GetEvalCtx()).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
+	contextopt.SessionVarsPropReader
+}
+
+// RequiredOptionalEvalProps implements the RequireOptionalEvalProps interface.
+func (b *builtinTiDBEncodeRecordKeySig) RequiredOptionalEvalProps() OptionalEvalPropKeySet {
+	return b.InfoSchemaPropReader.RequiredOptionalEvalProps() |
+		b.SessionVarsPropReader.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 {
+		if errors.ErrorEqual(err, plannererrors.ErrSpecificAccessDenied) {
+			sv, err2 := b.GetSessionVars(ctx)
+			if err2 != nil {
+				return "", isNull, err
+			}
+			tblName, isNull, err2 := b.args[1].EvalString(ctx, row)
+			if err2 != nil || isNull {
+				return "", isNull, err
+			}
+			return "", isNull, plannererrors.ErrTableaccessDenied.FastGenByArgs("SELECT", sv.User.AuthUsername, sv.User.AuthHostname, tblName)
+		}
+		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(ctx.GetEvalCtx()).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
+	contextopt.SessionVarsPropReader
+}
+
+// RequiredOptionalEvalProps implements the RequireOptionalEvalProps interface.
+func (b *builtinTiDBEncodeIndexKeySig) RequiredOptionalEvalProps() OptionalEvalPropKeySet {
+	return b.InfoSchemaPropReader.RequiredOptionalEvalProps() |
+		b.SessionVarsPropReader.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 {
+		if errors.ErrorEqual(err, plannererrors.ErrSpecificAccessDenied) {
+			sv, err2 := b.GetSessionVars(ctx)
+			if err2 != nil {
+				return "", isNull, err
+			}
+			tblName, isNull, err2 := b.args[1].EvalString(ctx, row)
+			if err2 != nil || isNull {
+				return "", isNull, err
+			}
+			return "", isNull, plannererrors.ErrTableaccessDenied.FastGenByArgs("SELECT", sv.User.AuthUsername, sv.User.AuthHostname, tblName)
+		}
+		return "", isNull, err
+	}
+	return hex.EncodeToString(idxKey), false, nil
+}
+
 type tidbDecodeKeyFunctionClass struct {
 	baseFunctionClass
 }
@@ -911,6 +1114,12 @@ func (c *tidbDecodeKeyFunctionClass) getFunction(ctx BuildContext, args []Expres
 // DecodeKeyFromString is used to decode key by expressions
 var DecodeKeyFromString func(types.Context, infoschema.MetaOnlyInfoSchema, string) string
 
+// EncodeRecordKeyFromRow is used to encode record key by expressions.
+var EncodeRecordKeyFromRow func(ctx EvalContext, isVer infoschema.MetaOnlyInfoSchema, args []Expression, row chunk.Row) ([]byte, bool, error)
+
+// EncodeIndexKeyFromRow is used to encode index key by expressions.
+var EncodeIndexKeyFromRow func(ctx EvalContext, isVer infoschema.MetaOnlyInfoSchema, 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 = 42,
+    shard_count = 43,
     deps = [
         "//pkg/config",
         "//pkg/domain",

pkg/expression/integration_test/integration_test.go

@@ -1106,6 +1106,45 @@ 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":{}}`)
+
+	tk.MustExec("create user 'alice'@'%';")
+	tk.MustExec("flush privileges;")
+	tk2 := testkit.NewTestKit(t, store)
+	err = tk2.Session().Auth(&auth.UserIdentity{Username: "alice", Hostname: "localhost"}, nil, nil, nil)
+	require.NoError(t, err)
+	err = tk2.QueryToErr("select tidb_mvcc_info('74800000000000006d5f728000000000000001');")
+	require.ErrorContains(t, err, "Access denied")
+	err = tk2.QueryToErr("select tidb_encode_record_key('test', 't1(p1)', 1);")
+	require.ErrorContains(t, err, "SELECT command denied")
+	err = tk2.QueryToErr("select tidb_encode_index_key('test', 't', 'i1', 1);")
+	require.ErrorContains(t, err, "SELECT command denied")
+	tk.MustExec("grant select on test.t1 to 'alice'@'%';")
+	tk2.MustQuery("select tidb_encode_record_key('test', 't1(p1)', 1);").Check(testkit.Rows("74800000000000006d5f728000000000000001"))
+}
+
 func TestIssue9710(t *testing.T) {
 	store := testkit.CreateMockStore(t)
 

pkg/parser/ast/functions.go

@@ -368,6 +368,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.

pkg/planner/core/core_init.go

@@ -78,6 +78,8 @@ func init() {
 	expression.EvalSimpleAst = evalAstExpr
 	expression.BuildSimpleExpr = buildSimpleExpr
 	expression.DecodeKeyFromString = decodeKeyFromString
+	expression.EncodeRecordKeyFromRow = encodeHandleFromRow
+	expression.EncodeIndexKeyFromRow = encodeIndexKeyFromRow
 	plannerutil.EvalAstExprWithPlanCtx = evalAstExprWithPlanCtx
 	plannerutil.RewriteAstExprWithPlanCtx = rewriteAstExprWithPlanCtx
 	DefaultDisabledLogicalRulesList = new(atomic.Value)