From dddec05805f44fe5f7e7eefc1c50a5a1caed876f Mon Sep 17 00:00:00 2001
From: angelapwen
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.payloads_for_trace(trace_id: int) → tuple{int AS span_id, string AS payload_type, jsonb AS payload_jsonb}
Returns the payload(s) of the requested trace.
+crdb_internal.pretty_key(raw_key: bytes, skip_fields: int) → string
This function is used only by CockroachDB’s developers for testing purposes.
crdb_internal.range_stats(key: bytes) → jsonb
This function is used to retrieve range statistics information as a JSON object.
diff --git a/pkg/cli/zip.go b/pkg/cli/zip.go index 14bd7f321d5a..2bdeace6e224 100644 --- a/pkg/cli/zip.go +++ b/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 diff --git a/pkg/cli/zip_cluster_wide.go b/pkg/cli/zip_cluster_wide.go index 62ac684611a7..58db78b626d8 100644 --- a/pkg/cli/zip_cluster_wide.go +++ b/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) } } diff --git a/pkg/cli/zip_per_node.go b/pkg/cli/zip_per_node.go index a4168a3542a5..243a383c1940 100644 --- a/pkg/cli/zip_per_node.go +++ b/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) } } diff --git a/pkg/cli/zip_test.go b/pkg/cli/zip_test.go index c3dd537bc179..4eb220da3b57 100644 --- a/pkg/cli/zip_test.go +++ b/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) } }() diff --git a/pkg/sql/logictest/testdata/logic_test/builtin_function b/pkg/sql/logictest/testdata/logic_test/builtin_function index 96ae4fbc2624..05806ebe9a5a 100644 --- a/pkg/sql/logictest/testdata/logic_test/builtin_function +++ b/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 diff --git a/pkg/sql/logictest/testdata/logic_test/contention_event b/pkg/sql/logictest/testdata/logic_test/contention_event index cc51e18fbfd4..c7019308911b 100644 --- a/pkg/sql/logictest/testdata/logic_test/contention_event +++ b/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 diff --git a/pkg/sql/parser/parse_test.go b/pkg/sql/parser/parse_test.go index ccaebfbae7aa..13dda0f380f5 100644 --- a/pkg/sql/parser/parse_test.go +++ b/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 diff --git a/pkg/sql/sem/builtins/generator_builtins.go b/pkg/sql/sem/builtins/generator_builtins.go index f4f7a23465a8..c3bfd379abdc 100644 --- a/pkg/sql/sem/builtins/generator_builtins.go +++ b/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 + } +}