parser: respect TiDB comment when DROP INDEX IF EXISTS

pingcap · Nov 26, 2021 · f20e4c0 · f20e4c0
1 parent 51b7a19
commit f20e4c0
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 2 deletions.
diff --git a/parser/ast/ddl.go b/parser/ast/ddl.go
@@ -15,6 +15,7 @@ package ast
 
 import (
 	"github.com/pingcap/errors"
+
 	"github.com/pingcap/tidb/parser/auth"
 	"github.com/pingcap/tidb/parser/format"
 	"github.com/pingcap/tidb/parser/model"
@@ -1689,7 +1690,13 @@ type DropIndexStmt struct {
 func (n *DropIndexStmt) Restore(ctx *format.RestoreCtx) error {
 	ctx.WriteKeyWord("DROP INDEX ")
 	if n.IfExists {
-		ctx.WriteKeyWord("IF EXISTS ")
+		if ctx.Flags.HasTiDBSpecialCommentFlag() {
+			ctx.WriteWithSpecialComments("", func() {
+				ctx.WriteKeyWord("IF EXISTS")
+			})
+		} else {
+			ctx.WriteKeyWord("IF EXISTS ")
+		}
 	}
 	ctx.WriteName(n.IndexName)
 	ctx.WriteKeyWord(" ON ")

diff --git a/parser/ast/ddl_test.go b/parser/ast/ddl_test.go
@@ -16,9 +16,10 @@ package ast_test
 import (
 	"testing"
 
+	"github.com/stretchr/testify/require"
+
 	. "github.com/pingcap/tidb/parser/ast"
 	"github.com/pingcap/tidb/parser/format"
-	"github.com/stretchr/testify/require"
 )
 
 func TestDDLVisitorCover(t *testing.T) {
@@ -623,3 +624,27 @@ func TestSequenceRestore(t *testing.T) {
 	}
 	runNodeRestoreTest(t, testCases, "%s", extractNodeFunc)
 }
+
+
+func TestDropIndexRestore(t *testing.T) {
+	t.Parallel()
+	sourceSQL := "drop index if exists idx on t"
+	cases := [] struct {
+		flags format.RestoreFlags
+		expectSQL string
+	}{
+		{format.DefaultRestoreFlags, "DROP INDEX IF EXISTS `idx` ON `t`"},
+		{format.DefaultRestoreFlags|format.RestoreTiDBSpecialComment, "DROP INDEX /*T! IF EXISTS */`idx` ON `t`"},
+	}
+
+	extractNodeFunc := func(node Node) Node {
+		return node
+	}
+
+	for _, ca := range cases {
+		testCases := []NodeRestoreTestCase{
+			{sourceSQL, ca.expectSQL},
+		}
+		runNodeRestoreTestWithFlags(t, testCases, "%s", extractNodeFunc, ca.flags)
+	}
+}