Golang kv.Client类代码示例

func aggFuncToPBExpr(sc *variable.StatementContext, client kv.Client, aggFunc expression.AggregationFunction) *tipb.Expr {
	pc := pbConverter{client: client, sc: sc}
	var tp tipb.ExprType
	switch aggFunc.GetName() {
	case ast.AggFuncCount:
		tp = tipb.ExprType_Count
	case ast.AggFuncFirstRow:
		tp = tipb.ExprType_First
	case ast.AggFuncGroupConcat:
		tp = tipb.ExprType_GroupConcat
	case ast.AggFuncMax:
		tp = tipb.ExprType_Max
	case ast.AggFuncMin:
		tp = tipb.ExprType_Min
	case ast.AggFuncSum:
		tp = tipb.ExprType_Sum
	case ast.AggFuncAvg:
		tp = tipb.ExprType_Avg
	}
	if !client.SupportRequestType(kv.ReqTypeSelect, int64(tp)) {
		return nil
	}

	children := make([]*tipb.Expr, 0, len(aggFunc.GetArgs()))
	for _, arg := range aggFunc.GetArgs() {
		pbArg := pc.exprToPB(arg)
		if pbArg == nil {
			return nil
		}
		children = append(children, pbArg)
	}
	return &tipb.Expr{Tp: tp, Children: children}
}

func (b *executorBuilder) columnToPBExpr(client kv.Client, column *expression.Column, tbl *model.TableInfo) *tipb.Expr {
	if !client.SupportRequestType(kv.ReqTypeSelect, int64(tipb.ExprType_ColumnRef)) {
		return nil
	}
	switch column.GetType().Tp {
	case mysql.TypeBit, mysql.TypeSet, mysql.TypeEnum, mysql.TypeDecimal, mysql.TypeGeometry,
		mysql.TypeDate, mysql.TypeNewDate, mysql.TypeDatetime, mysql.TypeTimestamp, mysql.TypeYear:
		return nil
	}

	id := int64(-1)
	for _, col := range tbl.Columns {
		if tbl.Name == column.TblName && col.Name == column.ColName {
			id = col.ID
			break
		}
	}
	// Zero Column ID is not a column from table, can not support for now.
	if id == 0 {
		return nil
	}
	// TODO：If the column ID isn't in fields, it means the column is from an outer table,
	// its value is available to use.
	if id == -1 {
		return nil
	}

	return &tipb.Expr{
		Tp:  tipb.ExprType_ColumnRef.Enum(),
		Val: codec.EncodeInt(nil, id)}
}

// Select do a select request, returns SelectResult.
// conncurrency: The max concurrency for underlying coprocessor request.
// keepOrder: If the result should returned in key order. For example if we need keep data in order by
//            scan index, we should set keepOrder to true.
func Select(client kv.Client, req *tipb.SelectRequest, concurrency int, keepOrder bool) (SelectResult, error) {
	// Convert tipb.*Request to kv.Request.
	kvReq, err := composeRequest(req, concurrency, keepOrder)
	if err != nil {
		return nil, errors.Trace(err)
	}
	resp := client.Send(kvReq)
	if resp == nil {
		return nil, errors.New("client returns nil response")
	}
	result := &selectResult{
		resp:    resp,
		results: make(chan PartialResult, 5),
		done:    make(chan error, 1),
	}
	// If Aggregates is not nil, we should set result fields latter.
	if len(req.Aggregates) == 0 && len(req.GroupBy) == 0 {
		if req.TableInfo != nil {
			result.fields = ProtoColumnsToFieldTypes(req.TableInfo.Columns)
		} else {
			result.fields = ProtoColumnsToFieldTypes(req.IndexInfo.Columns)
			length := len(req.IndexInfo.Columns)
			if req.IndexInfo.Columns[length-1].GetPkHandle() {
				// Returned index row do not contains extra PKHandle column.
				result.fields = result.fields[:length-1]
			}
			result.index = true
		}
	} else {
		result.aggregate = true
	}
	return result, nil
}

func bitwiseFuncToPBExpr(client kv.Client, expr *expression.ScalarFunction) *tipb.Expr {
	var tp tipb.ExprType
	switch expr.FuncName.L {
	case ast.And:
		tp = tipb.ExprType_BitAnd
	case ast.Or:
		tp = tipb.ExprType_BitOr
	case ast.Xor:
		tp = tipb.ExprType_BitXor
	case ast.LeftShift:
		tp = tipb.ExprType_LeftShift
	case ast.RightShift:
		tp = tipb.ExprType_RighShift
	case ast.BitNeg:
		return bitNegToPBExpr(client, expr)
	}
	if !client.SupportRequestType(kv.ReqTypeSelect, int64(tp)) {
		return nil
	}
	expr0 := exprToPB(client, expr.Args[0])
	if expr0 == nil {
		return nil
	}
	expr1 := exprToPB(client, expr.Args[1])
	if expr1 == nil {
		return nil
	}
	return &tipb.Expr{
		Tp:       tp,
		Children: []*tipb.Expr{expr0, expr1}}
}

func (b *executorBuilder) aggFuncToPBExpr(client kv.Client, af *ast.AggregateFuncExpr, tn *ast.TableName) *tipb.Expr {
	name := strings.ToLower(af.F)
	var tp tipb.ExprType
	switch name {
	case ast.AggFuncCount:
		tp = tipb.ExprType_Count
	case ast.AggFuncFirstRow:
		tp = tipb.ExprType_First
	case ast.AggFuncGroupConcat:
		tp = tipb.ExprType_GroupConcat
	case ast.AggFuncMax:
		tp = tipb.ExprType_Max
	case ast.AggFuncMin:
		tp = tipb.ExprType_Min
	case ast.AggFuncSum:
		tp = tipb.ExprType_Sum
	case ast.AggFuncAvg:
		tp = tipb.ExprType_Avg
	}
	if !client.SupportRequestType(kv.ReqTypeSelect, int64(tp)) {
		return nil
	}
	// convert it to pb
	children := make([]*tipb.Expr, 0, len(af.Args))
	for _, arg := range af.Args {
		pbArg := b.exprToPBExpr(client, arg, tn)
		if pbArg == nil {
			return nil
		}
		children = append(children, pbArg)
	}
	return &tipb.Expr{Tp: tp.Enum(), Children: children}
}

func (b *executorBuilder) columnNameToPBExpr(client kv.Client, column *ast.ColumnNameExpr, tn *ast.TableName) *tipb.Expr {
	if !client.SupportRequestType(kv.ReqTypeSelect, int64(tipb.ExprType_ColumnRef)) {
		return nil
	}
	// Zero Column ID is not a column from table, can not support for now.
	if column.Refer.Column.ID == 0 {
		return nil
	}
	switch column.Refer.Expr.GetType().Tp {
	case mysql.TypeBit, mysql.TypeSet, mysql.TypeEnum, mysql.TypeDecimal, mysql.TypeGeometry,
		mysql.TypeDate, mysql.TypeNewDate, mysql.TypeDatetime, mysql.TypeTimestamp, mysql.TypeYear:
		return nil
	}
	matched := false
	for _, f := range tn.GetResultFields() {
		if f.TableName == column.Refer.TableName && f.Column.ID == column.Refer.Column.ID {
			matched = true
			break
		}
	}
	if matched {
		pbExpr := new(tipb.Expr)
		pbExpr.Tp = tipb.ExprType_ColumnRef.Enum()
		pbExpr.Val = codec.EncodeInt(nil, column.Refer.Column.ID)
		return pbExpr
	}
	// If the column ID isn't in fields, it means the column is from an outer table,
	// its value is available to use.
	return b.datumToPBExpr(client, *column.Refer.Expr.GetDatum())
}

func logicalOpsToPBExpr(client kv.Client, expr *expression.ScalarFunction) *tipb.Expr {
	var tp tipb.ExprType
	switch expr.FuncName.L {
	case ast.AndAnd:
		tp = tipb.ExprType_And
	case ast.OrOr:
		tp = tipb.ExprType_Or
	case ast.LogicXor:
		tp = tipb.ExprType_Xor
	case ast.UnaryNot:
		return notToPBExpr(client, expr)
	}
	if !client.SupportRequestType(kv.ReqTypeSelect, int64(tp)) {
		return nil
	}

	expr0 := exprToPB(client, expr.Args[0])
	if expr0 == nil {
		return nil
	}
	expr1 := exprToPB(client, expr.Args[1])
	if expr1 == nil {
		return nil
	}
	return &tipb.Expr{
		Tp:       tp,
		Children: []*tipb.Expr{expr0, expr1}}
}

func arithmeticalOpsToPBExpr(client kv.Client, expr *expression.ScalarFunction) *tipb.Expr {
	var tp tipb.ExprType
	switch expr.FuncName.L {
	case ast.Plus:
		tp = tipb.ExprType_Plus
	case ast.Minus:
		tp = tipb.ExprType_Minus
	case ast.Mul:
		tp = tipb.ExprType_Mul
	case ast.Div:
		tp = tipb.ExprType_Div
	case ast.Mod:
		tp = tipb.ExprType_Mod
	case ast.IntDiv:
		tp = tipb.ExprType_IntDiv
	}
	if !client.SupportRequestType(kv.ReqTypeSelect, int64(tp)) {
		return nil
	}
	expr0 := exprToPB(client, expr.Args[0])
	if expr0 == nil {
		return nil
	}
	expr1 := exprToPB(client, expr.Args[1])
	if expr1 == nil {
		return nil
	}
	return &tipb.Expr{
		Tp:       tp,
		Children: []*tipb.Expr{expr0, expr1}}
}

func (b *executorBuilder) newAggFuncToPBExpr(client kv.Client, aggFunc expression.AggregationFunction,
	tbl *model.TableInfo) *tipb.Expr {
	var tp tipb.ExprType
	switch aggFunc.GetName() {
	case ast.AggFuncCount:
		tp = tipb.ExprType_Count
	case ast.AggFuncFirstRow:
		tp = tipb.ExprType_First
	case ast.AggFuncGroupConcat:
		tp = tipb.ExprType_GroupConcat
	case ast.AggFuncMax:
		tp = tipb.ExprType_Max
	case ast.AggFuncMin:
		tp = tipb.ExprType_Min
	case ast.AggFuncSum:
		tp = tipb.ExprType_Sum
	case ast.AggFuncAvg:
		tp = tipb.ExprType_Avg
	}
	if !client.SupportRequestType(kv.ReqTypeSelect, int64(tp)) {
		return nil
	}

	children := make([]*tipb.Expr, 0, len(aggFunc.GetArgs()))
	for _, arg := range aggFunc.GetArgs() {
		pbArg := b.newExprToPBExpr(client, arg, tbl)
		if pbArg == nil {
			return nil
		}
		children = append(children, pbArg)
	}
	return &tipb.Expr{Tp: tp.Enum(), Children: children}
}

func compareOpsToPBExpr(client kv.Client, expr *expression.ScalarFunction) *tipb.Expr {
	var tp tipb.ExprType
	switch expr.FuncName.L {
	case ast.LT:
		tp = tipb.ExprType_LT
	case ast.LE:
		tp = tipb.ExprType_LE
	case ast.EQ:
		tp = tipb.ExprType_EQ
	case ast.NE:
		tp = tipb.ExprType_NE
	case ast.GE:
		tp = tipb.ExprType_GE
	case ast.GT:
		tp = tipb.ExprType_GT
	case ast.NullEQ:
		tp = tipb.ExprType_NullEQ
	case ast.In:
		return inToPBExpr(client, expr)
	case ast.Like:
		// Only patterns like 'abc', '%abc', 'abc%', '%abc%' can be converted to *tipb.Expr for now.
		escape := expr.Args[2].(*expression.Constant).Value
		if escape.IsNull() || byte(escape.GetInt64()) != '\\' {
			return nil
		}
		pattern, ok := expr.Args[1].(*expression.Constant)
		if !ok || pattern.Value.Kind() != types.KindString {
			return nil
		}
		for i, b := range pattern.Value.GetString() {
			switch b {
			case '\\', '_':
				return nil
			case '%':
				if i != 0 && i != len(pattern.Value.GetString())-1 {
					return nil
				}
			}
		}
		tp = tipb.ExprType_Like
	}
	if !client.SupportRequestType(kv.ReqTypeSelect, int64(tp)) {
		return nil
	}
	expr0 := exprToPB(client, expr.Args[0])
	if expr0 == nil {
		return nil
	}
	expr1 := exprToPB(client, expr.Args[1])
	if expr1 == nil {
		return nil
	}
	return &tipb.Expr{
		Tp:       tp,
		Children: []*tipb.Expr{expr0, expr1}}
}

func (b *executorBuilder) scalarFuncToPBExpr(client kv.Client, expr *expression.ScalarFunction,
	tbl *model.TableInfo) *tipb.Expr {
	var tp tipb.ExprType
	switch expr.FuncName.L {
	case ast.LT:
		tp = tipb.ExprType_LT
	case ast.LE:
		tp = tipb.ExprType_LE
	case ast.EQ:
		tp = tipb.ExprType_EQ
	case ast.NE:
		tp = tipb.ExprType_NE
	case ast.GE:
		tp = tipb.ExprType_GE
	case ast.GT:
		tp = tipb.ExprType_GT
	case ast.NullEQ:
		tp = tipb.ExprType_NullEQ
	case ast.And:
		tp = tipb.ExprType_And
	case ast.Or:
		tp = tipb.ExprType_Or
	// It's the operation for unary operator.
	case ast.UnaryNot:
		tp = tipb.ExprType_Not
	default:
		return nil
	}

	if !client.SupportRequestType(kv.ReqTypeSelect, int64(tp)) {
		return nil
	}

	if len(expr.Args) == 1 {
		child := b.newExprToPBExpr(client, expr.Args[0], tbl)
		if child == nil {
			return nil
		}
		return &tipb.Expr{
			Tp:       tp.Enum(),
			Children: []*tipb.Expr{child}}
	}

	leftExpr := b.newExprToPBExpr(client, expr.Args[0], tbl)
	if leftExpr == nil {
		return nil
	}
	rightExpr := b.newExprToPBExpr(client, expr.Args[1], tbl)
	if rightExpr == nil {
		return nil
	}
	return &tipb.Expr{
		Tp:       tp.Enum(),
		Children: []*tipb.Expr{leftExpr, rightExpr}}
}

func notToPBExpr(client kv.Client, expr *expression.ScalarFunction) *tipb.Expr {
	if !client.SupportRequestType(kv.ReqTypeSelect, int64(tipb.ExprType_Not)) {
		return nil
	}

	child := exprToPB(client, expr.Args[0])
	if child == nil {
		return nil
	}
	return &tipb.Expr{
		Tp:       tipb.ExprType_Not,
		Children: []*tipb.Expr{child}}
}

func (b *executorBuilder) notToPBExpr(client kv.Client, expr *expression.ScalarFunction, tbl *model.TableInfo) *tipb.Expr {
	if !client.SupportRequestType(kv.ReqTypeSelect, int64(tipb.ExprType_Not)) {
		return nil
	}

	child := b.newExprToPBExpr(client, expr.Args[0], tbl)
	if child == nil {
		return nil
	}
	return &tipb.Expr{
		Tp:       tipb.ExprType_Not.Enum(),
		Children: []*tipb.Expr{child}}
}

// Select do a select request, returns SelectResult.
// conncurrency: The max concurrency for underlying coprocessor request.
// keepOrder: If the result should returned in key order. For example if we need keep data in order by
//            scan index, we should set keepOrder to true.
func Select(client kv.Client, req *tipb.SelectRequest, keyRanges []kv.KeyRange, concurrency int, keepOrder bool) (SelectResult, error) {
	var err error
	startTs := time.Now()
	defer func() {
		// Add metrics
		queryHistgram.Observe(time.Since(startTs).Seconds())
		if err != nil {
			queryCounter.WithLabelValues(queryFailed).Inc()
		} else {
			queryCounter.WithLabelValues(querySucc).Inc()
		}
	}()

	// Convert tipb.*Request to kv.Request.
	kvReq, err1 := composeRequest(req, keyRanges, concurrency, keepOrder)
	if err1 != nil {
		err = errors.Trace(err1)
		return nil, err
	}

	resp := client.Send(kvReq)
	if resp == nil {
		err = errors.New("client returns nil response")
		return nil, err
	}
	result := &selectResult{
		resp:    resp,
		results: make(chan PartialResult, 5),
		done:    make(chan error, 1),
		closed:  make(chan struct{}),
	}
	// If Aggregates is not nil, we should set result fields latter.
	if len(req.Aggregates) == 0 && len(req.GroupBy) == 0 {
		if req.TableInfo != nil {
			result.fields = ProtoColumnsToFieldTypes(req.TableInfo.Columns)
		} else {
			result.fields = ProtoColumnsToFieldTypes(req.IndexInfo.Columns)
			length := len(req.IndexInfo.Columns)
			if req.IndexInfo.Columns[length-1].GetPkHandle() {
				// Returned index row do not contains extra PKHandle column.
				result.fields = result.fields[:length-1]
			}
			result.index = true
		}
	} else {
		result.aggregate = true
	}
	return result, nil
}

func (b *executorBuilder) unaryToPBExpr(client kv.Client, expr *ast.UnaryOperationExpr, tn *ast.TableName) *tipb.Expr {
	switch expr.Op {
	case opcode.Not:
		if !client.SupportRequestType(kv.ReqTypeSelect, int64(tipb.ExprType_Not)) {
			return nil
		}
	default:
		return nil
	}
	child := b.exprToPBExpr(client, expr.V, tn)
	if child == nil {
		return nil
	}
	return &tipb.Expr{Tp: tipb.ExprType_Not.Enum(), Children: []*tipb.Expr{child}}
}