parser: respect TiDB comment when DROP INDEX IF EXISTS (#30173)

pingcap · Nov 30, 2021 · 13d0deb · 13d0deb
1 parent 0b2d9d1
commit 13d0deb
Show file tree

Hide file tree

Showing 2 changed files with 29 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,9 @@ type DropIndexStmt struct {
 func (n *DropIndexStmt) Restore(ctx *format.RestoreCtx) error {
 	ctx.WriteKeyWord("DROP INDEX ")
 	if n.IfExists {
-		ctx.WriteKeyWord("IF EXISTS ")
+		ctx.WriteWithSpecialComments("", func() {
+			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,26 @@ 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)
+	}
+}