lightning: fix check table empty logic (pingcap#36747)

close pingcap#36746
Dousir9 · Aug 1, 2022 · 4607cba · 4607cba
1 parent 1f6f7da
commit 4607cba
Show file tree

Hide file tree

Showing 3 changed files with 60 additions and 0 deletions.
diff --git a/br/pkg/lightning/restore/BUILD.bazel b/br/pkg/lightning/restore/BUILD.bazel
@@ -42,6 +42,7 @@ go_library(
         "//br/pkg/version",
         "//br/pkg/version/build",
         "//ddl",
+        "//errno",
         "//kv",
         "//meta/autoid",
         "//parser",
@@ -61,6 +62,7 @@ go_library(
         "//util/mock",
         "@com_github_coreos_go_semver//semver",
         "@com_github_docker_go_units//:go-units",
+        "@com_github_go_sql_driver_mysql//:mysql",
         "@com_github_google_uuid//:uuid",
         "@com_github_jedib0t_go_pretty_v6//table",
         "@com_github_jedib0t_go_pretty_v6//text",

diff --git a/br/pkg/lightning/restore/get_pre_info.go b/br/pkg/lightning/restore/get_pre_info.go
@@ -21,6 +21,7 @@ import (
 	"io"
 	"strings"
 
+	mysql_sql_driver "github.com/go-sql-driver/mysql"
 	"github.com/pingcap/errors"
 	"github.com/pingcap/failpoint"
 	"github.com/pingcap/tidb/br/pkg/lightning/backend"
@@ -38,6 +39,7 @@ import (
 	"github.com/pingcap/tidb/br/pkg/lightning/worker"
 	"github.com/pingcap/tidb/br/pkg/storage"
 	"github.com/pingcap/tidb/ddl"
+	"github.com/pingcap/tidb/errno"
 	"github.com/pingcap/tidb/meta/autoid"
 	"github.com/pingcap/tidb/parser"
 	"github.com/pingcap/tidb/parser/ast"
@@ -205,7 +207,14 @@ func (g *TargetInfoGetterImpl) IsTableEmpty(ctx context.Context, schemaName stri
 		&dump,
 	)
 
+	isNoSuchTableErr := false
+	rootErr := errors.Cause(err)
+	if mysqlErr, ok := rootErr.(*mysql_sql_driver.MySQLError); ok && mysqlErr.Number == errno.ErrNoSuchTable {
+		isNoSuchTableErr = true
+	}
 	switch {
+	case isNoSuchTableErr:
+		result = true
 	case errors.ErrorEqual(err, sql.ErrNoRows):
 		result = true
 	case err != nil:

diff --git a/br/pkg/lightning/restore/get_pre_info_test.go b/br/pkg/lightning/restore/get_pre_info_test.go
@@ -15,12 +15,17 @@ package restore
 
 import (
 	"context"
+	"database/sql"
 	"fmt"
 	"strings"
 	"testing"
 
+	"github.com/DATA-DOG/go-sqlmock"
+	mysql_sql_driver "github.com/go-sql-driver/mysql"
+	"github.com/pingcap/errors"
 	"github.com/pingcap/tidb/br/pkg/lightning/config"
 	"github.com/pingcap/tidb/br/pkg/lightning/restore/mock"
+	"github.com/pingcap/tidb/errno"
 	"github.com/pingcap/tidb/parser/model"
 	"github.com/pingcap/tidb/types"
 	"github.com/stretchr/testify/require"
@@ -495,3 +500,47 @@ func TestGetPreInfoEstimateSourceSize(t *testing.T) {
 	require.Equal(t, int64(len(testData)), sizeResult.SizeWithoutIndex)
 	require.False(t, sizeResult.HasUnsortedBigTables)
 }
+
+func TestGetPreInfoIsTableEmpty(t *testing.T) {
+	ctx := context.TODO()
+	db, mock, err := sqlmock.New()
+	require.NoError(t, err)
+	lnConfig := config.NewConfig()
+	lnConfig.TikvImporter.Backend = config.BackendLocal
+	targetGetter, err := NewTargetInfoGetterImpl(lnConfig, db)
+	require.NoError(t, err)
+
+	mock.ExpectQuery("SELECT 1 FROM `test_db`.`test_tbl` LIMIT 1").
+		WillReturnError(&mysql_sql_driver.MySQLError{
+			Number:  errno.ErrNoSuchTable,
+			Message: "Table 'test_db.test_tbl' doesn't exist",
+		})
+	pIsEmpty, err := targetGetter.IsTableEmpty(ctx, "test_db", "test_tbl")
+	require.NoError(t, err)
+	require.NotNil(t, pIsEmpty)
+	require.Equal(t, true, *pIsEmpty)
+
+	mock.ExpectQuery("SELECT 1 FROM `test_db`.`test_tbl` LIMIT 1").
+		WillReturnRows(
+			sqlmock.NewRows([]string{"1"}).
+				RowError(0, sql.ErrNoRows),
+		)
+	pIsEmpty, err = targetGetter.IsTableEmpty(ctx, "test_db", "test_tbl")
+	require.NoError(t, err)
+	require.NotNil(t, pIsEmpty)
+	require.Equal(t, true, *pIsEmpty)
+
+	mock.ExpectQuery("SELECT 1 FROM `test_db`.`test_tbl` LIMIT 1").
+		WillReturnRows(
+			sqlmock.NewRows([]string{"1"}).AddRow(1),
+		)
+	pIsEmpty, err = targetGetter.IsTableEmpty(ctx, "test_db", "test_tbl")
+	require.NoError(t, err)
+	require.NotNil(t, pIsEmpty)
+	require.Equal(t, false, *pIsEmpty)
+
+	mock.ExpectQuery("SELECT 1 FROM `test_db`.`test_tbl` LIMIT 1").
+		WillReturnError(errors.New("some dummy error"))
+	_, err = targetGetter.IsTableEmpty(ctx, "test_db", "test_tbl")
+	require.Error(t, err)
+}