planner: implements support for dumping multiple SQL statements in a single PLAN REPLAYER DUMP EXPLAIN command

pkg/domain/plan_replayer_dump.go

@@ -159,11 +159,12 @@ func (tne *tableNameExtractor) Leave(in ast.Node) (ast.Node, bool) {
 			tne.err = err
 			return in, true
 		}
-		if tne.is.TableExists(t.Schema, t.Name) {
-			tp := tableNamePair{DBName: t.Schema.L, TableName: t.Name.L, IsView: isView}
-			if tp.DBName == "" {
-				tp.DBName = tne.curDB.L
-			}
+		schema := t.Schema
+		if schema.L == "" {
+			schema = tne.curDB
+		}
+		if tne.is.TableExists(schema, t.Name) {
+			tp := tableNamePair{DBName: schema.L, TableName: t.Name.L, IsView: isView}
 			tne.names[tp] = struct{}{}
 		}
 	} else if s, ok := in.(*ast.SelectStmt); ok {
@@ -201,36 +202,31 @@ func (tne *tableNameExtractor) handleIsView(t *ast.TableName) (bool, error) {
 
 // DumpPlanReplayerInfo will dump the information about sqls.
 // The files will be organized into the following format:
-/*
- |-sql_meta.toml
- |-meta.txt
- |-schema
- |	 |-schema_meta.txt
- |	 |-db1.table1.schema.txt
- |	 |-db2.table2.schema.txt
- |	 |-....
- |-view
- | 	 |-db1.view1.view.txt
- |	 |-db2.view2.view.txt
- |	 |-....
- |-stats
- |   |-stats1.json
- |   |-stats2.json
- |   |-....
- |-statsMem
- |   |-stats1.txt
- |   |-stats2.txt
- |   |-....
- |-config.toml
- |-table_tiflash_replica.txt
- |-variables.toml
- |-bindings.sql
- |-sql
- |   |-sql1.sql
- |   |-sql2.sql
- |	 |-....
- |-explain.txt
-*/
+//
+// Single SQL dump:
+//
+//	|-sql_meta.toml
+//	|-meta.txt
+//	|-schema/...
+//	|-view/...
+//	|-stats/...
+//	|-statsMem/...
+//	|-config.toml
+//	|-table_tiflash_replica.txt
+//	|-variables.toml
+//	|-bindings.sql
+//	|-sql/sql0.sql
+//	|-explain.txt
+//
+// Multiple SQL dump (PLAN REPLAYER DUMP EXPLAIN ( "sql1", "sql2", ... )):
+//
+//	|-(same as above)
+//	|-sql/sql0.sql
+//	|-sql/sql1.sql
+//	|-...
+//	|-explain/explain0.txt
+//	|-explain/explain1.txt
+//	|-...
 func DumpPlanReplayerInfo(ctx context.Context, sctx sessionctx.Context,
 	task *PlanReplayerDumpTask,
 ) (err error) {
@@ -728,15 +724,32 @@ func dumpEncodedPlan(ctx sessionctx.Context, zw *zip.Writer, encodedPlan string)
 }
 
 func dumpExplain(ctx sessionctx.Context, zw *zip.Writer, isAnalyze bool, sqls []string, emptyAsNil bool) (debugTraces []any, err error) {
-	fw, err := zw.Create("explain.txt")
-	if err != nil {
-		return nil, errors.AddStack(err)
-	}
 	ctx.GetSessionVars().InPlanReplayer = true
 	defer func() {
 		ctx.GetSessionVars().InPlanReplayer = false
 	}()
+
+	// If there are multiple SQLs, write separate explain files
+	useSeparateFiles := len(sqls) > 1
+
+	// For single SQL, create explain.txt once before the loop
+	var fw io.Writer
+	if !useSeparateFiles && len(sqls) > 0 {
+		fw, err = zw.Create("explain.txt")
+		if err != nil {
+			return nil, errors.AddStack(err)
+		}
+	}
+
 	for i, sql := range sqls {
+		// For multiple SQLs, create a separate file for each
+		if useSeparateFiles {
+			fw, err = zw.Create(fmt.Sprintf("explain/explain%v.txt", i))
+			if err != nil {
+				return nil, errors.AddStack(err)
+			}
+		}
+
 		var recordSets []sqlexec.RecordSet
 		if isAnalyze {
 			// Explain analyze
@@ -763,7 +776,9 @@ func dumpExplain(ctx sessionctx.Context, zw *zip.Writer, isAnalyze bool, sqls []
 				return nil, err
 			}
 		}
-		if i < len(sqls)-1 {
+
+		// For single SQL, add separator between multiple explains in the same file
+		if !useSeparateFiles && i < len(sqls)-1 {
 			fmt.Fprintf(fw, "<--------->\n")
 		}
 	}

pkg/executor/builder.go

@@ -1181,7 +1181,19 @@ func (b *executorBuilder) buildPlanReplayer(v *plannercore.PlanReplayer) exec.Ex
 			HistoricalStatsTS: v.HistoricalStatsTS,
 		},
 	}
-	if v.ExecStmt != nil {
+	if len(v.StmtList) > 0 {
+		// Parse multiple SQL strings from StmtList
+		e.DumpInfo.ExecStmts = make([]ast.StmtNode, 0, len(v.StmtList))
+		for _, sqlStr := range v.StmtList {
+			node, err := b.ctx.GetRestrictedSQLExecutor().ParseWithParams(context.Background(), sqlStr)
+			if err != nil {
+				// If parsing fails, propagate the error immediately so the statement fails.
+				b.err = errors.Errorf("plan replayer: failed to parse SQL: %s, error: %v", sqlStr, err)
+				return nil
+			}
+			e.DumpInfo.ExecStmts = append(e.DumpInfo.ExecStmts, node)
+		}
+	} else if v.ExecStmt != nil {
 		e.DumpInfo.ExecStmts = []ast.StmtNode{v.ExecStmt}
 	} else {
 		e.BaseExecutor = exec.NewBaseExecutor(b.ctx, nil, v.ID())

pkg/executor/test/planreplayer/BUILD.bazel

@@ -8,7 +8,7 @@ go_test(
         "plan_replayer_test.go",
     ],
     flaky = True,
-    shard_count = 5,
+    shard_count = 7,
     deps = [
         "//pkg/config",
         "//pkg/meta/autoid",

pkg/executor/test/planreplayer/plan_replayer_test.go

@@ -226,3 +226,119 @@ func TestPlanReplayerDumpSingle(t *testing.T) {
 		require.True(t, checkFileName(file.Name), file.Name)
 	}
 }
+
+func TestPlanReplayerDumpMultipleError(t *testing.T) {
+	store := testkit.CreateMockStore(t)
+	tk := testkit.NewTestKit(t, store)
+	tk.MustExec("use test")
+	tk.MustExec("create table t(id int)")
+
+	// empty statement list should return error
+	tk.MustContainErrMsg("plan replayer dump explain ()", "[parser:1064]")
+
+	// one error statement
+	tk.MustContainErrMsg("plan replayer dump explain ('select x om t')", "[parser:1064]")
+
+	// multiple error statements
+	tk.MustContainErrMsg("plan replayer dump explain ('select x from t', 'select y om t')", "[parser:1064]")
+}
+
+func TestPlanReplayerDumpMultiple(t *testing.T) {
+	const numStmts = 50
+	const numTables = 5
+	ctx := context.Background()
+	tempDir := t.TempDir()
+	storage, err := extstore.NewExtStorage(ctx, "file://"+tempDir, "")
+	require.NoError(t, err)
+	extstore.SetGlobalExtStorageForTest(storage)
+	defer func() {
+		extstore.SetGlobalExtStorageForTest(nil)
+		storage.Close()
+	}()
+
+	store := testkit.CreateMockStore(t)
+	tk := testkit.NewTestKit(t, store)
+	// Prepare multiple databases and tables for multi-SQL dump.
+	dbs := []string{"test", "test_multi_db1", "test_multi_db2", "test_multi_db3", "test_multi_db4"}
+	for _, db := range dbs {
+		tk.MustExec(fmt.Sprintf("create database if not exists %s", db))
+	}
+	for _, db := range dbs {
+		tk.MustExec("use " + db)
+		for i := 1; i <= numTables; i++ {
+			tableName := fmt.Sprintf("t_dump_multi_%d", i)
+			tk.MustExec(fmt.Sprintf("drop table if exists %s", tableName))
+			tk.MustExec(fmt.Sprintf("create table %s(a int, b int)", tableName))
+			tk.MustExec(fmt.Sprintf("insert into %s values (1, 1)", tableName))
+			tk.MustExec(fmt.Sprintf("insert into %s values (2, 2)", tableName))
+			tk.MustExec(fmt.Sprintf("insert into %s values (3, 3)", tableName))
+			tk.MustExec(fmt.Sprintf("insert into %s values (4, 4)", tableName))
+			tk.MustExec(fmt.Sprintf("insert into %s values (5, 5)", tableName))
+			tk.MustExec(fmt.Sprintf("analyze table %s", tableName))
+		}
+	}
+	tk.MustExec("use test")
+
+	// Build multiple SQL statements using the tables across multiple databases with fully
+	// qualified names (db.table) so the plan replayer extractor finds them regardless
+	// of current DB / schema sync.
+	stmts := make([]string, numStmts)
+	pairMod := len(dbs) * numTables
+	for i := 0; i < numStmts; i++ {
+		// Make sure every (db, table) pair is covered at least once.
+		pairIdx := i % pairMod
+		db := dbs[pairIdx/numTables]
+		tbl := (pairIdx % numTables) + 1
+		switch i % 4 {
+		case 0:
+			stmts[i] = fmt.Sprintf("'select * from %s.t_dump_multi_%d'", db, tbl)
+		case 1:
+			stmts[i] = fmt.Sprintf("'select * from %s.t_dump_multi_%d where a=1'", db, tbl)
+		case 2:
+			stmts[i] = fmt.Sprintf("'select * from %s.t_dump_multi_%d where b>0'", db, tbl)
+		default:
+			// join two tables, potentially across databases
+			t2 := (tbl % numTables) + 1
+			otherDB := dbs[(i+1)%len(dbs)]
+			stmts[i] = fmt.Sprintf("'select * from %s.t_dump_multi_%d, %s.t_dump_multi_%d where %s.t_dump_multi_%d.a=%s.t_dump_multi_%d.a'",
+				db, tbl, otherDB, t2, db, tbl, otherDB, t2)
+		}
+	}
+	sqlCmd := "plan replayer dump explain (" + strings.Join(stmts, ", ") + ")"
+	res := tk.MustQuery(sqlCmd)
+	path := testdata.ConvertRowsToStrings(res.Rows())
+	require.Len(t, path, 1)
+
+	filePath := filepath.Join(replayer.GetPlanReplayerDirName(), path[0])
+	fileReader, err := storage.Open(ctx, filePath, nil)
+	require.NoError(t, err)
+	defer fileReader.Close()
+
+	content, err := io.ReadAll(fileReader)
+	require.NoError(t, err)
+
+	readerAt := bytes.NewReader(content)
+	zr, err := zip.NewReader(readerAt, int64(len(content)))
+	require.NoError(t, err)
+
+	names := make(map[string]struct{})
+	for _, f := range zr.File {
+		names[f.Name] = struct{}{}
+	}
+	for i := 0; i < numStmts; i++ {
+		require.Contains(t, names, fmt.Sprintf("sql/sql%d.sql", i))
+		require.Contains(t, names, fmt.Sprintf("explain/explain%d.txt", i))
+	}
+	require.NotContains(t, names, "explain.txt") // single explain.txt is not used for multi-SQL
+
+	// Check stats and schema files for all tables in all databases
+	for _, db := range dbs {
+		for i := 1; i <= numTables; i++ {
+			tableName := fmt.Sprintf("t_dump_multi_%d", i)
+			statsName := fmt.Sprintf("stats/%s.%s.json", db, tableName)
+			schemaName := fmt.Sprintf("schema/%s.%s.schema.txt", db, tableName)
+			require.Contains(t, names, statsName, "missing stats file for db=%s table=%s (expected %s)", db, tableName, statsName)
+			require.Contains(t, names, schemaName, "missing schema file for db=%s table=%s (expected %s)", db, tableName, schemaName)
+		}
+	}
+}

pkg/parser/ast/misc.go

@@ -301,6 +301,10 @@ type PlanReplayerStmt struct {
 	// 2. plan replayer dump explain <analyze> 'file'
 	File string
 
+	// StmtList is used for PLAN REPLAYER DUMP EXPLAIN [ANALYZE] ( "sql1", "sql2", ... )
+	// When non-nil, multiple SQL strings are dumped in one command.
+	StmtList []string
+
 	// Fields below are currently useless.
 
 	// Where is the where clause in select statement.
@@ -347,6 +351,17 @@ func (n *PlanReplayerStmt) Restore(ctx *format.RestoreCtx) error {
 	} else {
 		ctx.WriteKeyWord("EXPLAIN ")
 	}
+	if len(n.StmtList) > 0 {
+		ctx.WritePlain("(")
+		for i, s := range n.StmtList {
+			if i > 0 {
+				ctx.WritePlain(", ")
+			}
+			ctx.WriteString(s)
+		}
+		ctx.WritePlain(")")
+		return nil
+	}
 	if n.Stmt == nil {
 		if len(n.File) > 0 {
 			ctx.WriteString(n.File)

pkg/parser/ast/misc_test.go

@@ -343,6 +343,10 @@ func TestPlanReplayerStmtRestore(t *testing.T) {
 			"PLAN REPLAYER DUMP EXPLAIN ANALYZE 'test'"},
 		{"plan replayer dump with stats as of timestamp '12345' explain analyze 'test2'",
 			"PLAN REPLAYER DUMP WITH STATS AS OF TIMESTAMP _UTF8MB4'12345' EXPLAIN ANALYZE 'test2'"},
+		{"plan replayer dump explain ('SELECT * FROM t1', 'SELECT * FROM t2')",
+			"PLAN REPLAYER DUMP EXPLAIN ('SELECT * FROM t1', 'SELECT * FROM t2')"},
+		{"plan replayer dump explain analyze ('SELECT * FROM t1')",
+			"PLAN REPLAYER DUMP EXPLAIN ANALYZE ('SELECT * FROM t1')"},
 	}
 	extractNodeFunc := func(node ast.Node) ast.Node {
 		return node.(*ast.PlanReplayerStmt)