diff --git a/pkg/infoschema/infoschema_v2.go b/pkg/infoschema/infoschema_v2.go index 3a351191d0f62..8967389bb7aa4 100644 --- a/pkg/infoschema/infoschema_v2.go +++ b/pkg/infoschema/infoschema_v2.go @@ -97,6 +97,17 @@ type Data struct { // For information_schema/metrics_schema/performance_schema etc specials map[string]*schemaTables + + // pid2tid is used by FindTableInfoByPartitionID, it stores {partitionID, schemaVersion} => table ID + // Need full data in memory! + pid2tid *btree.BTreeG[partitionItem] +} + +type partitionItem struct { + partitionID int64 + schemaVersion int64 + tableID int64 + tomb bool } func (isd *Data) getVersionByTS(ts uint64) (int64, bool) { @@ -148,6 +159,7 @@ func NewData() *Data { // TODO: limit by size instead of by table count. tableCache: sieve.New[tableCacheKey, table.Table](1000), specials: make(map[string]*schemaTables), + pid2tid: btree.NewBTreeG[partitionItem](comparePartitionItem), } return ret } @@ -156,6 +168,11 @@ func (isd *Data) add(item tableItem, tbl table.Table) { isd.byID.Set(item) isd.byName.Set(item) isd.tableCache.Set(tableCacheKey{item.tableID, item.schemaVersion}, tbl) + if pi := tbl.Meta().GetPartitionInfo(); pi != nil { + for _, def := range pi.Definitions { + isd.pid2tid.Set(partitionItem{def.ID, item.schemaVersion, tbl.Meta().ID, false}) + } + } } func (isd *Data) addSpecialDB(di *model.DBInfo, tables *schemaTables) { @@ -207,6 +224,16 @@ func compareByName(a, b tableItem) bool { return a.schemaVersion < b.schemaVersion } +func comparePartitionItem(a, b partitionItem) bool { + if a.partitionID < b.partitionID { + return true + } + if a.partitionID > b.partitionID { + return false + } + return a.schemaVersion < b.schemaVersion +} + func compareSchemaItem(a, b schemaItem) bool { if a.Name() < b.Name() { return true @@ -464,23 +491,52 @@ func (is *infoschemaV2) SchemaExists(schema model.CIStr) bool { } func (is *infoschemaV2) FindTableByPartitionID(partitionID int64) (table.Table, *model.DBInfo, *model.PartitionDefinition) { - // TODO: This is quite inefficient! we need some better way or avoid this API. - dbInfos := is.AllSchemas() - for _, dbInfo := range dbInfos { - tbls := is.SchemaTables(dbInfo.Name) - for _, tbl := range tbls { - pi := tbl.Meta().GetPartitionInfo() - if pi == nil { - continue + var ok bool + var pi partitionItem + is.pid2tid.Descend(partitionItem{partitionID: partitionID, schemaVersion: math.MaxInt64}, + func(item partitionItem) bool { + if item.partitionID != partitionID { + return false } - for _, p := range pi.Definitions { - if p.ID == partitionID { - return tbl, dbInfo, &p - } + if item.schemaVersion > is.infoSchema.schemaMetaVersion { + // Skip the record. + return true } + if item.schemaVersion <= is.infoSchema.schemaMetaVersion { + ok = !item.tomb + pi = item + return false + } + return true + }) + if !ok { + return nil, nil, nil + } + + tbl, ok := is.TableByID(pi.tableID) + if !ok { + // something wrong? + return nil, nil, nil + } + + dbID := tbl.Meta().DBID + dbInfo, ok := is.SchemaByID(dbID) + if !ok { + // something wrong? + return nil, nil, nil + } + + partInfo := tbl.Meta().GetPartitionInfo() + var def *model.PartitionDefinition + for i := 0; i < len(partInfo.Definitions); i++ { + pdef := &partInfo.Definitions[i] + if pdef.ID == partitionID { + def = pdef + break } } - return nil, nil, nil + + return tbl, dbInfo, def } func (is *infoschemaV2) TableExists(schema, table model.CIStr) bool { @@ -777,6 +833,12 @@ func (b *Builder) applyDropTableV2(diff *model.SchemaDiff, dbInfo *model.DBInfo, // The old DBInfo still holds a reference to old table info, we need to remove it. b.infoSchema.deleteReferredForeignKeys(dbInfo.Name, table.Meta()) + if pi := table.Meta().GetPartitionInfo(); pi != nil { + for _, def := range pi.Definitions { + b.infoData.pid2tid.Set(partitionItem{def.ID, diff.Version, table.Meta().ID, true}) + } + } + b.infoData.remove(tableItem{ dbName: dbInfo.Name.L, dbID: dbInfo.ID, diff --git a/pkg/infoschema/test/infoschemav2test/BUILD.bazel b/pkg/infoschema/test/infoschemav2test/BUILD.bazel index 3dbafe928efe2..c1ca8337d21cc 100644 --- a/pkg/infoschema/test/infoschemav2test/BUILD.bazel +++ b/pkg/infoschema/test/infoschemav2test/BUILD.bazel @@ -12,6 +12,8 @@ go_test( "//pkg/domain", "//pkg/infoschema", "//pkg/parser/auth", + "//pkg/parser/model", + "//pkg/table", "//pkg/testkit", "//pkg/testkit/testsetup", "@com_github_stretchr_testify//require", diff --git a/pkg/infoschema/test/infoschemav2test/v2_test.go b/pkg/infoschema/test/infoschemav2test/v2_test.go index 2d76796890147..936f948cc29e0 100644 --- a/pkg/infoschema/test/infoschemav2test/v2_test.go +++ b/pkg/infoschema/test/infoschemav2test/v2_test.go @@ -20,6 +20,8 @@ import ( "github.com/pingcap/tidb/pkg/domain" "github.com/pingcap/tidb/pkg/infoschema" "github.com/pingcap/tidb/pkg/parser/auth" + "github.com/pingcap/tidb/pkg/parser/model" + "github.com/pingcap/tidb/pkg/table" "github.com/pingcap/tidb/pkg/testkit" "github.com/stretchr/testify/require" ) @@ -64,3 +66,83 @@ func TestSpecialSchemas(t *testing.T) { tk.MustExec("set @@global.tidb_schema_cache_size = default;") } + +func checkPIDNotExist(t *testing.T, dom *domain.Domain, pid int64) { + is := dom.InfoSchema() + ptbl, dbInfo, pdef := is.FindTableByPartitionID(pid) + require.Nil(t, ptbl) + require.Nil(t, dbInfo) + require.Nil(t, pdef) +} + +func getPIDForP3(t *testing.T, dom *domain.Domain) (int64, table.Table) { + is := dom.InfoSchema() + tbl, err := is.TableByName(model.NewCIStr("test"), model.NewCIStr("pt")) + require.NoError(t, err) + pi := tbl.Meta().GetPartitionInfo() + pid := pi.GetPartitionIDByName("p3") + ptbl, _, _ := is.FindTableByPartitionID(pid) + require.Equal(t, ptbl.Meta().ID, tbl.Meta().ID) + return pid, tbl +} + +func TestFindTableByPartitionID(t *testing.T) { + store, dom := testkit.CreateMockStoreAndDomain(t) + tk := testkit.NewTestKit(t, store) + tk.MustExec("use test") + tk.MustExec(`create table pt (id int) partition by range (id) ( +partition p0 values less than (10), +partition p1 values less than (20), +partition p2 values less than (30), +partition p3 values less than (40))`) + + pid, tbl := getPIDForP3(t, dom) + is := dom.InfoSchema() + tbl1, dbInfo, pdef := is.FindTableByPartitionID(pid) + require.Equal(t, tbl1.Meta().ID, tbl.Meta().ID) + require.Equal(t, dbInfo.Name.L, "test") + require.Equal(t, pdef.ID, pid) + + // Test FindTableByPartitionID after dropping a unrelated partition. + tk.MustExec("alter table pt drop partition p2") + is = dom.InfoSchema() + tbl2, dbInfo, pdef := is.FindTableByPartitionID(pid) + require.Equal(t, tbl2.Meta().ID, tbl.Meta().ID) + require.Equal(t, dbInfo.Name.L, "test") + require.Equal(t, pdef.ID, pid) + + // Test FindTableByPartitionID after dropping that partition. + tk.MustExec("alter table pt drop partition p3") + checkPIDNotExist(t, dom, pid) + + // Test FindTableByPartitionID after adding back the partition. + tk.MustExec("alter table pt add partition (partition p3 values less than (35))") + checkPIDNotExist(t, dom, pid) + pid, _ = getPIDForP3(t, dom) + + // Test FindTableByPartitionID after truncate partition. + tk.MustExec("alter table pt truncate partition p3") + checkPIDNotExist(t, dom, pid) + pid, _ = getPIDForP3(t, dom) + + // Test FindTableByPartitionID after reorganize partition. + tk.MustExec(`alter table pt reorganize partition p1,p3 INTO ( +PARTITION p3 VALUES LESS THAN (1970), +PARTITION p5 VALUES LESS THAN (1980))`) + checkPIDNotExist(t, dom, pid) + _, _ = getPIDForP3(t, dom) + + // Test FindTableByPartitionID after exchange partition. + tk.MustExec("create table nt (id int)") + is = dom.InfoSchema() + ntbl, err := is.TableByName(model.NewCIStr("test"), model.NewCIStr("nt")) + require.NoError(t, err) + + tk.MustExec("alter table pt exchange partition p3 with table nt") + is = dom.InfoSchema() + ptbl, err := is.TableByName(model.NewCIStr("test"), model.NewCIStr("pt")) + require.NoError(t, err) + pi := ptbl.Meta().GetPartitionInfo() + pid = pi.GetPartitionIDByName("p3") + require.Equal(t, pid, ntbl.Meta().ID) +}