Merge #156905

156905: sql: Record statement stats for statements executed in UDFs / SPs r=kyle-a-wong a=kyle-a-wong

Updates UDF and SP execution to support recording
statement statistics and insights. Now, each
sub statement executed in a UDF and SP will have
latencies and execution stats stored in the
sql stats subsystem.

Additionally, these statements will be appear
in the transaction statistics details page
in table containing "statement fingerprints for
this transaction".

Fixes: [CRDB-57025](https://cockroachlabs.atlassian.net/browse/CRDB-57025)
Epic: [CRDB-55081](https://cockroachlabs.atlassian.net/browse/CRDB-55081)
Release note (sql change): SQL Statements
executed in stored procedures and user defined
functions will now have sql stats recorded and
should appear in the sql activity and insights
pages in db console. Limitations: statements executed
in a UDF / SP won't be collected for active
statement diagnostics requests. Statement diagnostics
continue to only work for top level statement
executions

Co-authored-by: Kyle Wong <[email protected]>

Author: craig[bot]

pkg/sql/conn_executor.go

@@ -1311,8 +1311,6 @@ func (s *Server) newConnExecutor(
 
 	ex.applicationName.Store(ex.sessionData().ApplicationName)
 	ex.applicationStats = applicationStats
-	// We ignore statements and transactions run by the internal executor by
-	// passing a nil writer.
 	ex.statsCollector = sslocal.NewStatsCollector(
 		s.cfg.Settings,
 		applicationStats,
@@ -4020,7 +4018,7 @@ func (ex *connExecutor) initPlanner(ctx context.Context, p *planner) {
 	p.cancelChecker.Reset(ctx)
 
 	ex.initEvalCtx(ctx, &p.extendedEvalCtx, p)
-
+	p.statsCollector = ex.statsCollector
 	p.sessionDataMutatorIterator = ex.dataMutatorIterator
 	p.noticeSender = nil
 	p.preparedStatements = ex.getPrepStmtsAccessor()

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",
@@ -45,6 +46,7 @@ go_library(
         "//pkg/sql/sem/tree/treewindow",
         "//pkg/sql/sem/volatility",
         "//pkg/sql/sqlerrors",
+        "//pkg/sql/sqlstats",
         "//pkg/sql/sqltelemetry",
         "//pkg/sql/types",
         "//pkg/util/buildutil",

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

@@ -3719,6 +3719,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,8 +8,10 @@ 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/exec/explain"
 	"github.com/cockroachdb/cockroach/pkg/sql/opt/memo"
 	"github.com/cockroachdb/cockroach/pkg/sql/opt/norm"
 	"github.com/cockroachdb/cockroach/pkg/sql/opt/props"
@@ -24,6 +26,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree/treecmp"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/volatility"
 	"github.com/cockroachdb/cockroach/pkg/sql/sqlerrors"
+	"github.com/cockroachdb/cockroach/pkg/sql/sqlstats"
 	"github.com/cockroachdb/cockroach/pkg/sql/types"
 	"github.com/cockroachdb/cockroach/pkg/util/errorutil"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
@@ -695,6 +698,7 @@ func (b *Builder) buildExistsSubquery(
 			stmtProps,
 			nil, /* stmtStr */
 			make([]string, len(stmts)),
+			nil,  /* stmtASTs */
 			true, /* allowOuterWithRefs */
 			wrapRootExpr,
 			0, /* resultBufferID */
@@ -822,6 +826,7 @@ func (b *Builder) buildSubquery(
 			stmtProps,
 			nil, /* stmtStr */
 			make([]string, len(stmts)),
+			nil,  /* stmtASTs */
 			true, /* allowOuterWithRefs */
 			nil,  /* wrapRootExpr */
 			0,    /* resultBufferID */
@@ -900,7 +905,7 @@ func (b *Builder) buildSubquery(
 			if err != nil {
 				return err
 			}
-			err = fn(plan, "" /* stmtForDistSQLDiagram */, true /* isFinalPlan */)
+			err = fn(plan, nil /* routineStatsBuilder */, "" /* stmtForDistSQLDiagram */, true /* isFinalPlan */)
 			if err != nil {
 				return err
 			}
@@ -1017,6 +1022,7 @@ func (b *Builder) buildUDF(ctx *buildScalarCtx, scalar opt.ScalarExpr) (tree.Typ
 		udf.Def.BodyProps,
 		udf.Def.BodyStmts,
 		udf.Def.BodyTags,
+		udf.Def.BodyASTs,
 		false, /* allowOuterWithRefs */
 		nil,   /* wrapRootExpr */
 		udf.Def.ResultBufferID,
@@ -1091,6 +1097,7 @@ func (b *Builder) initRoutineExceptionHandler(
 			action.BodyProps,
 			action.BodyStmts,
 			action.BodyTags,
+			nil,   /* stmtASTs */
 			false, /* allowOuterWithRefs */
 			nil,   /* wrapRootExpr */
 			0,     /* resultBufferID */
@@ -1141,6 +1148,7 @@ func (b *Builder) buildRoutinePlanGenerator(
 	stmtProps []*physical.Required,
 	stmtStr []string,
 	stmtTags []string,
+	stmtASTs []tree.Statement,
 	allowOuterWithRefs bool,
 	wrapRootExpr wrapRootExprFn,
 	resultBufferID memo.RoutineResultBufferID,
@@ -1174,6 +1182,8 @@ func (b *Builder) buildRoutinePlanGenerator(
 	//
 	// Note: we put o outside of the function so we allocate it only once.
 	var o xform.Optimizer
+	var gistFactory explain.PlanGistFactory
+	var latencyRecorder = sqlstats.NewStatementLatencyRecorder()
 	originalMemo := b.mem
 	planGen := func(
 		ctx context.Context,
@@ -1206,8 +1216,14 @@ func (b *Builder) buildRoutinePlanGenerator(
 
 		dbName := b.evalCtx.SessionData().Database
 		appName := b.evalCtx.SessionData().ApplicationName
-
+		// TODO(yuzefovich): look into computing fingerprintFormat lazily.
+		fingerprintFormat := tree.FmtHideConstants | tree.FmtFlags(tree.QueryFormattingForFingerprintsMask.Get(&b.evalCtx.Settings.SV))
 		for i := range stmts {
+			latencyRecorder.Reset()
+			var builder *sqlstats.RecordedStatementStatsBuilder
+			var statsBuilderWithLatencies tree.RoutineStatsBuilder
+			sqlstats.RecordStatementPhase(latencyRecorder, sqlstats.StatementStarted)
+			sqlstats.RecordStatementPhase(latencyRecorder, sqlstats.StatementStartParsing)
 			stmt := stmts[i]
 			props := stmtProps[i]
 			var tag string
@@ -1216,6 +1232,22 @@ func (b *Builder) buildRoutinePlanGenerator(
 			if i < len(stmtTags) {
 				tag = stmtTags[i]
 			}
+			if i < len(stmtASTs) && stmtASTs[i] != nil {
+				fingerprint := tree.FormatStatementHideConstants(stmtASTs[i], fingerprintFormat)
+				fpId := appstatspb.ConstructStatementFingerprintID(fingerprint, b.evalCtx.TxnImplicit, dbName)
+				summary := tree.FormatStatementSummary(stmtASTs[i], fingerprintFormat)
+				stmtType := stmtASTs[i].StatementType()
+				builder = sqlstats.NewRecordedStatementStatsBuilder(
+					fpId, dbName, fingerprint, summary, stmtType, appName,
+				)
+
+				statsBuilderWithLatencies = &sqlstats.StatsBuilderWithLatencyRecorder{
+					StatsBuilder:    builder,
+					LatencyRecorder: latencyRecorder,
+				}
+			}
+			sqlstats.RecordStatementPhase(latencyRecorder, sqlstats.StatementEndParsing)
+			sqlstats.RecordStatementPhase(latencyRecorder, sqlstats.StatementStartPlanning)
 			o.Init(ctx, b.evalCtx, b.catalog)
 			f := o.Factory()
 
@@ -1287,9 +1319,14 @@ func (b *Builder) buildRoutinePlanGenerator(
 				tailCalls = make(map[opt.ScalarExpr]struct{})
 				memo.ExtractTailCalls(optimizedExpr, tailCalls)
 			}
-
 			// Build the memo into a plan.
 			ef := ref.(exec.Factory)
+			if builder != nil && !b.evalCtx.SessionData().DisablePlanGists {
+				gistFactory.Reset()
+				gistFactory.Init(ef)
+				ef = &gistFactory
+			}
+
 			eb := New(ctx, ef, &o, f.Memo(), b.catalog, optimizedExpr, b.semaCtx, b.evalCtx, false /* allowAutoCommit */, b.IsANSIDML)
 			eb.withExprs = withExprs
 			eb.disableTelemetry = true
@@ -1303,6 +1340,10 @@ func (b *Builder) buildRoutinePlanGenerator(
 				eb.addRoutineResultBuffer(resultBufferID, resultWriter)
 			}
 			plan, err := eb.Build()
+			if gistFactory.Initialized() {
+				planGist := gistFactory.PlanGist()
+				builder.PlanGist(planGist.String(), planGist.Hash())
+			}
 			if err != nil {
 				if errors.IsAssertionFailure(err) {
 					// Enhance the error with the EXPLAIN (OPT, VERBOSE) of the
@@ -1324,7 +1365,8 @@ func (b *Builder) buildRoutinePlanGenerator(
 				stmtForDistSQLDiagram = stmtStr[i]
 			}
 			incrementRoutineStmtCounter(b.evalCtx.StartedRoutineStatementCounters, dbName, appName, tag)
-			err = fn(plan, stmtForDistSQLDiagram, isFinalPlan)
+			sqlstats.RecordStatementPhase(latencyRecorder, sqlstats.StatementEndPlanning)
+			err = fn(plan, statsBuilderWithLatencies, stmtForDistSQLDiagram, isFinalPlan)
 			if err != nil {
 				return err
 			}

pkg/sql/opt/memo/expr.go

@@ -742,6 +742,11 @@ type UDFDefinition struct {
 	// at the same position in Body.
 	BodyProps []*physical.Required
 
+	// BodyASTs contains the AST representation of each statement in Body. The
+	// size of this slice matches that of Body, but it may contain nil entries
+	// for non tree.RoutineLangSQL types.
+	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/plpgsql.go

@@ -678,7 +678,7 @@ func (b *plpgsqlBuilder) buildPLpgSQLStatements(stmts []ast.Statement, s *scope)
 			if b.options.insideDataSource && b.setReturnType.Family() == types.TupleFamily {
 				retNextScope = b.ob.expandRoutineTupleIntoCols(retNextScope)
 			}
-			b.appendBodyStmtFromScope(&retCon, retNextScope, "" /* stmtTag */)
+			b.appendBodyStmtFromScope(&retCon, retNextScope, nil /* stmt */)
 			b.appendPlpgSQLStmts(&retCon, stmts[i+1:])
 			return b.callContinuation(&retCon, s)
 
@@ -720,7 +720,7 @@ func (b *plpgsqlBuilder) buildPLpgSQLStatements(stmts []ast.Statement, s *scope)
 			if !b.options.insideDataSource && b.setReturnType.Family() == types.TupleFamily {
 				retQueryScope = b.ob.combineRoutineColsIntoTuple(retQueryScope)
 			}
-			b.appendBodyStmtFromScope(&retCon, retQueryScope, t.SqlStmt.StatementTag())
+			b.appendBodyStmtFromScope(&retCon, retQueryScope, t.SqlStmt)
 			b.appendPlpgSQLStmts(&retCon, stmts[i+1:])
 			return b.callContinuation(&retCon, s)
 
@@ -949,7 +949,7 @@ func (b *plpgsqlBuilder) buildPLpgSQLStatements(stmts []ast.Statement, s *scope)
 			// crdb_internal.plpgsql_raise builtin function.
 			con := b.makeContinuation("_stmt_raise")
 			con.def.Volatility = volatility.Volatile
-			b.appendBodyStmtFromScope(&con, b.buildPLpgSQLRaise(con.s, b.getRaiseArgs(con.s, t)), "" /* stmtTag */)
+			b.appendBodyStmtFromScope(&con, b.buildPLpgSQLRaise(con.s, b.getRaiseArgs(con.s, t)), nil /* stmt */)
 			b.appendPlpgSQLStmts(&con, stmts[i+1:])
 			return b.callContinuation(&con, s)
 
@@ -970,7 +970,7 @@ func (b *plpgsqlBuilder) buildPLpgSQLStatements(stmts []ast.Statement, s *scope)
 			if len(t.Target) == 0 {
 				// When there is no INTO target, build the SQL statement into a body
 				// statement that is only executed for its side effects.
-				b.appendBodyStmtFromScope(&execCon, stmtScope, t.SqlStmt.StatementTag())
+				b.appendBodyStmtFromScope(&execCon, stmtScope, t.SqlStmt)
 				b.appendPlpgSQLStmts(&execCon, stmts[i+1:])
 				return b.callContinuation(&execCon, s)
 			}
@@ -1029,7 +1029,7 @@ func (b *plpgsqlBuilder) buildPLpgSQLStatements(stmts []ast.Statement, s *scope)
 			intoScope = b.callContinuation(&retCon, intoScope)
 
 			// Step 3: call the INTO continuation from the parent scope.
-			b.appendBodyStmtFromScope(&execCon, intoScope, t.SqlStmt.StatementTag())
+			b.appendBodyStmtFromScope(&execCon, intoScope, t.SqlStmt)
 			return b.callContinuation(&execCon, s)
 
 		case *ast.Open:
@@ -1069,7 +1069,7 @@ func (b *plpgsqlBuilder) buildPLpgSQLStatements(stmts []ast.Statement, s *scope)
 				// Cursors with mutations are invalid.
 				panic(cursorMutationErr)
 			}
-			b.appendBodyStmtFromScope(&openCon, openScope, query.StatementTag())
+			b.appendBodyStmtFromScope(&openCon, openScope, query)
 			b.appendPlpgSQLStmts(&openCon, stmts[i+1:])
 
 			// Build a statement to generate a unique name for the cursor if one
@@ -1079,7 +1079,7 @@ func (b *plpgsqlBuilder) buildPLpgSQLStatements(stmts []ast.Statement, s *scope)
 			nameCon := b.makeContinuation("_gen_cursor_name")
 			nameCon.def.Volatility = volatility.Volatile
 			nameScope := b.buildCursorNameGen(&nameCon, t.CurVar)
-			b.appendBodyStmtFromScope(&nameCon, b.callContinuation(&openCon, nameScope), "" /* stmtTag */)
+			b.appendBodyStmtFromScope(&nameCon, b.callContinuation(&openCon, nameScope), nil /* stmt */)
 			return b.callContinuation(&nameCon, s)
 
 		case *ast.Close:
@@ -1119,7 +1119,7 @@ func (b *plpgsqlBuilder) buildPLpgSQLStatements(stmts []ast.Statement, s *scope)
 			closeScope := closeCon.s.push()
 			b.ob.synthesizeColumn(closeScope, closeColName, types.Int, nil /* expr */, closeCall)
 			b.ob.constructProjectForScope(closeCon.s, closeScope)
-			b.appendBodyStmtFromScope(&closeCon, closeScope, "" /* stmtTag */)
+			b.appendBodyStmtFromScope(&closeCon, closeScope, nil /* stmt */)
 			b.appendPlpgSQLStmts(&closeCon, stmts[i+1:])
 			return b.callContinuation(&closeCon, s)
 
@@ -1143,7 +1143,7 @@ func (b *plpgsqlBuilder) buildPLpgSQLStatements(stmts []ast.Statement, s *scope)
 			fetchCon.def.Volatility = volatility.Volatile
 			fetchScope := b.buildFetch(fetchCon.s, t)
 			if t.IsMove {
-				b.appendBodyStmtFromScope(&fetchCon, fetchScope, "" /* stmtTag */)
+				b.appendBodyStmtFromScope(&fetchCon, fetchScope, nil /* stmt */)
 				b.appendPlpgSQLStmts(&fetchCon, stmts[i+1:])
 				return b.callContinuation(&fetchCon, s)
 			}
@@ -1174,7 +1174,7 @@ func (b *plpgsqlBuilder) buildPLpgSQLStatements(stmts []ast.Statement, s *scope)
 			intoScope = b.callContinuation(&retCon, intoScope)
 
 			// Add the built statement to the FETCH continuation.
-			b.appendBodyStmtFromScope(&fetchCon, intoScope, "" /* stmtTag */)
+			b.appendBodyStmtFromScope(&fetchCon, intoScope, nil /* stmt */)
 			return b.callContinuation(&fetchCon, s)
 
 		case *ast.Null:
@@ -1276,7 +1276,7 @@ func (b *plpgsqlBuilder) buildPLpgSQLStatements(stmts []ast.Statement, s *scope)
 			if len(target) == 0 {
 				// When there is no INTO target, build the nested procedure call into a
 				// body statement that is only executed for its side effects.
-				b.appendBodyStmtFromScope(&callCon, callScope, "" /* stmtTag */)
+				b.appendBodyStmtFromScope(&callCon, callScope, nil /* stmt */)
 				b.appendPlpgSQLStmts(&callCon, stmts[i+1:])
 				return b.callContinuation(&callCon, s)
 			}
@@ -1291,7 +1291,7 @@ func (b *plpgsqlBuilder) buildPLpgSQLStatements(stmts []ast.Statement, s *scope)
 			intoScope = b.callContinuation(&retCon, intoScope)
 
 			// Add the built statement to the CALL continuation.
-			b.appendBodyStmtFromScope(&callCon, intoScope, "" /* stmtTag */)
+			b.appendBodyStmtFromScope(&callCon, intoScope, nil /* stmt */)
 			return b.callContinuation(&callCon, s)
 
 		case *ast.DoBlock:
@@ -1307,7 +1307,7 @@ func (b *plpgsqlBuilder) buildPLpgSQLStatements(stmts []ast.Statement, s *scope)
 			doCon := b.makeContinuation("_stmt_do")
 			doCon.def.Volatility = volatility.Volatile
 			bodyScope := b.ob.buildPLpgSQLDoBody(t)
-			b.appendBodyStmtFromScope(&doCon, bodyScope, "" /* stmtTag */)
+			b.appendBodyStmtFromScope(&doCon, bodyScope, nil /* stmt */)
 			b.appendPlpgSQLStmts(&doCon, stmts[i+1:])
 			return b.callContinuation(&doCon, s)
 
@@ -1457,7 +1457,7 @@ func (b *plpgsqlBuilder) handleIntForLoop(
 	)
 	// Call recursively into the loop body continuation.
 	incScope = b.callContinuation(&loopCon, incScope)
-	b.appendBodyStmtFromScope(&incrementCon, incScope, "" /* stmtTag */)
+	b.appendBodyStmtFromScope(&incrementCon, incScope, nil /* stmt */)
 
 	// Notably, we call the loop body continuation here, rather than the
 	// increment continuation, because the counter should not be incremented
@@ -2021,7 +2021,7 @@ func (b *plpgsqlBuilder) buildEndOfFunctionRaise(con *continuation) {
 		pgcode.RoutineExceptionFunctionExecutedNoReturnStatement.String(), /* code */
 	)
 	con.def.Volatility = volatility.Volatile
-	b.appendBodyStmtFromScope(con, b.buildPLpgSQLRaise(con.s, args), "" /* stmtTag */)
+	b.appendBodyStmtFromScope(con, b.buildPLpgSQLRaise(con.s, args), nil /* stmt */)
 
 	// Build a dummy statement that returns NULL. It won't be executed, but
 	// ensures that the continuation routine's return type is correct.
@@ -2030,7 +2030,7 @@ func (b *plpgsqlBuilder) buildEndOfFunctionRaise(con *continuation) {
 	typedNull := b.ob.factory.ConstructNull(b.returnType)
 	b.ob.synthesizeColumn(eofScope, eofColName, b.returnType, nil /* expr */, typedNull)
 	b.ob.constructProjectForScope(con.s, eofScope)
-	b.appendBodyStmtFromScope(con, eofScope, "" /* stmtTag */)
+	b.appendBodyStmtFromScope(con, eofScope, nil /* stmt */)
 }
 
 // addOneRowCheck handles INTO STRICT, where a SQL statement is required to
@@ -2280,8 +2280,10 @@ func (b *plpgsqlBuilder) makeContinuationWithTyp(
 // appendBodyStmtFromScope is separate from makeContinuation to allow recursive
 // routine definitions, which need to push the continuation before it is
 // finished. The separation also allows for appending multiple body statements.
+// If the continuation does not map to an explicit statement in the body, stmt
+// may be nil.
 func (b *plpgsqlBuilder) appendBodyStmtFromScope(
-	con *continuation, bodyScope *scope, stmtTag string,
+	con *continuation, bodyScope *scope, stmt tree.Statement,
 ) {
 	// Set the volatility of the continuation routine to the least restrictive
 	// volatility level in the Relational properties of the body statements.
@@ -2291,8 +2293,11 @@ func (b *plpgsqlBuilder) appendBodyStmtFromScope(
 		con.def.Volatility = vol
 	}
 	con.def.Body = append(con.def.Body, bodyExpr)
-	con.def.BodyTags = append(con.def.BodyTags, stmtTag)
 	con.def.BodyProps = append(con.def.BodyProps, bodyScope.makePhysicalProps())
+	if stmt != nil {
+		con.def.BodyTags = append(con.def.BodyTags, stmt.StatementTag())
+		con.def.BodyASTs = append(con.def.BodyASTs, stmt)
+	}
 }
 
 // appendPlpgSQLStmts builds the given PLpgSQL statements into a relational
@@ -2302,7 +2307,7 @@ func (b *plpgsqlBuilder) appendPlpgSQLStmts(con *continuation, stmts []ast.State
 	// Make sure to push s before constructing the continuation scope to ensure
 	// that the parameter columns are not projected.
 	continuationScope := b.buildPLpgSQLStatements(stmts, con.s.push())
-	b.appendBodyStmtFromScope(con, continuationScope, "" /* stmtTag */)
+	b.appendBodyStmtFromScope(con, continuationScope, nil /* stmt */)
 }
 
 // callContinuation adds a column that projects the result of calling the