domain: fix play replay dump cannot save the table in the foreign key's reference (#56512)

close #56458

Author: hawkingrei

pkg/domain/extract.go

@@ -290,7 +290,10 @@ func (w *extractWorker) handleIsView(ctx context.Context, p *extractPlanPackage)
 	if tne.err != nil {
 		return tne.err
 	}
-	r := tne.getTablesAndViews()
+	r, err := tne.getTablesAndViews()
+	if err != nil {
+		return err
+	}
 	for t := range r {
 		p.tables[t] = struct{}{}
 	}

pkg/domain/plan_replayer_dump.go

@@ -97,7 +97,7 @@ type tableNameExtractor struct {
 	err      error
 }
 
-func (tne *tableNameExtractor) getTablesAndViews() map[tableNamePair]struct{} {
+func (tne *tableNameExtractor) getTablesAndViews() (map[tableNamePair]struct{}, error) {
 	r := make(map[tableNamePair]struct{})
 	for tablePair := range tne.names {
 		if tablePair.IsView {
@@ -109,8 +109,21 @@ func (tne *tableNameExtractor) getTablesAndViews() map[tableNamePair]struct{} {
 		if !ok {
 			r[tablePair] = struct{}{}
 		}
+		// if the table has a foreign key, we need to add the referenced table to the list
+		tblInfo, err := tne.is.TableByName(tne.ctx, model.NewCIStr(tablePair.DBName), model.NewCIStr(tablePair.TableName))
+		if err != nil {
+			return nil, err
+		}
+		for _, fk := range tblInfo.Meta().ForeignKeys {
+			key := tableNamePair{
+				DBName:    fk.RefSchema.L,
+				TableName: fk.RefTable.L,
+				IsView:    false,
+			}
+			r[key] = struct{}{}
+		}
 	}
-	return r
+	return r, nil
 }
 
 func (*tableNameExtractor) Enter(in ast.Node) (ast.Node, bool) {
@@ -776,7 +789,7 @@ func extractTableNames(ctx context.Context, sctx sessionctx.Context,
 	if tableExtractor.err != nil {
 		return nil, tableExtractor.err
 	}
-	return tableExtractor.getTablesAndViews(), nil
+	return tableExtractor.getTablesAndViews()
 }
 
 func getStatsForTable(do *Domain, pair tableNamePair, historyStatsTS uint64) (*util.JSONTable, []string, error) {

pkg/server/handler/optimizor/BUILD.bazel

@@ -44,7 +44,7 @@ go_test(
         "statistics_handler_test.go",
     ],
     flaky = True,
-    shard_count = 5,
+    shard_count = 6,
     deps = [
         ":optimizor",
         "//pkg/config",

pkg/server/handler/optimizor/plan_replayer_test.go

@@ -249,6 +249,55 @@ func prepareData4PlanReplayer(t *testing.T, client *testserverclient.TestServerC
 	return filename, filename3
 }
 
+func TestIssue56458(t *testing.T) {
+	store := testkit.CreateMockStore(t)
+	dom, err := session.GetDomain(store)
+	require.NoError(t, err)
+	// 1. setup and prepare plan replayer files by manual command and capture
+	server, client := prepareServerAndClientForTest(t, store, dom)
+	defer server.Close()
+
+	filename := prepareData4Issue56458(t, client, dom)
+	defer os.RemoveAll(replayer.GetPlanReplayerDirName())
+
+	// 2. check the contents of the plan replayer zip files.
+	var filesInReplayer []string
+	collectFileNameAndAssertFileSize := func(f *zip.File) {
+		// collect file name
+		filesInReplayer = append(filesInReplayer, f.Name)
+	}
+
+	// 2-1. check the plan replayer file from manual command
+	resp0, err := client.FetchStatus(filepath.Join("/plan_replayer/dump/", filename))
+	require.NoError(t, err)
+	defer func() {
+		require.NoError(t, resp0.Body.Close())
+	}()
+	body, err := io.ReadAll(resp0.Body)
+	require.NoError(t, err)
+	forEachFileInZipBytes(t, body, collectFileNameAndAssertFileSize)
+	slices.Sort(filesInReplayer)
+	require.Equal(t, []string{
+		"config.toml",
+		"debug_trace/debug_trace0.json",
+		"explain.txt",
+		"global_bindings.sql",
+		"meta.txt",
+		"schema/planreplayer.t.schema.txt",
+		"schema/planreplayer.v.schema.txt",
+		"schema/schema_meta.txt",
+		"session_bindings.sql",
+		"sql/sql0.sql",
+		"sql_meta.toml",
+		"stats/planreplayer.t.json",
+		"stats/planreplayer.v.json",
+		"statsMem/planreplayer.t.txt",
+		"statsMem/planreplayer.v.txt",
+		"table_tiflash_replica.txt",
+		"variables.toml",
+	}, filesInReplayer)
+}
+
 func TestIssue43192(t *testing.T) {
 	store := testkit.CreateMockStore(t)
 	dom, err := session.GetDomain(store)
@@ -344,6 +393,33 @@ func prepareData4Issue43192(t *testing.T, client *testserverclient.TestServerCli
 	return filename
 }
 
+func prepareData4Issue56458(t *testing.T, client *testserverclient.TestServerClient, dom *domain.Domain) string {
+	h := dom.StatsHandle()
+	db, err := sql.Open("mysql", client.GetDSN())
+	require.NoError(t, err, "Error connecting")
+	defer func() {
+		err := db.Close()
+		require.NoError(t, err)
+	}()
+	tk := testkit.NewDBTestKit(t, db)
+
+	tk.MustExec("create database planReplayer")
+	tk.MustExec("use planReplayer")
+	tk.MustExec("CREATE TABLE v(id INT PRIMARY KEY AUTO_INCREMENT);")
+	tk.MustExec("create table t(a int, b int, INDEX ia (a), INDEX ib (b), author_id int, FOREIGN KEY (author_id) REFERENCES v(id) ON DELETE CASCADE);")
+	err = h.HandleDDLEvent(<-h.DDLEventCh())
+	require.NoError(t, err)
+	tk.MustExec("create global binding for select a, b from t where a in (1, 2, 3) using select a, b from t use index (ib) where a in (1, 2, 3)")
+	rows := tk.MustQuery("plan replayer dump explain select a, b from t where a in (1, 2, 3)")
+	require.True(t, rows.Next(), "unexpected data")
+	var filename string
+	require.NoError(t, rows.Scan(&filename))
+	require.NoError(t, rows.Close())
+	rows = tk.MustQuery("select @@tidb_last_plan_replayer_token")
+	require.True(t, rows.Next(), "unexpected data")
+	return filename
+}
+
 func forEachFileInZipBytes(t *testing.T, b []byte, fn func(file *zip.File)) {
 	br := bytes.NewReader(b)
 	z, err := zip.NewReader(br, int64(len(b)))