From 334745d189a70fed8e8fa9ac4a4a383378b6d50e Mon Sep 17 00:00:00 2001
From: yangming <yangming6@xiaomi.com>
Date: Sat, 2 Apr 2022 11:06:44 +0800
Subject: [PATCH] =?UTF-8?q?=E4=BF=AE=E5=A4=8D=20select=20xxx=20from=20tb?=
 =?UTF-8?q?=20as=20t=20force=20index(u=5Fkey)=20where=20t.id=3Dx=20?=
 =?UTF-8?q?=E8=AF=AD=E6=B3=95=E9=94=99=E8=AF=AF?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 parser/ast/dml.go | 60 +++++++++++++++++++++++++++++++++++++----------
 1 file changed, 47 insertions(+), 13 deletions(-)

diff --git a/parser/ast/dml.go b/parser/ast/dml.go
index 254061ed..34153989 100755
--- a/parser/ast/dml.go
+++ b/parser/ast/dml.go
@@ -189,12 +189,15 @@ type TableName struct {
 }
 
 // Restore implements Node interface.
-func (n *TableName) Restore(ctx *format.RestoreCtx) error {
+func (n *TableName) restoreName(ctx *format.RestoreCtx) {
 	if n.Schema.String() != "" {
 		ctx.WriteName(n.Schema.String())
 		ctx.WritePlain(".")
 	}
 	ctx.WriteName(n.Name.String())
+}
+
+func (n *TableName) restorePartitions(ctx *format.RestoreCtx) {
 	if len(n.PartitionNames) > 0 {
 		ctx.WriteKeyWord(" PARTITION")
 		ctx.WritePlain("(")
@@ -206,6 +209,9 @@ func (n *TableName) Restore(ctx *format.RestoreCtx) error {
 		}
 		ctx.WritePlain(")")
 	}
+}
+
+func (n *TableName) restoreIndexHints(ctx *format.RestoreCtx) error {
 	for _, value := range n.IndexHints {
 		ctx.WritePlain(" ")
 		if err := value.Restore(ctx); err != nil {
@@ -216,6 +222,12 @@ func (n *TableName) Restore(ctx *format.RestoreCtx) error {
 	return nil
 }
 
+func (n *TableName) Restore(ctx *format.RestoreCtx) error {
+	n.restoreName(ctx)
+	n.restorePartitions(ctx)
+	return n.restoreIndexHints(ctx)
+}
+
 // IndexHintType is the type for index hint use, ignore or force.
 type IndexHintType int
 
@@ -383,18 +395,40 @@ func (n *TableSource) Restore(ctx *format.RestoreCtx) error {
 	case *SelectStmt, *UnionStmt:
 		needParen = true
 	}
-	if needParen {
-		ctx.WritePlain("(")
-	}
-	if err := n.Source.Restore(ctx); err != nil {
-		return errors.Annotate(err, "An error occurred while restore TableSource.Source")
-	}
-	if needParen {
-		ctx.WritePlain(")")
-	}
-	if asName := n.AsName.String(); asName != "" {
-		ctx.WriteKeyWord(" AS ")
-		ctx.WriteName(asName)
+
+	if tn, tnCase := n.Source.(*TableName); tnCase {
+		if needParen {
+			ctx.WritePlain("(")
+		}
+
+		tn.restoreName(ctx)
+		tn.restorePartitions(ctx)
+
+		if asName := n.AsName.String(); asName != "" {
+			ctx.WriteKeyWord(" AS ")
+			ctx.WriteName(asName)
+		}
+		if err := tn.restoreIndexHints(ctx); err != nil {
+			return errors.Annotate(err, "An error occurred while restore TableSource.Source.(*TableName).IndexHints")
+		}
+
+		if needParen {
+			ctx.WritePlain(")")
+		}
+	} else {
+		if needParen {
+			ctx.WritePlain("(")
+		}
+		if err := n.Source.Restore(ctx); err != nil {
+			return errors.Annotate(err, "An error occurred while restore TableSource.Source")
+		}
+		if needParen {
+			ctx.WritePlain(")")
+		}
+		if asName := n.AsName.String(); asName != "" {
+			ctx.WriteKeyWord(" AS ")
+			ctx.WriteName(asName)
+		}
 	}
 
 	return nil