domain: fix play replay dump cannot save the table in the foreign key…

…'s reference (#56512) (#56564)

close #56458

domain/extract.go

@@ -288,7 +288,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{}{}
 	}

domain/plan_replayer_dump.go

@@ -92,7 +92,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 {
@@ -104,8 +104,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(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 (tne *tableNameExtractor) Enter(in ast.Node) (ast.Node, bool) {
@@ -747,7 +760,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) (*handle.JSONTable, error) {

server/plan_replayer_test.go

@@ -27,9 +27,11 @@ import (
 	"github.com/go-sql-driver/mysql"
 	"github.com/gorilla/mux"
 	"github.com/pingcap/tidb/domain"
+	"github.com/pingcap/tidb/kv"
 	"github.com/pingcap/tidb/session"
 	"github.com/pingcap/tidb/statistics/handle"
 	"github.com/pingcap/tidb/testkit"
+	"github.com/pingcap/tidb/util/replayer"
 	"github.com/stretchr/testify/require"
 	"golang.org/x/exp/slices"
 )
@@ -241,3 +243,114 @@ func forEachFileInZipBytes(t *testing.T, b []byte, fn func(file *zip.File)) {
 		fn(f)
 	}
 }
+
+func prepareData4Issue56458(t *testing.T, client *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("create database planReplayer2")
+	tk.MustExec("use planReplayer")
+	tk.MustExec("create placement policy p " +
+		"LEARNERS=1 " +
+		"LEARNER_CONSTRAINTS=\"[+region=cn-west-1]\" " +
+		"FOLLOWERS=3 " +
+		"FOLLOWER_CONSTRAINTS=\"[+disk=ssd]\"")
+	tk.MustExec("CREATE TABLE v(id INT PRIMARY KEY AUTO_INCREMENT);")
+	err = h.HandleDDLEvent(<-h.DDLEventCh())
+	require.NoError(t, err)
+	tk.MustExec("create table planReplayer2.t(a int, b int, INDEX ia (a), INDEX ib (b), author_id int, FOREIGN KEY (author_id) REFERENCES planReplayer.v(id) ON DELETE CASCADE);")
+	err = h.HandleDDLEvent(<-h.DDLEventCh())
+	require.NoError(t, err)
+	tk.MustExec("create table t(a int, b int, INDEX ia (a), INDEX ib (b), author_id int, FOREIGN KEY (author_id) REFERENCES planReplayer2.t(a) ON DELETE CASCADE) placement policy p;")
+	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 prepareServerAndClientForTest(t *testing.T, store kv.Storage, dom *domain.Domain) (srv *Server, client *testServerClient) {
+	driver := NewTiDBDriver(store)
+	client = newTestServerClient()
+
+	cfg := newTestConfig()
+	cfg.Port = client.port
+	cfg.Status.StatusPort = client.statusPort
+	cfg.Status.ReportStatus = true
+
+	srv, err := NewServer(cfg, driver)
+	srv.SetDomain(dom)
+	require.NoError(t, err)
+	go func() {
+		err := srv.Run(nil)
+		require.NoError(t, err)
+	}()
+	<-RunInGoTestChan
+	client.port = getPortFromTCPAddr(srv.listener.Addr())
+	client.statusPort = getPortFromTCPAddr(srv.statusListener.Addr())
+	client.waitUntilServerOnline()
+	return
+}
+
+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/planreplayer2.t.schema.txt",
+		"schema/schema_meta.txt",
+		"session_bindings.sql",
+		"sql/sql0.sql",
+		"sql_meta.toml",
+		"stats/planreplayer.t.json",
+		"stats/planreplayer2.t.json",
+		"statsMem/planreplayer.t.txt",
+		"statsMem/planreplayer2.t.txt",
+		"table_tiflash_replica.txt",
+		"variables.toml",
+	}, filesInReplayer)
+}