This is an automated cherry-pick of pingcap#56457

Signed-off-by: ti-chi-bot <[email protected]>
ti-chi-bot · Oct 17, 2024 · 3993d4f · 3993d4f
1 parent 21d2679
commit 3993d4f
Show file tree

Hide file tree

Showing 3 changed files with 105 additions and 21 deletions.
diff --git a/pkg/domain/plan_replayer_dump.go b/pkg/domain/plan_replayer_dump.go
@@ -109,22 +109,38 @@ func (tne *tableNameExtractor) getTablesAndViews() (map[tableNamePair]struct{},
 			r[tablePair] = struct{}{}
 		}
 		// if the table has a foreign key, we need to add the referenced table to the list
+<<<<<<< HEAD
 		tblInfo, err := tne.is.TableByName(model.NewCIStr(tablePair.DBName), model.NewCIStr(tablePair.TableName))
+=======
+		err := findFK(tne.is, tablePair.DBName, tablePair.TableName, r)
+>>>>>>> 53a4c3eded2 (planner: play replay load restore the table with foreign key with right order. (#56457))
 		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, nil
 }
 
+func findFK(is infoschema.InfoSchema, dbName, tableName string, tableMap map[tableNamePair]struct{}) error {
+	tblInfo, err := is.TableByName(context.Background(), model.NewCIStr(dbName), model.NewCIStr(tableName))
+	if err != nil {
+		return err
+	}
+	for _, fk := range tblInfo.Meta().ForeignKeys {
+		key := tableNamePair{
+			DBName:    fk.RefSchema.L,
+			TableName: fk.RefTable.L,
+			IsView:    false,
+		}
+		tableMap[key] = struct{}{}
+		err := findFK(is, key.DBName, key.TableName, tableMap)
+		if err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
 func (*tableNameExtractor) Enter(in ast.Node) (ast.Node, bool) {
 	if _, ok := in.(*ast.TableName); ok {
 		return in, true

diff --git a/pkg/executor/plan_replayer.go b/pkg/executor/plan_replayer.go
@@ -482,17 +482,9 @@ func (e *PlanReplayerLoadInfo) Update(data []byte) error {
 	}
 
 	// build schema and table first
-	for _, zipFile := range z.File {
-		if zipFile.Name == fmt.Sprintf("schema/%v", domain.PlanReplayerSchemaMetaFile) {
-			continue
-		}
-		path := strings.Split(zipFile.Name, "/")
-		if len(path) == 2 && strings.Compare(path[0], "schema") == 0 && zipFile.Mode().IsRegular() {
-			err = createSchemaAndItems(e.Ctx, zipFile)
-			if err != nil {
-				return err
-			}
-		}
+	err = e.createTable(z)
+	if err != nil {
+		return err
 	}
 
 	// set tiflash replica if exists
@@ -530,3 +522,26 @@ func (e *PlanReplayerLoadInfo) Update(data []byte) error {
 	}
 	return nil
 }
+
+func (e *PlanReplayerLoadInfo) createTable(z *zip.Reader) error {
+	origin := e.Ctx.GetSessionVars().ForeignKeyChecks
+	// We need to disable foreign key check when we create schema and tables.
+	// because the order of creating schema and tables is not guaranteed.
+	e.Ctx.GetSessionVars().ForeignKeyChecks = false
+	defer func() {
+		e.Ctx.GetSessionVars().ForeignKeyChecks = origin
+	}()
+	for _, zipFile := range z.File {
+		if zipFile.Name == fmt.Sprintf("schema/%v", domain.PlanReplayerSchemaMetaFile) {
+			continue
+		}
+		path := strings.Split(zipFile.Name, "/")
+		if len(path) == 2 && strings.Compare(path[0], "schema") == 0 && zipFile.Mode().IsRegular() {
+			err := createSchemaAndItems(e.Ctx, zipFile)
+			if err != nil {
+				return err
+			}
+		}
+	}
+	return nil
+}
diff --git a/pkg/server/handler/optimizor/plan_replayer_test.go b/pkg/server/handler/optimizor/plan_replayer_test.go
@@ -214,6 +214,8 @@ func prepareData4PlanReplayer(t *testing.T, client *testserverclient.TestServerC
 	tk.MustExec("create database planReplayer")
 	tk.MustExec("use planReplayer")
 	tk.MustExec("create table t(a int)")
+	tk.MustExec("CREATE TABLE authors (id INT PRIMARY KEY AUTO_INCREMENT,name VARCHAR(100) NOT NULL,email VARCHAR(100) UNIQUE NOT NULL);")
+	tk.MustExec("CREATE TABLE books (id INT PRIMARY KEY AUTO_INCREMENT,title VARCHAR(200) NOT NULL,publication_date DATE NOT NULL,author_id INT,FOREIGN KEY (author_id) REFERENCES authors(id) ON DELETE CASCADE);")
 	err = h.HandleDDLEvent(<-h.DDLEventCh())
 	require.NoError(t, err)
 	tk.MustExec("insert into t values(1), (2), (3), (4)")
@@ -248,7 +250,7 @@ func prepareData4PlanReplayer(t *testing.T, client *testserverclient.TestServerC
 	return filename, filename3
 }
 
-func TestIssue56458(t *testing.T) {
+func TestPlanReplayerWithMultiForeignKey(t *testing.T) {
 	store := testkit.CreateMockStore(t)
 	dom, err := session.GetDomain(store)
 	require.NoError(t, err)
@@ -284,17 +286,60 @@ func TestIssue56458(t *testing.T) {
 		"meta.txt",
 		"schema/planreplayer.t.schema.txt",
 		"schema/planreplayer.v.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/planreplayer.v.json",
+		"stats/planreplayer2.t.json",
 		"statsMem/planreplayer.t.txt",
 		"statsMem/planreplayer.v.txt",
+		"statsMem/planreplayer2.t.txt",
 		"table_tiflash_replica.txt",
 		"variables.toml",
 	}, 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("drop table planReplayer2.t")
+	tk.MustExec("drop table planReplayer.v")
+	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 TestIssue43192(t *testing.T) {
@@ -403,11 +448,19 @@ func prepareData4Issue56458(t *testing.T, client *testserverclient.TestServerCli
 	tk := testkit.NewDBTestKit(t, db)
 
 	tk.MustExec("create database planReplayer")
+	tk.MustExec("create database planReplayer2")
 	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 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);")
+	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")