infoschema: add the SchemaNameByTableID API (#56156)

ref #55906
pingcap · Sep 25, 2024 · 8c12f12 · 8c12f12
1 parent 918b58d
commit 8c12f12
Show file tree

Hide file tree

Showing 5 changed files with 110 additions and 2 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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