From 7a1e3aa2b449f517301bfd9f2f86454a9b6740ff Mon Sep 17 00:00:00 2001
From: glorv <glorvs@163.com>
Date: Thu, 30 Dec 2021 10:35:50 +0800
Subject: [PATCH] lightning: don't ignore error if check new collation failed
 (#31089)

close pingcap/tidb#31088
---
 pkg/lightning/restore/restore.go   |  5 ++++-
 pkg/lightning/restore/tidb.go      |  8 ++++++--
 pkg/lightning/restore/tidb_test.go | 22 +++++++++++++++++++---
 3 files changed, 29 insertions(+), 6 deletions(-)

diff --git a/pkg/lightning/restore/restore.go b/pkg/lightning/restore/restore.go
index 79f132b1c..7cef2b467 100644
--- a/pkg/lightning/restore/restore.go
+++ b/pkg/lightning/restore/restore.go
@@ -1815,7 +1815,10 @@ func (rc *Controller) setGlobalVariables(ctx context.Context) error {
 		return nil
 	}
 	// set new collation flag base on tidb config
-	enabled := ObtainNewCollationEnabled(ctx, rc.tidbGlue.GetSQLExecutor())
+	enabled, err := ObtainNewCollationEnabled(ctx, rc.tidbGlue.GetSQLExecutor())
+	if err != nil {
+		return err
+	}
 	// we should enable/disable new collation here since in server mode, tidb config
 	// may be different in different tasks
 	collate.SetNewCollationEnabledForTest(enabled)
diff --git a/pkg/lightning/restore/tidb.go b/pkg/lightning/restore/tidb.go
index ee1252fd3..6b8d216dd 100644
--- a/pkg/lightning/restore/tidb.go
+++ b/pkg/lightning/restore/tidb.go
@@ -369,7 +369,7 @@ func ObtainImportantVariables(ctx context.Context, g glue.SQLExecutor, needTiDBV
 	return result
 }
 
-func ObtainNewCollationEnabled(ctx context.Context, g glue.SQLExecutor) bool {
+func ObtainNewCollationEnabled(ctx context.Context, g glue.SQLExecutor) (bool, error) {
 	newCollationEnabled := false
 	newCollationVal, err := g.ObtainStringWithLog(
 		ctx,
@@ -379,9 +379,13 @@ func ObtainNewCollationEnabled(ctx context.Context, g glue.SQLExecutor) bool {
 	)
 	if err == nil && newCollationVal == "True" {
 		newCollationEnabled = true
+	} else if errors.ErrorEqual(err, sql.ErrNoRows) {
+		// ignore if target variable is not found, this may happen if tidb < v4.0
+		newCollationEnabled = false
+		err = nil
 	}
 
-	return newCollationEnabled
+	return newCollationEnabled, errors.Trace(err)
 }
 
 // AlterAutoIncrement rebase the table auto increment id
diff --git a/pkg/lightning/restore/tidb_test.go b/pkg/lightning/restore/tidb_test.go
index bd1169ed2..42b5a10b9 100644
--- a/pkg/lightning/restore/tidb_test.go
+++ b/pkg/lightning/restore/tidb_test.go
@@ -16,6 +16,7 @@ package restore
 
 import (
 	"context"
+	"database/sql"
 	"testing"
 
 	"github.com/DATA-DOG/go-sqlmock"
@@ -505,8 +506,22 @@ func (s *tidbSuite) TestObtainNewCollationEnabled(c *C) {
 	ctx := context.Background()
 
 	s.mockDB.
-		ExpectQuery("\\QSELECT variable_value FROM mysql.tidb WHERE variable_name = 'new_collation_enabled'\\E")
-	version := ObtainNewCollationEnabled(ctx, s.tiGlue.GetSQLExecutor())
+		ExpectQuery("\\QSELECT variable_value FROM mysql.tidb WHERE variable_name = 'new_collation_enabled'\\E").
+		WillReturnError(errors.New("mock permission deny"))
+	s.mockDB.
+		ExpectQuery("\\QSELECT variable_value FROM mysql.tidb WHERE variable_name = 'new_collation_enabled'\\E").
+		WillReturnError(errors.New("mock permission deny"))
+	s.mockDB.
+		ExpectQuery("\\QSELECT variable_value FROM mysql.tidb WHERE variable_name = 'new_collation_enabled'\\E").
+		WillReturnError(errors.New("mock permission deny"))
+	_, err := ObtainNewCollationEnabled(ctx, s.tiGlue.GetSQLExecutor())
+	c.Assert(err, ErrorMatches, "obtain new collation enabled failed: mock permission deny")
+
+	s.mockDB.
+		ExpectQuery("\\QSELECT variable_value FROM mysql.tidb WHERE variable_name = 'new_collation_enabled'\\E").
+		WillReturnRows(sqlmock.NewRows([]string{"variable_value"}).RowError(0, sql.ErrNoRows))
+	version, err := ObtainNewCollationEnabled(ctx, s.tiGlue.GetSQLExecutor())
+	c.Assert(err, IsNil)
 	c.Assert(version, Equals, false)
 
 	kvMap := map[string]bool{
@@ -518,7 +533,8 @@ func (s *tidbSuite) TestObtainNewCollationEnabled(c *C) {
 			ExpectQuery("\\QSELECT variable_value FROM mysql.tidb WHERE variable_name = 'new_collation_enabled'\\E").
 			WillReturnRows(sqlmock.NewRows([]string{"variable_value"}).AddRow(k))
 
-		version := ObtainNewCollationEnabled(ctx, s.tiGlue.GetSQLExecutor())
+		version, err = ObtainNewCollationEnabled(ctx, s.tiGlue.GetSQLExecutor())
+		c.Assert(err, IsNil)
 		c.Assert(version, Equals, v)
 	}
 	s.mockDB.