planner: support dump file for trace plan statement (#30059)

executor/builder.go

@@ -975,12 +975,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,
+		baseExecutor:   newBaseExecutor(b.ctx, v.Schema(), v.ID()),
+		stmtNode:       v.StmtNode,
+		builder:        b,
+		format:         v.Format,
+		optimizerTrace: v.OptimizerTrace,
 	}
-	if t.format == plannercore.TraceFormatLog {
+	if t.format == plannercore.TraceFormatLog && !t.optimizerTrace {
 		return &SortExec{
 			baseExecutor: newBaseExecutor(b.ctx, v.Schema(), v.ID(), t),
 			ByItems: []*plannerutil.ByItems{

executor/executor.go

@@ -1689,7 +1689,6 @@ func ResetContextOfStmt(ctx sessionctx.Context, s ast.StmtNode) (err error) {
 	sc.CTEStorageMap = map[int]*CTEStorages{}
 	sc.IsStaleness = false
 	sc.LockTableIDs = make(map[int64]struct{})
-	sc.EnableOptimizeTrace = false
 	sc.LogicalOptimizeTrace = nil
 
 	sc.InitMemTracker(memory.LabelForSQLText, vars.MemQuotaQuery)

executor/trace.go

@@ -15,10 +15,16 @@
 package executor
 
 import (
+	"archive/zip"
 	"context"
+	"crypto/rand"
+	"encoding/base64"
 	"encoding/json"
 	"fmt"
+	"os"
+	"path/filepath"
 	"sort"
+	"strconv"
 	"time"
 
 	"github.com/opentracing/basictracer-go"
@@ -51,6 +57,8 @@ type TraceExec struct {
 
 	builder *executorBuilder
 	format  string
+	// optimizerTrace indicates 'trace plan statement'
+	optimizerTrace bool
 }
 
 // Next executes real query and collects span later.
@@ -71,6 +79,10 @@ func (e *TraceExec) Next(ctx context.Context, req *chunk.Chunk) error {
 		e.ctx.GetSessionVars().StmtCtx = stmtCtx
 	}()
 
+	if e.optimizerTrace {
+		return e.nextOptimizerPlanTrace(ctx, e.ctx, req)
+	}
+
 	switch e.format {
 	case core.TraceFormatLog:
 		return e.nextTraceLog(ctx, se, req)
@@ -79,6 +91,40 @@ 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)
+	}
+	e.executeChild(ctx, se.(sqlexec.SQLExecutor))
+	res, err := json.Marshal(se.GetSessionVars().StmtCtx.LogicalOptimizeTrace)
+	if err != nil {
+		return errors.AddStack(err)
+	}
+	_, err = traceZW.Write(res)
+	if err != nil {
+		return errors.AddStack(err)
+	}
+	req.AppendString(0, fileName)
+	e.exhausted = true
+	return nil
+}
+
 func (e *TraceExec) nextTraceLog(ctx context.Context, se sqlexec.SQLExecutor, req *chunk.Chunk) error {
 	recorder := basictracer.NewInMemoryRecorder()
 	tracer := basictracer.New(recorder)
@@ -142,8 +188,11 @@ func (e *TraceExec) executeChild(ctx context.Context, se sqlexec.SQLExecutor) {
 	vars := e.ctx.GetSessionVars()
 	origin := vars.InRestrictedSQL
 	vars.InRestrictedSQL = true
+	originOptimizeTrace := vars.EnableStmtOptimizeTrace
+	vars.EnableStmtOptimizeTrace = e.optimizerTrace
 	defer func() {
 		vars.InRestrictedSQL = origin
+		vars.EnableStmtOptimizeTrace = originOptimizeTrace
 	}()
 	rs, err := se.ExecuteStmt(ctx, e.stmtNode)
 	if err != nil {
@@ -252,3 +301,32 @@ func generateLogResult(allSpans []basictracer.RawSpan, chk *chunk.Chunk) {
 		}
 	}
 }
+
+func generateOptimizerTraceFile() (*os.File, string, error) {
+	dirPath := getOptimizerTraceDirName()
+	// Create path
+	err := os.MkdirAll(dirPath, os.ModePerm)
+	if err != nil {
+		return nil, "", errors.AddStack(err)
+	}
+	// Generate key and create zip file
+	time := time.Now().UnixNano()
+	b := make([]byte, 16)
+	_, err = rand.Read(b)
+	if err != nil {
+		return nil, "", errors.AddStack(err)
+	}
+	key := base64.URLEncoding.EncodeToString(b)
+	fileName := fmt.Sprintf("optimizer_trace_%v_%v.zip", key, time)
+	zf, err := os.Create(filepath.Join(dirPath, fileName))
+	if err != nil {
+		return nil, "", errors.AddStack(err)
+	}
+	return zf, fileName, nil
+}
+
+// getOptimizerTraceDirName returns optimizer trace directory path.
+// The path is related to the process id.
+func getOptimizerTraceDirName() string {
+	return filepath.Join(os.TempDir(), "optimizer_trace", strconv.Itoa(os.Getpid()))
+}

executor/trace_test.go

@@ -66,3 +66,13 @@ func rowsOrdered(rows [][]interface{}) bool {
 	}
 	return true
 }
+
+func (s *testSuite1) TestTracePlanStmt(c *C) {
+	tk := testkit.NewTestKit(c, s.store)
+	tk.MustExec("use test")
+	tk.MustExec("create table tp123(id int);")
+	rows := tk.MustQuery("trace plan select * from tp123").Rows()
+	c.Assert(rows, HasLen, 1)
+	c.Assert(rows[0], HasLen, 1)
+	c.Assert(rows[0][0].(string), Matches, ".*zip")
+}

planner/core/logical_plan_test.go

@@ -2107,7 +2107,7 @@ func (s *testPlanSuite) TestLogicalOptimizeWithTraceEnabled(c *C) {
 		err = Preprocess(s.ctx, stmt, WithPreprocessorReturn(&PreprocessorReturn{InfoSchema: s.is}))
 		c.Assert(err, IsNil, comment)
 		sctx := MockContext()
-		sctx.GetSessionVars().StmtCtx.EnableOptimizeTrace = true
+		sctx.GetSessionVars().EnableStmtOptimizeTrace = true
 		builder, _ := NewPlanBuilder().Init(sctx, s.is, &hint.BlockHintProcessor{})
 		domain.GetDomain(sctx).MockInfoCacheAndLoadInfoSchema(s.is)
 		ctx := context.TODO()

planner/core/optimizer.go

@@ -98,11 +98,11 @@ func (op *logicalOptimizeOp) withEnableOptimizeTracer(tracer *tracing.LogicalOpt
 	return op
 }
 
-func (op *logicalOptimizeOp) appendBeforeRuleOptimize(name string, before LogicalPlan) {
+func (op *logicalOptimizeOp) appendBeforeRuleOptimize(index int, name string, before LogicalPlan) {
 	if op.tracer == nil {
 		return
 	}
-	op.tracer.AppendRuleTracerBeforeRuleOptimize(name, before.buildLogicalPlanTrace())
+	op.tracer.AppendRuleTracerBeforeRuleOptimize(index, name, before.buildLogicalPlanTrace())
 }
 
 func (op *logicalOptimizeOp) appendStepToCurrent(id int, tp, reason, action string) {
@@ -377,12 +377,12 @@ func enableParallelApply(sctx sessionctx.Context, plan PhysicalPlan) PhysicalPla
 
 func logicalOptimize(ctx context.Context, flag uint64, logic LogicalPlan) (LogicalPlan, error) {
 	opt := defaultLogicalOptimizeOption()
-	stmtCtx := logic.SCtx().GetSessionVars().StmtCtx
-	if stmtCtx.EnableOptimizeTrace {
+	vars := logic.SCtx().GetSessionVars()
+	if vars.EnableStmtOptimizeTrace {
 		tracer := &tracing.LogicalOptimizeTracer{Steps: make([]*tracing.LogicalRuleOptimizeTracer, 0)}
 		opt = opt.withEnableOptimizeTracer(tracer)
 		defer func() {
-			stmtCtx.LogicalOptimizeTrace = tracer
+			vars.StmtCtx.LogicalOptimizeTrace = tracer
 		}()
 	}
 	var err error
@@ -393,7 +393,7 @@ func logicalOptimize(ctx context.Context, flag uint64, logic LogicalPlan) (Logic
 		if flag&(1<<uint(i)) == 0 || isLogicalRuleDisabled(rule) {
 			continue
 		}
-		opt.appendBeforeRuleOptimize(rule.name(), logic)
+		opt.appendBeforeRuleOptimize(i, rule.name(), logic)
 		logic, err = rule.optimize(ctx, logic, opt)
 		if err != nil {
 			return nil, err

planner/core/planbuilder.go

@@ -4048,7 +4048,15 @@ const (
 // underlying query and then constructs a schema, which will be used to constructs
 // rows result.
 func (b *PlanBuilder) buildTrace(trace *ast.TraceStmt) (Plan, error) {
-	p := &Trace{StmtNode: trace.Stmt, Format: trace.Format}
+	p := &Trace{StmtNode: trace.Stmt, Format: trace.Format, OptimizerTrace: trace.TracePlan}
+	// TODO: forbid trace plan if the statement isn't select read-only statement
+	if trace.TracePlan {
+		schema := newColumnsWithNames(1)
+		schema.Append(buildColumnWithName("", "Dump_link", mysql.TypeVarchar, 128))
+		p.SetSchema(schema.col2Schema())
+		p.names = schema.names
+		return p, nil
+	}
 	switch trace.Format {
 	case TraceFormatRow:
 		schema := newColumnsWithNames(3)
@@ -4380,7 +4388,7 @@ func buildShowSchema(s *ast.ShowStmt, isView bool, isSequence bool) (schema *exp
 func (b *PlanBuilder) buildPlanReplayer(pc *ast.PlanReplayerStmt) Plan {
 	p := &PlanReplayer{ExecStmt: pc.Stmt, Analyze: pc.Analyze, Load: pc.Load, File: pc.File}
 	schema := newColumnsWithNames(1)
-	schema.Append(buildColumnWithName("", "Dump_link", mysql.TypeVarchar, 128))
+	schema.Append(buildColumnWithName("", "File_token", mysql.TypeVarchar, 128))
 	p.SetSchema(schema.col2Schema())
 	p.names = schema.names
 	return p

planner/core/trace.go

@@ -24,4 +24,7 @@ type Trace struct {
 
 	StmtNode ast.StmtNode
 	Format   string
+
+	// OptimizerTrace indicates `trace plan <statement>` case
+	OptimizerTrace bool
 }

sessionctx/stmtctx/stmtctx.go

@@ -199,8 +199,6 @@ type StatementContext struct {
 	// InVerboseExplain indicates the statement is "explain format='verbose' ...".
 	InVerboseExplain bool
 
-	// EnableOptimizeTrace indicates whether the statement is enable optimize trace
-	EnableOptimizeTrace bool
 	// LogicalOptimizeTrace indicates the trace for optimize
 	LogicalOptimizeTrace *tracing.LogicalOptimizeTracer
 }

sessionctx/variable/session.go

@@ -958,6 +958,8 @@ type SessionVars struct {
 		data [2]stmtctx.StatementContext
 	}
 
+	// EnableStmtOptimizeTrace indicates whether enable optimizer trace by 'trace plan statement'
+	EnableStmtOptimizeTrace bool
 	// Rng stores the rand_seed1 and rand_seed2 for Rand() function
 	Rng *utilMath.MysqlRng
 }

util/tracing/opt_trace.go

@@ -16,24 +16,24 @@ package tracing
 
 // LogicalPlanTrace indicates for the LogicalPlan trace information
 type LogicalPlanTrace struct {
-	ID       int
-	TP       string
-	Children []*LogicalPlanTrace
+	ID       int                 `json:"id"`
+	TP       string              `json:"type"`
+	Children []*LogicalPlanTrace `json:"children"`
 
 	// ExplainInfo should be implemented by each implemented LogicalPlan
-	ExplainInfo string
+	ExplainInfo string `json:"info"`
 }
 
 // LogicalOptimizeTracer indicates the trace for the whole logicalOptimize processing
 type LogicalOptimizeTracer struct {
-	Steps []*LogicalRuleOptimizeTracer
+	Steps []*LogicalRuleOptimizeTracer `json:"steps"`
 	// curRuleTracer indicates the current rule Tracer during optimize by rule
 	curRuleTracer *LogicalRuleOptimizeTracer
 }
 
 // AppendRuleTracerBeforeRuleOptimize add plan tracer before optimize
-func (tracer *LogicalOptimizeTracer) AppendRuleTracerBeforeRuleOptimize(name string, before *LogicalPlanTrace) {
-	ruleTracer := buildLogicalRuleOptimizeTracerBeforeOptimize(name, before)
+func (tracer *LogicalOptimizeTracer) AppendRuleTracerBeforeRuleOptimize(index int, name string, before *LogicalPlanTrace) {
+	ruleTracer := buildLogicalRuleOptimizeTracerBeforeOptimize(index, name, before)
 	tracer.Steps = append(tracer.Steps, ruleTracer)
 	tracer.curRuleTracer = ruleTracer
 }
@@ -56,15 +56,17 @@ func (tracer *LogicalOptimizeTracer) TrackLogicalPlanAfterRuleOptimize(after *Lo
 // LogicalRuleOptimizeTracer indicates the trace for the LogicalPlan tree before and after
 // logical rule optimize
 type LogicalRuleOptimizeTracer struct {
-	Before   *LogicalPlanTrace
-	After    *LogicalPlanTrace
-	RuleName string
-	Steps    []LogicalRuleOptimizeTraceStep
+	Index    int                            `json:"index"`
+	Before   *LogicalPlanTrace              `json:"before"`
+	After    *LogicalPlanTrace              `json:"after"`
+	RuleName string                         `json:"name"`
+	Steps    []LogicalRuleOptimizeTraceStep `json:"steps"`
 }
 
 // buildLogicalRuleOptimizeTracerBeforeOptimize build rule tracer before rule optimize
-func buildLogicalRuleOptimizeTracerBeforeOptimize(name string, before *LogicalPlanTrace) *LogicalRuleOptimizeTracer {
+func buildLogicalRuleOptimizeTracerBeforeOptimize(index int, name string, before *LogicalPlanTrace) *LogicalRuleOptimizeTracer {
 	return &LogicalRuleOptimizeTracer{
+		Index:    index,
 		Before:   before,
 		RuleName: name,
 		Steps:    make([]LogicalRuleOptimizeTraceStep, 0),
@@ -74,8 +76,8 @@ func buildLogicalRuleOptimizeTracerBeforeOptimize(name string, before *LogicalPl
 // LogicalRuleOptimizeTraceStep indicates the trace for the detailed optimize changing in
 // logical rule optimize
 type LogicalRuleOptimizeTraceStep struct {
-	Action string
-	Reason string
-	ID     int
-	TP     string
+	Action string `json:"action"`
+	Reason string `json:"reason"`
+	ID     int    `json:"id"`
+	TP     string `json:"type"`
 }