Repository: pingcap/dumpling (go)

Issue:
dumpling costs too much time on getting schemas' tables
## Bug Report

Please answer these questions before submitting your issue. Thanks!

1. What did you do? If possible, provide a recipe for reproducing the error.

Run dumpling with `-T` args to dump only one table. The TiDB cluster has more than 1000 schemas and 100000 tables. Each `select` query will cost about 3s.

2. What did you expect to see?

Dumpling data quickly and successfully.

3. What did you see instead?

Dumpling is blocked at [listAllTables](https://github.com/pingcap/dumpling/blob/master/v4/export/prepare.go#L78).


Pull request description:
reduce select usage when query tables
<!--
Thank you for contributing to Dumpling! Please read the [CONTRIBUTING](https://github.com/pingcap/dumpling/blob/master/CONTRIBUTING.md) document **BEFORE** filing this PR.
-->

### What problem does this PR solve? <!--add issue link with summary if exists-->
Fix #109 

### What is changed and how it works?
Select all tables/views at one time.

### Check List <!--REMOVE the items that are not applicable-->

Tests <!-- At least one of them must be included. -->

 - Unit test
 - Integration test

Side effects

 - Possible performance regression
 - Increased code complexity
 
### Release note

<!-- bugfixes or new feature need a release note, must in the form of a list, such as

- support -T/--tables-list argument

or if no need to be included in the release note, just add the following line

- No release note
-->

- Fix the problem that dumpling get blocked when tidb cluster has a lot of tables.


Patch:
diff --git a/v4/export/prepare.go b/v4/export/prepare.go
index 3839f01..c1a6634 100644
--- a/v4/export/prepare.go
+++ b/v4/export/prepare.go
@@ -77,29 +77,12 @@ func prepareDumpingDatabases(conf *Config, db *sql.DB) ([]string, error) {
 
 func listAllTables(db *sql.DB, databaseNames []string) (DatabaseTables, error) {
 	log.Debug("list all the tables")
-	dbTables := DatabaseTables{}
-	for _, dbName := range databaseNames {
-		tables, err := ListAllTables(db, dbName)
-		if err != nil {
-			return nil, err
-		}
-		dbTables[dbName] = make([]*TableInfo, 0, len(tables))
-		dbTables = dbTables.AppendTables(dbName, tables...)
-	}
-	return dbTables, nil
+	return ListAllDatabasesTables(db, databaseNames, TableTypeBase)
 }
 
 func listAllViews(db *sql.DB, databaseNames []string) (DatabaseTables, error) {
 	log.Debug("list all the views")
-	dbTables := DatabaseTables{}
-	for _, dbName := range databaseNames {
-		views, err := ListAllViews(db, dbName)
-		if err != nil {
-			return nil, err
-		}
-		dbTables = dbTables.AppendViews(dbName, views...)
-	}
-	return dbTables, nil
+	return ListAllDatabasesTables(db, databaseNames, TableTypeView)
 }
 
 type databaseName = string
diff --git a/v4/export/sql.go b/v4/export/sql.go
index 2f258e5..14e0b45 100644
--- a/v4/export/sql.go
+++ b/v4/export/sql.go
@@ -77,6 +77,40 @@ func ListAllTables(db *sql.DB, database string) ([]string, error) {
 	return tables.data, nil
 }
 
+func ListAllDatabasesTables(db *sql.DB, databaseNames []string, tableType TableType) (DatabaseTables, error) {
+	var tableTypeStr string
+	switch tableType {
+	case TableTypeBase:
+		tableTypeStr = "BASE TABLE"
+	case TableTypeView:
+		tableTypeStr = "VIEW"
+	default:
+		return nil, errors.Errorf("unknown table type %v", tableType)
+	}
+
+	query := fmt.Sprintf("SELECT table_schema,table_name FROM information_schema.tables WHERE table_type = '%s'", tableTypeStr)
+	dbTables := DatabaseTables{}
+	for _, schema := range databaseNames {
+		dbTables[schema] = make([]*TableInfo, 0)
+	}
+
+	if err := simpleQueryWithArgs(db, func(rows *sql.Rows) error {
+		var schema, table string
+		if err := rows.Scan(&schema, &table); err != nil {
+			return withStack(err)
+		}
+
+		// only append tables to schemas in databaseNames
+		if _, ok := dbTables[schema]; ok {
+			dbTables[schema] = append(dbTables[schema], &TableInfo{table, tableType})
+		}
+		return nil
+	}, query); err != nil {
+		return nil, errors.WithMessage(err, query)
+	}
+	return dbTables, nil
+}
+
 func ListAllViews(db *sql.DB, database string) ([]string, error) {
 	var views oneStrColumnTable
 	const query = "SELECT table_name FROM information_schema.tables WHERE table_schema = ? and table_type = 'VIEW'"
@@ -426,7 +460,7 @@ func simpleQueryWithArgs(db *sql.DB, handleOneRow func(*sql.Rows) error, sql str
 			return withStack(err)
 		}
 	}
-	return nil
+	return rows.Err()
 }
 
 func pickupPossibleField(dbName, tableName string, db *sql.DB, conf *Config) (string, error) {


Test patch:
diff --git a/v4/export/prepare_test.go b/v4/export/prepare_test.go
index 395de9f..7233e92 100644
--- a/v4/export/prepare_test.go
+++ b/v4/export/prepare_test.go
@@ -64,19 +64,19 @@ func (s *testPrepareSuite) TestListAllTables(c *C) {
 		AppendViews("db3", "t6", "t7", "t8")
 
 	var dbNames []databaseName
+	rows := sqlmock.NewRows([]string{"table_schema", "table_name"})
 	for dbName, tableInfos := range data {
 		dbNames = append(dbNames, dbName)
 
-		rows := sqlmock.NewRows([]string{"Table_name"})
 		for _, tbInfo := range tableInfos {
 			if tbInfo.Type == TableTypeView {
 				continue
 			}
-			rows.AddRow(tbInfo.Name)
+			rows.AddRow(dbName, tbInfo.Name)
 		}
-		query := "SELECT table_name FROM information_schema.tables WHERE table_schema = (.*) and table_type = (.*)"
-		mock.ExpectQuery(query).WillReturnRows(rows)
 	}
+	query := "SELECT table_schema,table_name FROM information_schema.tables WHERE table_type = (.*)"
+	mock.ExpectQuery(query).WillReturnRows(rows)
 
 	tables, err := listAllTables(db, dbNames)
 	c.Assert(err, IsNil)
@@ -94,8 +94,8 @@ func (s *testPrepareSuite) TestListAllTables(c *C) {
 	data = NewDatabaseTables().
 		AppendTables("db", "t1").
 		AppendViews("db", "t2")
-	query := "SELECT table_name FROM information_schema.tables WHERE table_schema = (.*) and table_type = (.*)"
-	mock.ExpectQuery(query).WillReturnRows(sqlmock.NewRows([]string{"Table_name"}).AddRow("t2"))
+	query = "SELECT table_schema,table_name FROM information_schema.tables WHERE table_type = (.*)"
+	mock.ExpectQuery(query).WillReturnRows(sqlmock.NewRows([]string{"table_schema", "table_name"}).AddRow("db", "t2"))
 	tables, err = listAllViews(db, []string{"db"})
 	c.Assert(err, IsNil)
 	c.Assert(len(tables), Equals, 1)