sql,cli: add payloads_for_trace builtin

Previously it was quite cumbersome to view all payloads for a given
trace: we needed to join on the `node_inflight_trace_spans` vtable
to filter for span IDs that match a trace ID, then apply the
`payloads_for_span()` builtin to each span ID. This patch adds
syntactic sugar to the above query.

Instead of

```
WITH spans AS (
  SELECT span_id
  FROM crdb_internal.node_inflight_trace_spans
  WHERE trace_id = $TRACE_ID)
) SELECT *
  FROM spans, LATERAL crdb_internal.payloads_for_span(spans.span_id);
```

we can now simply use:
```
crdb_internal.payloads_for_trace($TRACE_ID);
```

and achieve the same result. The patch also adds all payloads for all
long-running spans to the `crdb_internal.node_inflight_trace_spans`
table of the debug.zip file.

Release note (sql change): Add `payloads_for_trace()` builtin so that
all payloads attached to all spans for a given trace ID will be
displayed, utilizing the `crdb_internal.payloads_for_span()`
builtin under the hood. All payloads for long-running spans are also
added to debug.zip in the `crdb_internal.node_inflight_trace_spans`
table dump.

Co-authored-by: Tobias Grieger <[email protected]>

Release justification: This patch is safe for release because it
adds syntactic sugar to an internal observability feature.

docs/generated/sql/functions.md

@@ -2647,6 +2647,8 @@ SELECT * FROM crdb_internal.check_consistency(true, ‘\x02’, ‘\x04’)</p>
 </span></td></tr>
 <tr><td><a name="crdb_internal.payloads_for_span"></a><code>crdb_internal.payloads_for_span(span_id: <a href="int.html">int</a>) &rarr; tuple{string AS payload_type, jsonb AS payload_jsonb}</code></td><td><span class="funcdesc"><p>Returns the payload(s) of the requested span.</p>
 </span></td></tr>
+<tr><td><a name="crdb_internal.payloads_for_trace"></a><code>crdb_internal.payloads_for_trace(trace_id: <a href="int.html">int</a>) &rarr; tuple{int AS span_id, string AS payload_type, jsonb AS payload_jsonb}</code></td><td><span class="funcdesc"><p>Returns the payload(s) of the requested trace.</p>
+</span></td></tr>
 <tr><td><a name="crdb_internal.pretty_key"></a><code>crdb_internal.pretty_key(raw_key: <a href="bytes.html">bytes</a>, skip_fields: <a href="int.html">int</a>) &rarr; <a href="string.html">string</a></code></td><td><span class="funcdesc"><p>This function is used only by CockroachDB’s developers for testing purposes.</p>
 </span></td></tr>
 <tr><td><a name="crdb_internal.range_stats"></a><code>crdb_internal.range_stats(key: <a href="bytes.html">bytes</a>) &rarr; jsonb</code></td><td><span class="funcdesc"><p>This function is used to retrieve range statistics information as a JSON object.</p>

pkg/cli/zip.go

@@ -37,11 +37,15 @@ type zipRequest struct {
 	pathName string
 }
 
-// Override for the default SELECT * when dumping one of the tables above.
-var customSelectClause = map[string]string{
-	"crdb.internal.node_inflight_trace_spans": "*, WHERE duration > 10*time.Second ORDER BY trace_id ASC, duration DESC",
-	"system.jobs":       "*, to_hex(payload) AS hex_payload, to_hex(progress) AS hex_progress",
-	"system.descriptor": "*, to_hex(descriptor) AS hex_descriptor",
+// Override for the default SELECT * FROM table when dumping one of the tables
+// in `debugZipTablesPerNode` or `debugZipTablesPerCluster`
+var customQuery = map[string]string{
+	"crdb_internal.node_inflight_trace_spans": "WITH spans AS (" +
+		"SELECT * FROM crdb_internal.node_inflight_trace_spans " +
+		"WHERE duration > INTERVAL '10' ORDER BY trace_id ASC, duration DESC" +
+		") SELECT * FROM spans, LATERAL crdb_internal.payloads_for_span(span_id)",
+	"system.jobs":       "SELECT *, to_hex(payload) AS hex_payload, to_hex(progress) AS hex_progress FROM system.jobs",
+	"system.descriptor": "SELECT *, to_hex(descriptor) AS hex_descriptor FROM system.descriptor",
 }
 
 type debugZipContext struct {
@@ -229,10 +233,8 @@ func maybeAddProfileSuffix(name string) string {
 //
 // An error is returned by this function if it is unable to write to
 // the output file or some other unrecoverable error is encountered.
-func (zc *debugZipContext) dumpTableDataForZip(
-	conn *sqlConn, base, table, selectClause string,
-) error {
-	query := fmt.Sprintf(`SET statement_timeout = '%s'; SELECT %s FROM %s`, zc.timeout, selectClause, table)
+func (zc *debugZipContext) dumpTableDataForZip(conn *sqlConn, base, table, query string) error {
+	fullQuery := fmt.Sprintf(`SET statement_timeout = '%s'; %s`, zc.timeout, query)
 	baseName := base + "/" + table
 
 	fmt.Printf("retrieving SQL data for %s... ", table)
@@ -250,7 +252,7 @@ func (zc *debugZipContext) dumpTableDataForZip(
 			}
 			// Pump the SQL rows directly into the zip writer, to avoid
 			// in-RAM buffering.
-			return runQueryAndFormatResults(conn, w, makeQuery(query))
+			return runQueryAndFormatResults(conn, w, makeQuery(fullQuery))
 		}(); err != nil {
 			if cErr := zc.z.createError(name, err); cErr != nil {
 				return cErr

pkg/cli/zip_cluster_wide.go

@@ -108,11 +108,11 @@ func (zc *debugZipContext) collectClusterData(
 	}
 
 	for _, table := range debugZipTablesPerCluster {
-		selectClause, ok := customSelectClause[table]
-		if !ok {
-			selectClause = "*"
+		query := fmt.Sprintf(`SELECT * FROM %s`, table)
+		if override, ok := customQuery[table]; ok {
+			query = override
 		}
-		if err := zc.dumpTableDataForZip(zc.firstNodeSQLConn, debugBase, table, selectClause); err != nil {
+		if err := zc.dumpTableDataForZip(zc.firstNodeSQLConn, debugBase, table, query); err != nil {
 			return nil, nil, errors.Wrapf(err, "fetching %s", table)
 		}
 	}

pkg/cli/zip_per_node.go

@@ -196,11 +196,11 @@ func (zc *debugZipContext) collectPerNodeData(
 	fmt.Printf("using SQL connection URL for node %s: %s\n", id, curSQLConn.url)
 
 	for _, table := range debugZipTablesPerNode {
-		selectClause, ok := customSelectClause[table]
-		if !ok {
-			selectClause = "*"
+		query := fmt.Sprintf(`SELECT * FROM %s`, table)
+		if override, ok := customQuery[table]; ok {
+			query = override
 		}
-		if err := zc.dumpTableDataForZip(curSQLConn, prefix, table, selectClause); err != nil {
+		if err := zc.dumpTableDataForZip(curSQLConn, prefix, table, query); err != nil {
 			return errors.Wrapf(err, "fetching %s", table)
 		}
 	}

pkg/cli/zip_test.go

@@ -417,8 +417,10 @@ func TestZipRetries(t *testing.T) {
 		}
 		if err := zc.dumpTableDataForZip(
 			sqlConn,
-			"test", `generate_series(1,15000) as t(x)`,
-			`if(x<11000,x,crdb_internal.force_retry('1h'))`); err != nil {
+			"test",
+			`generate_series(1,15000) as t(x)`,
+			`select if(x<11000,x,crdb_internal.force_retry('1h')) from generate_series(1,15000) as t(x)`,
+		); err != nil {
 			t.Fatal(err)
 		}
 	}()

pkg/sql/logictest/testdata/logic_test/builtin_function

@@ -2944,3 +2944,19 @@ SELECT * FROM crdb_internal.payloads_for_span(0)
 WHERE false
 ----
 payload_type  payload_jsonb
+
+subtest crdb_internal.payloads_for_trace
+
+# switch users -- this one has no permissions so expect errors
+user testuser
+
+query error pq: only users with the admin role are allowed to use crdb_internal.payloads_for_span
+SELECT * FROM crdb_internal.payloads_for_span(0)
+
+user root
+
+query TTT colnames
+SELECT * FROM crdb_internal.payloads_for_trace(0)
+WHERE false
+----
+span_id payload_type  payload_jsonb

pkg/sql/logictest/testdata/logic_test/contention_event

@@ -32,7 +32,7 @@ user root
 statement ok
 BEGIN;
 SET TRANSACTION PRIORITY HIGH;
-SELECT * FROM kv
+SELECT * FROM kv ORDER BY k ASC
 
 user testuser
 
@@ -60,5 +60,18 @@ WITH spans AS (
 ) SELECT count(*) > 0
   FROM payloads
   WHERE payload_type = 'roachpb.ContentionEvent'
+  AND crdb_internal.pretty_key(decode(payload_jsonb->>'key', 'base64'), 1) LIKE '/1/"k"/%'
+----
+true
+
+# crdb_internal.payloads_for_trace is syntactic sugar for much of the above statement.
+query B
+WITH payloads AS (
+  SELECT *
+  FROM crdb_internal.payloads_for_trace(crdb_internal.trace_id())
+) SELECT count(*) > 0
+  FROM payloads
+  WHERE payload_type = 'roachpb.ContentionEvent'
+  AND crdb_internal.pretty_key(decode(payload_jsonb->>'key', 'base64'), 1) LIKE '/1/"k"/%'
 ----
 true

pkg/sql/parser/parse_test.go

@@ -32,7 +32,7 @@ import (
 
 // TestParse verifies that we can parse the supplied SQL and regenerate the SQL
 // string from the syntax tree.  If the supplied SQL generates a different string
-// from the sytnax tree, use TestParse2 below.
+// from the syntax tree, use TestParse2 below.
 func TestParse(t *testing.T) {
 	testData := []struct {
 		sql string

pkg/sql/sem/builtins/generator_builtins.go

@@ -25,6 +25,8 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgerror"
 	"github.com/cockroachdb/cockroach/pkg/sql/protoreflect"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
+	"github.com/cockroachdb/cockroach/pkg/sql/sessiondata"
+	"github.com/cockroachdb/cockroach/pkg/sql/sqlutil"
 	"github.com/cockroachdb/cockroach/pkg/sql/types"
 	"github.com/cockroachdb/cockroach/pkg/util/arith"
 	"github.com/cockroachdb/cockroach/pkg/util/duration"
@@ -346,6 +348,22 @@ var generators = map[string]builtinDefinition{
 			tree.VolatilityVolatile,
 		),
 	),
+
+	"crdb_internal.payloads_for_trace": makeBuiltin(
+		tree.FunctionProperties{
+			Class:    tree.GeneratorClass,
+			Category: categorySystemInfo,
+		},
+		makeGeneratorOverload(
+			tree.ArgTypes{
+				{Name: "trace_id", Typ: types.Int},
+			},
+			payloadsForTraceGeneratorType,
+			makePayloadsForTraceGenerator,
+			"Returns the payload(s) of the requested trace.",
+			tree.VolatilityVolatile,
+		),
+	),
 }
 
 func makeGeneratorOverload(
@@ -1513,3 +1531,86 @@ func (p *payloadsForSpanGenerator) Values() (tree.Datums, error) {
 
 // Close implements the tree.ValueGenerator interface.
 func (p *payloadsForSpanGenerator) Close() {}
+
+var payloadsForTraceGeneratorLabels = []string{"span_id", "payload_type", "payload_jsonb"}
+
+var payloadsForTraceGeneratorType = types.MakeLabeledTuple(
+	[]*types.T{types.Int, types.String, types.Jsonb},
+	payloadsForTraceGeneratorLabels,
+)
+
+// payloadsForTraceGenerator is a value generator that iterates over all payloads
+// of a given Trace.
+type payloadsForTraceGenerator struct {
+	// Iterator over all internal rows of a query that retrieves all payloads
+	// of a trace.
+	it sqlutil.InternalRows
+}
+
+func makePayloadsForTraceGenerator(
+	ctx *tree.EvalContext, args tree.Datums,
+) (tree.ValueGenerator, error) {
+	// The user must be an admin to use this builtin.
+	isAdmin, err := ctx.SessionAccessor.HasAdminRole(ctx.Context)
+	if err != nil {
+		return nil, err
+	}
+	if !isAdmin {
+		return nil, pgerror.Newf(
+			pgcode.InsufficientPrivilege,
+			"only users with the admin role are allowed to use crdb_internal.payloads_for_trace",
+		)
+	}
+	traceID := uint64(*(args[0].(*tree.DInt)))
+
+	const query = `WITH spans AS(
+									SELECT span_id
+  	 							FROM crdb_internal.node_inflight_trace_spans
+ 		 							WHERE trace_id = $1
+									) SELECT * 
+										FROM spans, LATERAL crdb_internal.payloads_for_span(spans.span_id)`
+
+	ie := ctx.InternalExecutor.(sqlutil.InternalExecutor)
+	it, err := ie.QueryIteratorEx(
+		ctx.Ctx(),
+		"crdb_internal.payloads_for_trace",
+		ctx.Txn,
+		sessiondata.NoSessionDataOverride,
+		query,
+		traceID,
+	)
+	if err != nil {
+		return nil, err
+	}
+
+	return &payloadsForTraceGenerator{it: it}, nil
+}
+
+// ResolvedType implements the tree.ValueGenerator interface.
+func (p *payloadsForTraceGenerator) ResolvedType() *types.T {
+	return payloadsForSpanGeneratorType
+}
+
+// Start implements the tree.ValueGenerator interface.
+func (p *payloadsForTraceGenerator) Start(_ context.Context, _ *kv.Txn) error {
+	return nil
+}
+
+// Next implements the tree.ValueGenerator interface.
+func (p *payloadsForTraceGenerator) Next(ctx context.Context) (bool, error) {
+	return p.it.Next(ctx)
+}
+
+// Values implements the tree.ValueGenerator interface.
+func (p *payloadsForTraceGenerator) Values() (tree.Datums, error) {
+	return p.it.Cur(), nil
+}
+
+// Close implements the tree.ValueGenerator interface.
+func (p *payloadsForTraceGenerator) Close() {
+	err := p.it.Close()
+	if err != nil {
+		// TODO(angelapwen, yuzefovich): The iterator's error should be surfaced here.
+		return
+	}
+}