diff --git a/executor/plan_replayer.go b/executor/plan_replayer.go index f8a7ffe90fcca..b8fc52e45071d 100644 --- a/executor/plan_replayer.go +++ b/executor/plan_replayer.go @@ -62,6 +62,7 @@ type PlanReplayerCaptureInfo struct { type PlanReplayerDumpInfo struct { ExecStmts []ast.StmtNode Analyze bool + StartTS uint64 Path string File *os.File FileName string @@ -84,6 +85,15 @@ func (e *PlanReplayerExec) Next(ctx context.Context, req *chunk.Chunk) error { if err != nil { return err } + // Note: + // For the dumping for SQL file case (len(e.DumpInfo.Path) > 0), the DumpInfo.dump() is called in + // handleFileTransInConn(), which is after TxnManager.OnTxnEnd(), where we can't access the TxnManager anymore. + // So we must fetch the startTS now. + startTS, err := sessiontxn.GetTxnManager(e.ctx).GetStmtReadTS() + if err != nil { + return err + } + e.DumpInfo.StartTS = startTS if len(e.DumpInfo.Path) > 0 { err = e.prepare() if err != nil { @@ -163,12 +173,8 @@ func (e *PlanReplayerExec) createFile() error { func (e *PlanReplayerDumpInfo) dump(ctx context.Context) (err error) { fileName := e.FileName zf := e.File - startTS, err := sessiontxn.GetTxnManager(e.ctx).GetStmtReadTS() - if err != nil { - return err - } task := &domain.PlanReplayerDumpTask{ - StartTS: startTS, + StartTS: e.StartTS, FileName: fileName, Zf: zf, SessionVars: e.ctx.GetSessionVars(), diff --git a/server/server_test.go b/server/server_test.go index 0734b8f4b748d..2e6477d00ff56 100644 --- a/server/server_test.go +++ b/server/server_test.go @@ -15,7 +15,9 @@ package server import ( + "bufio" "bytes" + "context" "database/sql" "encoding/json" "fmt" @@ -39,7 +41,11 @@ import ( "github.com/pingcap/tidb/kv" tmysql "github.com/pingcap/tidb/parser/mysql" "github.com/pingcap/tidb/testkit" + "github.com/pingcap/tidb/testkit/testdata" "github.com/pingcap/tidb/testkit/testenv" + "github.com/pingcap/tidb/util/arena" + "github.com/pingcap/tidb/util/chunk" + "github.com/pingcap/tidb/util/replayer" "github.com/pingcap/tidb/util/versioninfo" "github.com/stretchr/testify/require" "go.uber.org/zap" @@ -2518,3 +2524,48 @@ func (cli *testServerClient) runTestInfoschemaClientErrors(t *testing.T) { } }) } + +func TestIssue46197(t *testing.T) { + store := testkit.CreateMockStore(t) + tk := testkit.NewTestKit(t, store) + tidbdrv := NewTiDBDriver(store) + cfg := newTestConfig() + cfg.Port, cfg.Status.StatusPort = 0, 0 + cfg.Status.ReportStatus = false + server, err := NewServer(cfg, tidbdrv) + require.NoError(t, err) + defer server.Close() + + // Mock the content of the SQL file in PacketIO buffer. + // First 4 bytes are the header, followed by the actual content. + // This acts like we are sending "select * from t1;" from the client when tidb requests the "a.txt" file. + var inBuffer bytes.Buffer + _, err = inBuffer.Write([]byte{0x11, 0x00, 0x00, 0x01}) + require.NoError(t, err) + _, err = inBuffer.Write([]byte("select * from t1;")) + require.NoError(t, err) + + // clientConn setup + brc := newBufferedReadConn(&bytesConn{b: inBuffer}) + pkt := newPacketIO(brc) + pkt.bufWriter = bufio.NewWriter(bytes.NewBuffer(nil)) + cc := &clientConn{ + server: server, + alloc: arena.NewAllocator(1024), + chunkAlloc: chunk.NewAllocator(), + pkt: pkt, + capability: tmysql.ClientLocalFiles, + } + ctx := context.Background() + cc.setCtx(&TiDBContext{Session: tk.Session(), stmts: make(map[int]*TiDBStatement)}) + + tk.MustExec("use test") + tk.MustExec("create table t1 (a int, b int)") + + // 3 is mysql.ComQuery, followed by the SQL text. + require.NoError(t, cc.dispatch(ctx, []byte("\u0003plan replayer dump explain 'a.txt'"))) + + // clean up + path := testdata.ConvertRowsToStrings(tk.MustQuery("select @@tidb_last_plan_replayer_token").Rows()) + require.NoError(t, os.Remove(filepath.Join(replayer.GetPlanReplayerDirName(), path[0]))) +}