builtins: add generator builtin for span payloads

docs/generated/sql/functions.md

@@ -2638,7 +2638,7 @@ SELECT * FROM crdb_internal.check_consistency(true, ‘\x02’, ‘\x04’)</p>
 </span></td></tr>
 <tr><td><a name="crdb_internal.num_inverted_index_entries"></a><code>crdb_internal.num_inverted_index_entries(val: jsonb, version: <a href="int.html">int</a>) &rarr; <a href="int.html">int</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.payloads_for_span"></a><code>crdb_internal.payloads_for_span(span ID: <a href="int.html">int</a>) &rarr; jsonb</code></td><td><span class="funcdesc"><p>Returns the payload(s) of the span whose ID is passed in the argument.</p>
+<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.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>

pkg/sql/logictest/testdata/logic_test/builtin_function

@@ -2913,3 +2913,34 @@ query T
 SELECT regexp_split_to_array('3aaa0AAa1', 'a+', 'i')
 ----
 {3,0,1}
+
+subtest crdb_internal.trace_id
+
+# switch users -- this one has no permissions so expect errors
+user testuser
+
+query error insufficient privilege
+SELECT * FROM crdb_internal.trace_id()
+
+user root
+
+query B
+SELECT count(*) = 1 FROM crdb_internal.trace_id()
+----
+true
+
+subtest crdb_internal.payloads_for_span
+
+# 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 TT colnames
+SELECT * FROM crdb_internal.payloads_for_span(0)
+WHERE false
+----
+payload_type  payload_jsonb

pkg/sql/logictest/testdata/logic_test/contention_event

@@ -9,7 +9,7 @@ GRANT ADMIN TO testuser
 
 statement ok
 CREATE TABLE kv (k VARCHAR PRIMARY KEY, v VARCHAR);
-ALTER TABLE kv SPLIT AT VALUES ('b'), ('d'), ('q'), ('z');
+ALTER TABLE kv SPLIT AT VALUES ('b'), ('d'), ('q'), ('z')
 
 query TT
 SELECT * FROM kv
@@ -18,7 +18,7 @@ SELECT * FROM kv
 user testuser
 
 statement ok
-BEGIN;
+BEGIN
 
 statement ok
 INSERT INTO kv VALUES('k', 'v')
@@ -32,7 +32,7 @@ user root
 statement ok
 BEGIN;
 SET TRANSACTION PRIORITY HIGH;
-SELECT * FROM kv;
+SELECT * FROM kv
 
 user testuser
 
@@ -49,16 +49,16 @@ user root
 #
 # NB: this needs the 5node-pretend59315 config because otherwise the span is not
 # tracked.
-#
 query B
 WITH spans AS (
-  SELECT span_id FROM crdb_internal.node_inflight_trace_spans
+  SELECT span_id
+  FROM crdb_internal.node_inflight_trace_spans
   WHERE trace_id = crdb_internal.trace_id()
-), payload_types AS (
-    SELECT jsonb_array_elements(crdb_internal.payloads_for_span(span_id))->>'@type' AS payload_type
-    FROM spans
+), payloads AS (
+  SELECT *
+  FROM spans, LATERAL crdb_internal.payloads_for_span(spans.span_id)
 ) SELECT count(*) > 0
-    FROM payload_types
-    WHERE payload_type = 'type.googleapis.com/cockroach.roachpb.ContentionEvent';
+  FROM payloads
+  WHERE payload_type = 'roachpb.ContentionEvent'
 ----
 true

pkg/sql/sem/builtins/builtins.go

@@ -73,7 +73,6 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/util/unaccent"
 	"github.com/cockroachdb/cockroach/pkg/util/uuid"
 	"github.com/cockroachdb/errors"
-	pbtypes "github.com/gogo/protobuf/types"
 	"github.com/knz/strtime"
 )
 
@@ -3593,7 +3592,9 @@ may increase either contention or retry errors, or both.`,
 					return nil, err
 				}
 				if !isAdmin {
-					return nil, pgerror.Newf(pgcode.InsufficientPrivilege, "user needs the admin role to view trace ID")
+					if err := checkPrivilegedUser(ctx); err != nil {
+						return nil, err
+					}
 				}
 
 				sp := tracing.SpanFromContext(ctx.Context)
@@ -3608,50 +3609,6 @@ may increase either contention or retry errors, or both.`,
 		},
 	),
 
-	"crdb_internal.payloads_for_span": makeBuiltin(
-		tree.FunctionProperties{Category: categorySystemInfo},
-		tree.Overload{
-			Types:      tree.ArgTypes{{"span ID", types.Int}},
-			ReturnType: tree.FixedReturnType(types.Jsonb),
-			Fn: func(ctx *tree.EvalContext, args tree.Datums) (tree.Datum, 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, "user needs the admin role to view payloads")
-				}
-
-				builder := json.NewArrayBuilder(len(args))
-
-				spanID := uint64(*(args[0].(*tree.DInt)))
-				span, found := ctx.Settings.Tracer.GetActiveSpanFromID(spanID)
-				// A span may not be found if its ID was surfaced previously but its
-				// corresponding trace has ended by the time this builtin was run.
-				if !found {
-					// Returns an empty JSON array.
-					return tree.NewDJSON(builder.Build()), nil
-				}
-
-				for _, rec := range span.GetRecording() {
-					rec.Structured(func(item *pbtypes.Any) {
-						payload, err := protoreflect.MessageToJSON(item, true /* emitDefaults */)
-						if err != nil {
-							return
-						}
-						if payload != nil {
-							builder.Add(payload)
-						}
-					})
-				}
-				return tree.NewDJSON(builder.Build()), nil
-			},
-			Info:       "Returns the payload(s) of the span whose ID is passed in the argument.",
-			Volatility: tree.VolatilityVolatile,
-		},
-	),
-
 	"crdb_internal.locality_value": makeBuiltin(
 		tree.FunctionProperties{Category: categorySystemInfo},
 		tree.Overload{

pkg/sql/sem/builtins/generator_builtins.go

@@ -13,6 +13,7 @@ package builtins
 import (
 	"bytes"
 	"context"
+	"strings"
 	"time"
 
 	"github.com/cockroachdb/cockroach/pkg/keys"
@@ -22,13 +23,16 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/lex"
 	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgcode"
 	"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/types"
 	"github.com/cockroachdb/cockroach/pkg/util/arith"
 	"github.com/cockroachdb/cockroach/pkg/util/duration"
 	"github.com/cockroachdb/cockroach/pkg/util/errorutil"
 	"github.com/cockroachdb/cockroach/pkg/util/json"
+	"github.com/cockroachdb/cockroach/pkg/util/tracing"
 	"github.com/cockroachdb/errors"
+	pbtypes "github.com/gogo/protobuf/types"
 )
 
 // See the comments at the start of generators.go for details about
@@ -326,6 +330,22 @@ var generators = map[string]builtinDefinition{
 			tree.VolatilityVolatile,
 		),
 	),
+
+	"crdb_internal.payloads_for_span": makeBuiltin(
+		tree.FunctionProperties{
+			Class:    tree.GeneratorClass,
+			Category: categorySystemInfo,
+		},
+		makeGeneratorOverload(
+			tree.ArgTypes{
+				{Name: "span_id", Typ: types.Int},
+			},
+			payloadsForSpanGeneratorType,
+			makePayloadsForSpanGenerator,
+			"Returns the payload(s) of the requested span.",
+			tree.VolatilityVolatile,
+		),
+	),
 }
 
 func makeGeneratorOverload(
@@ -1363,3 +1383,133 @@ func (rk *rangeKeyIterator) Values() (tree.Datums, error) {
 
 // Close implements the tree.ValueGenerator interface.
 func (rk *rangeKeyIterator) Close() {}
+
+var payloadsForSpanGeneratorLabels = []string{"payload_type", "payload_jsonb"}
+
+var payloadsForSpanGeneratorType = types.MakeLabeledTuple(
+	[]*types.T{types.String, types.Jsonb},
+	payloadsForSpanGeneratorLabels,
+)
+
+// payloadsForSpanGenerator is a value generator that iterates over all payloads
+// over all recordings for a given Span.
+type payloadsForSpanGenerator struct {
+	// The span to iterate over.
+	span *tracing.Span
+
+	// recordingIndex maintains the current position of the index of the iterator
+	// in the list of recordings surfaced by a given span. The payloads of the
+	// recording that this iterator points to are buffered in `payloads`
+	recordingIndex int
+
+	// payloads represents all payloads for a given recording currently accessed
+	// by the iterator, and accesses more in a streaming fashion.
+	payloads []json.JSON
+
+	// payloadIndex maintains the current position of the index of the iterator
+	// in the list of `payloads` associated with a given recording.
+	payloadIndex int
+}
+
+func makePayloadsForSpanGenerator(
+	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_span",
+		)
+	}
+	spanID := uint64(*(args[0].(*tree.DInt)))
+	span, found := ctx.Settings.Tracer.GetActiveSpanFromID(spanID)
+	if !found {
+		return nil, nil
+	}
+
+	return &payloadsForSpanGenerator{span: span}, nil
+}
+
+// ResolvedType implements the tree.ValueGenerator interface.
+func (p *payloadsForSpanGenerator) ResolvedType() *types.T {
+	return payloadsForSpanGeneratorType
+}
+
+// Start implements the tree.ValueGenerator interface.
+func (p *payloadsForSpanGenerator) Start(_ context.Context, _ *kv.Txn) error {
+	// The user of the generator first calls Next(), then Values(), so the index
+	// managing the iterator's position needs to start at -1 instead of 0.
+	p.recordingIndex = -1
+	p.payloadIndex = -1
+
+	return nil
+}
+
+// Next implements the tree.ValueGenerator interface.
+func (p *payloadsForSpanGenerator) Next(_ context.Context) (bool, error) {
+	p.payloadIndex++
+
+	// If payloadIndex is within payloads and there are some payloads, then we
+	// have more buffered payloads to return.
+	if p.payloads != nil && p.payloadIndex < len(p.payloads) {
+		return true, nil
+	}
+
+	// Otherwise either there are no payloads or we have exhausted the payloads in
+	// our current recording, and we need to access another set of payloads from
+	// another recording.
+	p.payloads = nil
+
+	// Keep searching recordings for one with a valid (non-nil) payload.
+	for p.payloads == nil {
+		p.recordingIndex++
+		// If there are no more recordings, then we cannot continue.
+		if !(p.recordingIndex < p.span.GetRecording().Len()) {
+			return false, nil
+		}
+		currRecording := p.span.GetRecording()[p.recordingIndex]
+		currRecording.Structured(func(item *pbtypes.Any) {
+			payload, err := protoreflect.MessageToJSON(item, true /* emitDefaults */)
+			if err != nil {
+				return
+			}
+			if payload != nil {
+				p.payloads = append(p.payloads, payload)
+			}
+		})
+	}
+
+	p.payloadIndex = 0
+	return true, nil
+}
+
+// Values implements the tree.ValueGenerator interface.
+func (p *payloadsForSpanGenerator) Values() (tree.Datums, error) {
+	payload := p.payloads[p.payloadIndex]
+	payloadTypeAsJSON, err := payload.FetchValKey("@type")
+	if err != nil {
+		return nil, err
+	}
+
+	// We trim the proto type prefix as well as the enclosing double quotes
+	// leftover from JSON value conversion.
+	payloadTypeAsString := strings.TrimSuffix(
+		strings.TrimPrefix(
+			payloadTypeAsJSON.String(),
+			"\"type.googleapis.com/cockroach.",
+		),
+		"\"",
+	)
+
+	return tree.Datums{
+		tree.NewDString(payloadTypeAsString),
+		tree.NewDJSON(payload),
+	}, nil
+}
+
+// Close implements the tree.ValueGenerator interface.
+func (p *payloadsForSpanGenerator) Close() {}

pkg/sql/sem/tree/eval.go

@@ -3940,7 +3940,7 @@ func EvalComparisonExprWithSubOperator(
 	return evalDatumsCmp(ctx, expr.Operator, expr.SubOperator, expr.Fn, left, datums)
 }
 
-// EvalArgsAndGetGenerator evaluates the arguments and instanciates a
+// EvalArgsAndGetGenerator evaluates the arguments and instantiates a
 // ValueGenerator for use by set projections.
 func (expr *FuncExpr) EvalArgsAndGetGenerator(ctx *EvalContext) (ValueGenerator, error) {
 	if expr.fn == nil || expr.fnProps.Class != GeneratorClass {