diff --git a/pkg/infoschema/infoschema.go b/pkg/infoschema/infoschema.go index 2f7b611c05f13..64e534ed55a3b 100644 --- a/pkg/infoschema/infoschema.go +++ b/pkg/infoschema/infoschema.go @@ -308,6 +308,19 @@ func (is *infoSchema) TableByID(_ stdctx.Context, id int64) (val table.Table, ok return slice[idx], true } +func (is *infoSchema) SchemaNameByTableID(tableID int64) (schemaName pmodel.CIStr, ok bool) { + tbl, ok := is.TableByID(stdctx.Background(), tableID) + if !ok { + return + } + db, ok := is.SchemaByID(tbl.Meta().DBID) + if !ok { + return + } + + return db.Name, true +} + // TableInfoByID implements InfoSchema.TableInfoByID func (is *infoSchema) TableInfoByID(id int64) (*model.TableInfo, bool) { tbl, ok := is.TableByID(stdctx.Background(), id) diff --git a/pkg/infoschema/infoschema_test.go b/pkg/infoschema/infoschema_test.go index d79af4fc72e1e..88d570940537c 100644 --- a/pkg/infoschema/infoschema_test.go +++ b/pkg/infoschema/infoschema_test.go @@ -244,6 +244,41 @@ func TestBasic(t *testing.T) { tbls, err := is.SchemaTableInfos(context.Background(), schema.Name) require.NoError(t, err) require.Equal(t, 1, len(tbls)) + + // Test SchemaNameByTableID + tests := []struct { + name string + tableID int64 + wantSchema pmodel.CIStr + wantOK bool + }{ + { + name: "valid table ID", + tableID: tbID, + wantSchema: dbName, + wantOK: true, + }, + { + name: "non-existent table ID", + tableID: tbID + 1, + wantSchema: pmodel.CIStr{}, + wantOK: false, + }, + { + name: "invalid table ID (negative)", + tableID: -1, + wantSchema: pmodel.CIStr{}, + wantOK: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + gotSchema, gotOK := is.SchemaNameByTableID(tt.tableID) + require.Equal(t, tt.wantOK, gotOK) + require.Equal(t, tt.wantSchema, gotSchema) + }) + } } func TestMockInfoSchema(t *testing.T) { diff --git a/pkg/infoschema/infoschema_v2.go b/pkg/infoschema/infoschema_v2.go index 88ba5fc7d88cc..1d4f2fd3caff0 100644 --- a/pkg/infoschema/infoschema_v2.go +++ b/pkg/infoschema/infoschema_v2.go @@ -589,6 +589,16 @@ func (is *infoschemaV2) CloneAndUpdateTS(startTS uint64) *infoschemaV2 { return &tmp } +func (is *infoschemaV2) searchTableItemByID(tableID int64) (tableItem, bool) { + eq := func(a, b *tableItem) bool { return a.tableID == b.tableID } + return search( + is.byID, + is.infoSchema.schemaMetaVersion, + tableItem{tableID: tableID, schemaVersion: math.MaxInt64}, + eq, + ) +} + // TableByID implements the InfoSchema interface. // As opposed to TableByName, TableByID will not refill cache when schema cache miss, // unless the caller changes the behavior by passing a context use WithRefillOption. @@ -604,8 +614,7 @@ func (is *infoschemaV2) TableByID(ctx context.Context, id int64) (val table.Tabl return tbl, true } - eq := func(a, b *tableItem) bool { return a.tableID == b.tableID } - itm, ok := search(is.byID, is.infoSchema.schemaMetaVersion, tableItem{tableID: id, schemaVersion: math.MaxInt64}, eq) + itm, ok := is.searchTableItemByID(id) if !ok { return nil, false } @@ -646,6 +655,19 @@ func (is *infoschemaV2) TableByID(ctx context.Context, id int64) (val table.Tabl return ret, true } +func (is *infoschemaV2) SchemaNameByTableID(tableID int64) (schemaName pmodel.CIStr, ok bool) { + if !tableIDIsValid(tableID) { + return + } + + itm, ok := is.searchTableItemByID(tableID) + if !ok { + return + } + + return itm.dbName, true +} + // TableItem is exported from tableItem. type TableItem struct { DBName pmodel.CIStr diff --git a/pkg/infoschema/infoschema_v2_test.go b/pkg/infoschema/infoschema_v2_test.go index 8ba6faeff1b70..f6ad9770691c2 100644 --- a/pkg/infoschema/infoschema_v2_test.go +++ b/pkg/infoschema/infoschema_v2_test.go @@ -121,6 +121,43 @@ func TestV2Basic(t *testing.T) { require.Equal(t, 0, len(tblInfos)) require.Equal(t, int64(2), is.SchemaMetaVersion()) + + // Test SchemaNameByTableID + schemaNameByTableIDTests := []struct { + name string + tableID int64 + wantSchema pmodel.CIStr + wantOK bool + }{ + { + name: "valid table ID", + tableID: tblInfo.ID, + wantSchema: schemaName, + wantOK: true, + }, + { + name: "non-existent table ID", + tableID: tblInfo.ID + 1, + wantSchema: pmodel.CIStr{}, + wantOK: false, + }, + { + name: "invalid table ID (negative)", + tableID: -1, + wantSchema: pmodel.CIStr{}, + wantOK: false, + }, + } + + for _, tt := range schemaNameByTableIDTests { + t.Run(tt.name, func(t *testing.T) { + gotSchema, gotOK := is.SchemaNameByTableID(tt.tableID) + + require.Equal(t, tt.wantOK, gotOK) + require.Equal(t, tt.wantSchema, gotSchema) + }) + } + // TODO: support FindTableByPartitionID. } diff --git a/pkg/infoschema/interface.go b/pkg/infoschema/interface.go index e6922e5e1ac53..da6f15f34956b 100644 --- a/pkg/infoschema/interface.go +++ b/pkg/infoschema/interface.go @@ -30,6 +30,7 @@ type InfoSchema interface { context.MetaOnlyInfoSchema TableByName(ctx stdctx.Context, schema, table pmodel.CIStr) (table.Table, error) TableByID(ctx stdctx.Context, id int64) (table.Table, bool) + SchemaNameByTableID(tableID int64) (pmodel.CIStr, bool) FindTableByPartitionID(partitionID int64) (table.Table, *model.DBInfo, *model.PartitionDefinition) ListTablesWithSpecialAttribute(filter specialAttributeFilter) []tableInfoResult base() *infoSchema