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

close #56458

Author: ti-chi-bot

pkg/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{}{}
 	}

pkg/domain/plan_replayer_dump.go

@@ -96,7 +96,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 {
@@ -108,8 +108,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) {
@@ -771,7 +784,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

@@ -43,7 +43,7 @@ go_test(
         "statistics_handler_test.go",
     ],
     flaky = True,
-    shard_count = 4,
+    shard_count = 6,
     deps = [
         ":optimizor",
         "//pkg/config",
@@ -60,6 +60,7 @@ go_test(
         "//pkg/store/mockstore/unistore",
         "//pkg/testkit",
         "//pkg/testkit/testsetup",
+        "//pkg/util/replayer",
         "//pkg/util/topsql/state",
         "@com_github_burntsushi_toml//:toml",
         "@com_github_go_sql_driver_mysql//:mysql",

pkg/server/handler/optimizor/plan_replayer_test.go

@@ -42,6 +42,7 @@ import (
 	"github.com/pingcap/tidb/pkg/session"
 	util2 "github.com/pingcap/tidb/pkg/statistics/handle/util"
 	"github.com/pingcap/tidb/pkg/testkit"
+	"github.com/pingcap/tidb/pkg/util/replayer"
 	"github.com/stretchr/testify/require"
 	"github.com/tikv/client-go/v2/oracle"
 )
@@ -246,6 +247,177 @@ 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)
+	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 := prepareData4Issue43192(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, expectedFilesInReplayer, filesInReplayer)
+
+	// 3. check plan replayer load
+	// 3-1. write the plan replayer file from manual command to a file
+	path := "/tmp/plan_replayer.zip"
+	fp, err := os.Create(path)
+	require.NoError(t, err)
+	require.NotNil(t, fp)
+	defer func() {
+		require.NoError(t, fp.Close())
+		require.NoError(t, os.Remove(path))
+	}()
+
+	_, err = io.Copy(fp, bytes.NewReader(body))
+	require.NoError(t, err)
+	require.NoError(t, fp.Sync())
+
+	// 3-2. connect to tidb and use PLAN REPLAYER LOAD to load this file
+	db, err := sql.Open("mysql", client.GetDSN(func(config *mysql.Config) {
+		config.AllowAllFiles = true
+	}))
+	require.NoError(t, err, "Error connecting")
+	defer func() {
+		err := db.Close()
+		require.NoError(t, err)
+	}()
+	tk := testkit.NewDBTestKit(t, db)
+	tk.MustExec("use planReplayer")
+	tk.MustExec("drop table planReplayer.t")
+	tk.MustExec(`plan replayer load "/tmp/plan_replayer.zip"`)
+
+	// 3-3. check whether binding takes effect
+	tk.MustExec(`select a, b from t where a in (1, 2, 3)`)
+	rows := tk.MustQuery("select @@last_plan_from_binding")
+	require.True(t, rows.Next(), "unexpected data")
+	var count int64
+	err = rows.Scan(&count)
+	require.NoError(t, err)
+	require.Equal(t, int64(1), count)
+}
+
+func prepareData4Issue43192(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 t(a int, b int, INDEX ia (a), INDEX ib (b));")
+	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 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)))