domain: support more params for extract plan task (#41975)

ref #41130

domain/domain.go

@@ -977,7 +977,7 @@ func NewDomain(store kv.Storage, ddlLease time.Duration, statsLease time.Duratio
 		infoCache:           infoschema.NewCache(16),
 		slowQuery:           newTopNSlowQueries(30, time.Hour*24*7, 500),
 		indexUsageSyncLease: idxUsageSyncLease,
-		dumpFileGcChecker:   &dumpFileGcChecker{gcLease: dumpFileGcLease, paths: []string{replayer.GetPlanReplayerDirName(), GetOptimizerTraceDirName()}},
+		dumpFileGcChecker:   &dumpFileGcChecker{gcLease: dumpFileGcLease, paths: []string{replayer.GetPlanReplayerDirName(), GetOptimizerTraceDirName(), GetExtractTaskDirName()}},
 		expiredTimeStamp4PC: types.NewTime(types.ZeroCoreTime, mysql.TypeTimestamp, types.DefaultFsp),
 		mdlCheckTableInfo: &mdlCheckTableInfo{
 			mu:         sync.Mutex{},

domain/extract.go

@@ -23,6 +23,7 @@ import (
 	"math/rand"
 	"os"
 	"path/filepath"
+	"strconv"
 	"strings"
 	"sync"
 	"time"
@@ -48,6 +49,8 @@ const (
 const (
 	// ExtractTaskType indicates type of extract task
 	ExtractTaskType = "taskType"
+	// ExtractPlanTaskSkipStats indicates skip stats for extract plan task
+	ExtractPlanTaskSkipStats = "SkipStats"
 )
 
 // ExtractType indicates type
@@ -102,6 +105,10 @@ type ExtractTask struct {
 	ExtractType     ExtractType
 	IsBackgroundJob bool
 
+	// Param for Extract Plan
+	SkipStats      bool
+	UseHistoryView bool
+
 	// variables for plan task type
 	Begin time.Time
 	End   time.Time
@@ -132,7 +139,7 @@ func (w *extractWorker) extractTask(ctx context.Context, task *ExtractTask) (str
 }
 
 func (w *extractWorker) extractPlanTask(ctx context.Context, task *ExtractTask) (string, error) {
-	if !config.GetGlobalConfig().Instance.StmtSummaryEnablePersistent {
+	if task.UseHistoryView && !config.GetGlobalConfig().Instance.StmtSummaryEnablePersistent {
 		return "", errors.New("tidb_stmt_summary_enable_persistent should be enabled for extract task")
 	}
 	records, err := w.collectRecords(ctx, task)
@@ -145,16 +152,20 @@ func (w *extractWorker) extractPlanTask(ctx context.Context, task *ExtractTask)
 		logutil.BgLogger().Error("package stmt summary records failed for extract plan task", zap.Error(err))
 		return "", err
 	}
-	return w.dumpExtractPlanPackage(p)
+	return w.dumpExtractPlanPackage(task, p)
 }
 
 func (w *extractWorker) collectRecords(ctx context.Context, task *ExtractTask) (map[stmtSummaryHistoryKey]*stmtSummaryHistoryRecord, error) {
 	w.Lock()
 	defer w.Unlock()
 	exec := w.sctx.(sqlexec.RestrictedSQLExecutor)
 	ctx1 := kv.WithInternalSourceType(ctx, kv.InternalTxnStats)
-	rows, _, err := exec.ExecRestrictedSQL(ctx1, nil, fmt.Sprintf("SELECT STMT_TYPE, DIGEST, PLAN_DIGEST,QUERY_SAMPLE_TEXT, BINARY_PLAN, TABLE_NAMES FROM INFORMATION_SCHEMA.STATEMENTS_SUMMARY_HISTORY WHERE SUMMARY_END_TIME > '%s' OR SUMMARY_BEGIN_TIME < '%s'",
-		task.Begin.Format(types.TimeFormat), task.End.Format(types.TimeFormat)))
+	sourceTable := "STATEMENTS_SUMMARY_HISTORY"
+	if !task.UseHistoryView {
+		sourceTable = "STATEMENTS_SUMMARY"
+	}
+	rows, _, err := exec.ExecRestrictedSQL(ctx1, nil, fmt.Sprintf("SELECT STMT_TYPE, DIGEST, PLAN_DIGEST,QUERY_SAMPLE_TEXT, BINARY_PLAN, TABLE_NAMES, SAMPLE_USER FROM INFORMATION_SCHEMA.%s WHERE SUMMARY_END_TIME > '%s' OR SUMMARY_BEGIN_TIME < '%s'",
+		sourceTable, task.Begin.Format(types.TimeFormat), task.End.Format(types.TimeFormat)))
 	if err != nil {
 		return nil, err
 	}
@@ -171,6 +182,7 @@ func (w *extractWorker) collectRecords(ctx context.Context, task *ExtractTask) (
 			digest:     record.digest,
 			planDigest: record.planDigest,
 		}
+		record.userName = row.GetString(6)
 		record.tables = make([]tableNamePair, 0)
 		setRecord, err := w.handleTableNames(tableNames, record)
 		if err != nil {
@@ -324,7 +336,7 @@ func (w *extractWorker) decodeBinaryPlan(ctx context.Context, bPlan string) (str
  |	 |-digest1.sql
  |	 |-...
 */
-func (w *extractWorker) dumpExtractPlanPackage(p *extractPlanPackage) (name string, err error) {
+func (w *extractWorker) dumpExtractPlanPackage(task *ExtractTask, p *extractPlanPackage) (name string, err error) {
 	f, name, err := GenerateExtractFile()
 	if err != nil {
 		return "", err
@@ -351,7 +363,7 @@ func (w *extractWorker) dumpExtractPlanPackage(p *extractPlanPackage) (name stri
 		return "", err
 	}
 	// dump extract plan task meta
-	if err = dumpExtractMeta(ExtractPlanType, zw); err != nil {
+	if err = dumpExtractMeta(task, zw); err != nil {
 		return "", err
 	}
 	// Dump Schema and View
@@ -371,8 +383,10 @@ func (w *extractWorker) dumpExtractPlanPackage(p *extractPlanPackage) (name stri
 		return "", err
 	}
 	// Dump stats
-	if err = dumpStats(zw, p.tables, GetDomain(w.sctx)); err != nil {
-		return "", err
+	if !task.SkipStats {
+		if err = dumpStats(zw, p.tables, GetDomain(w.sctx)); err != nil {
+			return "", err
+		}
 	}
 	// Dump sqls and plan
 	if err = dumpSQLRecords(p.records, zw); err != nil {
@@ -404,6 +418,7 @@ type singleSQLRecord struct {
 	SQL        string `json:"sql"`
 	Digest     string `json:"digest"`
 	BinaryPlan string `json:"binaryPlan"`
+	UserName   string `json:"userName"`
 }
 
 // dumpSQLRecord dumps sql records into one file for each record, the format is in json.
@@ -418,6 +433,7 @@ func dumpSQLRecord(record *stmtSummaryHistoryRecord, path string, zw *zip.Writer
 		SQL:        record.sql,
 		Digest:     record.digest,
 		BinaryPlan: record.binaryPlan,
+		UserName:   record.userName,
 	}
 	content, err := json.Marshal(singleSQLRecord)
 	if err != nil {
@@ -430,13 +446,18 @@ func dumpSQLRecord(record *stmtSummaryHistoryRecord, path string, zw *zip.Writer
 	return nil
 }
 
-func dumpExtractMeta(t ExtractType, zw *zip.Writer) error {
+func dumpExtractMeta(task *ExtractTask, zw *zip.Writer) error {
 	cf, err := zw.Create(ExtractMetaFile)
 	if err != nil {
 		return errors.AddStack(err)
 	}
 	varMap := make(map[string]string)
-	varMap[ExtractTaskType] = taskTypeToString(t)
+	varMap[ExtractTaskType] = taskTypeToString(task.ExtractType)
+	switch task.ExtractType {
+	case ExtractPlanType:
+		varMap[ExtractPlanTaskSkipStats] = strconv.FormatBool(task.SkipStats)
+	}
+
 	if err := toml.NewEncoder(cf).Encode(varMap); err != nil {
 		return errors.AddStack(err)
 	}
@@ -461,6 +482,7 @@ type stmtSummaryHistoryRecord struct {
 	planDigest string
 	sql        string
 	binaryPlan string
+	userName   string
 
 	plan string
 	skip bool

domain/extract_test.go

@@ -29,12 +29,23 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
+func TestExtractPlanWithoutHistoryView(t *testing.T) {
+	_, dom := testkit.CreateMockStoreAndDomain(t)
+	extractHandler := dom.GetExtractHandle()
+	task := domain.NewExtractPlanTask(time.Now(), time.Now())
+	task.UseHistoryView = false
+	_, err := extractHandler.ExtractTask(context.Background(), task)
+	require.NoError(t, err)
+}
+
 func TestExtractWithoutStmtSummaryPersistedEnabled(t *testing.T) {
 	setupStmtSummary()
 	closeStmtSummary()
 	_, dom := testkit.CreateMockStoreAndDomain(t)
 	extractHandler := dom.GetExtractHandle()
-	_, err := extractHandler.ExtractTask(context.Background(), domain.NewExtractPlanTask(time.Now(), time.Now()))
+	task := domain.NewExtractPlanTask(time.Now(), time.Now())
+	task.UseHistoryView = true
+	_, err := extractHandler.ExtractTask(context.Background(), task)
 	require.Error(t, err)
 }
 
@@ -61,7 +72,9 @@ func TestExtractHandlePlanTask(t *testing.T) {
 	time.Sleep(time.Second)
 	end := time.Now()
 	extractHandler := dom.GetExtractHandle()
-	name, err := extractHandler.ExtractTask(context.Background(), domain.NewExtractPlanTask(startTime, end))
+	task := domain.NewExtractPlanTask(startTime, end)
+	task.UseHistoryView = true
+	name, err := extractHandler.ExtractTask(context.Background(), task)
 	require.NoError(t, err)
 	require.True(t, len(name) > 0)
 }

domain/plan_replayer_test.go

@@ -46,7 +46,7 @@ func TestPlanReplayerGC(t *testing.T) {
 	require.True(t, os.IsNotExist(err))
 }
 
-func TestPlanReplayerParseTime(t *testing.T) {
+func TestDumpGCFileParseTime(t *testing.T) {
 	nowTime := time.Now()
 	name1 := fmt.Sprintf("replayer_single_xxxxxx_%v.zip", nowTime.UnixNano())
 	pt, err := parseTime(name1)
@@ -60,4 +60,49 @@ func TestPlanReplayerParseTime(t *testing.T) {
 	name3 := fmt.Sprintf("replayer_single_xxxxxx_%v._zip", nowTime.UnixNano())
 	_, err = parseTime(name3)
 	require.NotNil(t, err)
+
+	name4 := "extract_-brq6zKMarD9ayaifkHc4A==_1678168728477502000.zip"
+	_, err = parseTime(name4)
+	require.NoError(t, err)
+
+	var pName string
+	pName, err = replayer.GeneratePlanReplayerFileName(false, false, false)
+	require.NoError(t, err)
+	_, err = parseTime(pName)
+	require.NoError(t, err)
+
+	pName, err = replayer.GeneratePlanReplayerFileName(true, false, false)
+	require.NoError(t, err)
+	_, err = parseTime(pName)
+	require.NoError(t, err)
+
+	pName, err = replayer.GeneratePlanReplayerFileName(false, true, false)
+	require.NoError(t, err)
+	_, err = parseTime(pName)
+	require.NoError(t, err)
+
+	pName, err = replayer.GeneratePlanReplayerFileName(true, true, false)
+	require.NoError(t, err)
+	_, err = parseTime(pName)
+	require.NoError(t, err)
+
+	pName, err = replayer.GeneratePlanReplayerFileName(false, false, true)
+	require.NoError(t, err)
+	_, err = parseTime(pName)
+	require.NoError(t, err)
+
+	pName, err = replayer.GeneratePlanReplayerFileName(true, false, true)
+	require.NoError(t, err)
+	_, err = parseTime(pName)
+	require.NoError(t, err)
+
+	pName, err = replayer.GeneratePlanReplayerFileName(false, true, true)
+	require.NoError(t, err)
+	_, err = parseTime(pName)
+	require.NoError(t, err)
+
+	pName, err = replayer.GeneratePlanReplayerFileName(true, true, true)
+	require.NoError(t, err)
+	_, err = parseTime(pName)
+	require.NoError(t, err)
 }

server/extract.go

@@ -148,15 +148,26 @@ func buildExtractPlanTask(req *http.Request) (*domain.ExtractTask, bool, error)
 			return nil, false, err
 		}
 	}
-	isDumpStr := req.URL.Query().Get(pIsDump)
-	isDump, err := strconv.ParseBool(isDumpStr)
-	if err != nil {
-		isDump = false
-	}
+	isDump := extractBoolParam(pIsDump, false, req)
+
 	return &domain.ExtractTask{
 		ExtractType:     domain.ExtractPlanType,
 		IsBackgroundJob: false,
 		Begin:           begin,
 		End:             end,
+		SkipStats:       extractBoolParam(pIsSkipStats, false, req),
+		UseHistoryView:  extractBoolParam(pIsHistoryView, true, req),
 	}, isDump, nil
 }
+
+func extractBoolParam(param string, defaultValue bool, req *http.Request) bool {
+	str := req.URL.Query().Get(param)
+	if len(str) < 1 {
+		return defaultValue
+	}
+	v, err := strconv.ParseBool(str)
+	if err != nil {
+		return defaultValue
+	}
+	return v
+}

server/http_handler.go

@@ -91,8 +91,16 @@ const (
 	pDumpPartitionStats = "dumpPartitionStats"
 	pBegin              = "begin"
 	pEnd                = "end"
-	pType               = "type"
-	pIsDump             = "isDump"
+)
+
+// For extract task handler
+const (
+	pType   = "type"
+	pIsDump = "isDump"
+
+	// For extract plan task handler
+	pIsSkipStats   = "isSkipStats"
+	pIsHistoryView = "isHistoryView"
 )
 
 // For query string

util/replayer/replayer.go

@@ -50,6 +50,11 @@ func GeneratePlanReplayerFile(isCapture, isContinuesCapture, enableHistoricalSta
 	return zf, fileName, err
 }
 
+// GeneratePlanReplayerFileName generates plan replayer capture task name
+func GeneratePlanReplayerFileName(isCapture, isContinuesCapture, enableHistoricalStatsForCapture bool) (string, error) {
+	return generatePlanReplayerFileName(isCapture, isContinuesCapture, enableHistoricalStatsForCapture)
+}
+
 func generatePlanReplayerFileName(isCapture, isContinuesCapture, enableHistoricalStatsForCapture bool) (string, error) {
 	// Generate key and create zip file
 	time := time.Now().UnixNano()