sql: add builtin to request statement bundles

Makes cockroachdb#79547 less painful.

Previously, there was no way to request statement bundle builtins for
internal statements.
This commit introduces crdb_internal.request_statement_bundle builtin
which allows statement bundle being requested from SQL CLI. The new
builtin takes three parameters:
* statement fingerprint text
* minimum execution latency for the statement
* length of duration the statement bundle request will stay valid for
VIEWACTIVITY or ADMIN role option is required to use this builtin. If user has
VIEWACTIVITYREDACTED role option, user is not allowed to use this builtin.

Release note (sql change): new crdb_internal.request_statement_bundle builtin
allows statement bundle being requested from SQL CLI. The new builtin takes
three parameters:
* statement fingerprint text
* minimum execution latency for the statement
* length of duration the statement bundle request will stay valid for.
VIEWACTIVITY or ADMIN role option is required to use this builtin. If user has
VIEWACTIVITYREDACTED role option, user is not allowed to use this builtin.

docs/generated/sql/functions.md

@@ -3075,6 +3075,11 @@ SELECT * FROM crdb_internal.check_consistency(true, ‘\x02’, ‘\x04’)</p>
 </span></td></tr>
 <tr><td><a name="crdb_internal.repair_ttl_table_scheduled_job"></a><code>crdb_internal.repair_ttl_table_scheduled_job(oid: oid) &rarr; void</code></td><td><span class="funcdesc"><p>Repairs the scheduled job for a TTL table if it is missing.</p>
 </span></td></tr>
+<tr><td><a name="crdb_internal.request_statement_bundle"></a><code>crdb_internal.request_statement_bundle(stmtFingerprint: <a href="string.html">string</a>, minExecutionLatency: <a href="interval.html">interval</a>, expiresAfter: <a href="interval.html">interval</a>) &rarr; <a href="bool.html">bool</a></code></td><td><span class="funcdesc"><p>Used to request statement bundle for a given statement fingerprint
+that has execution latency greater than the ‘minExecutionLatency’. If the
+‘expiresAfter’ argument is empty, then the statement bundle request never
+expires until the statement bundle is collected</p>
+</span></td></tr>
 <tr><td><a name="crdb_internal.reset_index_usage_stats"></a><code>crdb_internal.reset_index_usage_stats() &rarr; <a href="bool.html">bool</a></code></td><td><span class="funcdesc"><p>This function is used to clear the collected index usage statistics.</p>
 </span></td></tr>
 <tr><td><a name="crdb_internal.reset_sql_stats"></a><code>crdb_internal.reset_sql_stats() &rarr; <a href="bool.html">bool</a></code></td><td><span class="funcdesc"><p>This function is used to clear the collected SQL statistics.</p>

pkg/server/status_test.go

@@ -2790,11 +2790,24 @@ func TestCreateStatementDiagnosticsReportWithViewActivityOptions(t *testing.T) {
 	defer s.Stopper().Stop(context.Background())
 	db := sqlutils.MakeSQLRunner(sqlDB)
 
+	ctx := context.Background()
+	ie := s.InternalExecutor().(*sql.InternalExecutor)
+
 	if err := getStatusJSONProtoWithAdminOption(s, "stmtdiagreports", &serverpb.CreateStatementDiagnosticsReportRequest{}, false); err != nil {
 		if !testutils.IsError(err, "status: 403") {
 			t.Fatalf("expected privilege error, got %v", err)
 		}
 	}
+	_, err := ie.ExecEx(
+		ctx,
+		"inserting-stmt-bundle-req",
+		nil, /* txn */
+		sessiondata.InternalExecutorOverride{
+			User: authenticatedUserNameNoAdmin(),
+		},
+		"SELECT crdb_internal.request_statement_bundle('SELECT _', 0::INTERVAL, 0::INTERVAL)",
+	)
+	require.Contains(t, err.Error(), "requesting statement bundle requires VIEWACTIVITY or ADMIN role option")
 
 	// Grant VIEWACTIVITY and all test should work.
 	db.Exec(t, fmt.Sprintf("ALTER USER %s VIEWACTIVITY", authenticatedUserNameNoAdmin().Normalized()))
@@ -2812,6 +2825,22 @@ func TestCreateStatementDiagnosticsReportWithViewActivityOptions(t *testing.T) {
 	if respGet.Reports[0].StatementFingerprint != req.StatementFingerprint {
 		t.Fatal("statement diagnostics request was not persisted")
 	}
+	_, err = ie.ExecEx(
+		ctx,
+		"inserting-stmt-bundle-req",
+		nil, /* txn */
+		sessiondata.InternalExecutorOverride{
+			User: authenticatedUserNameNoAdmin(),
+		},
+		"SELECT crdb_internal.request_statement_bundle('SELECT _', 0::INTERVAL, 0::INTERVAL)",
+	)
+	require.NoError(t, err)
+
+	db.CheckQueryResults(t, `
+      SELECT count(*)
+      FROM system.statement_diagnostics_requests
+      WHERE statement_fingerprint = 'SELECT _'
+`, [][]string{{"1"}})
 
 	// Grant VIEWACTIVITYREDACTED and all test should get permission errors.
 	db.Exec(t, fmt.Sprintf("ALTER USER %s VIEWACTIVITYREDACTED", authenticatedUserNameNoAdmin().Normalized()))
@@ -2826,6 +2855,17 @@ func TestCreateStatementDiagnosticsReportWithViewActivityOptions(t *testing.T) {
 			t.Fatalf("expected privilege error, got %v", err)
 		}
 	}
+
+	_, err = ie.ExecEx(
+		ctx,
+		"inserting-stmt-bundle-req",
+		nil, /* txn */
+		sessiondata.InternalExecutorOverride{
+			User: authenticatedUserNameNoAdmin(),
+		},
+		"SELECT crdb_internal.request_statement_bundle('SELECT _', 0::INTERVAL, 0::INTERVAL)",
+	)
+	require.Contains(t, err.Error(), "VIEWACTIVITYREDACTED role option cannot request statement bundle")
 }
 
 func TestStatementDiagnosticsCompleted(t *testing.T) {

pkg/sql/conn_executor.go

@@ -2665,22 +2665,23 @@ func (ex *connExecutor) asOfClauseWithSessionDefault(expr tree.AsOfClause) tree.
 func (ex *connExecutor) initEvalCtx(ctx context.Context, evalCtx *extendedEvalContext, p *planner) {
 	*evalCtx = extendedEvalContext{
 		Context: eval.Context{
-			Planner:                   p,
-			PrivilegedAccessor:        p,
-			SessionAccessor:           p,
-			JobExecContext:            p,
-			ClientNoticeSender:        p,
-			Sequence:                  p,
-			Tenant:                    p,
-			Regions:                   p,
-			JoinTokenCreator:          p,
-			PreparedStatementState:    &ex.extraTxnState.prepStmtsNamespace,
-			SessionDataStack:          ex.sessionDataStack,
-			ReCache:                   ex.server.reCache,
-			SQLStatsController:        ex.server.sqlStatsController,
-			IndexUsageStatsController: ex.server.indexUsageStatsController,
-			ConsistencyChecker:        p.execCfg.ConsistencyChecker,
-			RangeProber:               p.execCfg.RangeProber,
+			Planner:                        p,
+			PrivilegedAccessor:             p,
+			SessionAccessor:                p,
+			JobExecContext:                 p,
+			ClientNoticeSender:             p,
+			Sequence:                       p,
+			Tenant:                         p,
+			Regions:                        p,
+			JoinTokenCreator:               p,
+			PreparedStatementState:         &ex.extraTxnState.prepStmtsNamespace,
+			SessionDataStack:               ex.sessionDataStack,
+			ReCache:                        ex.server.reCache,
+			SQLStatsController:             ex.server.sqlStatsController,
+			IndexUsageStatsController:      ex.server.indexUsageStatsController,
+			ConsistencyChecker:             p.execCfg.ConsistencyChecker,
+			RangeProber:                    p.execCfg.RangeProber,
+			StmtDiagnosticsRequestInserter: ex.server.cfg.StmtDiagnosticsRecorder.InsertRequest,
 		},
 		Tracing:                &ex.sessionTracing,
 		MemMetrics:             &ex.memMetrics,

pkg/sql/planner.go

@@ -469,18 +469,19 @@ func internalExtendedEvalCtx(
 	}
 	ret := extendedEvalContext{
 		Context: eval.Context{
-			Txn:                       txn,
-			SessionDataStack:          sds,
-			TxnReadOnly:               false,
-			TxnImplicit:               true,
-			TxnIsSingleStmt:           true,
-			Context:                   ctx,
-			Mon:                       plannerMon,
-			TestingKnobs:              evalContextTestingKnobs,
-			StmtTimestamp:             stmtTimestamp,
-			TxnTimestamp:              txnTimestamp,
-			SQLStatsController:        sqlStatsController,
-			IndexUsageStatsController: indexUsageStatsController,
+			Txn:                            txn,
+			SessionDataStack:               sds,
+			TxnReadOnly:                    false,
+			TxnImplicit:                    true,
+			TxnIsSingleStmt:                true,
+			Context:                        ctx,
+			Mon:                            plannerMon,
+			TestingKnobs:                   evalContextTestingKnobs,
+			StmtTimestamp:                  stmtTimestamp,
+			TxnTimestamp:                   txnTimestamp,
+			SQLStatsController:             sqlStatsController,
+			IndexUsageStatsController:      indexUsageStatsController,
+			StmtDiagnosticsRequestInserter: execCfg.StmtDiagnosticsRecorder.InsertRequest,
 		},
 		Tracing:         &SessionTracing{},
 		Descs:           tables,

pkg/sql/sem/builtins/builtins.go

@@ -6968,6 +6968,69 @@ specified store on the node it's run from. One of 'mvccGC', 'merge', 'split',
 			Volatility: volatility.Volatile,
 		},
 	),
+
+	"crdb_internal.request_statement_bundle": makeBuiltin(
+		tree.FunctionProperties{
+			Category:         categorySystemInfo,
+			DistsqlBlocklist: true, // applicable only on the gateway
+		},
+		tree.Overload{
+			Types: tree.ArgTypes{
+				{"stmtFingerprint", types.String},
+				{"minExecutionLatency", types.Interval},
+				{"expiresAfter", types.Interval},
+			},
+			ReturnType: tree.FixedReturnType(types.Bool),
+			Fn: func(evalCtx *eval.Context, args tree.Datums) (tree.Datum, error) {
+				hasViewActivity, err := evalCtx.SessionAccessor.HasRoleOption(
+					evalCtx.Ctx(), roleoption.VIEWACTIVITY)
+				if err != nil {
+					return nil, err
+				}
+
+				if !hasViewActivity {
+					return nil, errors.New("requesting statement bundle requires " +
+						"VIEWACTIVITY or ADMIN role option")
+				}
+
+				isAdmin, err := evalCtx.SessionAccessor.HasAdminRole(evalCtx.Ctx())
+				if err != nil {
+					return nil, err
+				}
+
+				hasViewActivityRedacted, err := evalCtx.SessionAccessor.HasRoleOption(
+					evalCtx.Ctx(), roleoption.VIEWACTIVITYREDACTED)
+				if err != nil {
+					return nil, err
+				}
+
+				if !isAdmin && hasViewActivityRedacted {
+					return nil, errors.New("VIEWACTIVITYREDACTED role option cannot request " +
+						"statement bundle")
+				}
+
+				stmtFingerprint := string(tree.MustBeDString(args[0]))
+				minExecutionLatency := time.Duration(tree.MustBeDInterval(args[1]).Nanos())
+				expiresAfter := time.Duration(tree.MustBeDInterval(args[2]).Nanos())
+
+				if err := evalCtx.StmtDiagnosticsRequestInserter(
+					evalCtx.Ctx(),
+					stmtFingerprint,
+					minExecutionLatency,
+					expiresAfter,
+				); err != nil {
+					return nil, err
+				}
+
+				return tree.DBoolTrue, nil
+			},
+			Volatility: volatility.Volatile,
+			Info: `Used to request statement bundle for a given statement fingerprint
+that has execution latency greater than the 'minExecutionLatency'. If the
+'expiresAfter' argument is empty, then the statement bundle request never
+expires until the statement bundle is collected`,
+		},
+	),
 }
 
 var lengthImpls = func(incBitOverload bool) builtinDefinition {

pkg/sql/sem/eval/context.go

@@ -203,6 +203,11 @@ type Context struct {
 
 	// RangeProber is used in calls to crdb_internal.probe_ranges.
 	RangeProber RangeProber
+
+	// StmtDiagnosticsRequestInserter is used by the
+	// crdb_internal.request_statement_bundle builtin to insert a statement bundle
+	// request.
+	StmtDiagnosticsRequestInserter StmtDiagnosticsRequestInsertFunc
 }
 
 var _ tree.ParseTimeContext = &Context{}

pkg/sql/sem/eval/deps.go

@@ -490,6 +490,16 @@ type IndexUsageStatsController interface {
 	ResetIndexUsageStats(ctx context.Context) error
 }
 
+// StmtDiagnosticsRequestInsertFunc is an interface embedded in EvalCtx that can
+// be used by the builtins to insert a statement diagnostics request. This
+// interface is introduced to avoid circular dependency.
+type StmtDiagnosticsRequestInsertFunc func(
+	ctx context.Context,
+	stmtFingerprint string,
+	minExecutionLatency time.Duration,
+	expiresAfter time.Duration,
+) error
+
 // AsOfSystemTime represents the result from the evaluation of AS OF SYSTEM TIME
 // clause.
 type AsOfSystemTime struct {