protocol: support query attribute since mysql 8.0.23

pkg/expression/builtin.go

@@ -983,9 +983,10 @@ var funcs = map[string]functionClass{
 	ast.TiDBEncodeSQLDigest:  &tidbEncodeSQLDigestFunctionClass{baseFunctionClass{ast.TiDBEncodeSQLDigest, 1, 1}},
 
 	// TiDB Sequence function.
-	ast.NextVal: &nextValFunctionClass{baseFunctionClass{ast.NextVal, 1, 1}},
-	ast.LastVal: &lastValFunctionClass{baseFunctionClass{ast.LastVal, 1, 1}},
-	ast.SetVal:  &setValFunctionClass{baseFunctionClass{ast.SetVal, 2, 2}},
+	ast.NextVal:         &nextValFunctionClass{baseFunctionClass{ast.NextVal, 1, 1}},
+	ast.LastVal:         &lastValFunctionClass{baseFunctionClass{ast.LastVal, 1, 1}},
+	ast.SetVal:          &setValFunctionClass{baseFunctionClass{ast.SetVal, 2, 2}},
+	ast.QueryAttrString: &getQueryAttrFunctionClass{baseFunctionClass{ast.QueryAttrString, 1, 1}},
 }
 
 // IsFunctionSupported check if given function name is a builtin sql function.

pkg/expression/builtin_other.go

@@ -21,6 +21,7 @@ import (
 
 	"github.com/pingcap/errors"
 	"github.com/pingcap/tidb/pkg/expression/expropt"
+	"github.com/pingcap/tidb/pkg/param"
 	"github.com/pingcap/tidb/pkg/parser/ast"
 	"github.com/pingcap/tidb/pkg/parser/model"
 	"github.com/pingcap/tidb/pkg/parser/mysql"
@@ -46,6 +47,7 @@ var (
 	_ functionClass = &valuesFunctionClass{}
 	_ functionClass = &bitCountFunctionClass{}
 	_ functionClass = &getParamFunctionClass{}
+	_ functionClass = &getQueryAttrFunctionClass{}
 )
 
 var (
@@ -1676,3 +1678,60 @@ func (b *builtinGetParamStringSig) evalString(ctx EvalContext, row chunk.Row) (s
 	}
 	return str, false, nil
 }
+
+// getQueryAttrFunctionClass for plan cache of prepared statements
+type getQueryAttrFunctionClass struct {
+	baseFunctionClass
+}
+
+func (c *getQueryAttrFunctionClass) getFunction(ctx BuildContext, args []Expression) (builtinFunc, error) {
+	if err := c.verifyArgs(args); err != nil {
+		return nil, err
+	}
+	bf, err := newBaseBuiltinFuncWithTp(ctx, c.funcName, args, types.ETString, types.ETString)
+	if err != nil {
+		return nil, err
+	}
+	bf.tp.SetFlen(mysql.MaxFieldVarCharLength)
+	sig := &builtinGetQueryAttrStringSig{baseBuiltinFunc: bf}
+	return sig, nil
+}
+
+type builtinGetQueryAttrStringSig struct {
+	baseBuiltinFunc
+	expropt.SessionVarsPropReader
+}
+
+func (b *builtinGetQueryAttrStringSig) Clone() builtinFunc {
+	newSig := &builtinGetQueryAttrStringSig{}
+	newSig.cloneFrom(&b.baseBuiltinFunc)
+	return newSig
+}
+
+func (b *builtinGetQueryAttrStringSig) RequiredOptionalEvalProps() OptionalEvalPropKeySet {
+	return b.SessionVarsPropReader.RequiredOptionalEvalProps()
+}
+
+func (b *builtinGetQueryAttrStringSig) evalString(ctx EvalContext, row chunk.Row) (string, bool, error) {
+	sessionVars, err := b.GetSessionVars(ctx)
+	if err != nil {
+		return "", true, err
+	}
+
+	varName, isNull, err := b.args[0].EvalString(ctx, row)
+	if isNull || err != nil {
+		return "", true, err
+	}
+	attrs := sessionVars.QueryAttributes
+	if attrs == nil {
+		return "", true, nil
+	}
+	if v, ok := attrs[varName]; ok {
+		paramData, err := ExecBinaryParam(sessionVars.StmtCtx.TypeCtx(), []param.BinaryParam{v})
+		if err != nil {
+			return "", true, err
+		}
+		return paramData[0].EvalString(ctx, row)
+	}
+	return "", true, nil
+}

pkg/expression/function_traits.go

@@ -47,22 +47,23 @@ var UnCacheableFunctions = map[string]struct{}{
 
 // unFoldableFunctions stores functions which can not be folded duration constant folding stage.
 var unFoldableFunctions = map[string]struct{}{
-	ast.Sysdate:   {},
-	ast.FoundRows: {},
-	ast.Rand:      {},
-	ast.UUID:      {},
-	ast.Sleep:     {},
-	ast.RowFunc:   {},
-	ast.Values:    {},
-	ast.SetVar:    {},
-	ast.GetVar:    {},
-	ast.GetParam:  {},
-	ast.Benchmark: {},
-	ast.DayName:   {},
-	ast.NextVal:   {},
-	ast.LastVal:   {},
-	ast.SetVal:    {},
-	ast.AnyValue:  {},
+	ast.Sysdate:         {},
+	ast.FoundRows:       {},
+	ast.Rand:            {},
+	ast.UUID:            {},
+	ast.Sleep:           {},
+	ast.RowFunc:         {},
+	ast.Values:          {},
+	ast.SetVar:          {},
+	ast.GetVar:          {},
+	ast.GetParam:        {},
+	ast.Benchmark:       {},
+	ast.DayName:         {},
+	ast.NextVal:         {},
+	ast.LastVal:         {},
+	ast.SetVal:          {},
+	ast.AnyValue:        {},
+	ast.QueryAttrString: {},
 }
 
 // DisableFoldFunctions stores functions which prevent child scope functions from being constant folded.

pkg/param/binary_params.go

@@ -23,7 +23,7 @@ import (
 var ErrUnknownFieldType = dbterror.ClassServer.NewStd(errno.ErrUnknownFieldType)
 
 // BinaryParam stores the information decoded from the binary protocol
-// It can be further parsed into `expression.Expression` through the `ExecArgs` function in this package
+// It can be further parsed into `expression.Expression` through the expression.ExecBinaryParam function in expression package
 type BinaryParam struct {
 	Tp         byte
 	IsUnsigned bool

pkg/parser/ast/functions.go

@@ -377,9 +377,10 @@ const (
 	GetMvccInfo = "get_mvcc_info"
 
 	// Sequence function.
-	NextVal = "nextval"
-	LastVal = "lastval"
-	SetVal  = "setval"
+	NextVal         = "nextval"
+	LastVal         = "lastval"
+	SetVal          = "setval"
+	QueryAttrString = "mysql_query_attribute_string"
 )
 
 type FuncCallExprType int8

pkg/parser/mysql/const.go

@@ -171,7 +171,7 @@ const (
 	ClientDeprecateEOF                                  // CLIENT_DEPRECATE_EOF
 	ClientOptionalResultsetMetadata                     // CLIENT_OPTIONAL_RESULTSET_METADATA, Not supported: https://dev.mysql.com/doc/c-api/8.0/en/c-api-optional-metadata.html
 	ClientZstdCompressionAlgorithm                      // CLIENT_ZSTD_COMPRESSION_ALGORITHM
-	// 1 << 27 == CLIENT_QUERY_ATTRIBUTES
+	ClientQueryAttributes                               // CLIENT_QUERY_ATTRIBUTES
 	// 1 << 28 == MULTI_FACTOR_AUTHENTICATION
 	// 1 << 29 == CLIENT_CAPABILITY_EXTENSION
 	// 1 << 30 == CLIENT_SSL_VERIFY_SERVER_CERT
@@ -665,6 +665,8 @@ const (
 	CursorTypeReadOnly = 1 << iota
 	CursorTypeForUpdate
 	CursorTypeScrollable
+	// ParameterCountAvailable On when the client will send the parameter count even for 0 parameters.
+	ParameterCountAvailable
 )
 
 // ZlibCompressDefaultLevel is the zlib compression level for the compressed protocol

pkg/parser/parser_test.go

@@ -1979,6 +1979,7 @@ func TestBuiltin(t *testing.T) {
 		{`SELECT IS_USED_LOCK(@str);`, true, "SELECT IS_USED_LOCK(@`str`)"},
 		{`SELECT MASTER_POS_WAIT(@log_name, @log_pos), MASTER_POS_WAIT(@log_name, @log_pos, @timeout), MASTER_POS_WAIT(@log_name, @log_pos, @timeout, @channel_name);`, true, "SELECT MASTER_POS_WAIT(@`log_name`, @`log_pos`),MASTER_POS_WAIT(@`log_name`, @`log_pos`, @`timeout`),MASTER_POS_WAIT(@`log_name`, @`log_pos`, @`timeout`, @`channel_name`)"},
 		{`SELECT NAME_CONST('myname', 14);`, true, "SELECT NAME_CONST(_UTF8MB4'myname', 14)"},
+		{`SELECT MYSQL_QUERY_ATTRIBUTE_STRING(@str);`, true, "SELECT MYSQL_QUERY_ATTRIBUTE_STRING(@`str`)"},
 		{`SELECT RELEASE_ALL_LOCKS();`, true, "SELECT RELEASE_ALL_LOCKS()"},
 		{`SELECT UUID();`, true, "SELECT UUID()"},
 		{`SELECT UUID_SHORT()`, true, "SELECT UUID_SHORT()"},

pkg/server/conn.go

@@ -67,6 +67,7 @@ import (
 	"github.com/pingcap/tidb/pkg/kv"
 	"github.com/pingcap/tidb/pkg/meta/model"
 	"github.com/pingcap/tidb/pkg/metrics"
+	"github.com/pingcap/tidb/pkg/param"
 	"github.com/pingcap/tidb/pkg/parser"
 	"github.com/pingcap/tidb/pkg/parser/ast"
 	"github.com/pingcap/tidb/pkg/parser/auth"
@@ -1337,6 +1338,7 @@ func (cc *clientConn) dispatch(ctx context.Context, data []byte) error {
 
 		cc.server.releaseToken(token)
 		cc.lastActive = time.Now()
+		cc.ctx.GetSessionVars().QueryAttributes = nil
 	}()
 
 	vars := cc.ctx.GetSessionVars()
@@ -1373,8 +1375,14 @@ func (cc *clientConn) dispatch(ctx context.Context, data []byte) error {
 		// See http://dev.mysql.com/doc/internals/en/com-query.html
 		if len(data) > 0 && data[len(data)-1] == 0 {
 			data = data[:len(data)-1]
-			dataStr = string(hack.String(data))
 		}
+		pos, err := cc.parseQueryAttributes(ctx, data)
+		if err != nil {
+			return err
+		}
+		// fix lastPacket for display/log
+		cc.lastPacket = append([]byte{cc.lastPacket[0]}, data[pos:]...)
+		dataStr = string(hack.String(data[pos:]))
 		return cc.handleQuery(ctx, dataStr)
 	case mysql.ComFieldList:
 		return cc.handleFieldList(ctx, dataStr)
@@ -1698,6 +1706,61 @@ func (cc *clientConn) audit(eventType plugin.GeneralEvent) {
 	}
 }
 
+// parseQueryAttributes support query attributes since mysql 8.0.23
+// see https://dev.mysql.com/doc/refman/8.0/en/query-attributes.html
+// https://archive.fosdem.org/2021/schedule/event/mysql_protocl/attachments/slides/4274/export/events/attachments/mysql_protocl/slides/4274/FOSDEM21_MySQL_Protocols_Query_Attributes.pdf
+func (cc *clientConn) parseQueryAttributes(ctx context.Context, data []byte) (pos int, err error) {
+	if cc.capability&mysql.ClientQueryAttributes > 0 {
+		paraCount, _, np := util2.ParseLengthEncodedInt(data)
+		numParams := int(paraCount)
+		pos += np
+		_, _, np = util2.ParseLengthEncodedInt(data[pos:])
+		pos += np
+		ps := make([]param.BinaryParam, numParams)
+		names := make([]string, numParams)
+		if paraCount > 0 {
+			var (
+				nullBitmaps []byte
+				paramTypes  []byte
+			)
+			cc.initInputEncoder(ctx)
+			nullBitmapLen := (numParams + 7) >> 3
+			nullBitmaps = data[pos : pos+nullBitmapLen]
+			pos += nullBitmapLen
+			if data[pos] != 1 {
+				return 0, mysql.ErrMalformPacket
+			}
+
+			pos++
+			for i := 0; i < numParams; i++ {
+				paramTypes = append(paramTypes, data[pos:pos+2]...)
+				pos += 2
+				s, _, p, e := util2.ParseLengthEncodedBytes(data[pos:])
+				if e != nil {
+					return 0, mysql.ErrMalformPacket
+				}
+				names[i] = string(hack.String(s))
+				pos += p
+			}
+
+			boundParams := make([][]byte, numParams)
+			p := 0
+			if p, err = parseBinaryParams(ps, boundParams, nullBitmaps, paramTypes, data[pos:], cc.inputDecoder); err != nil {
+				return
+			}
+
+			pos += p
+			psWithName := make(map[string]param.BinaryParam, numParams)
+			for i := range names {
+				psWithName[names[i]] = ps[i]
+			}
+			cc.ctx.GetSessionVars().QueryAttributes = psWithName
+		}
+	}
+
+	return
+}
+
 // handleQuery executes the sql query string and writes result set or result ok to the client.
 // As the execution time of this function represents the performance of TiDB, we do time log and metrics here.
 // Some special queries like `load data` that does not return result, which is handled in handleFileTransInConn.

pkg/server/conn_stmt.go

@@ -54,6 +54,7 @@ import (
 	"github.com/pingcap/tidb/pkg/server/internal/dump"
 	"github.com/pingcap/tidb/pkg/server/internal/parse"
 	"github.com/pingcap/tidb/pkg/server/internal/resultset"
+	util2 "github.com/pingcap/tidb/pkg/server/internal/util"
 	"github.com/pingcap/tidb/pkg/sessionctx/variable"
 	"github.com/pingcap/tidb/pkg/sessiontxn"
 	storeerr "github.com/pingcap/tidb/pkg/store/driver/error"
@@ -180,6 +181,13 @@ func (cc *clientConn) handleStmtExecute(ctx context.Context, data []byte) (err e
 	)
 	cc.initInputEncoder(ctx)
 	numParams := stmt.NumParams()
+	clientHasQueryAttr := cc.capability&mysql.ClientQueryAttributes > 0
+	if clientHasQueryAttr && (numParams > 0 || flag&mysql.ParameterCountAvailable > 0) {
+		paraCount, _, np := util2.ParseLengthEncodedInt(data[pos:])
+		numParams = int(paraCount)
+		pos += np
+	}
+
 	args := make([]param.BinaryParam, numParams)
 	if numParams > 0 {
 		nullBitmapLen := (numParams + 7) >> 3
@@ -192,12 +200,28 @@ func (cc *clientConn) handleStmtExecute(ctx context.Context, data []byte) (err e
 		// new param bound flag
 		if data[pos] == 1 {
 			pos++
-			if len(data) < (pos + (numParams << 1)) {
-				return mysql.ErrMalformPacket
+			// For client that has query attribute ability, parameter's name will also be sent.
+			// However, it is useless for execute statement, so we ignore it here.
+			if clientHasQueryAttr {
+				for i := 0; i < numParams; i++ {
+					paramTypes = append(paramTypes, data[pos:pos+2]...)
+					pos += 2
+					// parse names
+					_, _, p, e := util2.ParseLengthEncodedBytes(data[pos:])
+					if e != nil {
+						return mysql.ErrMalformPacket
+					}
+					pos += p
+				}
+			} else {
+				if len(data) < (pos + (numParams << 1)) {
+					return mysql.ErrMalformPacket
+				}
+
+				paramTypes = data[pos : pos+(numParams<<1)]
+				pos += numParams << 1
 			}
 
-			paramTypes = data[pos : pos+(numParams<<1)]
-			pos += numParams << 1
 			paramValues = data[pos:]
 			// Just the first StmtExecute packet contain parameters type,
 			// we need save it for further use.
@@ -206,7 +230,7 @@ func (cc *clientConn) handleStmtExecute(ctx context.Context, data []byte) (err e
 			paramValues = data[pos+1:]
 		}
 
-		err = parseBinaryParams(args, stmt.BoundParams(), nullBitmaps, stmt.GetParamsType(), paramValues, cc.inputDecoder)
+		_, err = parseBinaryParams(args, stmt.BoundParams(), nullBitmaps, stmt.GetParamsType(), paramValues, cc.inputDecoder)
 		// This `.Reset` resets the arguments, so it's fine to just ignore the error (and the it'll be reset again in the following routine)
 		errReset := stmt.Reset()
 		if errReset != nil {

pkg/server/conn_stmt_params.go

@@ -27,8 +27,7 @@ import (
 var errUnknownFieldType = dbterror.ClassServer.NewStd(errno.ErrUnknownFieldType)
 
 // parseBinaryParams decodes the binary params according to the protocol
-func parseBinaryParams(params []param.BinaryParam, boundParams [][]byte, nullBitmap, paramTypes, paramValues []byte, enc *util2.InputDecoder) (err error) {
-	pos := 0
+func parseBinaryParams(params []param.BinaryParam, boundParams [][]byte, nullBitmap, paramTypes, paramValues []byte, enc *util2.InputDecoder) (pos int, err error) {
 	if enc == nil {
 		enc = util2.NewInputDecoder(charset.CharsetUTF8)
 	}
@@ -76,7 +75,7 @@ func parseBinaryParams(params []param.BinaryParam, boundParams [][]byte, nullBit
 		}
 
 		if (i<<1)+1 >= len(paramTypes) {
-			return mysql.ErrMalformPacket
+			return pos, mysql.ErrMalformPacket
 		}
 
 		tp := paramTypes[i<<1]