From 8210796ff1fc819056ea99af874f82f7e7211d4f Mon Sep 17 00:00:00 2001 From: angelapwen Date: Fri, 19 Feb 2021 10:47:38 +0100 Subject: [PATCH] builtins: add generator builtin for span payloads Previously we introduced a `crdb_internal.payloads_for_span` builtin that returned a JSONB array representing all payloads for a particular span. To improve the user experience of the builtin as well as prevent OOMs resulting from builtin usage, we replace it with a generator builtin that returns a table representing all payloads for a span, where each row represents a single payload. The new builtin, also called `crdb_internal.payloads_for_span`, has columns for the `payload_type` so that the user has the ability to easily filter on the type, and `payload_jsonb` so the user can use jsonb builtins to filter further. Release note (sql change): Update `crdb_internal.payloads_for_span` builtin to return a table instead of a JSONB array. Each row of the table represents one payload for the given span. It has columns for `payload_type` and `payload_jsonb`. --- docs/generated/sql/functions.md | 2 +- .../testdata/logic_test/builtin_function | 31 ++++ .../testdata/logic_test/contention_event | 20 +-- pkg/sql/sem/builtins/builtins.go | 49 +----- pkg/sql/sem/builtins/generator_builtins.go | 150 ++++++++++++++++++ pkg/sql/sem/tree/eval.go | 2 +- 6 files changed, 196 insertions(+), 58 deletions(-) diff --git a/docs/generated/sql/functions.md b/docs/generated/sql/functions.md index 529a7dcc8745..a08e6c9d38de 100644 --- a/docs/generated/sql/functions.md +++ b/docs/generated/sql/functions.md @@ -2638,7 +2638,7 @@ SELECT * FROM crdb_internal.check_consistency(true, ‘\x02’, ‘\x04’)

crdb_internal.num_inverted_index_entries(val: jsonb, version: int) → int

This function is used only by CockroachDB’s developers for testing purposes.

 -crdb_internal.payloads_for_span(span ID: int) → jsonb

Returns the payload(s) of the span whose ID is passed in the argument.

+crdb_internal.payloads_for_span(span_id: int) → tuple{string AS payload_type, jsonb AS payload_jsonb}

Returns the payload(s) of the requested span.

 crdb_internal.pretty_key(raw_key: bytes, skip_fields: int) → string

This function is used only by CockroachDB’s developers for testing purposes.

 diff --git a/pkg/sql/logictest/testdata/logic_test/builtin_function b/pkg/sql/logictest/testdata/logic_test/builtin_function index b708e533f7fd..96ae4fbc2624 100644 --- a/pkg/sql/logictest/testdata/logic_test/builtin_function +++ b/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 diff --git a/pkg/sql/logictest/testdata/logic_test/contention_event b/pkg/sql/logictest/testdata/logic_test/contention_event index 42f6c904223e..cc51e18fbfd4 100644 --- a/pkg/sql/logictest/testdata/logic_test/contention_event +++ b/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 diff --git a/pkg/sql/sem/builtins/builtins.go b/pkg/sql/sem/builtins/builtins.go index b25ea2527c24..e0f3684f6b32 100644 --- a/pkg/sql/sem/builtins/builtins.go +++ b/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{ diff --git a/pkg/sql/sem/builtins/generator_builtins.go b/pkg/sql/sem/builtins/generator_builtins.go index 9f121e1403db..f4f7a23465a8 100644 --- a/pkg/sql/sem/builtins/generator_builtins.go +++ b/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() {} diff --git a/pkg/sql/sem/tree/eval.go b/pkg/sql/sem/tree/eval.go index 34431f70499a..e40a64c0729a 100644 --- a/pkg/sql/sem/tree/eval.go +++ b/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 {