infoschema: add the SchemaNameByTableID API

Signed-off-by: Rustin170506 <[email protected]>

pkg/infoschema/infoschema.go

@@ -308,6 +308,26 @@ 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) {
+	if !tableIDIsValid(tableID) {
+		return
+	}
+
+	slice := is.sortedTablesBuckets[tableBucketIdx(tableID)]
+	idx := slice.searchTable(tableID)
+	if idx == -1 {
+		return
+	}
+
+	dbID := slice[idx].Meta().DBID
+	db, ok := is.SchemaByID(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)

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) {

pkg/infoschema/infoschema_v2.go

@@ -646,6 +646,25 @@ 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
+	}
+
+	eq := func(a, b *tableItem) bool { return a.tableID == b.tableID }
+	itm, ok := search(
+		is.byID,
+		is.infoSchema.schemaMetaVersion,
+		tableItem{tableID: tableID, schemaVersion: math.MaxInt64},
+		eq,
+	)
+	if !ok {
+		return
+	}
+
+	return itm.dbName, true
+}
+
 // TableItem is exported from tableItem.
 type TableItem struct {
 	DBName    pmodel.CIStr

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.
 }
 

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