executor: fix plan replayer for sql file input wrongly fetch startTS …

…after `OnTxnEnd` (#46201) (#46345)

close #46197

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(),

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])))
+}