*: support trace plan target='estimation' statement

executor/builder.go

@@ -978,11 +978,13 @@ func (b *executorBuilder) buildDDL(v *plannercore.DDL) Executor {
 // at build().
 func (b *executorBuilder) buildTrace(v *plannercore.Trace) Executor {
 	t := &TraceExec{
-		baseExecutor:   newBaseExecutor(b.ctx, v.Schema(), v.ID()),
-		stmtNode:       v.StmtNode,
-		builder:        b,
-		format:         v.Format,
-		optimizerTrace: v.OptimizerTrace,
+		baseExecutor: newBaseExecutor(b.ctx, v.Schema(), v.ID()),
+		stmtNode:     v.StmtNode,
+		builder:      b,
+		format:       v.Format,
+
+		optimizerTrace:       v.OptimizerTrace,
+		optimizerTraceTarget: v.OptimizerTraceTarget,
 	}
 	if t.format == plannercore.TraceFormatLog && !t.optimizerTrace {
 		return &SortExec{

executor/trace.go

@@ -24,6 +24,7 @@ import (
 	"os"
 	"path/filepath"
 	"sort"
+	"strings"
 	"time"
 
 	"github.com/opentracing/basictracer-go"
@@ -58,8 +59,10 @@ type TraceExec struct {
 
 	builder *executorBuilder
 	format  string
+
 	// optimizerTrace indicates 'trace plan statement'
-	optimizerTrace bool
+	optimizerTrace       bool
+	optimizerTraceTarget string
 }
 
 // Next executes real query and collects span later.
@@ -93,44 +96,74 @@ func (e *TraceExec) Next(ctx context.Context, req *chunk.Chunk) error {
 }
 
 func (e *TraceExec) nextOptimizerPlanTrace(ctx context.Context, se sessionctx.Context, req *chunk.Chunk) error {
-	zf, fileName, err := generateOptimizerTraceFile()
-	if err != nil {
-		return err
-	}
-	zw := zip.NewWriter(zf)
-	defer func() {
-		err := zw.Close()
-		if err != nil {
-			logutil.BgLogger().Warn("Closing zip writer failed", zap.Error(err))
-		}
-		err = zf.Close()
-		if err != nil {
-			logutil.BgLogger().Warn("Closing zip file failed", zap.Error(err))
-		}
-	}()
-	traceZW, err := zw.Create("trace.json")
-	if err != nil {
-		return errors.AddStack(err)
-	}
 	stmtCtx := se.GetSessionVars().StmtCtx
-	origin := stmtCtx.EnableOptimizeTrace
-	stmtCtx.EnableOptimizeTrace = true
+
+	var option *bool
+	var dumpToFile bool
+	if e.optimizerTraceTarget == core.TracePlanTargetEstimation {
+		option = &stmtCtx.EnableOptimizerCETrace
+		dumpToFile = false
+	} else {
+		option = &stmtCtx.EnableOptimizeTrace
+		dumpToFile = true
+	}
+
+	origin := *option
+	*option = true
 	defer func() {
-		stmtCtx.EnableOptimizeTrace = origin
+		*option = origin
 	}()
-	_, _, err = core.OptimizeAstNode(ctx, se, e.stmtNode, se.GetInfoSchema().(infoschema.InfoSchema))
+
+	_, _, err := core.OptimizeAstNode(ctx, se, e.stmtNode, se.GetInfoSchema().(infoschema.InfoSchema))
 	if err != nil {
 		return err
 	}
-	res, err := json.Marshal(se.GetSessionVars().StmtCtx.LogicalOptimizeTrace)
-	if err != nil {
-		return errors.AddStack(err)
+
+	var traceResult interface{}
+	if e.optimizerTraceTarget == core.TracePlanTargetEstimation {
+		traceResult = stmtCtx.OptimizerCETrace
+	} else {
+		traceResult = stmtCtx.LogicalOptimizeTrace
 	}
-	_, err = traceZW.Write(res)
+
+	writer := strings.Builder{}
+	jsonEncoder := json.NewEncoder(&writer)
+	// If we do not set this to false, ">", "<", "&"... will be escaped to "\u003c","\u003e", "\u0026"...
+	jsonEncoder.SetEscapeHTML(false)
+	err = jsonEncoder.Encode(traceResult)
 	if err != nil {
 		return errors.AddStack(err)
 	}
-	req.AppendString(0, fileName)
+	res := []byte(writer.String())
+
+	if dumpToFile {
+		zf, fileName, err := generateOptimizerTraceFile()
+		if err != nil {
+			return err
+		}
+		zw := zip.NewWriter(zf)
+		defer func() {
+			err := zw.Close()
+			if err != nil {
+				logutil.BgLogger().Warn("Closing zip writer failed", zap.Error(err))
+			}
+			err = zf.Close()
+			if err != nil {
+				logutil.BgLogger().Warn("Closing zip file failed", zap.Error(err))
+			}
+		}()
+		traceZW, err := zw.Create("trace.json")
+		if err != nil {
+			return errors.AddStack(err)
+		}
+		_, err = traceZW.Write(res)
+		if err != nil {
+			return errors.AddStack(err)
+		}
+		req.AppendString(0, fileName)
+	} else {
+		req.AppendBytes(0, res)
+	}
 	e.exhausted = true
 	return nil
 }

parser/ast/misc.go

@@ -120,16 +120,24 @@ func (n *AuthOption) Restore(ctx *format.RestoreCtx) error {
 type TraceStmt struct {
 	stmtNode
 
-	Stmt      StmtNode
-	Format    string
-	TracePlan bool
+	Stmt   StmtNode
+	Format string
+
+	TracePlan       bool
+	TracePlanTarget string
 }
 
 // Restore implements Node interface.
 func (n *TraceStmt) Restore(ctx *format.RestoreCtx) error {
 	ctx.WriteKeyWord("TRACE ")
 	if n.TracePlan {
 		ctx.WriteKeyWord("PLAN ")
+		if n.TracePlanTarget != "" {
+			ctx.WriteKeyWord("TARGET")
+			ctx.WritePlain(" = ")
+			ctx.WriteString(n.TracePlanTarget)
+			ctx.WritePlain(" ")
+		}
 	} else if n.Format != "row" {
 		ctx.WriteKeyWord("FORMAT")
 		ctx.WritePlain(" = ")

parser/misc.go

@@ -697,6 +697,7 @@ var tokenMap = map[string]int{
 	"SWITCHES":                 switchesSym,
 	"SYSTEM":                   system,
 	"SYSTEM_TIME":              systemTime,
+	"TARGET":                   target,
 	"TABLE_CHECKSUM":           tableChecksum,
 	"TABLE":                    tableKwd,
 	"TABLES":                   tables,