*: [DNM] read billing demo calibration surfaces

pkg/executor/adapter.go

@@ -366,11 +366,12 @@ type ExecStmt struct {
 	Ctx sessionctx.Context
 
 	// LowerPriority represents whether to lower the execution priority of a query.
-	LowerPriority     bool
-	isPreparedStmt    bool
-	isSelectForUpdate bool
-	retryCount        uint
-	retryStartTime    time.Time
+	LowerPriority        bool
+	isPreparedStmt       bool
+	isSelectForUpdate    bool
+	resolvedPreparedStmt ast.StmtNode
+	retryCount           uint
+	retryStartTime       time.Time
 
 	// Phase durations are splited into two parts: 1. trying to lock keys (but
 	// failed); 2. the final iteration of the retry loop. Here we use
@@ -770,11 +771,19 @@ func (a *ExecStmt) inheritContextFromExecuteStmt() {
 		a.Ctx.SetValue(sessionctx.QueryString, executePlan.Stmt.Text())
 		a.OutputNames = executePlan.OutputNames()
 		a.isPreparedStmt = true
+		a.resolvedPreparedStmt = executePlan.Stmt
 		a.Plan = executePlan.Plan
 		a.Ctx.GetSessionVars().StmtCtx.SetPlan(executePlan.Plan)
 	}
 }
 
+func (a *ExecStmt) readBillingDemoStmtNode() ast.StmtNode {
+	if a.resolvedPreparedStmt != nil {
+		return a.resolvedPreparedStmt
+	}
+	return a.StmtNode
+}
+
 func (a *ExecStmt) getSQLForProcessInfo() string {
 	sql := a.Text()
 	if simple, ok := a.Plan.(*plannercore.Simple); ok && simple.Statement != nil {
@@ -1689,10 +1698,14 @@ func (a *ExecStmt) FinishExecuteStmt(txnTS uint64, err error, hasMoreResults boo
 	}
 
 	a.finalizeStatementRUV2Metrics()
+	// Lazy SELECT metrics are emitted when the record set is closed. If a
+	// client abandons a result without closing/draining it, this first demo can
+	// miss that statement instead of guessing an incomplete status.
+	readBillingDemoStats := plannercore.RecordReadBillingDemoForStatement(a.Ctx, a.Plan, a.readBillingDemoStmtNode(), err)
 	a.updateNetworkTrafficStatsAndMetrics()
 	// `LowSlowQuery` and `SummaryStmt` must be called before recording `PrevStmt`.
 	a.LogSlowQuery(txnTS, succ, hasMoreResults)
-	a.SummaryStmt(succ)
+	a.SummaryStmt(succ, readBillingDemoStats)
 	a.observeStmtFinishedForTopProfiling()
 	a.UpdatePlanCacheRuntimeInfo()
 	if sessVars.StmtCtx.IsTiFlash.Load() {
@@ -2179,7 +2192,7 @@ func (digest planDigestAlias) planDigestDumpTriggerCheck(config *traceevent.Dump
 }
 
 // SummaryStmt collects statements for information_schema.statements_summary
-func (a *ExecStmt) SummaryStmt(succ bool) {
+func (a *ExecStmt) SummaryStmt(succ bool, readBillingDemoStats stmtsummary.ReadBillingDemoStatementStats) {
 	sessVars := a.Ctx.GetSessionVars()
 	var userString string
 	if sessVars.User != nil {
@@ -2283,6 +2296,8 @@ func (a *ExecStmt) SummaryStmt(succ bool) {
 	stmtExecInfo.KeyspaceID = keyspaceID
 	stmtExecInfo.RUDetail = ruDetail
 	stmtExecInfo.TotalRUV2 = calculateStatementTotalRUV2(sessVars.RUV2Metrics, sessVars.RUV2Weights(), ruDetail)
+	stmtExecInfo.ReadBillingDemoStats = readBillingDemoStats
+	stmtExecInfo.ReadBillingDemoBaseUnits = readBillingDemoStats.Totals
 	stmtExecInfo.ResourceGroupName = sessVars.StmtCtx.ResourceGroupName
 	stmtExecInfo.CPUUsages = sessVars.SQLCPUUsages.GetCPUUsages()
 	stmtExecInfo.PlanCacheUnqualified = sessVars.StmtCtx.PlanCacheUnqualified()

pkg/executor/builder.go

@@ -2616,10 +2616,18 @@ func (b *executorBuilder) buildMemTable(v *physicalop.PhysicalMemTable) exec.Exe
 		case strings.ToLower(infoschema.TableStatementsSummary),
 			strings.ToLower(infoschema.TableStatementsSummaryHistory),
 			strings.ToLower(infoschema.TableStatementsSummaryEvicted),
+			strings.ToLower(infoschema.TableStatementsSummaryReadBillingDemoBaseUnits),
+			strings.ToLower(infoschema.TableStatementsSummaryHistoryReadBillingDemoBaseUnits),
+			strings.ToLower(infoschema.TableStatementsSummaryReadBillingDemoStatus),
+			strings.ToLower(infoschema.TableStatementsSummaryHistoryReadBillingDemoStatus),
 			strings.ToLower(infoschema.TableTiDBStatementsStats),
 			strings.ToLower(infoschema.ClusterTableStatementsSummary),
 			strings.ToLower(infoschema.ClusterTableStatementsSummaryHistory),
 			strings.ToLower(infoschema.ClusterTableStatementsSummaryEvicted),
+			strings.ToLower(infoschema.ClusterTableStatementsSummaryReadBillingDemoBaseUnits),
+			strings.ToLower(infoschema.ClusterTableStatementsSummaryHistoryReadBillingDemoBaseUnits),
+			strings.ToLower(infoschema.ClusterTableStatementsSummaryReadBillingDemoStatus),
+			strings.ToLower(infoschema.ClusterTableStatementsSummaryHistoryReadBillingDemoStatus),
 			strings.ToLower(infoschema.ClusterTableTiDBStatementsStats):
 			var extractor *plannercore.StatementsSummaryExtractor
 			if v.Extractor != nil {

pkg/executor/compiler.go

@@ -153,6 +153,9 @@ func (c *Compiler) Compile(ctx context.Context, stmtNode ast.StmtNode) (_ *ExecS
 		OutputNames:   names,
 		Ti:            &TelemetryInfo{},
 	}
+	if preparedObj != nil && preparedObj.PreparedAst != nil {
+		stmt.resolvedPreparedStmt = preparedObj.PreparedAst.Stmt
+	}
 	// Use cached plan if possible.
 	if preparedObj != nil && plannercore.IsSafeToReusePointGetExecutor(c.Ctx, is, preparedObj) {
 		if exec, isExec := finalPlan.(*plannercore.Execute); isExec {