From 955bf39e80232b9ee778de5905c877d810fda8fe Mon Sep 17 00:00:00 2001 From: tangenta Date: Tue, 12 Mar 2024 11:10:48 +0800 Subject: [PATCH 01/21] expression: support tidb encode key function --- pkg/expression/BUILD.bazel | 1 + pkg/expression/builtin.go | 5 +- pkg/expression/builtin_info.go | 160 +++++++++++++++++++++++- pkg/expression/builtin_info_vec.go | 24 ++++ pkg/parser/ast/functions.go | 3 + pkg/planner/core/expression_rewriter.go | 110 ++++++++++++++++ 6 files changed, 300 insertions(+), 3 deletions(-) diff --git a/pkg/expression/BUILD.bazel b/pkg/expression/BUILD.bazel index bf5c92a855792..2e101559c1730 100644 --- a/pkg/expression/BUILD.bazel +++ b/pkg/expression/BUILD.bazel @@ -90,6 +90,7 @@ go_library( "//pkg/privilege", "//pkg/sessionctx/stmtctx", "//pkg/sessionctx/variable", + "//pkg/store/helper", "//pkg/types", "//pkg/types/parser_driver", "//pkg/util", diff --git a/pkg/expression/builtin.go b/pkg/expression/builtin.go index d4f1e2c71f169..02f6059ea40b5 100644 --- a/pkg/expression/builtin.go +++ b/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{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}}, diff --git a/pkg/expression/builtin_info.go b/pkg/expression/builtin_info.go index a52c249d47ba1..74636d92721dc 100644 --- a/pkg/expression/builtin_info.go +++ b/pkg/expression/builtin_info.go @@ -20,6 +20,7 @@ package expression import ( "context" + "encoding/hex" "encoding/json" "slices" "strings" @@ -27,12 +28,14 @@ import ( "github.com/pingcap/errors" "github.com/pingcap/tidb/pkg/expression/contextopt" + "github.com/pingcap/tidb/pkg/kv" "github.com/pingcap/tidb/pkg/parser" "github.com/pingcap/tidb/pkg/parser/ast" "github.com/pingcap/tidb/pkg/parser/model" "github.com/pingcap/tidb/pkg/parser/mysql" "github.com/pingcap/tidb/pkg/privilege" "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" "github.com/pingcap/tidb/pkg/util/chunk" @@ -59,6 +62,9 @@ var ( _ functionClass = &tidbVersionFunctionClass{} _ functionClass = &tidbIsDDLOwnerFunctionClass{} _ functionClass = &tidbDecodePlanFunctionClass{} + _ functionClass = &tidbMVCCInfoFunctionClass{} + _ functionClass = &tidbEncodeRecordKeyClass{} + _ functionClass = &tidbEncodeIndexKeyClass{} _ functionClass = &tidbDecodeKeyFunctionClass{} _ functionClass = &tidbDecodeSQLDigestsFunctionClass{} _ functionClass = &nextValFunctionClass{} @@ -80,6 +86,9 @@ var ( _ builtinFunc = &builtinVersionSig{} _ builtinFunc = &builtinTiDBVersionSig{} _ builtinFunc = &builtinRowCountSig{} + _ builtinFunc = &builtinTiDBMVCCInfoSig{} + _ builtinFunc = &builtinTiDBEncodeRecordKeySig{} + _ builtinFunc = &builtinTiDBEncodeIndexKeySig{} _ builtinFunc = &builtinTiDBDecodeKeySig{} _ builtinFunc = &builtinTiDBDecodeSQLDigestsSig{} _ builtinFunc = &builtinNextValSig{} @@ -831,6 +840,149 @@ 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 + } + var store helper.Storage + if ctxStore, ok := ctx.GetStore().(helper.Storage); ok { + store = ctxStore + } else { + return nil, errors.New("storage is not a helper.Storage") + } + sig := &builtinTiDBMVCCInfoSig{baseBuiltinFunc: bf, helper: helper.NewHelper(store)} + 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{bf} + return sig, nil +} + +type builtinTiDBEncodeRecordKeySig struct { + baseBuiltinFunc +} + +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) { + fn := ctx.Value(TiDBEncodeRecordFunctionKey) + if fn == nil { + return "", false, errors.New("missing encode record function") + } + encode := fn.(func(ctx EvalContext, args []Expression, row chunk.Row) (kv.Key, bool, error)) + recordKey, isNull, err := encode(ctx, 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{bf} + return sig, nil +} + +type builtinTiDBEncodeIndexKeySig struct { + baseBuiltinFunc +} + +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) { + fn := ctx.Value(TiDBEncodeIndexFunctionKey) + if fn == nil { + return "", false, errors.New("missing encode index function") + } + encode := fn.(func(ctx EvalContext, args []Expression, row chunk.Row) (kv.Key, bool, error)) + idxKey, isNull, err := encode(ctx, b.args, row) + if isNull || err != nil { + return "", isNull, err + } + return hex.EncodeToString(idxKey), false, nil +} + type tidbDecodeKeyFunctionClass struct { baseFunctionClass } @@ -878,8 +1030,12 @@ func (k TiDBDecodeKeyFunctionKeyType) String() string { return "tidb_decode_key" } -// TiDBDecodeKeyFunctionKey is used to identify the decoder function in context. -const TiDBDecodeKeyFunctionKey TiDBDecodeKeyFunctionKeyType = 0 +const ( + // TiDBDecodeKeyFunctionKey is used to identify the decoder function in context. + TiDBDecodeKeyFunctionKey TiDBDecodeKeyFunctionKeyType = 0 + TiDBEncodeRecordFunctionKey TiDBDecodeKeyFunctionKeyType = 1 + TiDBEncodeIndexFunctionKey TiDBDecodeKeyFunctionKeyType = 2 +) type tidbDecodeSQLDigestsFunctionClass struct { baseFunctionClass diff --git a/pkg/expression/builtin_info_vec.go b/pkg/expression/builtin_info_vec.go index 9c91b2c1048bc..cf6a40e672e1c 100644 --- a/pkg/expression/builtin_info_vec.go +++ b/pkg/expression/builtin_info_vec.go @@ -353,6 +353,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 } diff --git a/pkg/parser/ast/functions.go b/pkg/parser/ast/functions.go index db6625e14135a..674ae5e3484e2 100644 --- a/pkg/parser/ast/functions.go +++ b/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. diff --git a/pkg/planner/core/expression_rewriter.go b/pkg/planner/core/expression_rewriter.go index 4ccf0168f817a..20ad15e69b3fc 100644 --- a/pkg/planner/core/expression_rewriter.go +++ b/pkg/planner/core/expression_rewriter.go @@ -26,6 +26,7 @@ import ( "github.com/pingcap/tidb/pkg/expression" "github.com/pingcap/tidb/pkg/expression/aggregation" "github.com/pingcap/tidb/pkg/infoschema" + "github.com/pingcap/tidb/pkg/kv" "github.com/pingcap/tidb/pkg/parser/ast" "github.com/pingcap/tidb/pkg/parser/charset" "github.com/pingcap/tidb/pkg/parser/model" @@ -162,6 +163,8 @@ func buildSimpleExpr(ctx expression.BuildContext, node ast.ExprNode, opts ...exp rewriter.schema = expression.NewSchema() } rewriter.sctx.SetValue(expression.TiDBDecodeKeyFunctionKey, decodeKeyFromString) + rewriter.sctx.SetValue(expression.TiDBEncodeRecordFunctionKey, encodeHandleFromRow) + rewriter.sctx.SetValue(expression.TiDBEncodeIndexFunctionKey, encodeIndexKeyFromRow) expr, _, err := rewriteExprNode(rewriter, node, rewriter.asScalar) if err != nil { @@ -253,6 +256,8 @@ func (b *PlanBuilder) getExpressionRewriter(ctx context.Context, p LogicalPlan) planCtx: &exprRewriterPlanCtx{plan: p, builder: b, rollExpand: b.currentBlockExpand}, } rewriter.sctx.SetValue(expression.TiDBDecodeKeyFunctionKey, decodeKeyFromString) + rewriter.sctx.SetValue(expression.TiDBEncodeRecordFunctionKey, encodeHandleFromRow) + rewriter.sctx.SetValue(expression.TiDBEncodeIndexFunctionKey, encodeIndexKeyFromRow) b.rewriterPool = append(b.rewriterPool, rewriter) return } @@ -2572,6 +2577,111 @@ func hasCurrentDatetimeDefault(col *model.ColumnInfo) bool { return strings.ToLower(x) == ast.CurrentTimestamp } +func encodeHandleFromRow(ctx expression.EvalContext, args []expression.Expression, row chunk.Row) (kv.Key, bool, error) { + dbName, isNull, err := args[0].EvalString(ctx, row) + if err != nil || isNull { + return nil, isNull, err + } + tblName, isNull, err := args[1].EvalString(ctx, row) + if err != nil || isNull { + return nil, isNull, err + } + is := ctx.GetInfoSchema().(infoschema.InfoSchema) + if is == nil { + return nil, false, errors.New("missing information schema") + } + tbl, err := is.TableByName(model.NewCIStr(dbName), model.NewCIStr(tblName)) + if err != nil { + return nil, false, err + } + tblInfo := tbl.Meta() + var recordID kv.Handle + if !tblInfo.IsCommonHandle { + h, isNull, err := args[2].EvalInt(ctx, row) + if err != nil || isNull { + return nil, isNull, err + } + recordID = kv.IntHandle(h) + } else { + pkIdx := tables.FindPrimaryIndex(tblInfo) + if len(pkIdx.Columns) != row.Len()-2 { + return nil, false, errors.Errorf("column count mismatch, expected %d, got %d", len(pkIdx.Columns), row.Len()-2) + } + pkDts := make([]types.Datum, 0, len(pkIdx.Columns)) + for i, idxCol := range pkIdx.Columns { + dt, err := args[i+2].Eval(ctx, row) + ft := tblInfo.Columns[idxCol.Offset].FieldType + pkDt, err := dt.ConvertTo(ctx.TypeCtx(), &ft) + if err != nil { + return nil, false, err + } + pkDts = append(pkDts, pkDt) + } + tablecodec.TruncateIndexValues(tblInfo, pkIdx, pkDts) + var handleBytes []byte + handleBytes, err := codec.EncodeKey(ctx.Location(), nil, pkDts...) + ec := ctx.ErrCtx() + err = ec.HandleError(err) + if err != nil { + return nil, false, err + } + recordID, err = kv.NewCommonHandle(handleBytes) + if err != nil { + return nil, false, err + } + } + key := tablecodec.EncodeRecordKey(tbl.RecordPrefix(), recordID) + return key, false, nil +} + +func encodeIndexKeyFromRow(ctx expression.EvalContext, args []expression.Expression, row chunk.Row) (kv.Key, bool, error) { + dbName, isNull, err := args[0].EvalString(ctx, row) + if err != nil || isNull { + return nil, isNull, err + } + tblName, isNull, err := args[1].EvalString(ctx, row) + if err != nil || isNull { + return nil, isNull, err + } + idxName, isNull, err := args[2].EvalString(ctx, row) + if err != nil || isNull { + return nil, isNull, err + } + is := ctx.GetInfoSchema().(infoschema.InfoSchema) + if is == nil { + return nil, false, errors.New("missing information schema") + } + tbl, err := is.TableByName(model.NewCIStr(dbName), model.NewCIStr(tblName)) + if err != nil { + return nil, false, err + } + tblInfo := tbl.Meta() + idx := tblInfo.FindIndexByName(strings.ToLower(idxName)) + if idx == nil { + return nil, false, errors.New("index not found") + } + + if len(idx.Columns) != row.Len()-3 { + return nil, false, errors.Errorf("column count mismatch, expected %d, got %d", len(idx.Columns), row.Len()-3) + } + idxDts := make([]types.Datum, 0, len(idx.Columns)) + for i, idxCol := range idx.Columns { + dt, err := args[i+3].Eval(ctx, row) + if err != nil { + return nil, false, err + } + ft := tblInfo.Columns[idxCol.Offset].FieldType + idxDt, err := dt.ConvertTo(ctx.TypeCtx(), &ft) + if err != nil { + return nil, false, err + } + idxDts = append(idxDts, idxDt) + } + tablecodec.TruncateIndexValues(tblInfo, idx, idxDts) + idxKey, err := codec.EncodeKey(ctx.GetSessionVars().Location(), nil, idxDts...) + return idxKey, false, err +} + func decodeKeyFromString(ctx expression.EvalContext, s string) string { sc := ctx.GetSessionVars().StmtCtx key, err := hex.DecodeString(s) From 577ede5006f65d0ab171f2ce636f3d1e12d33634 Mon Sep 17 00:00:00 2001 From: tangenta Date: Tue, 12 Mar 2024 11:29:43 +0800 Subject: [PATCH 02/21] fix index encode key --- pkg/planner/core/expression_rewriter.go | 21 +++++++++++---------- 1 file changed, 11 insertions(+), 10 deletions(-) diff --git a/pkg/planner/core/expression_rewriter.go b/pkg/planner/core/expression_rewriter.go index 20ad15e69b3fc..0727eac9b0817 100644 --- a/pkg/planner/core/expression_rewriter.go +++ b/pkg/planner/core/expression_rewriter.go @@ -2604,8 +2604,8 @@ func encodeHandleFromRow(ctx expression.EvalContext, args []expression.Expressio recordID = kv.IntHandle(h) } else { pkIdx := tables.FindPrimaryIndex(tblInfo) - if len(pkIdx.Columns) != row.Len()-2 { - return nil, false, errors.Errorf("column count mismatch, expected %d, got %d", len(pkIdx.Columns), row.Len()-2) + if len(pkIdx.Columns) != len(args)-2 { + return nil, false, errors.Errorf("column count mismatch, expected %d, got %d", len(pkIdx.Columns), len(args)-2) } pkDts := make([]types.Datum, 0, len(pkIdx.Columns)) for i, idxCol := range pkIdx.Columns { @@ -2656,16 +2656,16 @@ func encodeIndexKeyFromRow(ctx expression.EvalContext, args []expression.Express return nil, false, err } tblInfo := tbl.Meta() - idx := tblInfo.FindIndexByName(strings.ToLower(idxName)) - if idx == nil { + idxInfo := tblInfo.FindIndexByName(strings.ToLower(idxName)) + if idxInfo == nil { return nil, false, errors.New("index not found") } - if len(idx.Columns) != row.Len()-3 { - return nil, false, errors.Errorf("column count mismatch, expected %d, got %d", len(idx.Columns), row.Len()-3) + if len(idxInfo.Columns) != len(args)-3 { + return nil, false, errors.Errorf("column count mismatch, expected %d, got %d", len(idxInfo.Columns), len(args)-3) } - idxDts := make([]types.Datum, 0, len(idx.Columns)) - for i, idxCol := range idx.Columns { + idxDts := make([]types.Datum, 0, len(idxInfo.Columns)) + for i, idxCol := range idxInfo.Columns { dt, err := args[i+3].Eval(ctx, row) if err != nil { return nil, false, err @@ -2677,8 +2677,9 @@ func encodeIndexKeyFromRow(ctx expression.EvalContext, args []expression.Express } idxDts = append(idxDts, idxDt) } - tablecodec.TruncateIndexValues(tblInfo, idx, idxDts) - idxKey, err := codec.EncodeKey(ctx.GetSessionVars().Location(), nil, idxDts...) + tablecodec.TruncateIndexValues(tblInfo, idxInfo, idxDts) + idx := tables.NewIndex(tblInfo.ID, tblInfo, idxInfo) + idxKey, _, err := idx.GenIndexKey(ctx.ErrCtx(), ctx.Location(), idxDts, nil, nil) return idxKey, false, err } From baad09de139c593f0ccae7dd651049bed205b2c4 Mon Sep 17 00:00:00 2001 From: tangenta Date: Tue, 12 Mar 2024 12:20:17 +0800 Subject: [PATCH 03/21] fix index encode key --- pkg/planner/core/expression_rewriter.go | 56 ++++++++++++++++++------- 1 file changed, 41 insertions(+), 15 deletions(-) diff --git a/pkg/planner/core/expression_rewriter.go b/pkg/planner/core/expression_rewriter.go index 0727eac9b0817..376a447f13ada 100644 --- a/pkg/planner/core/expression_rewriter.go +++ b/pkg/planner/core/expression_rewriter.go @@ -2594,26 +2594,37 @@ func encodeHandleFromRow(ctx expression.EvalContext, args []expression.Expressio if err != nil { return nil, false, err } - tblInfo := tbl.Meta() + recordID, err := buildHandle(ctx, tbl.Meta(), args[2:], row) + if err != nil { + return nil, false, err + } + key := tablecodec.EncodeRecordKey(tbl.RecordPrefix(), recordID) + return key, false, nil +} + +func buildHandle(ctx expression.EvalContext, tblInfo *model.TableInfo, pkArgs []expression.Expression, row chunk.Row) (kv.Handle, error) { var recordID kv.Handle if !tblInfo.IsCommonHandle { - h, isNull, err := args[2].EvalInt(ctx, row) - if err != nil || isNull { - return nil, isNull, err + h, _, err := pkArgs[0].EvalInt(ctx, row) + if err != nil { + return nil, err } recordID = kv.IntHandle(h) } else { pkIdx := tables.FindPrimaryIndex(tblInfo) - if len(pkIdx.Columns) != len(args)-2 { - return nil, false, errors.Errorf("column count mismatch, expected %d, got %d", len(pkIdx.Columns), len(args)-2) + if len(pkIdx.Columns) != len(pkArgs) { + return nil, errors.Errorf("pk column count mismatch, expected %d, got %d", len(pkIdx.Columns), pkArgs) } pkDts := make([]types.Datum, 0, len(pkIdx.Columns)) for i, idxCol := range pkIdx.Columns { - dt, err := args[i+2].Eval(ctx, row) + dt, err := pkArgs[i].Eval(ctx, row) + if err != nil { + return nil, err + } ft := tblInfo.Columns[idxCol.Offset].FieldType pkDt, err := dt.ConvertTo(ctx.TypeCtx(), &ft) if err != nil { - return nil, false, err + return nil, err } pkDts = append(pkDts, pkDt) } @@ -2623,15 +2634,14 @@ func encodeHandleFromRow(ctx expression.EvalContext, args []expression.Expressio ec := ctx.ErrCtx() err = ec.HandleError(err) if err != nil { - return nil, false, err + return nil, err } recordID, err = kv.NewCommonHandle(handleBytes) if err != nil { - return nil, false, err + return nil, err } } - key := tablecodec.EncodeRecordKey(tbl.RecordPrefix(), recordID) - return key, false, nil + return recordID, nil } func encodeIndexKeyFromRow(ctx expression.EvalContext, args []expression.Expression, row chunk.Row) (kv.Key, bool, error) { @@ -2661,9 +2671,24 @@ func encodeIndexKeyFromRow(ctx expression.EvalContext, args []expression.Express return nil, false, errors.New("index not found") } - if len(idxInfo.Columns) != len(args)-3 { - return nil, false, errors.Errorf("column count mismatch, expected %d, got %d", len(idxInfo.Columns), len(args)-3) + pkLen := 1 + var pkIdx *model.IndexInfo + if tblInfo.IsCommonHandle { + pkIdx = tables.FindPrimaryIndex(tblInfo) + pkLen = len(pkIdx.Columns) + } + + if len(idxInfo.Columns)+pkLen != len(args)-3 { + return nil, false, errors.Errorf( + "column count mismatch, expected %d (index length + pk/rowid length), got %d", + len(idxInfo.Columns), len(args)-3) } + + handle, err := buildHandle(ctx, tblInfo, args[3+len(idxInfo.Columns):], row) + if err != nil { + return nil, false, err + } + idxDts := make([]types.Datum, 0, len(idxInfo.Columns)) for i, idxCol := range idxInfo.Columns { dt, err := args[i+3].Eval(ctx, row) @@ -2679,7 +2704,8 @@ func encodeIndexKeyFromRow(ctx expression.EvalContext, args []expression.Express } tablecodec.TruncateIndexValues(tblInfo, idxInfo, idxDts) idx := tables.NewIndex(tblInfo.ID, tblInfo, idxInfo) - idxKey, _, err := idx.GenIndexKey(ctx.ErrCtx(), ctx.Location(), idxDts, nil, nil) + + idxKey, _, err := idx.GenIndexKey(ctx.ErrCtx(), ctx.Location(), idxDts, handle, nil) return idxKey, false, err } From bb6bf0c93712eaf48b0a4f1478d8fb789c0251d7 Mon Sep 17 00:00:00 2001 From: tangenta Date: Tue, 12 Mar 2024 12:22:35 +0800 Subject: [PATCH 04/21] fix linter --- pkg/expression/builtin_info.go | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/pkg/expression/builtin_info.go b/pkg/expression/builtin_info.go index 74636d92721dc..d42a3dfda82c7 100644 --- a/pkg/expression/builtin_info.go +++ b/pkg/expression/builtin_info.go @@ -849,10 +849,8 @@ func (c *tidbMVCCInfoFunctionClass) getFunction(ctx BuildContext, args []Express if err != nil { return nil, err } - var store helper.Storage - if ctxStore, ok := ctx.GetStore().(helper.Storage); ok { - store = ctxStore - } else { + store, ok := ctx.GetStore().(helper.Storage) + if !ok { return nil, errors.New("storage is not a helper.Storage") } sig := &builtinTiDBMVCCInfoSig{baseBuiltinFunc: bf, helper: helper.NewHelper(store)} From a4135e03b56bee2323d68e261ab2a53fac40db92 Mon Sep 17 00:00:00 2001 From: tangenta Date: Tue, 12 Mar 2024 12:37:42 +0800 Subject: [PATCH 05/21] fix linter --- pkg/expression/builtin_info.go | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/pkg/expression/builtin_info.go b/pkg/expression/builtin_info.go index d42a3dfda82c7..feec82524f8a4 100644 --- a/pkg/expression/builtin_info.go +++ b/pkg/expression/builtin_info.go @@ -1030,9 +1030,11 @@ func (k TiDBDecodeKeyFunctionKeyType) String() string { const ( // TiDBDecodeKeyFunctionKey is used to identify the decoder function in context. - TiDBDecodeKeyFunctionKey TiDBDecodeKeyFunctionKeyType = 0 + TiDBDecodeKeyFunctionKey TiDBDecodeKeyFunctionKeyType = 0 + // TiDBDecodeKeyFunctionKey is used to identify the encode record function in context. TiDBEncodeRecordFunctionKey TiDBDecodeKeyFunctionKeyType = 1 - TiDBEncodeIndexFunctionKey TiDBDecodeKeyFunctionKeyType = 2 + // TiDBDecodeKeyFunctionKey is used to identify the encode index function in context. + TiDBEncodeIndexFunctionKey TiDBDecodeKeyFunctionKeyType = 2 ) type tidbDecodeSQLDigestsFunctionClass struct { From b13c989aa680500b9e99af083020fef8e2c91814 Mon Sep 17 00:00:00 2001 From: tangenta Date: Tue, 12 Mar 2024 14:45:50 +0800 Subject: [PATCH 06/21] fix linter --- pkg/expression/builtin_info.go | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pkg/expression/builtin_info.go b/pkg/expression/builtin_info.go index feec82524f8a4..66bb5c2575d0b 100644 --- a/pkg/expression/builtin_info.go +++ b/pkg/expression/builtin_info.go @@ -1031,9 +1031,9 @@ func (k TiDBDecodeKeyFunctionKeyType) String() string { const ( // TiDBDecodeKeyFunctionKey is used to identify the decoder function in context. TiDBDecodeKeyFunctionKey TiDBDecodeKeyFunctionKeyType = 0 - // TiDBDecodeKeyFunctionKey is used to identify the encode record function in context. + // TiDBEncodeRecordFunctionKey is used to identify the encode record function in context. TiDBEncodeRecordFunctionKey TiDBDecodeKeyFunctionKeyType = 1 - // TiDBDecodeKeyFunctionKey is used to identify the encode index function in context. + // TiDBEncodeIndexFunctionKey is used to identify the encode index function in context. TiDBEncodeIndexFunctionKey TiDBDecodeKeyFunctionKeyType = 2 ) From 322daf73053ba28f4d6c7544f340d119f854e6b5 Mon Sep 17 00:00:00 2001 From: tangenta Date: Tue, 12 Mar 2024 22:36:33 +0800 Subject: [PATCH 07/21] update show builtin results --- tests/integrationtest/r/executor/show.result | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tests/integrationtest/r/executor/show.result b/tests/integrationtest/r/executor/show.result index db129c1c6f86d..70c7e17f2a480 100644 --- a/tests/integrationtest/r/executor/show.result +++ b/tests/integrationtest/r/executor/show.result @@ -839,8 +839,11 @@ tidb_decode_binary_plan tidb_decode_key tidb_decode_plan tidb_decode_sql_digests +tidb_encode_index_key +tidb_encode_record_key tidb_encode_sql_digest tidb_is_ddl_owner +tidb_mvcc_info tidb_parse_tso tidb_parse_tso_logical tidb_row_checksum From 56d480f61b247df93c488a7a5a5832953346e3bc Mon Sep 17 00:00:00 2001 From: tangenta Date: Mon, 25 Mar 2024 20:42:16 +0800 Subject: [PATCH 08/21] add partition table syntax --- .../integration_test/integration_test.go | 25 ++++++++++++++++++ pkg/planner/core/expression_rewriter.go | 26 ++++++++++++++++++- 2 files changed, 50 insertions(+), 1 deletion(-) diff --git a/pkg/expression/integration_test/integration_test.go b/pkg/expression/integration_test/integration_test.go index 34850e12c208d..b67c648dbf4cb 100644 --- a/pkg/expression/integration_test/integration_test.go +++ b/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) diff --git a/pkg/planner/core/expression_rewriter.go b/pkg/planner/core/expression_rewriter.go index 376a447f13ada..68843b0bce90b 100644 --- a/pkg/planner/core/expression_rewriter.go +++ b/pkg/planner/core/expression_rewriter.go @@ -2590,10 +2590,22 @@ func encodeHandleFromRow(ctx expression.EvalContext, args []expression.Expressio if is == nil { return nil, false, errors.New("missing information schema") } + tblName, partName := extractTablePartition(tblName) tbl, err := is.TableByName(model.NewCIStr(dbName), model.NewCIStr(tblName)) if err != nil { return nil, false, err } + if part := tbl.GetPartitionedTable(); part != nil { + pid, err := tables.FindPartitionByName(tbl.Meta(), partName) + if err != nil { + return nil, false, errors.Trace(err) + } + tbl = part.GetPartition(pid) + } else { + if len(partName) != 0 { + return nil, false, errors.New("not partition table") + } + } recordID, err := buildHandle(ctx, tbl.Meta(), args[2:], row) if err != nil { return nil, false, err @@ -2602,6 +2614,18 @@ func encodeHandleFromRow(ctx expression.EvalContext, args []expression.Expressio return key, false, nil } +func extractTablePartition(str string) (table, partition string) { + start := strings.IndexByte(str, '(') + if start == -1 { + return str, "" + } + end := strings.IndexByte(str, ')') + if end == -1 { + return str, "" + } + return str[:start], str[start+1 : end] +} + func buildHandle(ctx expression.EvalContext, tblInfo *model.TableInfo, pkArgs []expression.Expression, row chunk.Row) (kv.Handle, error) { var recordID kv.Handle if !tblInfo.IsCommonHandle { @@ -2681,7 +2705,7 @@ func encodeIndexKeyFromRow(ctx expression.EvalContext, args []expression.Express if len(idxInfo.Columns)+pkLen != len(args)-3 { return nil, false, errors.Errorf( "column count mismatch, expected %d (index length + pk/rowid length), got %d", - len(idxInfo.Columns), len(args)-3) + len(idxInfo.Columns)+pkLen, len(args)-3) } handle, err := buildHandle(ctx, tblInfo, args[3+len(idxInfo.Columns):], row) From 4aabb50ad4c25681f3be8dcd6b58b3116e914103 Mon Sep 17 00:00:00 2001 From: tangenta Date: Mon, 25 Mar 2024 21:08:02 +0800 Subject: [PATCH 09/21] fix build --- pkg/expression/builtin.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pkg/expression/builtin.go b/pkg/expression/builtin.go index 02f6059ea40b5..10a7db7a7b9a1 100644 --- a/pkg/expression/builtin.go +++ b/pkg/expression/builtin.go @@ -911,7 +911,7 @@ var funcs = map[string]functionClass{ // TiDB internal function. ast.TiDBDecodeKey: &tidbDecodeKeyFunctionClass{baseFunctionClass{ast.TiDBDecodeKey, 1, 1}}, - ast.TiDBMVCCInfo: &tidbMVCCInfoFunctionClass{baseFunctionClass{ast.TiDBMVCCInfo, 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. From 78c5c63b167d940136da7360aaa35b394a5e9e5e Mon Sep 17 00:00:00 2001 From: tangenta Date: Mon, 25 Mar 2024 21:14:32 +0800 Subject: [PATCH 10/21] update bazel --- pkg/expression/integration_test/BUILD.bazel | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pkg/expression/integration_test/BUILD.bazel b/pkg/expression/integration_test/BUILD.bazel index b7b517b8762fb..eab016ca24fa1 100644 --- a/pkg/expression/integration_test/BUILD.bazel +++ b/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", From 4aa711750eea938fb1a0a833e66493797ff0984d Mon Sep 17 00:00:00 2001 From: tangenta Date: Wed, 21 Aug 2024 15:04:13 +0800 Subject: [PATCH 11/21] remove unnecessary changes --- pkg/expression/builtin_info.go | 6 ------ 1 file changed, 6 deletions(-) diff --git a/pkg/expression/builtin_info.go b/pkg/expression/builtin_info.go index 5c93eec1736f3..c04d0566f6d51 100644 --- a/pkg/expression/builtin_info.go +++ b/pkg/expression/builtin_info.go @@ -902,12 +902,6 @@ func (b *builtinRowCountSig) evalInt(ctx EvalContext, row chunk.Row) (res int64, 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) { From a4e51f027d0490589e292090f8583f0f76e0de0f Mon Sep 17 00:00:00 2001 From: tangenta Date: Wed, 21 Aug 2024 15:10:15 +0800 Subject: [PATCH 12/21] update bazel --- pkg/expression/integration_test/BUILD.bazel | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pkg/expression/integration_test/BUILD.bazel b/pkg/expression/integration_test/BUILD.bazel index 6270c437829d9..93d8780b0fb1e 100644 --- a/pkg/expression/integration_test/BUILD.bazel +++ b/pkg/expression/integration_test/BUILD.bazel @@ -8,7 +8,7 @@ go_test( "main_test.go", ], flaky = True, - shard_count = 26, + shard_count = 28, deps = [ "//pkg/config", "//pkg/domain", From 03a9dc0df6b82269b1d1d5766a632cafccb8c271 Mon Sep 17 00:00:00 2001 From: tangenta Date: Mon, 2 Sep 2024 21:45:19 +0800 Subject: [PATCH 13/21] require super priv for tidb_mvcc_info and arbitrary priv for tidb_encode_xxx --- pkg/expression/builtin_info.go | 34 +++++++++++++++++-- .../integration_test/integration_test.go | 14 ++++++++ pkg/planner/core/expression_rewriter.go | 8 +++++ 3 files changed, 54 insertions(+), 2 deletions(-) diff --git a/pkg/expression/builtin_info.go b/pkg/expression/builtin_info.go index c10ce01c5b4d7..d505ea60af632 100644 --- a/pkg/expression/builtin_info.go +++ b/pkg/expression/builtin_info.go @@ -36,6 +36,7 @@ import ( "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" @@ -931,6 +932,9 @@ func (b *builtinTiDBMVCCInfoSig) Clone() builtinFunc { // 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 @@ -983,11 +987,13 @@ func (c *tidbEncodeRecordKeyClass) getFunction(ctx BuildContext, args []Expressi type builtinTiDBEncodeRecordKeySig struct { baseBuiltinFunc contextopt.InfoSchemaPropReader + contextopt.SessionVarsPropReader } // RequiredOptionalEvalProps implements the RequireOptionalEvalProps interface. func (b *builtinTiDBEncodeRecordKeySig) RequiredOptionalEvalProps() OptionalEvalPropKeySet { - return b.InfoSchemaPropReader.RequiredOptionalEvalProps() + return b.InfoSchemaPropReader.RequiredOptionalEvalProps() | + b.SessionVarsPropReader.RequiredOptionalEvalProps() } func (b *builtinTiDBEncodeRecordKeySig) Clone() builtinFunc { @@ -1007,6 +1013,17 @@ func (b *builtinTiDBEncodeRecordKeySig) evalString(ctx EvalContext, row chunk.Ro } 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 @@ -1036,11 +1053,13 @@ func (c *tidbEncodeIndexKeyClass) getFunction(ctx BuildContext, args []Expressio type builtinTiDBEncodeIndexKeySig struct { baseBuiltinFunc contextopt.InfoSchemaPropReader + contextopt.SessionVarsPropReader } // RequiredOptionalEvalProps implements the RequireOptionalEvalProps interface. func (b *builtinTiDBEncodeIndexKeySig) RequiredOptionalEvalProps() OptionalEvalPropKeySet { - return b.InfoSchemaPropReader.RequiredOptionalEvalProps() + return b.InfoSchemaPropReader.RequiredOptionalEvalProps() | + b.SessionVarsPropReader.RequiredOptionalEvalProps() } func (b *builtinTiDBEncodeIndexKeySig) Clone() builtinFunc { @@ -1060,6 +1079,17 @@ func (b *builtinTiDBEncodeIndexKeySig) evalString(ctx EvalContext, row chunk.Row } 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 diff --git a/pkg/expression/integration_test/integration_test.go b/pkg/expression/integration_test/integration_test.go index 8aeee538acde3..d3cc9bd0c95a0 100644 --- a/pkg/expression/integration_test/integration_test.go +++ b/pkg/expression/integration_test/integration_test.go @@ -1129,6 +1129,20 @@ func TestTiDBEncodeKey(t *testing.T) { 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) { diff --git a/pkg/planner/core/expression_rewriter.go b/pkg/planner/core/expression_rewriter.go index 3a04b649c2373..9a529cb795ce3 100644 --- a/pkg/planner/core/expression_rewriter.go +++ b/pkg/planner/core/expression_rewriter.go @@ -2664,6 +2664,10 @@ func encodeHandleFromRow( if err != nil { return nil, false, err } + if !ctx.RequestVerification(dbName, tblName, "", mysql.AllPrivMask) { + // The arguments will be filled by caller. + return nil, false, plannererrors.ErrSpecificAccessDenied + } if part := tbl.GetPartitionedTable(); part != nil { pid, err := tables.FindPartitionByName(tbl.Meta(), partName) if err != nil { @@ -2756,6 +2760,10 @@ func encodeIndexKeyFromRow( if err != nil || isNull { return nil, isNull, err } + if !ctx.RequestVerification(dbName, tblName, "", mysql.AllPrivMask) { + // The arguments will be filled by caller. + return nil, false, plannererrors.ErrSpecificAccessDenied + } idxName, isNull, err := args[2].EvalString(ctx, row) if err != nil || isNull { return nil, isNull, err From 4e07b96d212ad91e4ca8e1d272c9147a72c8b18b Mon Sep 17 00:00:00 2001 From: tangenta Date: Wed, 4 Sep 2024 20:53:57 +0800 Subject: [PATCH 14/21] move tidb codec funcs to a new file --- pkg/expression/builtin_info.go | 4 +- pkg/planner/core/core_init.go | 7 +- pkg/planner/core/expression_codec_fn.go | 461 ++++++++++++++++++++++++ pkg/planner/core/expression_rewriter.go | 414 --------------------- 4 files changed, 467 insertions(+), 419 deletions(-) create mode 100644 pkg/planner/core/expression_codec_fn.go diff --git a/pkg/expression/builtin_info.go b/pkg/expression/builtin_info.go index d505ea60af632..c8f7b5222417d 100644 --- a/pkg/expression/builtin_info.go +++ b/pkg/expression/builtin_info.go @@ -1115,10 +1115,10 @@ func (c *tidbDecodeKeyFunctionClass) getFunction(ctx BuildContext, args []Expres 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) +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, isVer infoschema.MetaOnlyInfoSchema, args []Expression, row chunk.Row) ([]byte, bool, error) +var EncodeIndexKeyFromRow func(ctx EvalContext, is infoschema.MetaOnlyInfoSchema, args []Expression, row chunk.Row) ([]byte, bool, error) type builtinTiDBDecodeKeySig struct { baseBuiltinFunc diff --git a/pkg/planner/core/core_init.go b/pkg/planner/core/core_init.go index 6bff1f65b7d4e..a0af6bad0cccb 100644 --- a/pkg/planner/core/core_init.go +++ b/pkg/planner/core/core_init.go @@ -77,9 +77,10 @@ func init() { base.InvalidTask = &RootTask{} // invalid if p is nil expression.EvalSimpleAst = evalAstExpr expression.BuildSimpleExpr = buildSimpleExpr - expression.DecodeKeyFromString = decodeKeyFromString - expression.EncodeRecordKeyFromRow = encodeHandleFromRow - expression.EncodeIndexKeyFromRow = encodeIndexKeyFromRow + helper := tidbCodecFuncHelper{} + expression.DecodeKeyFromString = helper.decodeKeyFromString + expression.EncodeRecordKeyFromRow = helper.encodeHandleFromRow + expression.EncodeIndexKeyFromRow = helper.encodeIndexKeyFromRow plannerutil.EvalAstExprWithPlanCtx = evalAstExprWithPlanCtx plannerutil.RewriteAstExprWithPlanCtx = rewriteAstExprWithPlanCtx DefaultDisabledLogicalRulesList = new(atomic.Value) diff --git a/pkg/planner/core/expression_codec_fn.go b/pkg/planner/core/expression_codec_fn.go new file mode 100644 index 0000000000000..d4c6b4308763d --- /dev/null +++ b/pkg/planner/core/expression_codec_fn.go @@ -0,0 +1,461 @@ +// Copyright 2024 PingCAP, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package core + +import ( + "context" + "encoding/hex" + "encoding/json" + "strconv" + "strings" + "time" + + "github.com/pingcap/errors" + "github.com/pingcap/tidb/pkg/expression" + "github.com/pingcap/tidb/pkg/infoschema" + infoschemactx "github.com/pingcap/tidb/pkg/infoschema/context" + "github.com/pingcap/tidb/pkg/kv" + "github.com/pingcap/tidb/pkg/parser/model" + "github.com/pingcap/tidb/pkg/parser/mysql" + "github.com/pingcap/tidb/pkg/table" + "github.com/pingcap/tidb/pkg/table/tables" + "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/codec" + "github.com/pingcap/tidb/pkg/util/dbterror/plannererrors" +) + +// tidbCodecFuncHelper contains some utililty functions for +// - tidb_decode_key(hex_string) +// - tidb_encode_record_key(database_name, table_name, handle/pk columns...) +// - tidb_encode_index_key(database_name, table_name, index_name, index columns..., handle/pk columns...) +// +// define an individual struct instead of a bunch of un-exported functions +// to avoid polluting the global scope of current package. +type tidbCodecFuncHelper struct{} + +func (h tidbCodecFuncHelper) encodeHandleFromRow( + ctx expression.EvalContext, + isVer infoschemactx.MetaOnlyInfoSchema, + args []expression.Expression, + row chunk.Row, +) ([]byte, bool, error) { + dbName, isNull, err := args[0].EvalString(ctx, row) + if err != nil || isNull { + return nil, isNull, err + } + tblName, isNull, err := args[1].EvalString(ctx, row) + if err != nil || isNull { + return nil, isNull, err + } + is := isVer.(infoschema.InfoSchema) + tbl, _, err := h.findCommonOrPartitionedTable(ctx, is, dbName, tblName) + if err != nil { + return nil, false, err + } + recordID, err := h.buildHandle(ctx, tbl.Meta(), args[2:], row) + if err != nil { + return nil, false, err + } + key := tablecodec.EncodeRecordKey(tbl.RecordPrefix(), recordID) + return key, false, nil +} + +func (h tidbCodecFuncHelper) findCommonOrPartitionedTable( + ctx expression.EvalContext, + is infoschema.InfoSchema, + dbName string, + tblName string, +) (table.Table, int64, error) { + tblName, partName := h.extractTablePartition(tblName) + tbl, err := is.TableByName(context.Background(), model.NewCIStr(dbName), model.NewCIStr(tblName)) + if err != nil { + return nil, 0, err + } + if !ctx.RequestVerification(dbName, tblName, "", mysql.AllPrivMask) { + // The arguments will be filled by caller. + return nil, 0, plannererrors.ErrSpecificAccessDenied + } + if part := tbl.GetPartitionedTable(); part != nil && len(partName) > 0 { + pid, err := tables.FindPartitionByName(tbl.Meta(), partName) + if err != nil { + return nil, 0, errors.Trace(err) + } + tbl = part.GetPartition(pid) + return tbl, pid, nil + } else { + if len(partName) != 0 { + return nil, 0, errors.New("not a partitioned table") + } + } + return tbl, tbl.Meta().ID, nil +} + +func (tidbCodecFuncHelper) extractTablePartition(str string) (table, partition string) { + start := strings.IndexByte(str, '(') + if start == -1 { + return str, "" + } + end := strings.IndexByte(str, ')') + if end == -1 { + return str, "" + } + return str[:start], str[start+1 : end] +} + +func (tidbCodecFuncHelper) buildHandle( + ctx expression.EvalContext, + tblInfo *model.TableInfo, + pkArgs []expression.Expression, + row chunk.Row, +) (kv.Handle, error) { + var recordID kv.Handle + if !tblInfo.IsCommonHandle { + h, _, err := pkArgs[0].EvalInt(ctx, row) + if err != nil { + return nil, err + } + recordID = kv.IntHandle(h) + } else { + pkIdx := tables.FindPrimaryIndex(tblInfo) + if len(pkIdx.Columns) != len(pkArgs) { + return nil, errors.Errorf("pk column count mismatch, expected %d, got %d", len(pkIdx.Columns), pkArgs) + } + pkDts := make([]types.Datum, 0, len(pkIdx.Columns)) + for i, idxCol := range pkIdx.Columns { + dt, err := pkArgs[i].Eval(ctx, row) + if err != nil { + return nil, err + } + ft := tblInfo.Columns[idxCol.Offset].FieldType + pkDt, err := dt.ConvertTo(ctx.TypeCtx(), &ft) + if err != nil { + return nil, err + } + pkDts = append(pkDts, pkDt) + } + tablecodec.TruncateIndexValues(tblInfo, pkIdx, pkDts) + var handleBytes []byte + handleBytes, err := codec.EncodeKey(ctx.Location(), nil, pkDts...) + ec := ctx.ErrCtx() + err = ec.HandleError(err) + if err != nil { + return nil, err + } + recordID, err = kv.NewCommonHandle(handleBytes) + if err != nil { + return nil, err + } + } + return recordID, nil +} + +func (h tidbCodecFuncHelper) encodeIndexKeyFromRow( + ctx expression.EvalContext, + isVer infoschemactx.MetaOnlyInfoSchema, + args []expression.Expression, + row chunk.Row, +) ([]byte, bool, error) { + dbName, isNull, err := args[0].EvalString(ctx, row) + if err != nil || isNull { + return nil, isNull, err + } + tblName, isNull, err := args[1].EvalString(ctx, row) + if err != nil || isNull { + return nil, isNull, err + } + idxName, isNull, err := args[2].EvalString(ctx, row) + if err != nil || isNull { + return nil, isNull, err + } + is := isVer.(infoschema.InfoSchema) + tbl, physicalID, err := h.findCommonOrPartitionedTable(ctx, is, dbName, tblName) + if err != nil { + return nil, false, err + } + tblInfo := tbl.Meta() + idxInfo := tblInfo.FindIndexByName(strings.ToLower(idxName)) + if idxInfo == nil { + return nil, false, errors.New("index not found") + } + + pkLen := 1 + var pkIdx *model.IndexInfo + if tblInfo.IsCommonHandle { + pkIdx = tables.FindPrimaryIndex(tblInfo) + pkLen = len(pkIdx.Columns) + } + + if len(idxInfo.Columns)+pkLen != len(args)-3 { + return nil, false, errors.Errorf( + "column count mismatch, expected %d (index length + pk/rowid length), got %d", + len(idxInfo.Columns)+pkLen, len(args)-3) + } + + handle, err := h.buildHandle(ctx, tblInfo, args[3+len(idxInfo.Columns):], row) + if err != nil { + return nil, false, err + } + + idxDts := make([]types.Datum, 0, len(idxInfo.Columns)) + for i, idxCol := range idxInfo.Columns { + dt, err := args[i+3].Eval(ctx, row) + if err != nil { + return nil, false, err + } + ft := tblInfo.Columns[idxCol.Offset].FieldType + idxDt, err := dt.ConvertTo(ctx.TypeCtx(), &ft) + if err != nil { + return nil, false, err + } + idxDts = append(idxDts, idxDt) + } + tablecodec.TruncateIndexValues(tblInfo, idxInfo, idxDts) + idx := tables.NewIndex(physicalID, tblInfo, idxInfo) + + idxKey, _, err := idx.GenIndexKey(ctx.ErrCtx(), ctx.Location(), idxDts, handle, nil) + return idxKey, false, err +} + +func (h tidbCodecFuncHelper) decodeKeyFromString( + tc types.Context, isVer infoschemactx.MetaOnlyInfoSchema, s string) string { + key, err := hex.DecodeString(s) + if err != nil { + tc.AppendWarning(errors.NewNoStackErrorf("invalid key: %X", key)) + return s + } + // Auto decode byte if needed. + _, bs, err := codec.DecodeBytes(key, nil) + if err == nil { + key = bs + } + tableID := tablecodec.DecodeTableID(key) + if tableID <= 0 { + tc.AppendWarning(errors.NewNoStackErrorf("invalid key: %X", key)) + return s + } + + is, ok := isVer.(infoschema.InfoSchema) + if !ok { + tc.AppendWarning(errors.NewNoStackErrorf("infoschema not found when decoding key: %X", key)) + return s + } + tbl, _ := infoschema.FindTableByTblOrPartID(is, tableID) + loc := tc.Location() + if tablecodec.IsRecordKey(key) { + ret, err := h.decodeRecordKey(key, tableID, tbl, loc) + if err != nil { + tc.AppendWarning(err) + return s + } + return ret + } else if tablecodec.IsIndexKey(key) { + ret, err := h.decodeIndexKey(key, tableID, tbl, loc) + if err != nil { + tc.AppendWarning(err) + return s + } + return ret + } else if tablecodec.IsTableKey(key) { + ret, err := h.decodeTableKey(key, tableID, tbl) + if err != nil { + tc.AppendWarning(err) + return s + } + return ret + } + tc.AppendWarning(errors.NewNoStackErrorf("invalid key: %X", key)) + return s +} + +func (h tidbCodecFuncHelper) decodeRecordKey( + key []byte, tableID int64, tbl table.Table, loc *time.Location) (string, error) { + _, handle, err := tablecodec.DecodeRecordKey(key) + if err != nil { + return "", errors.Trace(err) + } + if handle.IsInt() { + ret := make(map[string]any) + if tbl != nil && tbl.Meta().Partition != nil { + ret["partition_id"] = tableID + tableID = tbl.Meta().ID + } + ret["table_id"] = strconv.FormatInt(tableID, 10) + // When the clustered index is enabled, we should show the PK name. + if tbl != nil && tbl.Meta().HasClusteredIndex() { + ret[tbl.Meta().GetPkName().String()] = handle.IntValue() + } else { + ret["_tidb_rowid"] = handle.IntValue() + } + retStr, err := json.Marshal(ret) + if err != nil { + return "", errors.Trace(err) + } + return string(retStr), nil + } + if tbl != nil { + tblInfo := tbl.Meta() + idxInfo := tables.FindPrimaryIndex(tblInfo) + if idxInfo == nil { + return "", errors.Trace(errors.Errorf("primary key not found when decoding record key: %X", key)) + } + cols := make(map[int64]*types.FieldType, len(tblInfo.Columns)) + for _, col := range tblInfo.Columns { + cols[col.ID] = &(col.FieldType) + } + handleColIDs := make([]int64, 0, len(idxInfo.Columns)) + for _, col := range idxInfo.Columns { + handleColIDs = append(handleColIDs, tblInfo.Columns[col.Offset].ID) + } + + if len(handleColIDs) != handle.NumCols() { + return "", errors.Trace(errors.Errorf("primary key length not match handle columns number in key")) + } + datumMap, err := tablecodec.DecodeHandleToDatumMap(handle, handleColIDs, cols, loc, nil) + if err != nil { + return "", errors.Trace(err) + } + ret := make(map[string]any) + if tbl.Meta().Partition != nil { + ret["partition_id"] = tableID + tableID = tbl.Meta().ID + } + ret["table_id"] = tableID + handleRet := make(map[string]any) + for colID := range datumMap { + dt := datumMap[colID] + dtStr, err := h.datumToJSONObject(&dt) + if err != nil { + return "", errors.Trace(err) + } + found := false + for _, colInfo := range tblInfo.Columns { + if colInfo.ID == colID { + found = true + handleRet[colInfo.Name.L] = dtStr + break + } + } + if !found { + return "", errors.Trace(errors.Errorf("column not found when decoding record key: %X", key)) + } + } + ret["handle"] = handleRet + retStr, err := json.Marshal(ret) + if err != nil { + return "", errors.Trace(err) + } + return string(retStr), nil + } + ret := make(map[string]any) + ret["table_id"] = tableID + ret["handle"] = handle.String() + retStr, err := json.Marshal(ret) + if err != nil { + return "", errors.Trace(err) + } + return string(retStr), nil +} + +func (h tidbCodecFuncHelper) decodeIndexKey( + key []byte, tableID int64, tbl table.Table, loc *time.Location) (string, error) { + if tbl != nil { + _, indexID, _, err := tablecodec.DecodeKeyHead(key) + if err != nil { + return "", errors.Trace(errors.Errorf("invalid record/index key: %X", key)) + } + tblInfo := tbl.Meta() + var targetIndex *model.IndexInfo + for _, idx := range tblInfo.Indices { + if idx.ID == indexID { + targetIndex = idx + break + } + } + if targetIndex == nil { + return "", errors.Trace(errors.Errorf("index not found when decoding index key: %X", key)) + } + colInfos := tables.BuildRowcodecColInfoForIndexColumns(targetIndex, tblInfo) + tps := tables.BuildFieldTypesForIndexColumns(targetIndex, tblInfo) + values, err := tablecodec.DecodeIndexKV(key, []byte{0}, len(colInfos), tablecodec.HandleNotNeeded, colInfos) + if err != nil { + return "", errors.Trace(err) + } + ds := make([]types.Datum, 0, len(colInfos)) + for i := 0; i < len(colInfos); i++ { + d, err := tablecodec.DecodeColumnValue(values[i], tps[i], loc) + if err != nil { + return "", errors.Trace(err) + } + ds = append(ds, d) + } + ret := make(map[string]any) + if tbl.Meta().Partition != nil { + ret["partition_id"] = tableID + tableID = tbl.Meta().ID + } + ret["table_id"] = tableID + ret["index_id"] = indexID + idxValMap := make(map[string]any, len(targetIndex.Columns)) + for i := 0; i < len(targetIndex.Columns); i++ { + dtStr, err := h.datumToJSONObject(&ds[i]) + if err != nil { + return "", errors.Trace(err) + } + idxValMap[targetIndex.Columns[i].Name.L] = dtStr + } + ret["index_vals"] = idxValMap + retStr, err := json.Marshal(ret) + if err != nil { + return "", errors.Trace(err) + } + return string(retStr), nil + } + _, indexID, indexValues, err := tablecodec.DecodeIndexKey(key) + if err != nil { + return "", errors.Trace(errors.Errorf("invalid index key: %X", key)) + } + ret := make(map[string]any) + ret["table_id"] = tableID + ret["index_id"] = indexID + ret["index_vals"] = strings.Join(indexValues, ", ") + retStr, err := json.Marshal(ret) + if err != nil { + return "", errors.Trace(err) + } + return string(retStr), nil +} + +func (tidbCodecFuncHelper) decodeTableKey(_ []byte, tableID int64, tbl table.Table) (string, error) { + ret := map[string]int64{} + if tbl != nil && tbl.Meta().GetPartitionInfo() != nil { + ret["partition_id"] = tableID + tableID = tbl.Meta().ID + } + ret["table_id"] = tableID + retStr, err := json.Marshal(ret) + if err != nil { + return "", errors.Trace(err) + } + return string(retStr), nil +} + +func (tidbCodecFuncHelper) datumToJSONObject(d *types.Datum) (any, error) { + if d.IsNull() { + return nil, nil + } + return d.ToString() +} diff --git a/pkg/planner/core/expression_rewriter.go b/pkg/planner/core/expression_rewriter.go index 9a529cb795ce3..956e3362d791d 100644 --- a/pkg/planner/core/expression_rewriter.go +++ b/pkg/planner/core/expression_rewriter.go @@ -16,12 +16,9 @@ package core import ( "context" - "encoding/hex" - "encoding/json" "fmt" "strconv" "strings" - "time" "github.com/pingcap/errors" "github.com/pingcap/tidb/pkg/expression" @@ -29,8 +26,6 @@ import ( exprctx "github.com/pingcap/tidb/pkg/expression/context" "github.com/pingcap/tidb/pkg/expression/contextopt" "github.com/pingcap/tidb/pkg/infoschema" - infoschemactx "github.com/pingcap/tidb/pkg/infoschema/context" - "github.com/pingcap/tidb/pkg/kv" "github.com/pingcap/tidb/pkg/parser/ast" "github.com/pingcap/tidb/pkg/parser/charset" "github.com/pingcap/tidb/pkg/parser/model" @@ -42,12 +37,9 @@ import ( "github.com/pingcap/tidb/pkg/planner/util/coreusage" "github.com/pingcap/tidb/pkg/sessionctx/variable" "github.com/pingcap/tidb/pkg/table" - "github.com/pingcap/tidb/pkg/table/tables" - "github.com/pingcap/tidb/pkg/tablecodec" "github.com/pingcap/tidb/pkg/types" driver "github.com/pingcap/tidb/pkg/types/parser_driver" "github.com/pingcap/tidb/pkg/util/chunk" - "github.com/pingcap/tidb/pkg/util/codec" "github.com/pingcap/tidb/pkg/util/collate" "github.com/pingcap/tidb/pkg/util/dbterror/plannererrors" "github.com/pingcap/tidb/pkg/util/hint" @@ -2640,409 +2632,3 @@ func hasCurrentDatetimeDefault(col *model.ColumnInfo) bool { } return strings.ToLower(x) == ast.CurrentTimestamp } - -func encodeHandleFromRow( - ctx expression.EvalContext, - isVer infoschemactx.MetaOnlyInfoSchema, - args []expression.Expression, - row chunk.Row, -) ([]byte, bool, error) { - dbName, isNull, err := args[0].EvalString(ctx, row) - if err != nil || isNull { - return nil, isNull, err - } - tblName, isNull, err := args[1].EvalString(ctx, row) - if err != nil || isNull { - return nil, isNull, err - } - is := isVer.(infoschema.InfoSchema) - if is == nil { - return nil, false, errors.New("missing information schema") - } - tblName, partName := extractTablePartition(tblName) - tbl, err := is.TableByName(context.Background(), model.NewCIStr(dbName), model.NewCIStr(tblName)) - if err != nil { - return nil, false, err - } - if !ctx.RequestVerification(dbName, tblName, "", mysql.AllPrivMask) { - // The arguments will be filled by caller. - return nil, false, plannererrors.ErrSpecificAccessDenied - } - if part := tbl.GetPartitionedTable(); part != nil { - pid, err := tables.FindPartitionByName(tbl.Meta(), partName) - if err != nil { - return nil, false, errors.Trace(err) - } - tbl = part.GetPartition(pid) - } else { - if len(partName) != 0 { - return nil, false, errors.New("not partition table") - } - } - recordID, err := buildHandle(ctx, tbl.Meta(), args[2:], row) - if err != nil { - return nil, false, err - } - key := tablecodec.EncodeRecordKey(tbl.RecordPrefix(), recordID) - return key, false, nil -} - -func extractTablePartition(str string) (table, partition string) { - start := strings.IndexByte(str, '(') - if start == -1 { - return str, "" - } - end := strings.IndexByte(str, ')') - if end == -1 { - return str, "" - } - return str[:start], str[start+1 : end] -} - -func buildHandle( - ctx expression.EvalContext, - tblInfo *model.TableInfo, - pkArgs []expression.Expression, - row chunk.Row, -) (kv.Handle, error) { - var recordID kv.Handle - if !tblInfo.IsCommonHandle { - h, _, err := pkArgs[0].EvalInt(ctx, row) - if err != nil { - return nil, err - } - recordID = kv.IntHandle(h) - } else { - pkIdx := tables.FindPrimaryIndex(tblInfo) - if len(pkIdx.Columns) != len(pkArgs) { - return nil, errors.Errorf("pk column count mismatch, expected %d, got %d", len(pkIdx.Columns), pkArgs) - } - pkDts := make([]types.Datum, 0, len(pkIdx.Columns)) - for i, idxCol := range pkIdx.Columns { - dt, err := pkArgs[i].Eval(ctx, row) - if err != nil { - return nil, err - } - ft := tblInfo.Columns[idxCol.Offset].FieldType - pkDt, err := dt.ConvertTo(ctx.TypeCtx(), &ft) - if err != nil { - return nil, err - } - pkDts = append(pkDts, pkDt) - } - tablecodec.TruncateIndexValues(tblInfo, pkIdx, pkDts) - var handleBytes []byte - handleBytes, err := codec.EncodeKey(ctx.Location(), nil, pkDts...) - ec := ctx.ErrCtx() - err = ec.HandleError(err) - if err != nil { - return nil, err - } - recordID, err = kv.NewCommonHandle(handleBytes) - if err != nil { - return nil, err - } - } - return recordID, nil -} - -func encodeIndexKeyFromRow( - ctx expression.EvalContext, - isVer infoschemactx.MetaOnlyInfoSchema, - args []expression.Expression, - row chunk.Row, -) ([]byte, bool, error) { - dbName, isNull, err := args[0].EvalString(ctx, row) - if err != nil || isNull { - return nil, isNull, err - } - tblName, isNull, err := args[1].EvalString(ctx, row) - if err != nil || isNull { - return nil, isNull, err - } - if !ctx.RequestVerification(dbName, tblName, "", mysql.AllPrivMask) { - // The arguments will be filled by caller. - return nil, false, plannererrors.ErrSpecificAccessDenied - } - idxName, isNull, err := args[2].EvalString(ctx, row) - if err != nil || isNull { - return nil, isNull, err - } - is := isVer.(infoschema.InfoSchema) - if is == nil { - return nil, false, errors.New("missing information schema") - } - tbl, err := is.TableByName(context.Background(), model.NewCIStr(dbName), model.NewCIStr(tblName)) - if err != nil { - return nil, false, err - } - tblInfo := tbl.Meta() - idxInfo := tblInfo.FindIndexByName(strings.ToLower(idxName)) - if idxInfo == nil { - return nil, false, errors.New("index not found") - } - - pkLen := 1 - var pkIdx *model.IndexInfo - if tblInfo.IsCommonHandle { - pkIdx = tables.FindPrimaryIndex(tblInfo) - pkLen = len(pkIdx.Columns) - } - - if len(idxInfo.Columns)+pkLen != len(args)-3 { - return nil, false, errors.Errorf( - "column count mismatch, expected %d (index length + pk/rowid length), got %d", - len(idxInfo.Columns)+pkLen, len(args)-3) - } - - handle, err := buildHandle(ctx, tblInfo, args[3+len(idxInfo.Columns):], row) - if err != nil { - return nil, false, err - } - - idxDts := make([]types.Datum, 0, len(idxInfo.Columns)) - for i, idxCol := range idxInfo.Columns { - dt, err := args[i+3].Eval(ctx, row) - if err != nil { - return nil, false, err - } - ft := tblInfo.Columns[idxCol.Offset].FieldType - idxDt, err := dt.ConvertTo(ctx.TypeCtx(), &ft) - if err != nil { - return nil, false, err - } - idxDts = append(idxDts, idxDt) - } - tablecodec.TruncateIndexValues(tblInfo, idxInfo, idxDts) - idx := tables.NewIndex(tblInfo.ID, tblInfo, idxInfo) - - idxKey, _, err := idx.GenIndexKey(ctx.ErrCtx(), ctx.Location(), idxDts, handle, nil) - return idxKey, false, err -} - -func decodeKeyFromString(tc types.Context, isVer infoschemactx.MetaOnlyInfoSchema, s string) string { - key, err := hex.DecodeString(s) - if err != nil { - tc.AppendWarning(errors.NewNoStackErrorf("invalid key: %X", key)) - return s - } - // Auto decode byte if needed. - _, bs, err := codec.DecodeBytes(key, nil) - if err == nil { - key = bs - } - tableID := tablecodec.DecodeTableID(key) - if tableID <= 0 { - tc.AppendWarning(errors.NewNoStackErrorf("invalid key: %X", key)) - return s - } - - is, ok := isVer.(infoschema.InfoSchema) - if !ok { - tc.AppendWarning(errors.NewNoStackErrorf("infoschema not found when decoding key: %X", key)) - return s - } - tbl, _ := infoschema.FindTableByTblOrPartID(is, tableID) - loc := tc.Location() - if tablecodec.IsRecordKey(key) { - ret, err := decodeRecordKey(key, tableID, tbl, loc) - if err != nil { - tc.AppendWarning(err) - return s - } - return ret - } else if tablecodec.IsIndexKey(key) { - ret, err := decodeIndexKey(key, tableID, tbl, loc) - if err != nil { - tc.AppendWarning(err) - return s - } - return ret - } else if tablecodec.IsTableKey(key) { - ret, err := decodeTableKey(key, tableID, tbl) - if err != nil { - tc.AppendWarning(err) - return s - } - return ret - } - tc.AppendWarning(errors.NewNoStackErrorf("invalid key: %X", key)) - return s -} - -func decodeRecordKey(key []byte, tableID int64, tbl table.Table, loc *time.Location) (string, error) { - _, handle, err := tablecodec.DecodeRecordKey(key) - if err != nil { - return "", errors.Trace(err) - } - if handle.IsInt() { - ret := make(map[string]any) - if tbl != nil && tbl.Meta().Partition != nil { - ret["partition_id"] = tableID - tableID = tbl.Meta().ID - } - ret["table_id"] = strconv.FormatInt(tableID, 10) - // When the clustered index is enabled, we should show the PK name. - if tbl != nil && tbl.Meta().HasClusteredIndex() { - ret[tbl.Meta().GetPkName().String()] = handle.IntValue() - } else { - ret["_tidb_rowid"] = handle.IntValue() - } - retStr, err := json.Marshal(ret) - if err != nil { - return "", errors.Trace(err) - } - return string(retStr), nil - } - if tbl != nil { - tblInfo := tbl.Meta() - idxInfo := tables.FindPrimaryIndex(tblInfo) - if idxInfo == nil { - return "", errors.Trace(errors.Errorf("primary key not found when decoding record key: %X", key)) - } - cols := make(map[int64]*types.FieldType, len(tblInfo.Columns)) - for _, col := range tblInfo.Columns { - cols[col.ID] = &(col.FieldType) - } - handleColIDs := make([]int64, 0, len(idxInfo.Columns)) - for _, col := range idxInfo.Columns { - handleColIDs = append(handleColIDs, tblInfo.Columns[col.Offset].ID) - } - - if len(handleColIDs) != handle.NumCols() { - return "", errors.Trace(errors.Errorf("primary key length not match handle columns number in key")) - } - datumMap, err := tablecodec.DecodeHandleToDatumMap(handle, handleColIDs, cols, loc, nil) - if err != nil { - return "", errors.Trace(err) - } - ret := make(map[string]any) - if tbl.Meta().Partition != nil { - ret["partition_id"] = tableID - tableID = tbl.Meta().ID - } - ret["table_id"] = tableID - handleRet := make(map[string]any) - for colID := range datumMap { - dt := datumMap[colID] - dtStr, err := datumToJSONObject(&dt) - if err != nil { - return "", errors.Trace(err) - } - found := false - for _, colInfo := range tblInfo.Columns { - if colInfo.ID == colID { - found = true - handleRet[colInfo.Name.L] = dtStr - break - } - } - if !found { - return "", errors.Trace(errors.Errorf("column not found when decoding record key: %X", key)) - } - } - ret["handle"] = handleRet - retStr, err := json.Marshal(ret) - if err != nil { - return "", errors.Trace(err) - } - return string(retStr), nil - } - ret := make(map[string]any) - ret["table_id"] = tableID - ret["handle"] = handle.String() - retStr, err := json.Marshal(ret) - if err != nil { - return "", errors.Trace(err) - } - return string(retStr), nil -} - -func decodeIndexKey(key []byte, tableID int64, tbl table.Table, loc *time.Location) (string, error) { - if tbl != nil { - _, indexID, _, err := tablecodec.DecodeKeyHead(key) - if err != nil { - return "", errors.Trace(errors.Errorf("invalid record/index key: %X", key)) - } - tblInfo := tbl.Meta() - var targetIndex *model.IndexInfo - for _, idx := range tblInfo.Indices { - if idx.ID == indexID { - targetIndex = idx - break - } - } - if targetIndex == nil { - return "", errors.Trace(errors.Errorf("index not found when decoding index key: %X", key)) - } - colInfos := tables.BuildRowcodecColInfoForIndexColumns(targetIndex, tblInfo) - tps := tables.BuildFieldTypesForIndexColumns(targetIndex, tblInfo) - values, err := tablecodec.DecodeIndexKV(key, []byte{0}, len(colInfos), tablecodec.HandleNotNeeded, colInfos) - if err != nil { - return "", errors.Trace(err) - } - ds := make([]types.Datum, 0, len(colInfos)) - for i := 0; i < len(colInfos); i++ { - d, err := tablecodec.DecodeColumnValue(values[i], tps[i], loc) - if err != nil { - return "", errors.Trace(err) - } - ds = append(ds, d) - } - ret := make(map[string]any) - if tbl.Meta().Partition != nil { - ret["partition_id"] = tableID - tableID = tbl.Meta().ID - } - ret["table_id"] = tableID - ret["index_id"] = indexID - idxValMap := make(map[string]any, len(targetIndex.Columns)) - for i := 0; i < len(targetIndex.Columns); i++ { - dtStr, err := datumToJSONObject(&ds[i]) - if err != nil { - return "", errors.Trace(err) - } - idxValMap[targetIndex.Columns[i].Name.L] = dtStr - } - ret["index_vals"] = idxValMap - retStr, err := json.Marshal(ret) - if err != nil { - return "", errors.Trace(err) - } - return string(retStr), nil - } - _, indexID, indexValues, err := tablecodec.DecodeIndexKey(key) - if err != nil { - return "", errors.Trace(errors.Errorf("invalid index key: %X", key)) - } - ret := make(map[string]any) - ret["table_id"] = tableID - ret["index_id"] = indexID - ret["index_vals"] = strings.Join(indexValues, ", ") - retStr, err := json.Marshal(ret) - if err != nil { - return "", errors.Trace(err) - } - return string(retStr), nil -} - -func decodeTableKey(_ []byte, tableID int64, tbl table.Table) (string, error) { - ret := map[string]int64{} - if tbl != nil && tbl.Meta().GetPartitionInfo() != nil { - ret["partition_id"] = tableID - tableID = tbl.Meta().ID - } - ret["table_id"] = tableID - retStr, err := json.Marshal(ret) - if err != nil { - return "", errors.Trace(err) - } - return string(retStr), nil -} - -func datumToJSONObject(d *types.Datum) (any, error) { - if d.IsNull() { - return nil, nil - } - return d.ToString() -} From 59307539b64bb9977e0e11f19b6bc62301c8b24e Mon Sep 17 00:00:00 2001 From: tangenta Date: Wed, 4 Sep 2024 21:02:44 +0800 Subject: [PATCH 15/21] update bazel --- pkg/planner/core/BUILD.bazel | 1 + 1 file changed, 1 insertion(+) diff --git a/pkg/planner/core/BUILD.bazel b/pkg/planner/core/BUILD.bazel index 960ad2f493817..6a7fa154f4ae5 100644 --- a/pkg/planner/core/BUILD.bazel +++ b/pkg/planner/core/BUILD.bazel @@ -11,6 +11,7 @@ go_library( "encode.go", "exhaust_physical_plans.go", "explain.go", + "expression_codec_fn.go", "expression_rewriter.go", "find_best_task.go", "flat_plan.go", From c6ab4446ddee26d83a4b09d91f3b80461c73a62b Mon Sep 17 00:00:00 2001 From: tangenta Date: Wed, 4 Sep 2024 21:06:10 +0800 Subject: [PATCH 16/21] add comment --- pkg/planner/core/expression_codec_fn.go | 1 + 1 file changed, 1 insertion(+) diff --git a/pkg/planner/core/expression_codec_fn.go b/pkg/planner/core/expression_codec_fn.go index 4e834cf04a544..417702a0d60c8 100644 --- a/pkg/planner/core/expression_codec_fn.go +++ b/pkg/planner/core/expression_codec_fn.go @@ -225,6 +225,7 @@ func (h tidbCodecFuncHelper) encodeIndexKeyFromRow( idxDts = append(idxDts, idxDt) } tablecodec.TruncateIndexValues(tblInfo, idxInfo, idxDts) + // Use physicalID instead of tblInfo.ID here to handle the partition case. idx := tables.NewIndex(physicalID, tblInfo, idxInfo) idxKey, _, err := idx.GenIndexKey(ctx.ErrCtx(), ctx.Location(), idxDts, handle, nil) From f8dca63f4824c05ea5df471f88e494974e0cfb35 Mon Sep 17 00:00:00 2001 From: tangenta Date: Wed, 4 Sep 2024 21:47:34 +0800 Subject: [PATCH 17/21] support showing temp index for index keys --- pkg/expression/builtin_info.go | 20 +++++++++++++++++++- 1 file changed, 19 insertions(+), 1 deletion(-) diff --git a/pkg/expression/builtin_info.go b/pkg/expression/builtin_info.go index c8f7b5222417d..176c342339026 100644 --- a/pkg/expression/builtin_info.go +++ b/pkg/expression/builtin_info.go @@ -27,6 +27,7 @@ import ( "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" @@ -34,6 +35,7 @@ import ( "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" @@ -956,7 +958,23 @@ func (b *builtinTiDBMVCCInfoSig) evalString(ctx EvalContext, row chunk.Row) (str if err != nil { return "", false, err } - js, err := json.Marshal(resp) + 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 } From dc3af961089fbfd9d1aaba75072e7e887f744d94 Mon Sep 17 00:00:00 2001 From: tangenta Date: Wed, 4 Sep 2024 21:50:09 +0800 Subject: [PATCH 18/21] fix linter --- pkg/planner/core/expression_codec_fn.go | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/pkg/planner/core/expression_codec_fn.go b/pkg/planner/core/expression_codec_fn.go index 417702a0d60c8..3631d4b62a911 100644 --- a/pkg/planner/core/expression_codec_fn.go +++ b/pkg/planner/core/expression_codec_fn.go @@ -90,13 +90,15 @@ func (h tidbCodecFuncHelper) findCommonOrPartitionedTable( // The arguments will be filled by caller. return nil, 0, plannererrors.ErrSpecificAccessDenied } - if part := tbl.GetPartitionedTable(); part != nil && len(partName) > 0 { - pid, err := tables.FindPartitionByName(tbl.Meta(), partName) - if err != nil { - return nil, 0, errors.Trace(err) + if part := tbl.GetPartitionedTable(); part != nil { + if len(partName) > 0 { + pid, err := tables.FindPartitionByName(tbl.Meta(), partName) + if err != nil { + return nil, 0, errors.Trace(err) + } + tbl = part.GetPartition(pid) + return tbl, pid, nil } - tbl = part.GetPartition(pid) - return tbl, pid, nil } else { if len(partName) != 0 { return nil, 0, errors.New("not a partitioned table") From 9aa48284e0c7b6dd3fba260779838933a6d2d715 Mon Sep 17 00:00:00 2001 From: tangenta Date: Wed, 4 Sep 2024 21:52:22 +0800 Subject: [PATCH 19/21] refine code --- pkg/planner/core/expression_codec_fn.go | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/pkg/planner/core/expression_codec_fn.go b/pkg/planner/core/expression_codec_fn.go index 3631d4b62a911..461aca426e321 100644 --- a/pkg/planner/core/expression_codec_fn.go +++ b/pkg/planner/core/expression_codec_fn.go @@ -90,8 +90,8 @@ func (h tidbCodecFuncHelper) findCommonOrPartitionedTable( // The arguments will be filled by caller. return nil, 0, plannererrors.ErrSpecificAccessDenied } - if part := tbl.GetPartitionedTable(); part != nil { - if len(partName) > 0 { + if len(partName) > 0 { + if part := tbl.GetPartitionedTable(); part != nil { pid, err := tables.FindPartitionByName(tbl.Meta(), partName) if err != nil { return nil, 0, errors.Trace(err) @@ -99,10 +99,7 @@ func (h tidbCodecFuncHelper) findCommonOrPartitionedTable( tbl = part.GetPartition(pid) return tbl, pid, nil } - } else { - if len(partName) != 0 { - return nil, 0, errors.New("not a partitioned table") - } + return nil, 0, errors.New("not a partitioned table") } return tbl, tbl.Meta().ID, nil } From a0dab565b2922985ee87025c38d8774aca0d2611 Mon Sep 17 00:00:00 2001 From: tangenta Date: Wed, 4 Sep 2024 21:55:23 +0800 Subject: [PATCH 20/21] update bazel --- pkg/expression/BUILD.bazel | 2 ++ 1 file changed, 2 insertions(+) diff --git a/pkg/expression/BUILD.bazel b/pkg/expression/BUILD.bazel index ae4c16a94fc66..bfb562f84bd2c 100644 --- a/pkg/expression/BUILD.bazel +++ b/pkg/expression/BUILD.bazel @@ -94,6 +94,7 @@ go_library( "//pkg/sessionctx/stmtctx", "//pkg/sessionctx/variable", "//pkg/store/helper", + "//pkg/tablecodec", "//pkg/types", "//pkg/types/parser_driver", "//pkg/util", @@ -127,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", From bd283836f6d1c8b550a7e583800f8b9665678b37 Mon Sep 17 00:00:00 2001 From: tangenta Date: Thu, 5 Sep 2024 11:23:03 +0800 Subject: [PATCH 21/21] add test and fix unstable test --- pkg/executor/infoschema_reader.go | 2 +- .../addindextest3/functional_test.go | 27 +++++++++++++++++++ 2 files changed, 28 insertions(+), 1 deletion(-) diff --git a/pkg/executor/infoschema_reader.go b/pkg/executor/infoschema_reader.go index a232c32d71ad6..05403fd4c31ba 100644 --- a/pkg/executor/infoschema_reader.go +++ b/pkg/executor/infoschema_reader.go @@ -749,7 +749,7 @@ func (e *memtableRetriever) setDataFromTables(ctx context.Context, sctx sessionc if !ex.HasTableSchema(t.DBName.L) { return true } - if checker != nil && !checker.RequestVerification(sctx.GetSessionVars().ActiveRoles, t.DBName.L, t.TableName.L, "", mysql.SelectPriv) { + if checker != nil && !checker.RequestVerification(sctx.GetSessionVars().ActiveRoles, t.DBName.L, t.TableName.L, "", mysql.AllPrivMask) { return true } diff --git a/tests/realtikvtest/addindextest3/functional_test.go b/tests/realtikvtest/addindextest3/functional_test.go index 0e9af59522793..ce8e7ba06cf27 100644 --- a/tests/realtikvtest/addindextest3/functional_test.go +++ b/tests/realtikvtest/addindextest3/functional_test.go @@ -17,6 +17,7 @@ package addindextest import ( "context" "fmt" + "strings" "sync" "testing" @@ -116,3 +117,29 @@ func TestMockMemoryUsedUp(t *testing.T) { tk.MustExec("insert into t values (1,1,1,1), (2,2,2,2), (3,3,3,3);") tk.MustGetErrMsg("alter table t add index i(c), add index i2(c2);", "[ddl:8247]Ingest failed: memory used up") } + +func TestTiDBEncodeKeyTempIndexKey(t *testing.T) { + store := realtikvtest.CreateMockStoreAndSetup(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);") + runDML := false + testfailpoint.EnableCall(t, "github.com/pingcap/tidb/pkg/ddl/onJobRunBefore", func(job *model.Job) { + if !runDML && job.Type == model.ActionAddIndex && job.SchemaState == model.StateWriteOnly { + tk2 := testkit.NewTestKit(t, store) + tk2.MustExec("use test") + tk2.MustExec("insert into t values (2, 2);") + runDML = true + } + }) + tk.MustExec("create index idx on t(b);") + require.True(t, runDML) + + rows := tk.MustQuery("select tidb_mvcc_info(tidb_encode_index_key('test', 't', 'idx', 1, 1));").Rows() + rs := rows[0][0].(string) + require.Equal(t, 1, strings.Count(rs, "writes"), rs) + rows = tk.MustQuery("select tidb_mvcc_info(tidb_encode_index_key('test', 't', 'idx', 2, 2));").Rows() + rs = rows[0][0].(string) + require.Equal(t, 2, strings.Count(rs, "writes"), rs) +}