sql,*: record UDF stats

Updates the UDF routine generators to pass down
the necessary information in order to generate
and record RecordedStmtStats structs

Epic: CRDB-55081
Release note: None

Author: kyle-a-wong

pkg/sql/conn_executor.go

@@ -3849,6 +3849,7 @@ func (ex *connExecutor) initEvalCtx(ctx context.Context, evalCtx *extendedEvalCo
 		validateDbZoneConfig: &ex.extraTxnState.validateDbZoneConfig,
 		persistedSQLStats:    ex.server.persistedSQLStats,
 		localSQLStats:        ex.server.localSqlStats,
+		sqlStatsIngester:     ex.server.sqlStatsIngester,
 		indexUsageStats:      ex.indexUsageStats,
 		statementPreparer:    ex,
 	}

pkg/sql/opt/exec/execbuilder/BUILD.bazel

@@ -17,6 +17,7 @@ go_library(
     deps = [
         "//pkg/kv/kvserver/concurrency/isolation",
         "//pkg/server/telemetry",
+        "//pkg/sql/appstatspb",
         "//pkg/sql/catalog/colinfo",
         "//pkg/sql/catalog/descpb",
         "//pkg/sql/lexbase",

pkg/sql/opt/exec/execbuilder/relational.go

@@ -3731,6 +3731,7 @@ func (b *Builder) buildCall(c *memo.CallExpr) (_ execPlan, outputCols colOrdMap,
 		udf.Def.BodyProps,
 		udf.Def.BodyStmts,
 		udf.Def.BodyTags,
+		udf.Def.BodyASTs,
 		false, /* allowOuterWithRefs */
 		nil,   /* wrapRootExpr */
 		0,     /* resultBufferID */

pkg/sql/opt/exec/execbuilder/scalar.go

@@ -8,6 +8,7 @@ package execbuilder
 import (
 	"context"
 
+	"github.com/cockroachdb/cockroach/pkg/sql/appstatspb"
 	"github.com/cockroachdb/cockroach/pkg/sql/opt"
 	"github.com/cockroachdb/cockroach/pkg/sql/opt/exec"
 	"github.com/cockroachdb/cockroach/pkg/sql/opt/memo"
@@ -689,16 +690,7 @@ func (b *Builder) buildExistsSubquery(
 
 		// Create a plan generator that can plan the single statement
 		// representing the subquery, and wrap the routine in a COALESCE.
-		planGen := b.buildRoutinePlanGenerator(
-			params,
-			stmts,
-			stmtProps,
-			nil, /* stmtStr */
-			make([]string, len(stmts)),
-			true, /* allowOuterWithRefs */
-			wrapRootExpr,
-			0, /* resultBufferID */
-		)
+		planGen := b.buildRoutinePlanGenerator(params, stmts, stmtProps, nil, make([]string, len(stmts)), nil, true, wrapRootExpr, 0)
 		return tree.NewTypedCoalesceExpr(tree.TypedExprs{
 			tree.NewTypedRoutineExpr(
 				"exists",
@@ -816,16 +808,7 @@ func (b *Builder) buildSubquery(
 
 		// Create a tree.RoutinePlanFn that can plan the single statement
 		// representing the subquery.
-		planGen := b.buildRoutinePlanGenerator(
-			params,
-			stmts,
-			stmtProps,
-			nil, /* stmtStr */
-			make([]string, len(stmts)),
-			true, /* allowOuterWithRefs */
-			nil,  /* wrapRootExpr */
-			0,    /* resultBufferID */
-		)
+		planGen := b.buildRoutinePlanGenerator(params, stmts, stmtProps, nil, make([]string, len(stmts)), nil, true, nil, 0)
 		_, tailCall := b.tailCalls[subquery]
 		return tree.NewTypedRoutineExpr(
 			"subquery",
@@ -900,7 +883,7 @@ func (b *Builder) buildSubquery(
 			if err != nil {
 				return err
 			}
-			err = fn(plan, "" /* stmtForDistSQLDiagram */, true /* isFinalPlan */)
+			err = fn(plan, tree.RoutineBodyStmt{} /* stmtForDistSQLDiagram */, true /* isFinalPlan */)
 			if err != nil {
 				return err
 			}
@@ -1011,16 +994,7 @@ func (b *Builder) buildUDF(ctx *buildScalarCtx, scalar opt.ScalarExpr) (tree.Typ
 	}
 
 	// Create a tree.RoutinePlanFn that can plan the statements in the UDF body.
-	planGen := b.buildRoutinePlanGenerator(
-		udf.Def.Params,
-		udf.Def.Body,
-		udf.Def.BodyProps,
-		udf.Def.BodyStmts,
-		udf.Def.BodyTags,
-		false, /* allowOuterWithRefs */
-		nil,   /* wrapRootExpr */
-		udf.Def.ResultBufferID,
-	)
+	planGen := b.buildRoutinePlanGenerator(udf.Def.Params, udf.Def.Body, udf.Def.BodyProps, udf.Def.BodyStmts, udf.Def.BodyTags, udf.Def.BodyASTs, false, nil, udf.Def.ResultBufferID)
 
 	// Enable stepping for volatile functions so that statements within the UDF
 	// see mutations made by the invoking statement and by previously executed
@@ -1085,16 +1059,7 @@ func (b *Builder) initRoutineExceptionHandler(
 		Actions: make([]*tree.RoutineExpr, len(exceptionBlock.Actions)),
 	}
 	for i, action := range exceptionBlock.Actions {
-		actionPlanGen := b.buildRoutinePlanGenerator(
-			action.Params,
-			action.Body,
-			action.BodyProps,
-			action.BodyStmts,
-			action.BodyTags,
-			false, /* allowOuterWithRefs */
-			nil,   /* wrapRootExpr */
-			0,     /* resultBufferID */
-		)
+		actionPlanGen := b.buildRoutinePlanGenerator(action.Params, action.Body, action.BodyProps, action.BodyStmts, action.BodyTags, nil, false, nil, 0)
 		// Build a routine with no arguments for the exception handler. The actual
 		// arguments will be supplied when (if) the handler is invoked.
 		exceptionHandler.Actions[i] = tree.NewTypedRoutineExpr(
@@ -1141,6 +1106,7 @@ func (b *Builder) buildRoutinePlanGenerator(
 	stmtProps []*physical.Required,
 	stmtStr []string,
 	stmtTags []string,
+	stmtASTs []tree.Statement,
 	allowOuterWithRefs bool,
 	wrapRootExpr wrapRootExprFn,
 	resultBufferID memo.RoutineResultBufferID,
@@ -1207,15 +1173,42 @@ func (b *Builder) buildRoutinePlanGenerator(
 		dbName := b.evalCtx.SessionData().Database
 		appName := b.evalCtx.SessionData().ApplicationName
 
+		format := tree.FmtHideConstants | tree.FmtFlags(tree.QueryFormattingForFingerprintsMask.Get(&b.evalCtx.Settings.SV))
 		for i := range stmts {
+			var bodyStmt tree.RoutineBodyStmt
 			stmt := stmts[i]
 			props := stmtProps[i]
 			var tag string
+			var fingerprint string
+			var summary string
+			var fpId appstatspb.StmtFingerprintID
+			var stmtString string
+			var stmtType tree.StatementType
 			// Theoretically, stmts and stmtTags should have the same length,
 			// but just to avoid an out-of-bounds panic, we have this check.
 			if i < len(stmtTags) {
 				tag = stmtTags[i]
 			}
+			if i < len(stmtStr) {
+				stmtString = stmtStr[i]
+			}
+			if i < len(stmtASTs) {
+				fingerprint = tree.FormatStatementHideConstants(stmtASTs[i], format)
+				summary = tree.FormatStatementSummary(stmtASTs[i], format)
+				fpId = appstatspb.ConstructStatementFingerprintID(fingerprint, b.evalCtx.TxnImplicit, dbName)
+				stmtType = stmtASTs[i].StatementType()
+			}
+
+			bodyStmt = tree.RoutineBodyStmt{
+				FingerprintStr: fingerprint,
+				SummaryStr:     summary,
+				StmtType:       stmtType,
+				StmtString:     stmtString,
+				DbName:         dbName,
+				AppName:        appName,
+				FingerprintId:  uint64(fpId),
+			}
+
 			o.Init(ctx, b.evalCtx, b.catalog)
 			f := o.Factory()
 
@@ -1319,12 +1312,8 @@ func (b *Builder) buildRoutinePlanGenerator(
 					return expectedLazyRoutineError("subquery")
 				}
 			}
-			var stmtForDistSQLDiagram string
-			if i < len(stmtStr) {
-				stmtForDistSQLDiagram = stmtStr[i]
-			}
 			incrementRoutineStmtCounter(b.evalCtx.StartedRoutineStatementCounters, dbName, appName, tag)
-			err = fn(plan, stmtForDistSQLDiagram, isFinalPlan)
+			err = fn(plan, bodyStmt, isFinalPlan)
 			if err != nil {
 				return err
 			}

pkg/sql/opt/memo/expr.go

@@ -742,6 +742,8 @@ type UDFDefinition struct {
 	// at the same position in Body.
 	BodyProps []*physical.Required
 
+	BodyASTs []tree.Statement
+
 	// BodyStmts, if set, is the string representation of each statement in
 	// Body. It is only populated when verbose tracing is enabled.
 	BodyStmts []string

pkg/sql/opt/optbuilder/routine.go

@@ -420,10 +420,15 @@ func (b *Builder) buildRoutine(
 	var bodyProps []*physical.Required
 	var bodyStmts []string
 	var bodyTags []string
+	var bodyASTs []tree.Statement
 	switch o.Language {
 	case tree.RoutineLangSQL:
 		// Parse the function body.
 		stmts, err := parser.Parse(o.Body)
+		bodyASTs = make([]tree.Statement, len(stmts))
+		for i := range stmts {
+			bodyASTs[i] = stmts[i].AST
+		}
 		if err != nil {
 			panic(err)
 		}
@@ -539,6 +544,7 @@ func (b *Builder) buildRoutine(
 				BodyProps:          bodyProps,
 				BodyStmts:          bodyStmts,
 				BodyTags:           bodyTags,
+				BodyASTs:           bodyASTs,
 				Params:             params,
 				ResultBufferID:     resultBufferID,
 			},

pkg/sql/planner.go

@@ -109,6 +109,8 @@ type extendedEvalContext struct {
 
 	localSQLStats *sslocal.SQLStats
 
+	sqlStatsIngester *sslocal.SQLStatsIngester
+
 	indexUsageStats *idxusage.LocalIndexUsageStats
 
 	SchemaChangerState *SchemaChangerState
@@ -536,12 +538,14 @@ func internalExtendedEvalCtx(
 	var indexUsageStatsController eval.IndexUsageStatsController
 	var sqlStatsProvider *persistedsqlstats.PersistedSQLStats
 	var localSqlStatsProvider *sslocal.SQLStats
+	var sqlStatsIngester *sslocal.SQLStatsIngester
 	if ief := execCfg.InternalDB; ief != nil {
 		if ief.server != nil {
 			indexUsageStats = ief.server.indexUsageStats
 			schemaTelemetryController = ief.server.schemaTelemetryController
 			sqlStatsProvider = ief.server.persistedSQLStats
 			localSqlStatsProvider = ief.server.localSqlStats
+			sqlStatsIngester = ief.server.sqlStatsIngester
 		} else {
 			// If the indexUsageStats is nil from the sql.Server, we create a dummy
 			// index usage stats collector. The sql.Server in the ExecutorConfig
@@ -553,6 +557,7 @@ func internalExtendedEvalCtx(
 			indexUsageStatsController = &idxusage.Controller{}
 			sqlStatsProvider = &persistedsqlstats.PersistedSQLStats{}
 			localSqlStatsProvider = &sslocal.SQLStats{}
+			sqlStatsIngester = &sslocal.SQLStatsIngester{}
 		}
 	}
 	ret := extendedEvalContext{
@@ -578,6 +583,7 @@ func internalExtendedEvalCtx(
 		indexUsageStats:   indexUsageStats,
 		persistedSQLStats: sqlStatsProvider,
 		localSQLStats:     localSqlStatsProvider,
+		sqlStatsIngester:  sqlStatsIngester,
 		jobs:              newTxnJobsCollection(),
 	}
 	ret.copyFromExecCfg(execCfg)

pkg/sql/routine.go

@@ -11,14 +11,18 @@ import (
 
 	"github.com/cockroachdb/cockroach/pkg/kv"
 	"github.com/cockroachdb/cockroach/pkg/server/telemetry"
+	"github.com/cockroachdb/cockroach/pkg/sql/appstatspb"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/colinfo"
+	"github.com/cockroachdb/cockroach/pkg/sql/clusterunique"
 	"github.com/cockroachdb/cockroach/pkg/sql/isql"
+	"github.com/cockroachdb/cockroach/pkg/sql/opt/exec/explain"
 	"github.com/cockroachdb/cockroach/pkg/sql/opt/memo"
 	"github.com/cockroachdb/cockroach/pkg/sql/parser/statements"
 	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgcode"
 	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgerror"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/eval"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
+	"github.com/cockroachdb/cockroach/pkg/sql/sqlstats"
 	"github.com/cockroachdb/cockroach/pkg/sql/sqltelemetry"
 	"github.com/cockroachdb/cockroach/pkg/sql/types"
 	"github.com/cockroachdb/cockroach/pkg/storage/enginepb"
@@ -325,7 +329,7 @@ func (g *routineGenerator) startInternal(ctx context.Context, txn *kv.Txn) (err
 	rrw := NewRowResultWriter(&g.rch)
 	var cursorHelper *plpgsqlCursorHelper
 	err = g.expr.ForEachPlan(ctx, ef, rrw, g.args,
-		func(plan tree.RoutinePlan, stmtForDistSQLDiagram string, isFinalPlan bool) error {
+		func(plan tree.RoutinePlan, stmt tree.RoutineBodyStmt, isFinalPlan bool) error {
 			stmtIdx++
 			opName := "routine-stmt-" + g.expr.Name + "-" + strconv.Itoa(stmtIdx)
 			ctx, sp := tracing.ChildSpan(ctx, opName)
@@ -368,10 +372,31 @@ func (g *routineGenerator) startInternal(ctx context.Context, txn *kv.Txn) (err
 
 			// Run the plan.
 			params := runParams{ctx, g.p.ExtendedEvalContext(), g.p}
-			queryStats, err := runPlanInsidePlan(ctx, params, plan.(*planComponents), w, g, stmtForDistSQLDiagram)
+			var builder sqlstats.RecordedStatementStatsBuilder[sqlstats.StatementLatencyRecorder, topLevelQueryStats, PlanInfo]
+			if stmt.FingerprintId != 0 {
+				builder = sqlstats.NewRecordedStatementStatsBuilder[sqlstats.StatementLatencyRecorder, topLevelQueryStats, PlanInfo](
+					appstatspb.StmtFingerprintID(stmt.FingerprintId),
+					g.p.extendedEvalCtx.SessionID,
+					stmt.DbName,
+					stmt.FingerprintStr,
+					stmt.SummaryStr,
+					clusterunique.ID{},
+					stmt.StmtType, stmt.AppName,
+					PlanInfo{
+						planFlags: g.p.curPlan.flags, // TODO: Is this right?
+						planGist:  explain.PlanGist{},
+					})
+				defer func() {
+					g.p.ExtendedEvalContext().sqlStatsIngester.RecordStatement(builder.Build())
+				}()
+			}
+			queryStats, err := runPlanInsidePlan(ctx, params, plan.(*planComponents), w, g, stmt.StmtString)
+
 			if err != nil {
+				builder = builder.StatementError(err)
 				return err
 			}
+			builder = builder.QueryLevelStats(queryStats)
 			forwardInnerQueryStats(g.p.routineMetadataForwarder, queryStats)
 			if openCursor {
 				return cursorHelper.createCursor(g.p)

pkg/sql/sem/tree/routine.go

@@ -29,13 +29,23 @@ type RoutinePlanGenerator func(
 	fn RoutinePlanGeneratedFunc,
 ) error
 
+type RoutineBodyStmt struct {
+	FingerprintStr string
+	SummaryStr     string
+	StmtType       StatementType
+	StmtString     string
+	DbName         string
+	AppName        string
+	FingerprintId  uint64
+}
+
 // RoutinePlanGeneratedFunc is the function type that is called for each plan
 // enumerated by a RoutinePlanGenerator.
 // - stmtForDistSQLDiagram, if set, will be used when generating DistSQL diagram
 // to specify the SQL stmt corresponding to the plan.
 // - isFinalPlan is true if no more plans will be generated after the current
 // plan.
-type RoutinePlanGeneratedFunc func(plan RoutinePlan, stmtForDistSQLDiagram string, isFinalPlan bool) error
+type RoutinePlanGeneratedFunc func(plan RoutinePlan, stmt RoutineBodyStmt, isFinalPlan bool) error
 
 // RoutinePlan represents a plan for a statement in a routine. It currently maps
 // to exec.Plan. We use the empty interface here rather than exec.Plan to avoid

pkg/sql/sqlstats/insights/BUILD.bazel

@@ -28,6 +28,7 @@ go_library(
         "//pkg/util/metric",
         "//pkg/util/quantile",
         "//pkg/util/syncutil",
+        "//pkg/util/uint128",
         "//pkg/util/uuid",
         "@com_github_cockroachdb_redact//:redact",
         "@com_github_prometheus_client_model//go",