infoschema: fix panic for TableByID when table id is negative (#52016)

close #52012

br/pkg/restore/db_test.go

@@ -406,7 +406,7 @@ func TestGetExistedUserDBs(t *testing.T) {
 			{Name: model.NewCIStr("d1")},
 			{
 				Name:   model.NewCIStr("test"),
-				Tables: []*model.TableInfo{{Name: model.NewCIStr("t1"), State: model.StatePublic}},
+				Tables: []*model.TableInfo{{ID: 1, Name: model.NewCIStr("t1"), State: model.StatePublic}},
 				State:  model.StatePublic,
 			},
 		},

pkg/infoschema/builder.go

@@ -35,6 +35,7 @@ import (
 	"github.com/pingcap/tidb/pkg/table"
 	"github.com/pingcap/tidb/pkg/table/tables"
 	"github.com/pingcap/tidb/pkg/util/domainutil"
+	"github.com/pingcap/tidb/pkg/util/intest"
 )
 
 // Builder builds a new InfoSchema.
@@ -970,9 +971,10 @@ func NewBuilder(r autoid.Requirement, factory func() (pools.Resource, error), in
 }
 
 func tableBucketIdx(tableID int64) int {
+	intest.Assert(tableID > 0)
 	return int(tableID % bucketCount)
 }
 
 func tableIDIsValid(tableID int64) bool {
-	return tableID != 0
+	return tableID > 0
 }

pkg/infoschema/infoschema.go

@@ -104,7 +104,17 @@ func MockInfoSchema(tbList []*model.TableInfo) InfoSchema {
 		tables: make(map[string]table.Table),
 	}
 	result.schemaMap["test"] = tableNames
+	var tableIDs map[int64]struct{}
 	for _, tb := range tbList {
+		intest.AssertFunc(func() bool {
+			if tableIDs == nil {
+				tableIDs = make(map[int64]struct{})
+			}
+			_, ok := tableIDs[tb.ID]
+			intest.Assert(!ok)
+			tableIDs[tb.ID] = struct{}{}
+			return true
+		})
 		tb.DBID = dbInfo.ID
 		tbl := table.MockTableFromMeta(tb)
 		tableNames.tables[tb.Name.L] = tbl
@@ -238,6 +248,10 @@ func SchemaByTable(is InfoSchema, tableInfo *model.TableInfo) (val *model.DBInfo
 }
 
 func (is *infoSchema) TableByID(id int64) (val table.Table, ok bool) {
+	if !tableIDIsValid(id) {
+		return nil, false
+	}
+
 	slice := is.sortedTablesBuckets[tableBucketIdx(id)]
 	idx := slice.searchTable(id)
 	if idx == -1 {
@@ -713,6 +727,10 @@ func (ts *SessionExtendedInfoSchema) FindTableInfoByPartitionID(
 
 // TableByID implements InfoSchema.TableByID
 func (ts *SessionExtendedInfoSchema) TableByID(id int64) (table.Table, bool) {
+	if !tableIDIsValid(id) {
+		return nil, false
+	}
+
 	if ts.LocalTemporaryTables != nil {
 		if tbl, ok := ts.LocalTemporaryTables.TableByID(id); ok {
 			return tbl, true

pkg/infoschema/infoschema_test.go

@@ -172,6 +172,14 @@ func TestBasic(t *testing.T) {
 	require.False(t, ok)
 	require.Nil(t, gotTblInfo)
 
+	tb, ok = is.TableByID(-12345)
+	require.False(t, ok)
+	require.Nil(t, tb)
+
+	gotTblInfo, ok = is.TableInfoByID(-12345)
+	require.False(t, ok)
+	require.Nil(t, gotTblInfo)
+
 	tb, err = is.TableByName(dbName, tbName)
 	require.NoError(t, err)
 	require.NotNil(t, tb)
@@ -187,6 +195,17 @@ func TestBasic(t *testing.T) {
 	require.Error(t, err)
 	require.Nil(t, gotTblInfo)
 
+	// negative id should always be seen as not exists
+	tb, ok = is.TableByID(-1)
+	require.False(t, ok)
+	require.Nil(t, tb)
+	schema, ok = is.SchemaByID(-1)
+	require.False(t, ok)
+	require.Nil(t, schema)
+	gotTblInfo, ok = is.TableInfoByID(-1)
+	require.Nil(t, gotTblInfo)
+	require.False(t, ok)
+
 	tbs := is.SchemaTables(dbName)
 	require.Len(t, tbs, 1)
 
@@ -855,6 +874,14 @@ func TestLocalTemporaryTables(t *testing.T) {
 	info, ok = is.SchemaByID(tb22.Meta().DBID)
 	require.False(t, ok)
 	require.Nil(t, info)
+
+	// negative id should always be seen as not exists
+	tbl, ok = is.TableByID(-1)
+	require.False(t, ok)
+	require.Nil(t, tbl)
+	info, ok = is.SchemaByID(-1)
+	require.False(t, ok)
+	require.Nil(t, info)
 }
 
 // TestInfoSchemaCreateTableLike tests the table's column ID and index ID for memory database.

pkg/infoschema/infoschema_v2.go

@@ -278,6 +278,10 @@ func (is *infoschemaV2) base() *infoSchema {
 }
 
 func (is *infoschemaV2) TableByID(id int64) (val table.Table, ok bool) {
+	if !tableIDIsValid(id) {
+		return
+	}
+
 	// Get from the cache.
 	key := tableCacheKey{id, is.infoSchema.schemaMetaVersion}
 	tbl, found := is.tableCache.Get(key)

pkg/infoschema/infoschema_v2_test.go

@@ -81,6 +81,17 @@ func TestV2Basic(t *testing.T) {
 	require.True(t, ok)
 	require.Same(t, gotTblInfo, getTableInfo.Meta())
 
+	// negative id should always be seen as not exists
+	getTableInfo, ok = is.TableByID(-1)
+	require.False(t, ok)
+	require.Nil(t, getTableInfo)
+	gotTblInfo, ok = is.TableInfoByID(-1)
+	require.False(t, ok)
+	require.Nil(t, gotTblInfo)
+	getDBInfo, ok = is.SchemaByID(-1)
+	require.False(t, ok)
+	require.Nil(t, getDBInfo)
+
 	gotTblInfo, ok = is.TableInfoByID(1234567)
 	require.False(t, ok)
 	require.Nil(t, gotTblInfo)

pkg/planner/core/logical_plans_test.go

@@ -83,7 +83,7 @@ func createPlannerSuite() (s *plannerSuite) {
 		MockListPartitionTable(),
 		MockStateNoneColumnTable(),
 	}
-	id := int64(0)
+	id := int64(1)
 	for _, tblInfo := range tblInfos {
 		tblInfo.ID = id
 		id += 1

pkg/planner/core/mock.go

@@ -267,6 +267,7 @@ func MockSignedTable() *model.TableInfo {
 	col5.SetFlag(mysql.NotNullFlag)
 	col6.SetFlag(mysql.NoDefaultValueFlag)
 	table := &model.TableInfo{
+		ID:         1,
 		Columns:    []*model.ColumnInfo{pkColumn, col0, col1, col2, col3, colStr1, colStr2, colStr3, col4, col5, col6, col7},
 		Indices:    indices,
 		Name:       model.NewCIStr("t"),
@@ -336,6 +337,7 @@ func MockUnsignedTable() *model.TableInfo {
 	col0.SetFlag(mysql.NotNullFlag)
 	col1.SetFlag(mysql.UnsignedFlag)
 	table := &model.TableInfo{
+		ID:         2,
 		Columns:    []*model.ColumnInfo{pkColumn, col0, col1},
 		Indices:    indices,
 		Name:       model.NewCIStr("t2"),
@@ -365,6 +367,7 @@ func MockNoPKTable() *model.TableInfo {
 	col0.SetFlag(mysql.NotNullFlag)
 	col1.SetFlag(mysql.UnsignedFlag)
 	table := &model.TableInfo{
+		ID:         3,
 		Columns:    []*model.ColumnInfo{col0, col1},
 		Name:       model.NewCIStr("t3"),
 		PKIsHandle: true,
@@ -395,6 +398,7 @@ func MockView() *model.TableInfo {
 	}
 	view := &model.ViewInfo{SelectStmt: selectStmt, Security: model.SecurityDefiner, Definer: &auth.UserIdentity{Username: "root", Hostname: ""}, Cols: []model.CIStr{col0.Name, col1.Name, col2.Name}}
 	table := &model.TableInfo{
+		ID:      4,
 		Name:    model.NewCIStr("v"),
 		Columns: []*model.ColumnInfo{col0, col1, col2},
 		View:    view,
@@ -462,6 +466,7 @@ func MockRangePartitionTable() *model.TableInfo {
 		},
 	}
 	tableInfo := MockSignedTable()
+	tableInfo.ID = 5
 	tableInfo.Name = model.NewCIStr("pt1")
 	cols := make([]*model.ColumnInfo, 0, len(tableInfo.Columns))
 	cols = append(cols, tableInfo.Columns...)
@@ -497,6 +502,7 @@ func MockHashPartitionTable() *model.TableInfo {
 		},
 	}
 	tableInfo := MockSignedTable()
+	tableInfo.ID = 6
 	tableInfo.Name = model.NewCIStr("pt2")
 	cols := make([]*model.ColumnInfo, 0, len(tableInfo.Columns))
 	cols = append(cols, tableInfo.Columns...)
@@ -543,6 +549,7 @@ func MockListPartitionTable() *model.TableInfo {
 		},
 	}
 	tableInfo := MockSignedTable()
+	tableInfo.ID = 7
 	tableInfo.Name = model.NewCIStr("pt3")
 	cols := make([]*model.ColumnInfo, 0, len(tableInfo.Columns))
 	cols = append(cols, tableInfo.Columns...)
@@ -610,6 +617,7 @@ func MockStateNoneColumnTable() *model.TableInfo {
 	col0.SetFlag(mysql.NotNullFlag)
 	col1.SetFlag(mysql.UnsignedFlag)
 	table := &model.TableInfo{
+		ID:         8,
 		Columns:    []*model.ColumnInfo{pkColumn, col0, col1},
 		Indices:    indices,
 		Name:       model.NewCIStr("T_StateNoneColumn"),

pkg/ttl/ttlworker/job_manager_test.go

@@ -260,8 +260,8 @@ func TestReadyForLockHBTimeoutJobTables(t *testing.T) {
 			tables := m.readyForLockHBTimeoutJobTables(se.Now())
 			if c.shouldSchedule {
 				assert.Len(t, tables, 1)
-				assert.Equal(t, int64(0), tables[0].ID)
-				assert.Equal(t, int64(0), tables[0].TableInfo.ID)
+				assert.Equal(t, tbl.ID, tables[0].ID)
+				assert.Equal(t, tbl.ID, tables[0].TableInfo.ID)
 			} else {
 				assert.Len(t, tables, 0)
 			}

pkg/ttl/ttlworker/session_test.go

@@ -18,6 +18,7 @@ import (
 	"context"
 	"errors"
 	"strings"
+	"sync/atomic"
 	"testing"
 	"time"
 
@@ -36,8 +37,11 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
+var idAllocator atomic.Int64
+
 func newMockTTLTbl(t *testing.T, name string) *cache.PhysicalTable {
 	tblInfo := &model.TableInfo{
+		ID:   idAllocator.Add(1),
 		Name: model.NewCIStr(name),
 		Columns: []*model.ColumnInfo{
 			{