*: support cancel query like 'select * from information_schema.tables' (#57766) (#57806)

ti-chi-bot · web-flow · commit afb68c65e4d0 · 2024-12-02T07:00:44.000Z
close #57749
diff --git a/pkg/ddl/placement_policy.go b/pkg/ddl/placement_policy.go
@@ -458,7 +458,7 @@ func CheckPlacementPolicyNotInUseFromMeta(t *meta.Mutator, policy *model.PolicyI
 			return dbterror.ErrPlacementPolicyInUse.GenWithStackByArgs(policy.Name)
 		}
 
-		tables, err := t.ListTables(dbInfo.ID)
+		tables, err := t.ListTables(context.Background(), dbInfo.ID)
 		if err != nil {
 			return err
 		}
diff --git a/pkg/ddl/placement_policy_test.go b/pkg/ddl/placement_policy_test.go
@@ -984,7 +984,7 @@ func testGetPolicyDependency(storage kv.Storage, name string) []int64 {
 			return err
 		}
 		for _, db := range dbs {
-			tbls, err := t.ListTables(db.ID)
+			tbls, err := t.ListTables(context.Background(), db.ID)
 			if err != nil {
 				return err
 			}
diff --git a/pkg/ddl/schema.go b/pkg/ddl/schema.go
@@ -180,7 +180,7 @@ func (w *worker) onDropSchema(jobCtx *jobContext, job *model.Job) (ver int64, _
 			return ver, errors.Trace(err)
 		}
 		var tables []*model.TableInfo
-		tables, err = metaMut.ListTables(job.SchemaID)
+		tables, err = metaMut.ListTables(jobCtx.stepCtx, job.SchemaID)
 		if err != nil {
 			return ver, errors.Trace(err)
 		}
@@ -205,7 +205,7 @@ func (w *worker) onDropSchema(jobCtx *jobContext, job *model.Job) (ver int64, _
 	case model.StateDeleteOnly:
 		dbInfo.State = model.StateNone
 		var tables []*model.TableInfo
-		tables, err = metaMut.ListTables(job.SchemaID)
+		tables, err = metaMut.ListTables(jobCtx.stepCtx, job.SchemaID)
 		if err != nil {
 			return ver, errors.Trace(err)
 		}
@@ -294,7 +294,7 @@ func (w *worker) onRecoverSchema(jobCtx *jobContext, job *model.Job) (ver int64,
 			sid := recoverSchemaInfo.DBInfo.ID
 			snap := w.store.GetSnapshot(kv.NewVersion(recoverSchemaInfo.SnapshotTS))
 			snapMeta := meta.NewReader(snap)
-			tables, err2 := snapMeta.ListTables(sid)
+			tables, err2 := snapMeta.ListTables(jobCtx.stepCtx, sid)
 			if err2 != nil {
 				job.State = model.JobStateCancelled
 				return ver, errors.Trace(err2)
diff --git a/pkg/ddl/table.go b/pkg/ddl/table.go
@@ -1239,7 +1239,7 @@ func checkConstraintNamesNotExists(t *meta.Mutator, schemaID int64, constraints
 	if len(constraints) == 0 {
 		return nil
 	}
-	tbInfos, err := t.ListTables(schemaID)
+	tbInfos, err := t.ListTables(context.Background(), schemaID)
 	if err != nil {
 		return err
 	}
diff --git a/pkg/domain/domain.go b/pkg/domain/domain.go
@@ -521,7 +521,7 @@ func (*Domain) fetchSchemasWithTables(ctx context.Context, schemas []*model.DBIn
 			di.TableName2ID = name2ID
 			tables = specialTableInfos
 		} else {
-			tables, err = m.ListTables(di.ID)
+			tables, err = m.ListTables(ctx, di.ID)
 			if err != nil {
 				return err
 			}
diff --git a/pkg/executor/infoschema_reader.go b/pkg/executor/infoschema_reader.go
@@ -834,6 +834,9 @@ func (e *memtableRetriever) setDataFromTables(ctx context.Context, sctx sessionc
 		if err != nil {
 			return errors.Trace(err)
 		}
+		if ctx.Err() != nil {
+			return errors.Trace(ctx.Err())
+		}
 	}
 	e.rows = rows
 	return nil
@@ -1204,6 +1207,10 @@ func (e *memtableRetriever) setDataFromPartitions(ctx context.Context, sctx sess
 		}
 		createTime := types.NewTime(types.FromGoTime(table.GetUpdateTime()), createTimeTp, types.DefaultFsp)
 
+		if ctx.Err() != nil {
+			return errors.Trace(ctx.Err())
+		}
+
 		var rowCount, dataLength, indexLength uint64
 		if useStatsCache {
 			if table.GetPartitionInfo() == nil {
diff --git a/pkg/infoschema/infoschema_v2.go b/pkg/infoschema/infoschema_v2.go
@@ -844,7 +844,7 @@ retry:
 	// the meta region leader is slow.
 	snapshot.SetOption(kv.TiKVClientReadTimeout, uint64(3000)) // 3000ms.
 	m := meta.NewReader(snapshot)
-	tblInfos, err := m.ListTables(dbInfo.ID)
+	tblInfos, err := m.ListTables(ctx, dbInfo.ID)
 	if err != nil {
 		if meta.ErrDBNotExists.Equal(err) {
 			return nil, nil
diff --git a/pkg/meta/meta.go b/pkg/meta/meta.go
@@ -16,6 +16,7 @@ package meta
 
 import (
 	"bytes"
+	"context"
 	"encoding/binary"
 	"encoding/json"
 	"fmt"
@@ -1131,7 +1132,7 @@ func GetTableInfoWithAttributes(m *Mutator, dbID int64, filterAttrs ...string) (
 }
 
 // ListTables shows all tables in database.
-func (m *Mutator) ListTables(dbID int64) ([]*model.TableInfo, error) {
+func (m *Mutator) ListTables(ctx context.Context, dbID int64) ([]*model.TableInfo, error) {
 	res, err := m.GetMetasByDBID(dbID)
 	if err != nil {
 		return nil, errors.Trace(err)
@@ -1144,6 +1145,9 @@ func (m *Mutator) ListTables(dbID int64) ([]*model.TableInfo, error) {
 		if !strings.HasPrefix(tableKey, mTablePrefix) {
 			continue
 		}
+		if ctx.Err() != nil {
+			return nil, errors.Trace(ctx.Err())
+		}
 
 		tbInfo := &model.TableInfo{}
 		err = json.Unmarshal(r.Value, tbInfo)
diff --git a/pkg/meta/meta_test.go b/pkg/meta/meta_test.go
@@ -289,7 +289,7 @@ func TestMeta(t *testing.T) {
 	tableNames, err := m.ListSimpleTables(1)
 	require.NoError(t, err)
 	require.Equal(t, []*model.TableNameInfo{tblName, tblName2}, tableNames)
-	tables, err := m.ListTables(1)
+	tables, err := m.ListTables(context.Background(), 1)
 	require.NoError(t, err)
 	require.Equal(t, []*model.TableInfo{tbInfo, tbInfo2}, tables)
 	{
@@ -327,7 +327,7 @@ func TestMeta(t *testing.T) {
 	tableNames, err = m.ListSimpleTables(1)
 	require.NoError(t, err)
 	require.Equal(t, []*model.TableNameInfo{tblName}, tableNames)
-	tables, err = m.ListTables(1)
+	tables, err = m.ListTables(context.Background(), 1)
 	require.NoError(t, err)
 	require.Equal(t, []*model.TableInfo{tbInfo}, tables)
 	{
diff --git a/pkg/meta/reader.go b/pkg/meta/reader.go
@@ -15,6 +15,8 @@
 package meta
 
 import (
+	"context"
+
 	"github.com/pingcap/tidb/pkg/kv"
 	"github.com/pingcap/tidb/pkg/meta/model"
 	"github.com/pingcap/tidb/pkg/structure"
@@ -25,7 +27,7 @@ type Reader interface {
 	GetDatabase(dbID int64) (*model.DBInfo, error)
 	ListDatabases() ([]*model.DBInfo, error)
 	GetTable(dbID int64, tableID int64) (*model.TableInfo, error)
-	ListTables(dbID int64) ([]*model.TableInfo, error)
+	ListTables(ctx context.Context, dbID int64) ([]*model.TableInfo, error)
 	ListSimpleTables(dbID int64) ([]*model.TableNameInfo, error)
 	IterTables(dbID int64, fn func(info *model.TableInfo) error) error
 	GetAutoIDAccessors(dbID, tableID int64) AutoIDAccessors
diff --git a/pkg/planner/core/memtable_infoschema_extractor.go b/pkg/planner/core/memtable_infoschema_extractor.go
@@ -707,6 +707,9 @@ func listTablesForEachSchema(
 		if err != nil {
 			return nil, nil, errors.Trace(err)
 		}
+		if ctx.Err() != nil {
+			return nil, nil, errors.Trace(err)
+		}
 		tables = filterSchemaObjectByRegexp(e, ec.table, tables, extractStrTableInfo)
 		for _, t := range tables {
 			schemaSlice = append(schemaSlice, s)

Original file line number	Diff line number	Diff line change
`@@ -458,7 +458,7 @@ func CheckPlacementPolicyNotInUseFromMeta(t meta.Mutator, policy model.PolicyI`
`458`	`458`	`return dbterror.ErrPlacementPolicyInUse.GenWithStackByArgs(policy.Name)`
`459`	`459`	`}`
`460`	`460`
`461`		`- tables, err := t.ListTables(dbInfo.ID)`
	`461`	`+ tables, err := t.ListTables(context.Background(), dbInfo.ID)`
`462`	`462`	`if err != nil {`
`463`	`463`	`return err`
`464`	`464`	`}`
Original file line number	Diff line number	Diff line change
`@@ -984,7 +984,7 @@ func testGetPolicyDependency(storage kv.Storage, name string) []int64 {`
`984`	`984`	`return err`
`985`	`985`	`}`
`986`	`986`	`for _, db := range dbs {`
`987`		`- tbls, err := t.ListTables(db.ID)`
	`987`	`+ tbls, err := t.ListTables(context.Background(), db.ID)`
`988`	`988`	`if err != nil {`
`989`	`989`	`return err`
`990`	`990`	`}`
Original file line number	Diff line number	Diff line change
`@@ -1239,7 +1239,7 @@ func checkConstraintNamesNotExists(t *meta.Mutator, schemaID int64, constraints`
`1239`	`1239`	`if len(constraints) == 0 {`
`1240`	`1240`	`return nil`
`1241`	`1241`	`}`
`1242`		`- tbInfos, err := t.ListTables(schemaID)`
	`1242`	`+ tbInfos, err := t.ListTables(context.Background(), schemaID)`
`1243`	`1243`	`if err != nil {`
`1244`	`1244`	`return err`
`1245`	`1245`	`}`
Original file line number	Diff line number	Diff line change
`@@ -521,7 +521,7 @@ func (Domain) fetchSchemasWithTables(ctx context.Context, schemas []model.DBIn`
`521`	`521`	`di.TableName2ID = name2ID`
`522`	`522`	`tables = specialTableInfos`
`523`	`523`	`} else {`
`524`		`- tables, err = m.ListTables(di.ID)`
	`524`	`+ tables, err = m.ListTables(ctx, di.ID)`
`525`	`525`	`if err != nil {`
`526`	`526`	`return err`
`527`	`527`	`}`
Original file line number	Diff line number	Diff line change
`@@ -834,6 +834,9 @@ func (e *memtableRetriever) setDataFromTables(ctx context.Context, sctx sessionc`
`834`	`834`	`if err != nil {`
`835`	`835`	`return errors.Trace(err)`
`836`	`836`	`}`
	`837`	`+ if ctx.Err() != nil {`
	`838`	`+ return errors.Trace(ctx.Err())`
	`839`	`+ }`
`837`	`840`	`}`
`838`	`841`	`e.rows = rows`
`839`	`842`	`return nil`
`@@ -1204,6 +1207,10 @@ func (e *memtableRetriever) setDataFromPartitions(ctx context.Context, sctx sess`
`1204`	`1207`	`}`
`1205`	`1208`	`createTime := types.NewTime(types.FromGoTime(table.GetUpdateTime()), createTimeTp, types.DefaultFsp)`
`1206`	`1209`
	`1210`	`+ if ctx.Err() != nil {`
	`1211`	`+ return errors.Trace(ctx.Err())`
	`1212`	`+ }`
	`1213`	`+`
`1207`	`1214`	`var rowCount, dataLength, indexLength uint64`
`1208`	`1215`	`if useStatsCache {`
`1209`	`1216`	`if table.GetPartitionInfo() == nil {`