pkg/lightning: fix incorrect table counter (#996) (#1005)

pkg/lightning/restore/restore_test.go

@@ -1271,6 +1271,7 @@ type restoreSchemaSuite struct {
 	ctx        context.Context
 	rc         *Controller
 	controller *gomock.Controller
+	tableInfos []*model.TableInfo
 }
 
 func (s *restoreSchemaSuite) SetUpSuite(c *C) {
@@ -1286,14 +1287,28 @@ func (s *restoreSchemaSuite) SetUpSuite(c *C) {
 	c.Assert(err, IsNil)
 	// restore table schema files
 	fakeTableFilesCount := 8
+
+	p := parser.New()
+	p.SetSQLMode(mysql.ModeANSIQuotes)
+	se := tmock.NewContext()
+
+	tableInfos := make([]*model.TableInfo, 0, fakeTableFilesCount)
 	for i := 1; i <= fakeTableFilesCount; i++ {
 		fakeTableName := fmt.Sprintf("tbl%d", i)
 		// please follow the `mydump.defaultFileRouteRules`, matches files like '{schema}.{table}-schema.sql'
 		fakeFileName := fmt.Sprintf("%s.%s-schema.sql", fakeDBName, fakeTableName)
-		fakeFileContent := []byte(fmt.Sprintf("CREATE TABLE %s(i TINYINT);", fakeTableName))
-		err = store.WriteFile(ctx, fakeFileName, fakeFileContent)
+		fakeFileContent := fmt.Sprintf("CREATE TABLE %s(i TINYINT);", fakeTableName)
+		err = store.WriteFile(ctx, fakeFileName, []byte(fakeFileContent))
+		c.Assert(err, IsNil)
+
+		node, err := p.ParseOneStmt(fakeFileContent, "", "")
+		c.Assert(err, IsNil)
+		core, err := ddl.MockTableInfo(se, node.(*ast.CreateTableStmt), 0xabcdef)
 		c.Assert(err, IsNil)
+		core.State = model.StatePublic
+		tableInfos = append(tableInfos, core)
 	}
+	s.tableInfos = tableInfos
 	// restore view schema files
 	fakeViewFilesCount := 8
 	for i := 1; i <= fakeViewFilesCount; i++ {
@@ -1323,11 +1338,11 @@ func (s *restoreSchemaSuite) SetUpSuite(c *C) {
 func (s *restoreSchemaSuite) SetUpTest(c *C) {
 	s.controller, s.ctx = gomock.WithContext(context.Background(), c)
 	mockBackend := mock.NewMockBackend(s.controller)
-	// We don't care the execute results of those
 	mockBackend.EXPECT().
 		FetchRemoteTableModels(gomock.Any(), gomock.Any()).
 		AnyTimes().
-		Return(make([]*model.TableInfo, 0), nil)
+		Return(s.tableInfos, nil)
+	mockBackend.EXPECT().Close()
 	s.rc.backend = backend.MakeBackend(mockBackend)
 	mockSQLExecutor := mock.NewMockSQLExecutor(s.controller)
 	mockSQLExecutor.EXPECT().
@@ -1361,6 +1376,7 @@ func (s *restoreSchemaSuite) SetUpTest(c *C) {
 		GetParser().
 		AnyTimes().
 		Return(parser)
+	mockSQLExecutor.EXPECT().Close()
 	s.rc.tidbGlue = mockTiDBGlue
 }
 

pkg/lightning/restore/tidb.go

@@ -244,14 +244,20 @@ func LoadSchemaInfo(
 			return nil, err
 		}
 
+		tableMap := make(map[string]*model.TableInfo, len(tables))
+		for _, tbl := range tables {
+			tableMap[tbl.Name.L] = tbl
+		}
+
 		dbInfo := &checkpoints.TidbDBInfo{
 			Name:   schema.Name,
 			Tables: make(map[string]*checkpoints.TidbTableInfo),
 		}
 
-		for _, tbl := range tables {
-			tableName := tbl.Name.String()
-			if tbl.State != model.StatePublic {
+		for _, tbl := range schema.Tables {
+			tblInfo := tableMap[strings.ToLower(tbl.Name)]
+			tableName := tblInfo.Name.String()
+			if tblInfo.State != model.StatePublic {
 				err := errors.Errorf("table [%s.%s] state is not public", schema.Name, tableName)
 				metric.RecordTableCount(metric.TableStatePending, err)
 				return nil, err
@@ -261,10 +267,10 @@ func LoadSchemaInfo(
 				return nil, errors.Trace(err)
 			}
 			tableInfo := &checkpoints.TidbTableInfo{
-				ID:   tbl.ID,
+				ID:   tblInfo.ID,
 				DB:   schema.Name,
 				Name: tableName,
-				Core: tbl,
+				Core: tblInfo,
 			}
 			dbInfo.Tables[tableName] = tableInfo
 		}

pkg/lightning/restore/tidb_test.go

@@ -27,9 +27,9 @@ import (
 	"github.com/pingcap/tidb/ddl"
 	"github.com/pingcap/tidb/util/mock"
 
-	"github.com/pingcap/br/pkg/lightning/glue"
-
 	"github.com/pingcap/br/pkg/lightning/checkpoints"
+	"github.com/pingcap/br/pkg/lightning/glue"
+	"github.com/pingcap/br/pkg/lightning/metric"
 	"github.com/pingcap/br/pkg/lightning/mydump"
 )
 
@@ -284,10 +284,14 @@ func (s *tidbSuite) TestDropTable(c *C) {
 func (s *tidbSuite) TestLoadSchemaInfo(c *C) {
 	ctx := context.Background()
 
+	tableCntBefore := metric.ReadCounter(metric.TableCounter.WithLabelValues(metric.TableStatePending, metric.TableResultSuccess))
+
 	// Prepare the mock reply.
 	nodes, _, err := s.timgr.parser.Parse(
 		"CREATE TABLE `t1` (`a` INT PRIMARY KEY);"+
-			"CREATE TABLE `t2` (`b` VARCHAR(20), `c` BOOL, KEY (`b`, `c`))",
+			"CREATE TABLE `t2` (`b` VARCHAR(20), `c` BOOL, KEY (`b`, `c`));"+
+			// an extra table that not exists in dbMetas
+			"CREATE TABLE `t3` (`d` VARCHAR(20), `e` BOOL);",
 		"", "")
 	c.Assert(err, IsNil)
 	tableInfos := make([]*model.TableInfo, 0, len(nodes))
@@ -300,7 +304,23 @@ func (s *tidbSuite) TestLoadSchemaInfo(c *C) {
 		tableInfos = append(tableInfos, info)
 	}
 
-	loaded, err := LoadSchemaInfo(ctx, []*mydump.MDDatabaseMeta{{Name: "db"}}, func(ctx context.Context, schema string) ([]*model.TableInfo, error) {
+	dbMetas := []*mydump.MDDatabaseMeta{
+		{
+			Name: "db",
+			Tables: []*mydump.MDTableMeta{
+				{
+					DB:   "db",
+					Name: "t1",
+				},
+				{
+					DB:   "db",
+					Name: "t2",
+				},
+			},
+		},
+	}
+
+	loaded, err := LoadSchemaInfo(ctx, dbMetas, func(ctx context.Context, schema string) ([]*model.TableInfo, error) {
 		c.Assert(schema, Equals, "db")
 		return tableInfos, nil
 	})
@@ -324,6 +344,10 @@ func (s *tidbSuite) TestLoadSchemaInfo(c *C) {
 			},
 		},
 	})
+
+	tableCntAfter := metric.ReadCounter(metric.TableCounter.WithLabelValues(metric.TableStatePending, metric.TableResultSuccess))
+
+	c.Assert(tableCntAfter-tableCntBefore, Equals, 2.0)
 }
 
 func (s *tidbSuite) TestLoadSchemaInfoMissing(c *C) {