expression: support tidb encode key function (#51678)

close #51683

pkg/expression/BUILD.bazel

@@ -93,6 +93,8 @@ go_library(
         "//pkg/planner/cascades/base",
         "//pkg/sessionctx/stmtctx",
         "//pkg/sessionctx/variable",
+        "//pkg/store/helper",
+        "//pkg/tablecodec",
         "//pkg/types",
         "//pkg/types/parser_driver",
         "//pkg/util",
@@ -126,6 +128,7 @@ go_library(
         "@com_github_google_uuid//:uuid",
         "@com_github_pingcap_errors//:errors",
         "@com_github_pingcap_failpoint//:failpoint",
+        "@com_github_pingcap_kvproto//pkg/kvrpcpb",
         "@com_github_pingcap_tipb//go-tipb",
         "@com_github_qri_io_jsonschema//:jsonschema",
         "@com_github_tikv_client_go_v2//oracle",

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,20 +20,25 @@ package expression
 
 import (
 	"context"
+	"encoding/hex"
 	"encoding/json"
 	"slices"
 	"strings"
 	"time"
 
 	"github.com/pingcap/errors"
+	"github.com/pingcap/kvproto/pkg/kvrpcpb"
 	"github.com/pingcap/tidb/pkg/expression/contextopt"
 	infoschema "github.com/pingcap/tidb/pkg/infoschema/context"
 	"github.com/pingcap/tidb/pkg/parser"
 	"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/tablecodec"
 	"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 +62,9 @@ var (
 	_ functionClass = &tidbVersionFunctionClass{}
 	_ functionClass = &tidbIsDDLOwnerFunctionClass{}
 	_ functionClass = &tidbDecodePlanFunctionClass{}
+	_ functionClass = &tidbMVCCInfoFunctionClass{}
+	_ functionClass = &tidbEncodeRecordKeyClass{}
+	_ functionClass = &tidbEncodeIndexKeyClass{}
 	_ functionClass = &tidbDecodeKeyFunctionClass{}
 	_ functionClass = &tidbDecodeSQLDigestsFunctionClass{}
 	_ functionClass = &nextValFunctionClass{}
@@ -78,6 +86,9 @@ var (
 	_ builtinFunc = &builtinVersionSig{}
 	_ builtinFunc = &builtinTiDBVersionSig{}
 	_ builtinFunc = &builtinRowCountSig{}
+	_ builtinFunc = &builtinTiDBMVCCInfoSig{}
+	_ builtinFunc = &builtinTiDBEncodeRecordKeySig{}
+	_ builtinFunc = &builtinTiDBEncodeIndexKeySig{}
 	_ builtinFunc = &builtinTiDBDecodeKeySig{}
 	_ builtinFunc = &builtinTiDBDecodeSQLDigestsSig{}
 	_ builtinFunc = &builtinNextValSig{}
@@ -892,6 +903,216 @@ 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
+	}
+	type mvccInfoResult struct {
+		Key  string                        `json:"key"`
+		Resp *kvrpcpb.MvccGetByKeyResponse `json:"mvcc"`
+	}
+	mvccInfo := []*mvccInfoResult{{s, resp}}
+	if tablecodec.IsIndexKey(encodedKey) && !tablecodec.IsTempIndexKey(encodedKey) {
+		tablecodec.IndexKey2TempIndexKey(encodedKey)
+		hexStr := hex.EncodeToString(encodedKey)
+		res, err := h.GetMvccByEncodedKey(encodedKey)
+		if err != nil {
+			return "", false, err
+		}
+		if res.Info != nil && (len(res.Info.Writes) > 0 || len(res.Info.Values) > 0 || res.Info.Lock != nil) {
+			mvccInfo = append(mvccInfo, &mvccInfoResult{hexStr, res})
+		}
+	}
+	js, err := json.Marshal(mvccInfo)
+	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 +1132,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, is infoschema.MetaOnlyInfoSchema, args []Expression, row chunk.Row) ([]byte, bool, error)
+
+// EncodeIndexKeyFromRow is used to encode index key by expressions.
+var EncodeIndexKeyFromRow func(ctx EvalContext, is 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

@@ -1107,6 +1107,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)