From 92285cf769c9b130c19dc6a76c32b735b82800e3 Mon Sep 17 00:00:00 2001 From: vicgao Date: Sun, 4 Aug 2024 13:17:02 +0800 Subject: [PATCH 01/12] support query attribute since mysql 8.0.23 --- pkg/expression/builtin.go | 7 ++- pkg/expression/builtin_other.go | 59 ++++++++++++++++++ pkg/expression/function_traits.go | 33 +++++----- pkg/param/binary_params.go | 2 +- pkg/parser/ast/functions.go | 7 ++- pkg/parser/mysql/const.go | 4 +- pkg/parser/parser_test.go | 1 + pkg/server/conn.go | 65 +++++++++++++++++++- pkg/server/conn_stmt.go | 4 +- pkg/server/conn_stmt_params.go | 5 +- pkg/server/conn_stmt_params_test.go | 51 ++++++++++++++- pkg/server/conn_test.go | 32 ++++++++++ pkg/server/server.go | 2 +- pkg/sessionctx/variable/BUILD.bazel | 1 + pkg/sessionctx/variable/session.go | 7 +++ tests/integrationtest/r/executor/show.result | 1 + 16 files changed, 249 insertions(+), 32 deletions(-) diff --git a/pkg/expression/builtin.go b/pkg/expression/builtin.go index 7daac7fe0bec1..e7697ec37f8d2 100644 --- a/pkg/expression/builtin.go +++ b/pkg/expression/builtin.go @@ -992,9 +992,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. diff --git a/pkg/expression/builtin_other.go b/pkg/expression/builtin_other.go index d910879d72027..c5c465a81b3ca 100644 --- a/pkg/expression/builtin_other.go +++ b/pkg/expression/builtin_other.go @@ -16,6 +16,7 @@ package expression import ( "fmt" + "github.com/pingcap/tidb/pkg/param" "strings" "time" @@ -46,6 +47,7 @@ var ( _ functionClass = &valuesFunctionClass{} _ functionClass = &bitCountFunctionClass{} _ functionClass = &getParamFunctionClass{} + _ functionClass = &getQueryAttrFunctionClass{} ) var ( @@ -1804,3 +1806,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 + contextopt.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 +} diff --git a/pkg/expression/function_traits.go b/pkg/expression/function_traits.go index 978b8eba79eaf..82253a85a8ebe 100644 --- a/pkg/expression/function_traits.go +++ b/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. diff --git a/pkg/param/binary_params.go b/pkg/param/binary_params.go index 0a3d5197bb49a..27c448844888f 100644 --- a/pkg/param/binary_params.go +++ b/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 diff --git a/pkg/parser/ast/functions.go b/pkg/parser/ast/functions.go index 7b803db6f94e9..5318cc7f14509 100644 --- a/pkg/parser/ast/functions.go +++ b/pkg/parser/ast/functions.go @@ -369,9 +369,10 @@ const ( TiDBDecodeBase64Key = "tidb_decode_base64_key" // Sequence function. - NextVal = "nextval" - LastVal = "lastval" - SetVal = "setval" + NextVal = "nextval" + LastVal = "lastval" + SetVal = "setval" + QueryAttrString = "mysql_query_attribute_string" ) type FuncCallExprType int8 diff --git a/pkg/parser/mysql/const.go b/pkg/parser/mysql/const.go index 591c089069435..0fb3be0ab00a3 100644 --- a/pkg/parser/mysql/const.go +++ b/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 diff --git a/pkg/parser/parser_test.go b/pkg/parser/parser_test.go index a2a487b08dc7d..08cb52fec4d6e 100644 --- a/pkg/parser/parser_test.go +++ b/pkg/parser/parser_test.go @@ -1989,6 +1989,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()"}, diff --git a/pkg/server/conn.go b/pkg/server/conn.go index 94e488dc57bc3..6785ce2e5a3cb 100644 --- a/pkg/server/conn.go +++ b/pkg/server/conn.go @@ -66,6 +66,7 @@ import ( "github.com/pingcap/tidb/pkg/infoschema" "github.com/pingcap/tidb/pkg/kv" "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" @@ -1335,6 +1336,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() @@ -1371,8 +1373,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) @@ -1699,6 +1707,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. diff --git a/pkg/server/conn_stmt.go b/pkg/server/conn_stmt.go index d0188c2dddfd8..db321c8259afe 100644 --- a/pkg/server/conn_stmt.go +++ b/pkg/server/conn_stmt.go @@ -54,6 +54,8 @@ 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/sessionctx/vardef" "github.com/pingcap/tidb/pkg/sessiontxn" storeerr "github.com/pingcap/tidb/pkg/store/driver/error" @@ -375,7 +377,7 @@ func (cc *clientConn) executeWithCursor(ctx context.Context, stmt PreparedStatem rowContainer.GetMemTracker().SetLabel(memory.LabelForCursorFetch) rowContainer.GetDiskTracker().AttachTo(vars.DiskTracker) rowContainer.GetDiskTracker().SetLabel(memory.LabelForCursorFetch) - if vardef.EnableTmpStorageOnOOM.Load() { + if variable.EnableTmpStorageOnOOM.Load() { failpoint.Inject("testCursorFetchSpill", func(val failpoint.Value) { if val, ok := val.(bool); val && ok { actionSpill := rowContainer.ActionSpillForTest() diff --git a/pkg/server/conn_stmt_params.go b/pkg/server/conn_stmt_params.go index 29fbfd9b058c4..798e8e9f211c8 100644 --- a/pkg/server/conn_stmt_params.go +++ b/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] diff --git a/pkg/server/conn_stmt_params_test.go b/pkg/server/conn_stmt_params_test.go index bb808920771cc..e199261c279c9 100644 --- a/pkg/server/conn_stmt_params_test.go +++ b/pkg/server/conn_stmt_params_test.go @@ -35,12 +35,12 @@ import ( "github.com/stretchr/testify/require" ) -// decodeAndParse uses the `parseBinaryParams` and `parse.ExecArgs` to parse the params passed through binary protocol +// decodeAndParse uses the `parseBinaryParams` and `expression.ExecBinaryParam` to parse the params passed through binary protocol // It helps to test the integration of these two functions func decodeAndParse(typectx types.Context, args []expression.Expression, boundParams [][]byte, nullBitmap, paramTypes, paramValues []byte, enc *util.InputDecoder) (err error) { binParams := make([]param.BinaryParam, len(args)) - err = parseBinaryParams(binParams, boundParams, nullBitmap, paramTypes, paramValues, enc) + _, err = parseBinaryParams(binParams, boundParams, nullBitmap, paramTypes, paramValues, enc) if err != nil { return err } @@ -319,6 +319,20 @@ func buildDatetimeParam(year uint16, month uint8, day uint8, hour uint8, minute return result } +func buildDatetimeParamWithClientQueryAttr(year uint16, month uint8, day uint8, hour uint8, min uint8, sec uint8, msec uint32) []byte { + endian := binary.LittleEndian + + result := []byte{mysql.TypeDatetime, 0x0, 0xfb, 11} + result = endian.AppendUint16(result, year) + result = append(result, month) + result = append(result, day) + result = append(result, hour) + result = append(result, min) + result = append(result, sec) + result = endian.AppendUint32(result, msec) + return result +} + func expectedDatetimeExecuteResult(t *testing.T, c *mockConn, time types.Time, warnCount int) []byte { return getExpectOutput(t, c, func(conn *clientConn) { var err error @@ -380,3 +394,36 @@ func TestDateTimeTypes(t *testing.T) { require.Equal(t, expected, out.Bytes()) } + +func TestQueryAttrExecute(t *testing.T) { + store, dom := testkit.CreateMockStoreAndDomain(t) + srv := CreateMockServer(t, store) + srv.SetDomain(dom) + defer srv.Close() + + appendUint32 := binary.LittleEndian.AppendUint32 + ctx := context.Background() + c := CreateMockConn(t, srv).(*mockConn) + c.capability = mysql.ClientProtocol41 | mysql.ClientDeprecateEOF | mysql.ClientQueryAttributes + + tk := testkit.NewTestKitWithSession(t, store, c.Context().Session) + tk.MustExec("use test") + stmt, _, _, err := c.Context().Prepare("select ? as t") + require.NoError(t, err) + + expectedTimeDatum, err := types.ParseDatetime(types.DefaultStmtNoWarningContext, "2023-11-09 14:23:45.000100") + require.NoError(t, err) + expected := expectedDatetimeExecuteResult(t, c, expectedTimeDatum, 1) + + // execute the statement with datetime parameter + req := append( + appendUint32([]byte{mysql.ComStmtExecute}, uint32(stmt.ID())), + 0x0, 0x1, 0x0, 0x0, 0x0, 0x1, + 0x0, 0x1, + ) + req = append(req, buildDatetimeParamWithClientQueryAttr(2023, 11, 9, 14, 23, 45, 100)...) + out := c.GetOutput() + require.NoError(t, c.Dispatch(ctx, req)) + + require.Equal(t, expected, out.Bytes()) +} diff --git a/pkg/server/conn_test.go b/pkg/server/conn_test.go index 395b6acbb98e5..c6a25e980695b 100644 --- a/pkg/server/conn_test.go +++ b/pkg/server/conn_test.go @@ -2227,3 +2227,35 @@ func TestIssue54335(t *testing.T) { _ = cc.handleQuery(context.Background(), "select /*+ MAX_EXECUTION_TIME(1)*/ * FROM testTable2;") } } + +func TestQueryAttributes(t *testing.T) { + inputs := []dispatchInput{ + { + com: mysql.ComQuery, + in: append([]byte{0x0, 0x01}, append([]byte("SELECT mysql_query_attribute_string('a');"), 0x0)...), + err: nil, + out: []byte{0x1, 0x0, 0x0, 0x0, 0x1, 0x37, 0x0, 0x0, 0x1, + 0x3, 0x64, 0x65, 0x66, 0x0, 0x0, 0x0, 0x21, 0x6d, 0x79, + 0x73, 0x71, 0x6c, 0x5f, 0x71, 0x75, 0x65, 0x72, 0x79, 0x5f, + 0x61, 0x74, 0x74, 0x72, 0x69, 0x62, 0x75, 0x74, 0x65, 0x5f, + 0x73, 0x74, 0x72, 0x69, 0x6e, 0x67, 0x28, 0x27, 0x61, 0x27, 0x29, + 0x0, 0xc, 0x2e, 0x0, 0xfc, 0xff, 0x3, 0x0, 0xfd, 0x0, 0x0, 0x1f, 0x0, + 0x0, 0x1, 0x0, 0x0, 0x2, 0xfe, 0x1, 0x0, 0x0, 0x3, 0xfb, 0x1, + 0x0, 0x0, 0x4, 0xfe}, + }, + { + com: mysql.ComQuery, + in: append([]byte{0x1, 0x1, 0x0, 0x1, 0xfe, 0x0, 0x1, 0x61, 0x1, 0x62}, append([]byte("SELECT mysql_query_attribute_string('a');"), 0x0)...), + err: nil, + out: []byte{0x1, 0x0, 0x0, 0x5, 0x1, 0x37, 0x0, 0x0, 0x6, + 0x3, 0x64, 0x65, 0x66, 0x0, 0x0, 0x0, 0x21, 0x6d, 0x79, + 0x73, 0x71, 0x6c, 0x5f, 0x71, 0x75, 0x65, 0x72, 0x79, 0x5f, + 0x61, 0x74, 0x74, 0x72, 0x69, 0x62, 0x75, 0x74, 0x65, 0x5f, + 0x73, 0x74, 0x72, 0x69, 0x6e, 0x67, 0x28, 0x27, 0x61, 0x27, + 0x29, 0x0, 0xc, 0x2e, 0x0, 0xfc, 0xff, 0x3, 0x0, 0xfd, 0x0, + 0x0, 0x1f, 0x0, 0x0, 0x1, 0x0, 0x0, 0x7, 0xfe, 0x2, 0x0, + 0x0, 0x8, 0x1, 0x62, 0x1, 0x0, 0x0, 0x9, 0xfe}, + }, + } + testDispatch(t, inputs, mysql.ClientQueryAttributes) +} diff --git a/pkg/server/server.go b/pkg/server/server.go index 3775779bc3154..9a8c541650473 100644 --- a/pkg/server/server.go +++ b/pkg/server/server.go @@ -111,7 +111,7 @@ const defaultCapability = mysql.ClientLongPassword | mysql.ClientLongFlag | mysql.ClientTransactions | mysql.ClientSecureConnection | mysql.ClientFoundRows | mysql.ClientMultiStatements | mysql.ClientMultiResults | mysql.ClientLocalFiles | mysql.ClientConnectAtts | mysql.ClientPluginAuth | mysql.ClientInteractive | - mysql.ClientDeprecateEOF | mysql.ClientCompress | mysql.ClientZstdCompressionAlgorithm + mysql.ClientDeprecateEOF | mysql.ClientCompress | mysql.ClientZstdCompressionAlgorithm | mysql.ClientQueryAttributes // Server is the MySQL protocol server type Server struct { diff --git a/pkg/sessionctx/variable/BUILD.bazel b/pkg/sessionctx/variable/BUILD.bazel index 29d5e6050db4d..ed4d7ce3e3114 100644 --- a/pkg/sessionctx/variable/BUILD.bazel +++ b/pkg/sessionctx/variable/BUILD.bazel @@ -26,6 +26,7 @@ go_library( "//pkg/kv", "//pkg/meta/model", "//pkg/metrics", + "//pkg/param", "//pkg/parser", "//pkg/parser/ast", "//pkg/parser/auth", diff --git a/pkg/sessionctx/variable/session.go b/pkg/sessionctx/variable/session.go index 8d685273ff9b5..c0abe681ec415 100644 --- a/pkg/sessionctx/variable/session.go +++ b/pkg/sessionctx/variable/session.go @@ -38,6 +38,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" @@ -1087,6 +1088,9 @@ type SessionVars struct { // OptimizerEnableNAAJ enables TiDB to use null-aware anti join. OptimizerEnableNAAJ bool + // EnableTablePartition enables table partition feature. + EnableTablePartition string + // EnableCascadesPlanner enables the cascades planner. EnableCascadesPlanner bool @@ -1687,6 +1691,9 @@ type SessionVars struct { // `select for update` statements which do acquire pessimsitic locks. SharedLockPromotion bool + // QueryAttributes record query attributed for each statement. + QueryAttributes map[string]param.BinaryParam + // ScatterRegion will scatter the regions for DDLs when it is "table" or "global", "" indicates not trigger scatter. ScatterRegion string diff --git a/tests/integrationtest/r/executor/show.result b/tests/integrationtest/r/executor/show.result index 219a543d797d0..77b556db78fdb 100644 --- a/tests/integrationtest/r/executor/show.result +++ b/tests/integrationtest/r/executor/show.result @@ -763,6 +763,7 @@ mod month monthname mul +mysql_query_attribute_string name_const ne nextval From b4df1e6dfc5035b950d15d299304b5e7827cd666 Mon Sep 17 00:00:00 2001 From: vicgao Date: Sun, 1 Sep 2024 21:33:04 +0800 Subject: [PATCH 02/12] code review. --- pkg/expression/builtin_other.go | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/pkg/expression/builtin_other.go b/pkg/expression/builtin_other.go index c5c465a81b3ca..c7053c8710559 100644 --- a/pkg/expression/builtin_other.go +++ b/pkg/expression/builtin_other.go @@ -16,11 +16,12 @@ package expression import ( "fmt" - "github.com/pingcap/tidb/pkg/param" "strings" "time" "github.com/pingcap/errors" + "github.com/pingcap/tidb/pkg/expression/contextopt" + "github.com/pingcap/tidb/pkg/param" "github.com/pingcap/tidb/pkg/expression/expropt" "github.com/pingcap/tidb/pkg/parser/ast" "github.com/pingcap/tidb/pkg/parser/mysql" From 3f7c090c3e01944b925a0a2730dec49304fc3210 Mon Sep 17 00:00:00 2001 From: vicgao Date: Wed, 25 Sep 2024 18:54:51 +0800 Subject: [PATCH 03/12] code review. --- pkg/expression/builtin_other.go | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/pkg/expression/builtin_other.go b/pkg/expression/builtin_other.go index c7053c8710559..d9d28a2961a95 100644 --- a/pkg/expression/builtin_other.go +++ b/pkg/expression/builtin_other.go @@ -20,9 +20,8 @@ import ( "time" "github.com/pingcap/errors" - "github.com/pingcap/tidb/pkg/expression/contextopt" - "github.com/pingcap/tidb/pkg/param" "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/mysql" "github.com/pingcap/tidb/pkg/types" @@ -1828,7 +1827,7 @@ func (c *getQueryAttrFunctionClass) getFunction(ctx BuildContext, args []Express type builtinGetQueryAttrStringSig struct { baseBuiltinFunc - contextopt.SessionVarsPropReader + expropt.SessionVarsPropReader } func (b *builtinGetQueryAttrStringSig) Clone() builtinFunc { From 5993b04175e7745cdcbf5c38a3bf6e182fa95083 Mon Sep 17 00:00:00 2001 From: vicgao Date: Mon, 24 Feb 2025 19:11:17 +0800 Subject: [PATCH 04/12] resolve conflict --- pkg/server/conn_stmt.go | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/pkg/server/conn_stmt.go b/pkg/server/conn_stmt.go index db321c8259afe..31d901a069b20 100644 --- a/pkg/server/conn_stmt.go +++ b/pkg/server/conn_stmt.go @@ -54,8 +54,6 @@ 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/sessionctx/vardef" "github.com/pingcap/tidb/pkg/sessiontxn" storeerr "github.com/pingcap/tidb/pkg/store/driver/error" @@ -208,7 +206,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 { @@ -377,7 +375,7 @@ func (cc *clientConn) executeWithCursor(ctx context.Context, stmt PreparedStatem rowContainer.GetMemTracker().SetLabel(memory.LabelForCursorFetch) rowContainer.GetDiskTracker().AttachTo(vars.DiskTracker) rowContainer.GetDiskTracker().SetLabel(memory.LabelForCursorFetch) - if variable.EnableTmpStorageOnOOM.Load() { + if vardef.EnableTmpStorageOnOOM.Load() { failpoint.Inject("testCursorFetchSpill", func(val failpoint.Value) { if val, ok := val.(bool); val && ok { actionSpill := rowContainer.ActionSpillForTest() From df6d2538470f9bd1502266d6d4f5402715100cb9 Mon Sep 17 00:00:00 2001 From: vicgao Date: Mon, 24 Feb 2025 21:35:26 +0800 Subject: [PATCH 05/12] handle execute statement. --- pkg/expression/builtin.go | 8 +++++--- pkg/expression/builtin_other.go | 1 + pkg/param/binary_params.go | 2 +- pkg/server/conn.go | 2 ++ pkg/server/conn_stmt.go | 32 ++++++++++++++++++++++++----- pkg/server/conn_stmt_params.go | 3 ++- pkg/server/conn_stmt_params_test.go | 4 ++-- 7 files changed, 40 insertions(+), 12 deletions(-) diff --git a/pkg/expression/builtin.go b/pkg/expression/builtin.go index e7697ec37f8d2..d8a1934cd4ccb 100644 --- a/pkg/expression/builtin.go +++ b/pkg/expression/builtin.go @@ -992,9 +992,11 @@ 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}}, + + // TiDB Query ast.QueryAttrString: &getQueryAttrFunctionClass{baseFunctionClass{ast.QueryAttrString, 1, 1}}, } diff --git a/pkg/expression/builtin_other.go b/pkg/expression/builtin_other.go index d9d28a2961a95..73b8d5b4eaf7d 100644 --- a/pkg/expression/builtin_other.go +++ b/pkg/expression/builtin_other.go @@ -1840,6 +1840,7 @@ func (b *builtinGetQueryAttrStringSig) RequiredOptionalEvalProps() OptionalEvalP return b.SessionVarsPropReader.RequiredOptionalEvalProps() } +// This implements `mysql_query_attribute_string(str)` func (b *builtinGetQueryAttrStringSig) evalString(ctx EvalContext, row chunk.Row) (string, bool, error) { sessionVars, err := b.GetSessionVars(ctx) if err != nil { diff --git a/pkg/param/binary_params.go b/pkg/param/binary_params.go index 27c448844888f..cd63874e51896 100644 --- a/pkg/param/binary_params.go +++ b/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 expression.ExecBinaryParam function in expression package +// It can be further parsed into `expression.Expression` through the expression.ExecBinaryParam function in the expression package type BinaryParam struct { Tp byte IsUnsigned bool diff --git a/pkg/server/conn.go b/pkg/server/conn.go index 6785ce2e5a3cb..970c096fe6ccb 100644 --- a/pkg/server/conn.go +++ b/pkg/server/conn.go @@ -1710,6 +1710,8 @@ 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 +// https://dev.mysql.com/doc/dev/mysql-server/latest/page_protocol_com_query.html +// https://dev.mysql.com/doc/dev/mysql-server/latest/page_protocol_com_stmt_execute.html func (cc *clientConn) parseQueryAttributes(ctx context.Context, data []byte) (pos int, err error) { if cc.capability&mysql.ClientQueryAttributes > 0 { paraCount, _, np := util2.ParseLengthEncodedInt(data) diff --git a/pkg/server/conn_stmt.go b/pkg/server/conn_stmt.go index 31d901a069b20..c5fd1f33414f2 100644 --- a/pkg/server/conn_stmt.go +++ b/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/vardef" "github.com/pingcap/tidb/pkg/sessiontxn" storeerr "github.com/pingcap/tidb/pkg/store/driver/error" @@ -180,6 +181,12 @@ func (cc *clientConn) handleStmtExecute(ctx context.Context, data []byte) (err e ) cc.initInputEncoder(ctx) numParams := stmt.NumParams() + clientHasQueryAttr := cc.ctx.GetSessionVars().ClientCapability&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 +199,27 @@ 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. diff --git a/pkg/server/conn_stmt_params.go b/pkg/server/conn_stmt_params.go index 798e8e9f211c8..659de1a8cf324 100644 --- a/pkg/server/conn_stmt_params.go +++ b/pkg/server/conn_stmt_params.go @@ -31,12 +31,13 @@ func parseBinaryParams(params []param.BinaryParam, boundParams [][]byte, nullBit if enc == nil { enc = util2.NewInputDecoder(charset.CharsetUTF8) } + boundParamsSize := len(boundParams) for i := 0; i < len(params); i++ { // if params had received via ComStmtSendLongData, use them directly. // ref https://dev.mysql.com/doc/internals/en/com-stmt-send-long-data.html // see clientConn#handleStmtSendLongData - if boundParams[i] != nil { + if i < boundParamsSize && boundParams[i] != nil { params[i] = param.BinaryParam{ Tp: mysql.TypeBlob, Val: boundParams[i], diff --git a/pkg/server/conn_stmt_params_test.go b/pkg/server/conn_stmt_params_test.go index e199261c279c9..0a90a201b42e5 100644 --- a/pkg/server/conn_stmt_params_test.go +++ b/pkg/server/conn_stmt_params_test.go @@ -319,7 +319,7 @@ func buildDatetimeParam(year uint16, month uint8, day uint8, hour uint8, minute return result } -func buildDatetimeParamWithClientQueryAttr(year uint16, month uint8, day uint8, hour uint8, min uint8, sec uint8, msec uint32) []byte { +func buildDatetimeParamWithClientQueryAttr(year uint16, month uint8, day uint8, hour uint8, minute uint8, sec uint8, msec uint32) []byte { endian := binary.LittleEndian result := []byte{mysql.TypeDatetime, 0x0, 0xfb, 11} @@ -327,7 +327,7 @@ func buildDatetimeParamWithClientQueryAttr(year uint16, month uint8, day uint8, result = append(result, month) result = append(result, day) result = append(result, hour) - result = append(result, min) + result = append(result, minute) result = append(result, sec) result = endian.AppendUint32(result, msec) return result From 03ee920d278fb4815813dfb51f290bd9b1c15617 Mon Sep 17 00:00:00 2001 From: vicgao Date: Tue, 25 Feb 2025 19:04:34 +0800 Subject: [PATCH 06/12] handle execute statement. --- pkg/server/conn_stmt.go | 28 ++++++++++++++++++++++++---- 1 file changed, 24 insertions(+), 4 deletions(-) diff --git a/pkg/server/conn_stmt.go b/pkg/server/conn_stmt.go index c5fd1f33414f2..ecb0bbffb03ff 100644 --- a/pkg/server/conn_stmt.go +++ b/pkg/server/conn_stmt.go @@ -60,6 +60,7 @@ import ( storeerr "github.com/pingcap/tidb/pkg/store/driver/error" "github.com/pingcap/tidb/pkg/util/chunk" "github.com/pingcap/tidb/pkg/util/execdetails" + "github.com/pingcap/tidb/pkg/util/hack" "github.com/pingcap/tidb/pkg/util/logutil" "github.com/pingcap/tidb/pkg/util/memory" "github.com/pingcap/tidb/pkg/util/redact" @@ -180,7 +181,8 @@ func (cc *clientConn) handleStmtExecute(ctx context.Context, data []byte) (err e paramValues []byte ) cc.initInputEncoder(ctx) - numParams := stmt.NumParams() + stmtNumParams := stmt.NumParams() + numParams := stmtNumParams clientHasQueryAttr := cc.ctx.GetSessionVars().ClientCapability&mysql.ClientQueryAttributes > 0 if clientHasQueryAttr && (numParams > 0 || flag&mysql.ParameterCountAvailable > 0) { paraCount, _, np := util2.ParseLengthEncodedInt(data[pos:]) @@ -195,21 +197,28 @@ func (cc *clientConn) handleStmtExecute(ctx context.Context, data []byte) (err e } nullBitmaps = data[pos : pos+nullBitmapLen] pos += nullBitmapLen + var attributeNames []string // new param bound flag if data[pos] == 1 { pos++ - // 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. + // For client that has query attribute ability, query attributes' name will also be sent. if clientHasQueryAttr { + if numParams > stmtNumParams { + attributeNames = make([]string, 0, numParams-stmt.NumParams()) + } for i := 0; i < numParams; i++ { paramTypes = append(paramTypes, data[pos:pos+2]...) pos += 2 // parse names - _, _, p, e := util2.ParseLengthEncodedBytes(data[pos:]) + pName, _, p, e := util2.ParseLengthEncodedBytes(data[pos:]) if e != nil { return mysql.ErrMalformPacket } + // Only the names of the parameters for query attributes will be sent. + if len(pName) > 0 { + attributeNames = append(attributeNames, string(hack.String(pName))) + } pos += p } } else { @@ -229,6 +238,17 @@ func (cc *clientConn) handleStmtExecute(ctx context.Context, data []byte) (err e } _, err = parseBinaryParams(args, stmt.BoundParams(), nullBitmaps, stmt.GetParamsType(), paramValues, cc.inputDecoder) + if len(attributeNames) != 0 { + if len(attributeNames) != len(args)-stmtNumParams { + return mysql.ErrMalformPacket + } + psWithName := make(map[string]param.BinaryParam, numParams) + for i := range attributeNames { + psWithName[attributeNames[i]] = args[i+stmtNumParams] + } + cc.ctx.GetSessionVars().QueryAttributes = psWithName + args = args[:stmtNumParams] + } // 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 { From 2d565adb26ebb5dfa36f4dece33e67f7069fbef2 Mon Sep 17 00:00:00 2001 From: vicgao Date: Tue, 25 Feb 2025 19:42:08 +0800 Subject: [PATCH 07/12] fix tests. --- pkg/expression/function_traits_test.go | 1 + 1 file changed, 1 insertion(+) diff --git a/pkg/expression/function_traits_test.go b/pkg/expression/function_traits_test.go index da0ff30872614..6fbd4e097fe7c 100644 --- a/pkg/expression/function_traits_test.go +++ b/pkg/expression/function_traits_test.go @@ -194,6 +194,7 @@ func TestIllegalFunctions4GeneratedColumns(t *testing.T) { "month", "monthname", "mul", + "mysql_query_attribute_string", "ne", "nextval", "not", From 13471df2dd78d009f9dc4d4fc001a5a38e4dea0c Mon Sep 17 00:00:00 2001 From: vicgao Date: Tue, 25 Feb 2025 19:44:29 +0800 Subject: [PATCH 08/12] fix comments. --- pkg/server/conn.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pkg/server/conn.go b/pkg/server/conn.go index 970c096fe6ccb..38735091b1109 100644 --- a/pkg/server/conn.go +++ b/pkg/server/conn.go @@ -1709,7 +1709,7 @@ 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 +// https://archive.fosdem.org/2021/schedule/event/mysql_protocl/ // https://dev.mysql.com/doc/dev/mysql-server/latest/page_protocol_com_query.html // https://dev.mysql.com/doc/dev/mysql-server/latest/page_protocol_com_stmt_execute.html func (cc *clientConn) parseQueryAttributes(ctx context.Context, data []byte) (pos int, err error) { From 8347fe6dd77393087aa5c8a0552cdaae84bc8d12 Mon Sep 17 00:00:00 2001 From: vicgao Date: Tue, 25 Feb 2025 19:46:12 +0800 Subject: [PATCH 09/12] fix comments. --- 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 d8a1934cd4ccb..3aebb0424d27e 100644 --- a/pkg/expression/builtin.go +++ b/pkg/expression/builtin.go @@ -996,7 +996,7 @@ var funcs = map[string]functionClass{ ast.LastVal: &lastValFunctionClass{baseFunctionClass{ast.LastVal, 1, 1}}, ast.SetVal: &setValFunctionClass{baseFunctionClass{ast.SetVal, 2, 2}}, - // TiDB Query + // TiDB Query Attribute function. ast.QueryAttrString: &getQueryAttrFunctionClass{baseFunctionClass{ast.QueryAttrString, 1, 1}}, } From 170fcbe9f3a3e393f9304486c56e6ca89de98a1a Mon Sep 17 00:00:00 2001 From: vicgao Date: Tue, 25 Feb 2025 20:02:35 +0800 Subject: [PATCH 10/12] fix comments. --- pkg/server/conn.go | 85 +++++++++++++++++++++++----------------------- 1 file changed, 43 insertions(+), 42 deletions(-) diff --git a/pkg/server/conn.go b/pkg/server/conn.go index 38735091b1109..99f4e9701158c 100644 --- a/pkg/server/conn.go +++ b/pkg/server/conn.go @@ -1713,54 +1713,55 @@ func (cc *clientConn) audit(eventType plugin.GeneralEvent) { // https://dev.mysql.com/doc/dev/mysql-server/latest/page_protocol_com_query.html // https://dev.mysql.com/doc/dev/mysql-server/latest/page_protocol_com_stmt_execute.html 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 - } + if cc.capability&mysql.ClientQueryAttributes == 0 { + return + } - 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 + 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 - } + 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 + pos += p + psWithName := make(map[string]param.BinaryParam, numParams) + for i := range names { + psWithName[names[i]] = ps[i] } + cc.ctx.GetSessionVars().QueryAttributes = psWithName } - return } From ac2851ccfbce9f77903c65b031dc530569dce1f5 Mon Sep 17 00:00:00 2001 From: vicgao Date: Tue, 25 Feb 2025 20:05:56 +0800 Subject: [PATCH 11/12] fix comments. --- pkg/parser/ast/functions.go | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/pkg/parser/ast/functions.go b/pkg/parser/ast/functions.go index 5318cc7f14509..293af96744bf7 100644 --- a/pkg/parser/ast/functions.go +++ b/pkg/parser/ast/functions.go @@ -369,9 +369,11 @@ const ( TiDBDecodeBase64Key = "tidb_decode_base64_key" // Sequence function. - NextVal = "nextval" - LastVal = "lastval" - SetVal = "setval" + NextVal = "nextval" + LastVal = "lastval" + SetVal = "setval" + + // TiDB Query Attribute function QueryAttrString = "mysql_query_attribute_string" ) From 19f2ec858d9a26eb5f91f59fdc77cf1726100c7f Mon Sep 17 00:00:00 2001 From: vicgao Date: Tue, 25 Feb 2025 20:19:40 +0800 Subject: [PATCH 12/12] code format --- pkg/expression/builtin_threadunsafe_generated.go | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/pkg/expression/builtin_threadunsafe_generated.go b/pkg/expression/builtin_threadunsafe_generated.go index abaae4eacc80b..e5f8fbc36d953 100644 --- a/pkg/expression/builtin_threadunsafe_generated.go +++ b/pkg/expression/builtin_threadunsafe_generated.go @@ -306,6 +306,11 @@ func (s *builtinValuesVectorFloat32Sig) SafeToShareAcrossSession() bool { return false } +// SafeToShareAcrossSession implements BuiltinFunc.SafeToShareAcrossSession. +func (s *builtinGetQueryAttrStringSig) SafeToShareAcrossSession() bool { + return false +} + // SafeToShareAcrossSession implements BuiltinFunc.SafeToShareAcrossSession. func (s *builtinRegexpLikeFuncSig) SafeToShareAcrossSession() bool { return false