tidb hint_processor 源码
tidb hint_processor 代码
文件路径：/util/hint/hint_processor.go
// Copyright 2020 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 hint
import (
"fmt"
"strconv"
"strings"
"github.com/pingcap/tidb/errno"
"github.com/pingcap/tidb/parser"
"github.com/pingcap/tidb/parser/ast"
"github.com/pingcap/tidb/parser/format"
"github.com/pingcap/tidb/parser/model"
"github.com/pingcap/tidb/sessionctx"
"github.com/pingcap/tidb/util/dbterror"
"github.com/pingcap/tidb/util/logutil"
"go.uber.org/zap"
)
var supportedHintNameForInsertStmt = map[string]struct{}{}
var errWarnConflictingHint = dbterror.ClassUtil.NewStd(errno.ErrWarnConflictingHint)
func init() {
supportedHintNameForInsertStmt["memory_quota"] = struct{}{}
}
// HintsSet contains all hints of a query.
type HintsSet struct {
tableHints [][]*ast.TableOptimizerHint // Slice offset is the traversal order of `SelectStmt` in the ast.
indexHints [][]*ast.IndexHint // Slice offset is the traversal order of `TableName` in the ast.
}
// GetFirstTableHints gets the first table hints.
func (hs *HintsSet) GetFirstTableHints() []*ast.TableOptimizerHint {
if len(hs.tableHints) > 0 {
return hs.tableHints[0]
}
return nil
}
// ContainTableHint checks whether the table hint set contains a hint.
func (hs *HintsSet) ContainTableHint(hint string) bool {
for _, tableHintsForBlock := range hs.tableHints {
for _, tableHint := range tableHintsForBlock {
if tableHint.HintName.String() == hint {
return true
}
}
}
return false
}
// setTableHints4StmtNode sets table hints for select/update/delete.
func setTableHints4StmtNode(node ast.Node, hints []*ast.TableOptimizerHint) {
switch x := node.(type) {
case *ast.SelectStmt:
x.TableHints = hints
case *ast.UpdateStmt:
x.TableHints = hints
case *ast.DeleteStmt:
x.TableHints = hints
}
}
// ExtractTableHintsFromStmtNode extracts table hints from this node.
func ExtractTableHintsFromStmtNode(node ast.Node, sctx sessionctx.Context) []*ast.TableOptimizerHint {
switch x := node.(type) {
case *ast.SelectStmt:
return x.TableHints
case *ast.UpdateStmt:
return x.TableHints
case *ast.DeleteStmt:
return x.TableHints
case *ast.InsertStmt:
// check duplicated hints
checkInsertStmtHintDuplicated(node, sctx)
return x.TableHints
case *ast.ExplainStmt:
return ExtractTableHintsFromStmtNode(x.Stmt, sctx)
default:
return nil
}
}
// checkInsertStmtHintDuplicated check whether existed the duplicated hints in both insertStmt and its selectStmt.
// If existed, it would send a warning message.
func checkInsertStmtHintDuplicated(node ast.Node, sctx sessionctx.Context) {
switch x := node.(type) {
case *ast.InsertStmt:
if len(x.TableHints) > 0 {
var supportedHint *ast.TableOptimizerHint
for _, hint := range x.TableHints {
if _, ok := supportedHintNameForInsertStmt[hint.HintName.L]; ok {
supportedHint = hint
break
}
}
if supportedHint != nil {
var duplicatedHint *ast.TableOptimizerHint
for _, hint := range ExtractTableHintsFromStmtNode(x.Select, nil) {
if hint.HintName.L == supportedHint.HintName.L {
duplicatedHint = hint
break
}
}
if duplicatedHint != nil {
hint := fmt.Sprintf("%s(`%v`)", duplicatedHint.HintName.O, duplicatedHint.HintData)
sctx.GetSessionVars().StmtCtx.AppendWarning(errWarnConflictingHint.FastGenByArgs(hint))
}
}
}
default:
return
}
}
// RestoreOptimizerHints restores these hints.
func RestoreOptimizerHints(hints []*ast.TableOptimizerHint) string {
hintsStr := make([]string, 0, len(hints))
hintsMap := make(map[string]struct{}, len(hints))
for _, hint := range hints {
hintStr := RestoreTableOptimizerHint(hint)
if _, ok := hintsMap[hintStr]; ok {
continue
}
hintsMap[hintStr] = struct{}{}
hintsStr = append(hintsStr, hintStr)
}
return strings.Join(hintsStr, ", ")
}
// RestoreTableOptimizerHint returns string format of TableOptimizerHint.
func RestoreTableOptimizerHint(hint *ast.TableOptimizerHint) string {
var sb strings.Builder
ctx := format.NewRestoreCtx(format.DefaultRestoreFlags, &sb)
err := hint.Restore(ctx)
// There won't be any error for optimizer hint.
if err != nil {
logutil.BgLogger().Debug("restore TableOptimizerHint failed", zap.Error(err))
}
return strings.ToLower(sb.String())
}
// RestoreIndexHint returns string format of IndexHint.
func RestoreIndexHint(hint *ast.IndexHint) (string, error) {
var sb strings.Builder
ctx := format.NewRestoreCtx(format.DefaultRestoreFlags, &sb)
err := hint.Restore(ctx)
if err != nil {
logutil.BgLogger().Debug("restore IndexHint failed", zap.Error(err))
return "", err
}
return strings.ToLower(sb.String()), nil
}
// Restore returns the string format of HintsSet.
func (hs *HintsSet) Restore() (string, error) {
hintsStr := make([]string, 0, len(hs.tableHints)+len(hs.indexHints))
for _, tblHints := range hs.tableHints {
for _, tblHint := range tblHints {
hintsStr = append(hintsStr, RestoreTableOptimizerHint(tblHint))
}
}
for _, idxHints := range hs.indexHints {
for _, idxHint := range idxHints {
str, err := RestoreIndexHint(idxHint)
if err != nil {
return "", err
}
hintsStr = append(hintsStr, str)
}
}
return strings.Join(hintsStr, ", "), nil
}
type hintProcessor struct {
*HintsSet
// bindHint2Ast indicates the behavior of the processor, `true` for bind hint to ast, `false` for extract hint from ast.
bindHint2Ast bool
tableCounter int
indexCounter int
blockCounter int
}
func (hp *hintProcessor) Enter(in ast.Node) (ast.Node, bool) {
switch v := in.(type) {
case *ast.SelectStmt, *ast.UpdateStmt, *ast.DeleteStmt:
if hp.bindHint2Ast {
if hp.tableCounter < len(hp.tableHints) {
setTableHints4StmtNode(in, hp.tableHints[hp.tableCounter])
} else {
setTableHints4StmtNode(in, nil)
}
hp.tableCounter++
} else {
hp.tableHints = append(hp.tableHints, ExtractTableHintsFromStmtNode(in, nil))
}
hp.blockCounter++
case *ast.TableName:
// Insert cases.
if hp.blockCounter == 0 {
return in, false
}
if hp.bindHint2Ast {
if hp.indexCounter < len(hp.indexHints) {
v.IndexHints = hp.indexHints[hp.indexCounter]
} else {
v.IndexHints = nil
}
hp.indexCounter++
} else {
hp.indexHints = append(hp.indexHints, v.IndexHints)
}
}
return in, false
}
func (hp *hintProcessor) Leave(in ast.Node) (ast.Node, bool) {
switch in.(type) {
case *ast.SelectStmt, *ast.UpdateStmt, *ast.DeleteStmt:
hp.blockCounter--
}
return in, true
}
// CollectHint collects hints for a statement.
func CollectHint(in ast.StmtNode) *HintsSet {
hp := hintProcessor{HintsSet: &HintsSet{tableHints: make([][]*ast.TableOptimizerHint, 0, 4), indexHints: make([][]*ast.IndexHint, 0, 4)}}
in.Accept(&hp)
return hp.HintsSet
}
// BindHint will add hints for stmt according to the hints in `hintsSet`.
func BindHint(stmt ast.StmtNode, hintsSet *HintsSet) ast.StmtNode {
hp := hintProcessor{HintsSet: hintsSet, bindHint2Ast: true}
stmt.Accept(&hp)
return stmt
}
// ParseHintsSet parses a SQL string, then collects and normalizes the HintsSet.
func ParseHintsSet(p *parser.Parser, sql, charset, collation, db string) (*HintsSet, ast.StmtNode, []error, error) {
stmtNodes, warns, err := p.ParseSQL(sql,
parser.CharsetConnection(charset),
parser.CollationConnection(collation))
if err != nil {
return nil, nil, nil, err
}
if len(stmtNodes) != 1 {
return nil, nil, nil, fmt.Errorf("bind_sql must be a single statement: %s", sql)
}
hs := CollectHint(stmtNodes[0])
processor := &BlockHintProcessor{}
stmtNodes[0].Accept(processor)
topNodeType := nodeType4Stmt(stmtNodes[0])
for i, tblHints := range hs.tableHints {
newHints := make([]*ast.TableOptimizerHint, 0, len(tblHints))
curOffset := i + 1
if topNodeType == TypeDelete || topNodeType == TypeUpdate {
curOffset = curOffset - 1
}
for _, tblHint := range tblHints {
if tblHint.HintName.L == hintQBName {
continue
}
offset := processor.GetHintOffset(tblHint.QBName, curOffset)
if offset < 0 || !processor.checkTableQBName(tblHint.Tables) {
hintStr := RestoreTableOptimizerHint(tblHint)
return nil, nil, nil, fmt.Errorf("Unknown query block name in hint %s", hintStr)
}
tblHint.QBName, err = GenerateQBName(topNodeType, offset)
if err != nil {
return nil, nil, nil, err
}
for i, tbl := range tblHint.Tables {
if tbl.DBName.String() == "" {
tblHint.Tables[i].DBName = model.NewCIStr(db)
}
}
newHints = append(newHints, tblHint)
}
hs.tableHints[i] = newHints
}
return hs, stmtNodes[0], extractHintWarns(warns), nil
}
func extractHintWarns(warns []error) []error {
for _, w := range warns {
if parser.ErrParse.Equal(w) ||
parser.ErrWarnOptimizerHintUnsupportedHint.Equal(w) ||
parser.ErrWarnOptimizerHintInvalidToken.Equal(w) ||
parser.ErrWarnMemoryQuotaOverflow.Equal(w) ||
parser.ErrWarnOptimizerHintParseError.Equal(w) ||
parser.ErrWarnOptimizerHintInvalidInteger.Equal(w) {
// Just one warning is enough, however we use a slice here to stop golint complaining
// "error should be the last type when returning multiple items" for `ParseHintsSet`.
return []error{w}
}
}
return nil
}
// BlockHintProcessor processes hints at different level of sql statement.
type BlockHintProcessor struct {
QbNameMap map[string]int // Map from query block name to select stmt offset.
QbHints map[int][]*ast.TableOptimizerHint // Group all hints at same query block.
Ctx sessionctx.Context
selectStmtOffset int
}
// MaxSelectStmtOffset returns the current stmt offset.
func (p *BlockHintProcessor) MaxSelectStmtOffset() int {
return p.selectStmtOffset
}
// Enter implements Visitor interface.
func (p *BlockHintProcessor) Enter(in ast.Node) (ast.Node, bool) {
switch node := in.(type) {
case *ast.UpdateStmt:
p.checkQueryBlockHints(node.TableHints, 0)
case *ast.DeleteStmt:
p.checkQueryBlockHints(node.TableHints, 0)
case *ast.SelectStmt:
p.selectStmtOffset++
node.QueryBlockOffset = p.selectStmtOffset
p.checkQueryBlockHints(node.TableHints, node.QueryBlockOffset)
}
return in, false
}
// Leave implements Visitor interface.
func (*BlockHintProcessor) Leave(in ast.Node) (ast.Node, bool) {
return in, true
}
const hintQBName = "qb_name"
// checkQueryBlockHints checks the validity of query blocks and records the map of query block name to select offset.
func (p *BlockHintProcessor) checkQueryBlockHints(hints []*ast.TableOptimizerHint, offset int) {
var qbName string
for _, hint := range hints {
if hint.HintName.L != hintQBName {
continue
}
if qbName != "" {
if p.Ctx != nil {
p.Ctx.GetSessionVars().StmtCtx.AppendWarning(fmt.Errorf("There are more than two query names in same query block,, using the first one %s", qbName))
}
} else {
qbName = hint.QBName.L
}
}
if qbName == "" {
return
}
if p.QbNameMap == nil {
p.QbNameMap = make(map[string]int)
}
if _, ok := p.QbNameMap[qbName]; ok {
if p.Ctx != nil {
p.Ctx.GetSessionVars().StmtCtx.AppendWarning(fmt.Errorf("Duplicate query block name %s, only the first one is effective", qbName))
}
} else {
p.QbNameMap[qbName] = offset
}
}
const (
defaultUpdateBlockName = "upd_1"
defaultDeleteBlockName = "del_1"
defaultSelectBlockPrefix = "sel_"
)
// NodeType indicates if the node is for SELECT / UPDATE / DELETE.
type NodeType int
const (
// TypeUpdate for Update.
TypeUpdate NodeType = iota
// TypeDelete for DELETE.
TypeDelete
// TypeSelect for SELECT.
TypeSelect
// TypeInvalid for unexpected statements.
TypeInvalid
)
// nodeType4Stmt returns the NodeType for a statement. The type is used for SQL bind.
func nodeType4Stmt(node ast.StmtNode) NodeType {
switch node.(type) {
// This type is used by SQL bind, we only handle SQL bind for INSERT INTO SELECT, so we treat InsertStmt as TypeSelect.
case *ast.SelectStmt, *ast.InsertStmt:
return TypeSelect
case *ast.UpdateStmt:
return TypeUpdate
case *ast.DeleteStmt:
return TypeDelete
}
return TypeInvalid
}
// getBlockName finds the offset of query block name. It uses 0 as offset for top level update or delete,
// -1 for invalid block name.
func (p *BlockHintProcessor) getBlockOffset(blockName model.CIStr) int {
if p.QbNameMap != nil {
level, ok := p.QbNameMap[blockName.L]
if ok {
return level
}
}
// Handle the default query block name.
if blockName.L == defaultUpdateBlockName || blockName.L == defaultDeleteBlockName {
return 0
}
if strings.HasPrefix(blockName.L, defaultSelectBlockPrefix) {
suffix := blockName.L[len(defaultSelectBlockPrefix):]
level, err := strconv.ParseInt(suffix, 10, 64)
if err != nil || level > int64(p.selectStmtOffset) {
return -1
}
return int(level)
}
return -1
}
// GetHintOffset gets the offset of stmt that the hints take effects.
func (p *BlockHintProcessor) GetHintOffset(qbName model.CIStr, currentOffset int) int {
if qbName.L != "" {
return p.getBlockOffset(qbName)
}
return currentOffset
}
func (p *BlockHintProcessor) checkTableQBName(tables []ast.HintTable) bool {
for _, table := range tables {
if table.QBName.L != "" && p.getBlockOffset(table.QBName) < 0 {
return false
}
}
return true
}
// GetCurrentStmtHints extracts all hints that take effects at current stmt.
func (p *BlockHintProcessor) GetCurrentStmtHints(hints []*ast.TableOptimizerHint, currentOffset int) []*ast.TableOptimizerHint {
if p.QbHints == nil {
p.QbHints = make(map[int][]*ast.TableOptimizerHint)
}
for _, hint := range hints {
if hint.HintName.L == hintQBName {
continue
}
offset := p.GetHintOffset(hint.QBName, currentOffset)
if offset < 0 || !p.checkTableQBName(hint.Tables) {
hintStr := RestoreTableOptimizerHint(hint)
p.Ctx.GetSessionVars().StmtCtx.AppendWarning(fmt.Errorf("Hint %s is ignored due to unknown query block name", hintStr))
continue
}
p.QbHints[offset] = append(p.QbHints[offset], hint)
}
return p.QbHints[currentOffset]
}
// GenerateQBName builds QBName from offset.
func GenerateQBName(nodeType NodeType, blockOffset int) (model.CIStr, error) {
if blockOffset == 0 {
if nodeType == TypeDelete {
return model.NewCIStr(defaultDeleteBlockName), nil
}
if nodeType == TypeUpdate {
return model.NewCIStr(defaultUpdateBlockName), nil
}
return model.NewCIStr(""), fmt.Errorf("Unexpected NodeType %d when block offset is 0", nodeType)
}
return model.NewCIStr(fmt.Sprintf("%s%d", defaultSelectBlockPrefix, blockOffset)), nil
}
相关信息
相关文章
0
赞
热门推荐
-
2、 - 优质文章
-
3、 gate.io
-
6、 golang
-
7、 openharmony