planner: fix unstable test TestDAGPlanBuilderUnionScan (#37895)

build/nogo_config.json

@@ -313,7 +313,6 @@
   "nilness": {
     "exclude_files": {
       "/external/": "no need to vet third party code",
-      "planner/core/physical_plan_test.go": "please fix it",
       ".*_generated\\.go$": "ignore generated code",
       "/cgo/": "ignore cgo"
     }

planner/core/physical_plan_test.go

@@ -363,36 +363,37 @@ func TestDAGPlanBuilderUnion(t *testing.T) {
 
 func TestDAGPlanBuilderUnionScan(t *testing.T) {
 	store := testkit.CreateMockStore(t)
-
 	tk := testkit.NewTestKit(t, store)
 	tk.MustExec("use test")
+	tk.MustExec("drop table if exists t")
+	tk.MustExec("create table t(a int, b int, c int)")
 
 	var input []string
 	var output []struct {
 		SQL  string
 		Best string
 	}
+	planSuiteData := core.GetPlanSuiteData()
+	planSuiteData.LoadTestCases(t, &input, &output)
+
 	p := parser.New()
-	is := infoschema.MockInfoSchema([]*model.TableInfo{core.MockSignedTable(), core.MockUnsignedTable()})
 	for i, tt := range input {
+		tk.MustExec("begin;")
+		tk.MustExec("insert into t values(2, 2, 2);")
+
 		comment := fmt.Sprintf("input: %s", tt)
 		stmt, err := p.ParseOneStmt(tt, "", "")
 		require.NoError(t, err, comment)
-		require.NoError(t, sessiontxn.NewTxn(context.Background(), tk.Session()))
-
-		// Make txn not read only.
-		txn, err := tk.Session().Txn(true)
-		require.NoError(t, err)
-		err = txn.Set(kv.Key("AAA"), []byte("BBB"))
-		require.NoError(t, err)
-		tk.Session().StmtCommit()
-		p, _, err := planner.Optimize(context.TODO(), tk.Session(), stmt, is)
+		dom := domain.GetDomain(tk.Session())
+		require.NoError(t, dom.Reload())
+		plan, _, err := planner.Optimize(context.TODO(), tk.Session(), stmt, dom.InfoSchema())
 		require.NoError(t, err)
 		testdata.OnRecord(func() {
 			output[i].SQL = tt
-			output[i].Best = core.ToString(p)
+			output[i].Best = core.ToString(plan)
 		})
-		require.Equal(t, output[i].Best, core.ToString(p), fmt.Sprintf("input: %s", tt))
+		require.Equal(t, output[i].Best, core.ToString(plan), fmt.Sprintf("input: %s", tt))
+		tk.MustExec("rollback;")
 	}
 }
 

planner/core/testdata/plan_suite_out.json

@@ -831,7 +831,40 @@
   },
   {
     "Name": "TestDAGPlanBuilderUnionScan",
-    "Cases": null
+    "Cases": [
+      {
+        "SQL": "select * from t",
+        "Best": "TableReader(Table(t))->UnionScan([])->Projection"
+      },
+      {
+        "SQL": "select * from t where b = 1",
+        "Best": "TableReader(Table(t)->Sel([eq(test.t.b, 1)]))->UnionScan([eq(test.t.b, 1)])->Projection"
+      },
+      {
+        "SQL": "select * from t where a = 1",
+        "Best": "TableReader(Table(t)->Sel([eq(test.t.a, 1)]))->UnionScan([eq(test.t.a, 1)])->Projection"
+      },
+      {
+        "SQL": "select * from t where a = 1 order by a",
+        "Best": "TableReader(Table(t)->Sel([eq(test.t.a, 1)]))->UnionScan([eq(test.t.a, 1)])->Projection->Sort"
+      },
+      {
+        "SQL": "select * from t where a = 1 order by b",
+        "Best": "TableReader(Table(t)->Sel([eq(test.t.a, 1)]))->UnionScan([eq(test.t.a, 1)])->Projection->Sort"
+      },
+      {
+        "SQL": "select * from t where a = 1 limit 1",
+        "Best": "TableReader(Table(t)->Sel([eq(test.t.a, 1)]))->UnionScan([eq(test.t.a, 1)])->Limit"
+      },
+      {
+        "SQL": "select * from t where c = 1",
+        "Best": "TableReader(Table(t)->Sel([eq(test.t.c, 1)]))->UnionScan([eq(test.t.c, 1)])->Projection"
+      },
+      {
+        "SQL": "select c from t where c = 1",
+        "Best": "TableReader(Table(t)->Sel([eq(test.t.c, 1)]))->UnionScan([eq(test.t.c, 1)])->Projection"
+      }
+    ]
   },
   {
     "Name": "TestDAGPlanBuilderAgg",