move attribute to infoschema/context

br/pkg/task/BUILD.bazel

@@ -53,6 +53,7 @@ go_library(
         "//pkg/ddl",
         "//pkg/domain",
         "//pkg/infoschema",
+        "//pkg/infoschema/context",
         "//pkg/kv",
         "//pkg/meta/model",
         "//pkg/parser/model",

br/pkg/task/restore_data.go

@@ -22,7 +22,7 @@ import (
 	"github.com/pingcap/tidb/br/pkg/summary"
 	"github.com/pingcap/tidb/br/pkg/utils"
 	tidbconfig "github.com/pingcap/tidb/pkg/config"
-	"github.com/pingcap/tidb/pkg/infoschema"
+	infoschemacontext "github.com/pingcap/tidb/pkg/infoschema/context"
 	"github.com/pingcap/tidb/pkg/kv"
 	pd "github.com/tikv/pd/client"
 	"go.uber.org/zap"
@@ -181,7 +181,7 @@ func resetTiFlashReplicas(ctx context.Context, g glue.Glue, storage kv.Storage,
 
 	expectTiFlashStoreCount := uint64(0)
 	needTiFlash := false
-	tableInfoRes := info.ListTablesWithSpecialAttribute(infoschema.TiFlashAttribute)
+	tableInfoRes := info.ListTablesWithSpecialAttribute(infoschemacontext.TiFlashAttribute)
 	for _, s := range tableInfoRes {
 		for _, t := range s.TableInfos {
 			if t.TiFlashReplica != nil {

pkg/ddl/BUILD.bazel

@@ -98,6 +98,7 @@ go_library(
         "//pkg/expression/context",
         "//pkg/expression/contextstatic",
         "//pkg/infoschema",
+        "//pkg/infoschema/context",
         "//pkg/kv",
         "//pkg/lightning/backend",
         "//pkg/lightning/backend/external",

pkg/ddl/ddl_tiflash_api.go

@@ -34,6 +34,7 @@ import (
 	ddlutil "github.com/pingcap/tidb/pkg/ddl/util"
 	"github.com/pingcap/tidb/pkg/domain/infosync"
 	"github.com/pingcap/tidb/pkg/infoschema"
+	infoschemacontext "github.com/pingcap/tidb/pkg/infoschema/context"
 	"github.com/pingcap/tidb/pkg/meta/model"
 	"github.com/pingcap/tidb/pkg/sessionctx"
 	"github.com/pingcap/tidb/pkg/table"
@@ -451,7 +452,7 @@ func (d *ddl) refreshTiFlashTicker(ctx sessionctx.Context, pollTiFlashContext *T
 	var tableList = make([]TiFlashReplicaStatus, 0)
 
 	// Collect TiFlash Replica info, for every table.
-	ch := schema.ListTablesWithSpecialAttribute(infoschema.TiFlashAttribute)
+	ch := schema.ListTablesWithSpecialAttribute(infoschemacontext.TiFlashAttribute)
 	for _, v := range ch {
 		for _, tblInfo := range v.TableInfos {
 			LoadTiFlashReplicaInfo(tblInfo, &tableList)

pkg/ddl/executor.go

@@ -39,6 +39,7 @@ import (
 	"github.com/pingcap/tidb/pkg/errctx"
 	"github.com/pingcap/tidb/pkg/expression"
 	"github.com/pingcap/tidb/pkg/infoschema"
+	infoschemacontext "github.com/pingcap/tidb/pkg/infoschema/context"
 	"github.com/pingcap/tidb/pkg/kv"
 	"github.com/pingcap/tidb/pkg/meta"
 	"github.com/pingcap/tidb/pkg/meta/autoid"
@@ -413,7 +414,7 @@ func (e *executor) getPendingTiFlashTableCount(originVersion int64, pendingCount
 		return originVersion, pendingCount
 	}
 	cnt := uint32(0)
-	dbs := is.ListTablesWithSpecialAttribute(infoschema.TiFlashAttribute)
+	dbs := is.ListTablesWithSpecialAttribute(infoschemacontext.TiFlashAttribute)
 	for _, db := range dbs {
 		if util.IsMemOrSysDB(db.DBName.L) {
 			continue

pkg/ddl/placement_policy.go

@@ -23,6 +23,7 @@ import (
 	"github.com/pingcap/tidb/pkg/ddl/placement"
 	"github.com/pingcap/tidb/pkg/domain/infosync"
 	"github.com/pingcap/tidb/pkg/infoschema"
+	infoschemacontext "github.com/pingcap/tidb/pkg/infoschema/context"
 	"github.com/pingcap/tidb/pkg/meta"
 	"github.com/pingcap/tidb/pkg/meta/model"
 	"github.com/pingcap/tidb/pkg/parser/ast"
@@ -370,7 +371,7 @@ func CheckPlacementPolicyNotInUseFromInfoSchema(is infoschema.InfoSchema, policy
 		}
 	}
 
-	schemaTables := is.ListTablesWithSpecialAttribute(infoschema.AllPlacementPolicyAttribute)
+	schemaTables := is.ListTablesWithSpecialAttribute(infoschemacontext.AllPlacementPolicyAttribute)
 	for _, schemaTable := range schemaTables {
 		for _, tblInfo := range schemaTable.TableInfos {
 			if err := checkPlacementPolicyNotUsedByTable(tblInfo, policy); err != nil {

pkg/executor/infoschema_reader.go

@@ -40,6 +40,7 @@ import (
 	"github.com/pingcap/tidb/pkg/executor/internal/pdhelper"
 	"github.com/pingcap/tidb/pkg/expression"
 	"github.com/pingcap/tidb/pkg/infoschema"
+	infoschemacontext "github.com/pingcap/tidb/pkg/infoschema/context"
 	"github.com/pingcap/tidb/pkg/keyspace"
 	"github.com/pingcap/tidb/pkg/kv"
 	"github.com/pingcap/tidb/pkg/meta/autoid"
@@ -2657,7 +2658,7 @@ func (e *memtableRetriever) dataForTableTiFlashReplica(_ context.Context, sctx s
 		rows          [][]types.Datum
 		tiFlashStores map[int64]pd.StoreInfo
 	)
-	rs := e.is.ListTablesWithSpecialAttribute(infoschema.TiFlashAttribute)
+	rs := e.is.ListTablesWithSpecialAttribute(infoschemacontext.TiFlashAttribute)
 	for _, schema := range rs {
 		for _, tbl := range schema.TableInfos {
 			if checker != nil && !checker.RequestVerification(sctx.GetSessionVars().ActiveRoles, schema.DBName.L, tbl.Name.L, "", mysql.AllPrivMask) {

pkg/executor/join/join_row_table_test.go

@@ -44,4 +44,4 @@ func TestBitMaskInUint32(t *testing.T) {
 
 func TestUintptrCanHoldPointer(t *testing.T) {
 	require.Equal(t, true, sizeOfUintptr >= sizeOfUnsafePointer)
-}
+}

pkg/infoschema/builder_misc.go

@@ -18,6 +18,7 @@ import (
 	"fmt"
 
 	"github.com/pingcap/errors"
+	infoschemacontext "github.com/pingcap/tidb/pkg/infoschema/context"
 	"github.com/pingcap/tidb/pkg/meta"
 	"github.com/pingcap/tidb/pkg/meta/model"
 )
@@ -115,7 +116,7 @@ func (b *Builder) initMisc(dbInfos []*model.DBInfo, policies []*model.PolicyInfo
 
 	// Maintain foreign key reference information.
 	if b.enableV2 {
-		rs := b.ListTablesWithSpecialAttribute(ForeignKeysAttribute)
+		rs := b.ListTablesWithSpecialAttribute(infoschemacontext.ForeignKeysAttribute)
 		for _, db := range rs {
 			for _, tbl := range db.TableInfos {
 				info.addReferredForeignKeys(db.DBName, tbl)

pkg/infoschema/context/infoschema.go

@@ -25,6 +25,70 @@ import (
 // SpecialAttributeFilter is used to filter tables with special attributes.
 type SpecialAttributeFilter func(*model.TableInfo) bool
 
+// TTLAttribute is the TTL attribute filter used by ListTablesWithSpecialAttribute.
+var TTLAttribute SpecialAttributeFilter = func(t *model.TableInfo) bool {
+	return t.State == model.StatePublic && t.TTLInfo != nil
+}
+
+// TiFlashAttribute is the TiFlashReplica attribute filter used by ListTablesWithSpecialAttribute.
+var TiFlashAttribute SpecialAttributeFilter = func(t *model.TableInfo) bool {
+	return t.TiFlashReplica != nil
+}
+
+// PlacementPolicyAttribute is the Placement Policy attribute filter used by ListTablesWithSpecialAttribute.
+var PlacementPolicyAttribute SpecialAttributeFilter = func(t *model.TableInfo) bool {
+	if t.PlacementPolicyRef != nil {
+		return true
+	}
+	if parInfo := t.GetPartitionInfo(); parInfo != nil {
+		for _, def := range parInfo.Definitions {
+			if def.PlacementPolicyRef != nil {
+				return true
+			}
+		}
+	}
+	return false
+}
+
+// AllPlacementPolicyAttribute is the Placement Policy attribute filter used by ListTablesWithSpecialAttribute.
+// Different from PlacementPolicyAttribute, Partition.Enable flag will be ignored.
+var AllPlacementPolicyAttribute SpecialAttributeFilter = func(t *model.TableInfo) bool {
+	if t.PlacementPolicyRef != nil {
+		return true
+	}
+	if t.Partition != nil {
+		for _, def := range t.Partition.Definitions {
+			if def.PlacementPolicyRef != nil {
+				return true
+			}
+		}
+	}
+	return false
+}
+
+// TableLockAttribute is the Table Lock attribute filter used by ListTablesWithSpecialAttribute.
+var TableLockAttribute SpecialAttributeFilter = func(t *model.TableInfo) bool {
+	return t.Lock != nil
+}
+
+// ForeignKeysAttribute is the ForeignKeys attribute filter used by ListTablesWithSpecialAttribute.
+var ForeignKeysAttribute SpecialAttributeFilter = func(t *model.TableInfo) bool {
+	return len(t.ForeignKeys) > 0
+}
+
+// PartitionAttribute is the Partition attribute filter used by ListTablesWithSpecialAttribute.
+var PartitionAttribute SpecialAttributeFilter = func(t *model.TableInfo) bool {
+	return t.GetPartitionInfo() != nil
+}
+
+// HasSpecialAttributes checks if a table has any special attributes.
+func HasSpecialAttributes(t *model.TableInfo) bool {
+	return TTLAttribute(t) || TiFlashAttribute(t) || PlacementPolicyAttribute(t) || PartitionAttribute(t) || TableLockAttribute(t) || ForeignKeysAttribute(t)
+}
+
+// AllSpecialAttribute marks a model.TableInfo with any special attributes.
+var AllSpecialAttribute SpecialAttributeFilter = HasSpecialAttributes
+
 // TableInfoResult is used to store the result of ListTablesWithSpecialAttribute.
 type TableInfoResult struct {
 	DBName     pmodel.CIStr

pkg/infoschema/infoschema_v2.go

@@ -214,7 +214,7 @@ func (isd *Data) add(item tableItem, tbl table.Table) {
 			isd.pid2tid.Set(partitionItem{def.ID, item.schemaVersion, tbl.Meta().ID, false})
 		}
 	}
-	if hasSpecialAttributes(ti) {
+	if infoschemacontext.HasSpecialAttributes(ti) {
 		isd.tableInfoResident.Set(tableInfoItem{
 			dbName:        item.dbName,
 			tableID:       item.tableID,
@@ -1379,7 +1379,7 @@ func (b *bundleInfoBuilder) updateInfoSchemaBundlesV2(is *infoschemaV2) {
 
 	// do full update bundles
 	is.ruleBundleMap = make(map[int64]*placement.Bundle)
-	tmp := is.ListTablesWithSpecialAttribute(PlacementPolicyAttribute)
+	tmp := is.ListTablesWithSpecialAttribute(infoschemacontext.PlacementPolicyAttribute)
 	for _, v := range tmp {
 		for _, tbl := range v.TableInfos {
 			b.updateTableBundles(is, tbl.ID)
@@ -1392,7 +1392,7 @@ func (b *bundleInfoBuilder) completeUpdateTablesV2(is *infoschemaV2) {
 		return
 	}
 
-	dbs := is.ListTablesWithSpecialAttribute(AllSpecialAttribute)
+	dbs := is.ListTablesWithSpecialAttribute(infoschemacontext.AllSpecialAttribute)
 	for _, db := range dbs {
 		for _, tbl := range db.TableInfos {
 			tblInfo := tbl
@@ -1413,69 +1413,6 @@ func (b *bundleInfoBuilder) completeUpdateTablesV2(is *infoschemaV2) {
 	}
 }
 
-// TTLAttribute is the TTL attribute filter used by ListTablesWithSpecialAttribute.
-var TTLAttribute infoschemacontext.SpecialAttributeFilter = func(t *model.TableInfo) bool {
-	return t.State == model.StatePublic && t.TTLInfo != nil
-}
-
-// TiFlashAttribute is the TiFlashReplica attribute filter used by ListTablesWithSpecialAttribute.
-var TiFlashAttribute infoschemacontext.SpecialAttributeFilter = func(t *model.TableInfo) bool {
-	return t.TiFlashReplica != nil
-}
-
-// PlacementPolicyAttribute is the Placement Policy attribute filter used by ListTablesWithSpecialAttribute.
-var PlacementPolicyAttribute infoschemacontext.SpecialAttributeFilter = func(t *model.TableInfo) bool {
-	if t.PlacementPolicyRef != nil {
-		return true
-	}
-	if parInfo := t.GetPartitionInfo(); parInfo != nil {
-		for _, def := range parInfo.Definitions {
-			if def.PlacementPolicyRef != nil {
-				return true
-			}
-		}
-	}
-	return false
-}
-
-// AllPlacementPolicyAttribute is the Placement Policy attribute filter used by ListTablesWithSpecialAttribute.
-// Different from PlacementPolicyAttribute, Partition.Enable flag will be ignored.
-var AllPlacementPolicyAttribute infoschemacontext.SpecialAttributeFilter = func(t *model.TableInfo) bool {
-	if t.PlacementPolicyRef != nil {
-		return true
-	}
-	if t.Partition != nil {
-		for _, def := range t.Partition.Definitions {
-			if def.PlacementPolicyRef != nil {
-				return true
-			}
-		}
-	}
-	return false
-}
-
-// TableLockAttribute is the Table Lock attribute filter used by ListTablesWithSpecialAttribute.
-var TableLockAttribute infoschemacontext.SpecialAttributeFilter = func(t *model.TableInfo) bool {
-	return t.Lock != nil
-}
-
-// ForeignKeysAttribute is the ForeignKeys attribute filter used by ListTablesWithSpecialAttribute.
-var ForeignKeysAttribute infoschemacontext.SpecialAttributeFilter = func(t *model.TableInfo) bool {
-	return len(t.ForeignKeys) > 0
-}
-
-// PartitionAttribute is the Partition attribute filter used by ListTablesWithSpecialAttribute.
-var PartitionAttribute infoschemacontext.SpecialAttributeFilter = func(t *model.TableInfo) bool {
-	return t.GetPartitionInfo() != nil
-}
-
-func hasSpecialAttributes(t *model.TableInfo) bool {
-	return TTLAttribute(t) || TiFlashAttribute(t) || PlacementPolicyAttribute(t) || PartitionAttribute(t) || TableLockAttribute(t) || ForeignKeysAttribute(t)
-}
-
-// AllSpecialAttribute marks a model.TableInfo with any special attributes.
-var AllSpecialAttribute infoschemacontext.SpecialAttributeFilter = hasSpecialAttributes
-
 func (is *infoschemaV2) ListTablesWithSpecialAttribute(filter infoschemacontext.SpecialAttributeFilter) []infoschemacontext.TableInfoResult {
 	ret := make([]infoschemacontext.TableInfoResult, 0, 10)
 	var currDB string

pkg/infoschema/infoschema_v2_test.go

@@ -525,44 +525,44 @@ func TestSpecialAttributeCorrectnessInSchemaChange(t *testing.T) {
 		DDLState: model.StatePublic,
 	}
 	// add partition
-	tblInfo1 := updateTableSpecialAttribute(t, dbInfo, tblInfo, builder, r, model.ActionAddTablePartition, 3, PartitionAttribute, true)
+	tblInfo1 := updateTableSpecialAttribute(t, dbInfo, tblInfo, builder, r, model.ActionAddTablePartition, 3, infoschemacontext.PartitionAttribute, true)
 	require.Equal(t, tblInfo.Partition, tblInfo1.Partition)
 	// drop partition
 	tblInfo.Partition.Definitions = tblInfo.Partition.Definitions[:1]
-	tblInfo1 = updateTableSpecialAttribute(t, dbInfo, tblInfo, builder, r, model.ActionDropTablePartition, 4, PartitionAttribute, true)
+	tblInfo1 = updateTableSpecialAttribute(t, dbInfo, tblInfo, builder, r, model.ActionDropTablePartition, 4, infoschemacontext.PartitionAttribute, true)
 	require.Equal(t, tblInfo.Partition, tblInfo1.Partition)
 
 	// test placement policy correctness in schema change
 	tblInfo.PlacementPolicyRef = &model.PolicyRefInfo{
 		ID:   1,
 		Name: pmodel.NewCIStr("p3"),
 	}
-	tblInfo1 = updateTableSpecialAttribute(t, dbInfo, tblInfo, builder, r, model.ActionAlterTablePlacement, 5, PlacementPolicyAttribute, true)
+	tblInfo1 = updateTableSpecialAttribute(t, dbInfo, tblInfo, builder, r, model.ActionAlterTablePlacement, 5, infoschemacontext.PlacementPolicyAttribute, true)
 	require.Equal(t, tblInfo.PlacementPolicyRef, tblInfo1.PlacementPolicyRef)
 	tblInfo.PlacementPolicyRef = nil
-	updateTableSpecialAttribute(t, dbInfo, tblInfo, builder, r, model.ActionAlterTablePlacement, 6, PlacementPolicyAttribute, false)
+	updateTableSpecialAttribute(t, dbInfo, tblInfo, builder, r, model.ActionAlterTablePlacement, 6, infoschemacontext.PlacementPolicyAttribute, false)
 
 	// test tiflash replica correctness in schema change
 	tblInfo.TiFlashReplica = &model.TiFlashReplicaInfo{
 		Count:          1,
 		Available:      true,
 		LocationLabels: []string{"zone"},
 	}
-	tblInfo1 = updateTableSpecialAttribute(t, dbInfo, tblInfo, builder, r, model.ActionSetTiFlashReplica, 7, TiFlashAttribute, true)
+	tblInfo1 = updateTableSpecialAttribute(t, dbInfo, tblInfo, builder, r, model.ActionSetTiFlashReplica, 7, infoschemacontext.TiFlashAttribute, true)
 	require.Equal(t, tblInfo.TiFlashReplica, tblInfo1.TiFlashReplica)
 	tblInfo.TiFlashReplica = nil
-	updateTableSpecialAttribute(t, dbInfo, tblInfo, builder, r, model.ActionSetTiFlashReplica, 8, TiFlashAttribute, false)
+	updateTableSpecialAttribute(t, dbInfo, tblInfo, builder, r, model.ActionSetTiFlashReplica, 8, infoschemacontext.TiFlashAttribute, false)
 
 	// test table lock correctness in schema change
 	tblInfo.Lock = &model.TableLockInfo{
 		Tp:    pmodel.TableLockRead,
 		State: model.TableLockStatePublic,
 		TS:    1,
 	}
-	tblInfo1 = updateTableSpecialAttribute(t, dbInfo, tblInfo, builder, r, model.ActionLockTable, 9, TableLockAttribute, true)
+	tblInfo1 = updateTableSpecialAttribute(t, dbInfo, tblInfo, builder, r, model.ActionLockTable, 9, infoschemacontext.TableLockAttribute, true)
 	require.Equal(t, tblInfo.Lock, tblInfo1.Lock)
 	tblInfo.Lock = nil
-	updateTableSpecialAttribute(t, dbInfo, tblInfo, builder, r, model.ActionUnlockTable, 10, TableLockAttribute, false)
+	updateTableSpecialAttribute(t, dbInfo, tblInfo, builder, r, model.ActionUnlockTable, 10, infoschemacontext.TableLockAttribute, false)
 
 	// test foreign key correctness in schema change
 	tblInfo.ForeignKeys = []*model.FKInfo{{
@@ -574,10 +574,10 @@ func TestSpecialAttributeCorrectnessInSchemaChange(t *testing.T) {
 		Cols:      []pmodel.CIStr{pmodel.NewCIStr("t_a")},
 		State:     model.StateWriteOnly,
 	}}
-	tblInfo1 = updateTableSpecialAttribute(t, dbInfo, tblInfo, builder, r, model.ActionAddForeignKey, 11, ForeignKeysAttribute, true)
+	tblInfo1 = updateTableSpecialAttribute(t, dbInfo, tblInfo, builder, r, model.ActionAddForeignKey, 11, infoschemacontext.ForeignKeysAttribute, true)
 	require.Equal(t, tblInfo.ForeignKeys, tblInfo1.ForeignKeys)
 	tblInfo.ForeignKeys = nil
-	updateTableSpecialAttribute(t, dbInfo, tblInfo, builder, r, model.ActionDropForeignKey, 12, ForeignKeysAttribute, false)
+	updateTableSpecialAttribute(t, dbInfo, tblInfo, builder, r, model.ActionDropForeignKey, 12, infoschemacontext.ForeignKeysAttribute, false)
 }
 
 func TestDataStructFieldsCorrectnessInSchemaChange(t *testing.T) {
@@ -636,7 +636,7 @@ func TestDataStructFieldsCorrectnessInSchemaChange(t *testing.T) {
 		Enable:   true,
 		DDLState: model.StatePublic,
 	}
-	tblInfo1 := updateTableSpecialAttribute(t, dbInfo, tblInfo, builder, r, model.ActionAddTablePartition, 3, PartitionAttribute, true)
+	tblInfo1 := updateTableSpecialAttribute(t, dbInfo, tblInfo, builder, r, model.ActionAddTablePartition, 3, infoschemacontext.PartitionAttribute, true)
 	require.Equal(t, tblInfo.Partition, tblInfo1.Partition)
 	require.Equal(t, v2.Data.pid2tid.Len(), 2)
 	tblInfoItem, ok := v2.Data.pid2tid.Get(partitionItem{partitionID: 2, schemaVersion: 3})
@@ -645,7 +645,7 @@ func TestDataStructFieldsCorrectnessInSchemaChange(t *testing.T) {
 
 	// verify partition related fields drop partition
 	tblInfo.Partition.Definitions = tblInfo.Partition.Definitions[:1]
-	tblInfo1 = updateTableSpecialAttribute(t, dbInfo, tblInfo, builder, r, model.ActionDropTablePartition, 4, PartitionAttribute, true)
+	tblInfo1 = updateTableSpecialAttribute(t, dbInfo, tblInfo, builder, r, model.ActionDropTablePartition, 4, infoschemacontext.PartitionAttribute, true)
 	require.Equal(t, tblInfo.Partition, tblInfo1.Partition)
 	require.Equal(t, v2.Data.pid2tid.Len(), 4)
 	tblInfoItem, ok = v2.Data.pid2tid.Get(partitionItem{partitionID: 1, schemaVersion: 4})

pkg/infoschema/test/infoschemav2test/BUILD.bazel

@@ -13,6 +13,7 @@ go_test(
         "//pkg/domain",
         "//pkg/domain/infosync",
         "//pkg/infoschema",
+        "//pkg/infoschema/context",
         "//pkg/parser/auth",
         "//pkg/parser/model",
         "//pkg/sessionctx/variable",