From 130568716ef49b71746d1d14ffe2b2a216961c35 Mon Sep 17 00:00:00 2001
From: crazycs <crazycs520@gmail.com>
Date: Thu, 22 Dec 2022 23:02:55 +0800
Subject: [PATCH] ddl: fix issue of add foreign key too slow in big table
 (#40112)

close pingcap/tidb#40111
---
 ddl/fktest/foreign_key_test.go | 19 ++++++++++++++++++-
 ddl/foreign_key.go             |  9 +++++++--
 2 files changed, 25 insertions(+), 3 deletions(-)

diff --git a/ddl/fktest/foreign_key_test.go b/ddl/fktest/foreign_key_test.go
index f64de90ca4955..349d9b39ea682 100644
--- a/ddl/fktest/foreign_key_test.go
+++ b/ddl/fktest/foreign_key_test.go
@@ -19,6 +19,7 @@ import (
 	"context"
 	"fmt"
 	"testing"
+	"time"
 
 	"github.com/pingcap/tidb/domain"
 	"github.com/pingcap/tidb/infoschema"
@@ -1563,7 +1564,7 @@ func getLatestSchemaDiff(t *testing.T, tk *testkit.TestKit) *model.SchemaDiff {
 	return diff
 }
 
-func TestTestMultiSchemaAddForeignKey(t *testing.T) {
+func TestMultiSchemaAddForeignKey(t *testing.T) {
 	store, _ := testkit.CreateMockStoreAndDomain(t)
 	tk := testkit.NewTestKit(t, store)
 	tk.MustExec("set @@foreign_key_checks=1;")
@@ -1591,3 +1592,19 @@ func TestTestMultiSchemaAddForeignKey(t *testing.T) {
 	tk.MustExec("create table t2 (a int, b int, index idx0(a,b), index idx1(a), index idx2(b));")
 	tk.MustExec("alter table t2 drop index idx1, add foreign key (a) references t1(id), add foreign key (b) references t1(id)")
 }
+
+func TestAddForeignKeyInBigTable(t *testing.T) {
+	store := testkit.CreateMockStore(t)
+	tk := testkit.NewTestKit(t, store)
+	tk.MustExec("set @@foreign_key_checks=1;")
+	tk.MustExec("use test")
+	tk.MustExec("create table employee (id bigint auto_increment key, pid bigint)")
+	tk.MustExec("insert into employee (id) values (1),(2),(3),(4),(5),(6),(7),(8)")
+	for i := 0; i < 14; i++ {
+		tk.MustExec("insert into employee (pid) select pid from employee")
+	}
+	tk.MustExec("update employee set pid=id-1 where id>1")
+	start := time.Now()
+	tk.MustExec("alter table employee add foreign key fk_1(pid) references employee(id)")
+	require.Less(t, time.Since(start), time.Minute)
+}
diff --git a/ddl/foreign_key.go b/ddl/foreign_key.go
index cdf3695cdc017..e2506770e27c4 100644
--- a/ddl/foreign_key.go
+++ b/ddl/foreign_key.go
@@ -674,7 +674,12 @@ func checkForeignKeyConstrain(w *worker, schema, table string, fkInfo *model.FKI
 	if err != nil {
 		return errors.Trace(err)
 	}
-	defer w.sessPool.put(sctx)
+	originValue := sctx.GetSessionVars().OptimizerEnableNAAJ
+	sctx.GetSessionVars().OptimizerEnableNAAJ = true
+	defer func() {
+		sctx.GetSessionVars().OptimizerEnableNAAJ = originValue
+		w.sessPool.put(sctx)
+	}()
 
 	var buf strings.Builder
 	buf.WriteString("select 1 from %n.%n where ")
@@ -709,7 +714,7 @@ func checkForeignKeyConstrain(w *worker, schema, table string, fkInfo *model.FKI
 	}
 	buf.WriteString(" from %n.%n ) limit 1")
 	paramsList = append(paramsList, fkInfo.RefSchema.L, fkInfo.RefTable.L)
-	rows, _, err := sctx.(sqlexec.RestrictedSQLExecutor).ExecRestrictedSQL(w.ctx, nil, buf.String(), paramsList...)
+	rows, _, err := sctx.(sqlexec.RestrictedSQLExecutor).ExecRestrictedSQL(w.ctx, []sqlexec.OptionFuncAlias{sqlexec.ExecOptionUseCurSession}, buf.String(), paramsList...)
 	if err != nil {
 		return errors.Trace(err)
 	}