diff --git a/pkg/sql/conn_executor_exec.go b/pkg/sql/conn_executor_exec.go index 798fa2bc5b9b..6663e701a272 100644 --- a/pkg/sql/conn_executor_exec.go +++ b/pkg/sql/conn_executor_exec.go @@ -18,8 +18,11 @@ import ( "context" "fmt" "strconv" + "strings" "time" + "github.com/pkg/errors" + "github.com/cockroachdb/cockroach/pkg/roachpb" "github.com/cockroachdb/cockroach/pkg/sql/coltypes" "github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgerror" @@ -30,6 +33,7 @@ import ( "github.com/cockroachdb/cockroach/pkg/util/hlc" "github.com/cockroachdb/cockroach/pkg/util/log" "github.com/cockroachdb/cockroach/pkg/util/timeutil" + "github.com/cockroachdb/cockroach/pkg/util/tracing" ) var errSavepointNotUsed = pgerror.NewErrorf( @@ -899,6 +903,9 @@ func (ex *connExecutor) runObserverStatement( return ex.runShowTransactionState(ctx, res) case *tree.ShowSyntax: return ex.runShowSyntax(ctx, sqlStmt.Statement, res) + case *tree.SetTracing: + ex.runSetTracing(ctx, sqlStmt, res) + return nil default: res.SetError(pgerror.NewErrorf(pgerror.CodeInternalError, "programming error: unrecognized observer statement type %T", stmt.AST)) @@ -941,6 +948,57 @@ func (ex *connExecutor) runShowTransactionState( return res.AddRow(ctx, tree.Datums{tree.NewDString(state)}) } +func (ex *connExecutor) runSetTracing( + ctx context.Context, n *tree.SetTracing, res RestrictedCommandResult, +) { + if len(n.Values) == 0 { + res.SetError(fmt.Errorf("set tracing missing argument")) + return + } + + modes := make([]string, len(n.Values)) + for i, v := range n.Values { + v = unresolvedNameToStrVal(v) + strVal, ok := v.(*tree.StrVal) + if !ok { + res.SetError(fmt.Errorf("expected string for set tracing argument, not %T", v)) + return + } + modes[i] = strVal.RawString() + } + + if err := ex.enableTracing(modes); err != nil { + res.SetError(err) + } +} + +func (ex *connExecutor) enableTracing(modes []string) error { + traceKV := false + recordingType := tracing.SnowballRecording + enableMode := true + + for _, s := range modes { + switch strings.ToLower(s) { + case "on": + enableMode = true + case "off": + enableMode = false + case "kv": + traceKV = true + case "local": + recordingType = tracing.SingleNodeRecording + case "cluster": + recordingType = tracing.SnowballRecording + default: + return errors.Errorf("set tracing: unknown mode %q", s) + } + } + if !enableMode { + return stopTracing(&ex.dataMutator) + } + return ex.dataMutator.StartSessionTracing(recordingType, traceKV) +} + // addActiveQuery adds a running query to the list of running queries. // // It returns a cleanup function that needs to be run when the query is no diff --git a/pkg/sql/logictest/testdata/logic_test/show_trace b/pkg/sql/logictest/testdata/logic_test/show_trace index 4cc57b258ba7..be0667e3744a 100644 --- a/pkg/sql/logictest/testdata/logic_test/show_trace +++ b/pkg/sql/logictest/testdata/logic_test/show_trace @@ -12,18 +12,15 @@ SET tracing = on; BEGIN; SELECT 1; COMMIT; SELECT 2; SET tracing = off; query ITT SELECT span, message, operation FROM [SHOW TRACE FOR SESSION] WHERE message LIKE '%SPAN START%' OR message LIKE '%pos%executing%'; ---- -0 === SPAN START: sql txn === sql txn -1 === SPAN START: session recording === session recording -1 [NoTxn pos:5] executing ExecStmt: BEGIN TRANSACTION session recording +0 === SPAN START: session recording === session recording +0 [NoTxn pos:5] executing ExecStmt: BEGIN TRANSACTION session recording +1 === SPAN START: sql txn === sql txn +1 [Open pos:6] executing ExecStmt: SELECT 1 sql txn +1 [Open pos:7] executing ExecStmt: COMMIT TRANSACTION sql txn +0 [NoTxn pos:8] executing ExecStmt: SELECT 2 session recording 2 === SPAN START: sql txn === sql txn -2 [Open pos:6] executing ExecStmt: SELECT 1 sql txn -2 [Open pos:7] executing ExecStmt: COMMIT TRANSACTION sql txn -1 [NoTxn pos:8] executing ExecStmt: SELECT 2 session recording -3 === SPAN START: sql txn === sql txn -3 [Open pos:8] executing ExecStmt: SELECT 2 sql txn -1 [NoTxn pos:9] executing ExecStmt: SET tracing = off session recording -4 === SPAN START: sql txn === sql txn -4 [Open pos:9] executing ExecStmt: SET tracing = off sql txn +2 [Open pos:8] executing ExecStmt: SELECT 2 sql txn +0 [NoTxn pos:9] executing ExecStmt: SET TRACING = off session recording # Same, with SHOW TRACE FOR. # This also tests that sub-spans are reported properly. @@ -827,3 +824,28 @@ query T SELECT message FROM [ SHOW TRACE FOR SELECT * FROM t.enginestats ] WHERE message LIKE '%InternalDelete%' ---- engine stats: {InternalDeleteSkippedCount:0 TimeBoundNumSSTs:0} + +# Check that we can run set tracing regardless of the current tracing state. +# This is convenient; sometimes it's unclear, for example, if you previously +# stopped tracing or not, so issuing a set tracing=off should just work. +subtest idempotent + +statement ok +SET tracing = on; SET tracing = on; + +statement ok +SET tracing = off; SET tracing = off; + +# Check that we can run set tracing in the aborted state (this is implemented by +# making set tracing an ObserverStmt). This is very convenient for clients that +# start tracing, then might get an error, then want to stop tracing. +subtest aborted_txn + +query error pq: foo +BEGIN; SELECT crdb_internal.force_error('', 'foo') + +statement ok +SET tracing = off + +statement ok +ROLLBACK diff --git a/pkg/sql/parser/parse_test.go b/pkg/sql/parser/parse_test.go index 91bb1c427cf7..8c1a7f603d15 100644 --- a/pkg/sql/parser/parse_test.go +++ b/pkg/sql/parser/parse_test.go @@ -785,6 +785,7 @@ func TestParse(t *testing.T) { {`SET TRANSACTION PRIORITY NORMAL`}, {`SET TRANSACTION PRIORITY HIGH`}, {`SET TRANSACTION ISOLATION LEVEL SNAPSHOT, PRIORITY HIGH`}, + {`SET TRACING = off`}, {`SET SESSION CHARACTERISTICS AS TRANSACTION ISOLATION LEVEL SERIALIZABLE`}, {`SET SESSION CHARACTERISTICS AS TRANSACTION ISOLATION LEVEL SNAPSHOT`}, {`SET CLUSTER SETTING a = 3`}, diff --git a/pkg/sql/parser/sql.y b/pkg/sql/parser/sql.y index c28a2054a3b6..aea28c1a8331 100644 --- a/pkg/sql/parser/sql.y +++ b/pkg/sql/parser/sql.y @@ -525,6 +525,7 @@ func newNameFromStr(s string) *tree.Name { %token TABLE TABLES TEMP TEMPLATE TEMPORARY TESTING_RANGES EXPERIMENTAL_RANGES TESTING_RELOCATE EXPERIMENTAL_RELOCATE TEXT THAN THEN %token TIME TIMETZ TIMESTAMP TIMESTAMPTZ TO TRAILING TRACE TRANSACTION TREAT TRIM TRUE %token TRUNCATE TYPE +%token TRACING %token UNBOUNDED UNCOMMITTED UNION UNIQUE UNKNOWN %token UPDATE UPSERT USE USER USERS USING UUID @@ -2516,6 +2517,7 @@ set_exprs_internal: // SET [SESSION] { TO | = } // SET [SESSION] TIME ZONE // SET [SESSION] CHARACTERISTICS AS TRANSACTION ISOLATION LEVEL { SNAPSHOT | SERIALIZABLE } +// SET [SESSION] TRACING { TO | = } { off | cluster | on | kv | local } [,...] // // %SeeAlso: SHOW SESSION, RESET, DISCARD, SHOW, SET CLUSTER SETTING, SET TRANSACTION, // WEBDOCS/set-vars.html @@ -2564,8 +2566,23 @@ generic_set: } | var_name '=' var_list { - $$.val = &tree.SetVar{Name: strings.Join($1.strs(), "."), Values: $3.exprs()} + if $1.strs()[0] == "tracing" { + $$.val = &tree.SetTracing{Values: $3.exprs()} + } else { + $$.val = &tree.SetVar{Name: strings.Join($1.strs(), "."), Values: $3.exprs()} + } + } +/* + * | TRACING '=' var_list + { + $$.val = &tree.SetTracing{Values: $3.exprs()} + } +| TRACING TO var_list + { + $$.val = &tree.SetTracing{Values: $3.exprs()} } +*/ + set_rest_more: // Generic SET syntaxes: diff --git a/pkg/sql/sem/tree/set.go b/pkg/sql/sem/tree/set.go index b80c6eb83311..db2a10ea786d 100644 --- a/pkg/sql/sem/tree/set.go +++ b/pkg/sql/sem/tree/set.go @@ -82,3 +82,14 @@ func (node *SetSessionCharacteristics) Format(ctx *FmtCtx) { ctx.WriteString("SET SESSION CHARACTERISTICS AS TRANSACTION") node.Modes.Format(ctx) } + +// SetTracing represents a SET TRACING statement. +type SetTracing struct { + Values Exprs +} + +// Format implements the NodeFormatter interface. +func (node *SetTracing) Format(ctx *FmtCtx) { + ctx.WriteString("SET TRACING = ") + ctx.FormatNode(&node.Values) +} diff --git a/pkg/sql/sem/tree/stmt.go b/pkg/sql/sem/tree/stmt.go index 6c412407db9f..863f144a9dc6 100644 --- a/pkg/sql/sem/tree/stmt.go +++ b/pkg/sql/sem/tree/stmt.go @@ -151,11 +151,17 @@ func (l *StatementList) Format(ctx *FmtCtx) { } } -// ObserverStatement is a marker interface to be implemented by -// statements which are always valid and do not modify neither the -// session state, the txn state or the database. They are to be -// allowed at any point, including when a current transaction is in a -// transient state like retry_wait. +// ObserverStatement is a marker interface for statements which are allowed to +// run regardless of the current transaction state: statements other than +// rollback are generally rejected if the session is in a failed transaction +// state, but it's convenient to allow some statements (e.g. "show syntax; set +// tracing"). +// Such statements are not expected to modify the database, the transaction or +// session state (other than special cases such as enabling/disabling tracing). +// +// These statements short-circuit the regular execution - they don't get planned +// (there's no corresponding planNodes). The connExecutor recognizes them and +// handles them. type ObserverStatement interface { observerStatement() } @@ -584,6 +590,17 @@ func (*SetTransaction) StatementTag() string { return "SET TRANSACTION" } func (*SetTransaction) hiddenFromStats() {} +// StatementType implements the Statement interface. +func (*SetTracing) StatementType() StatementType { return Ack } + +// StatementTag returns a short string identifying the type of statement. +func (*SetTracing) StatementTag() string { return "SET TRACING" } + +func (*SetTracing) hiddenFromStats() {} + +// observerStatement implements the ObserverStatement interface. +func (*SetTracing) observerStatement() {} + // StatementType implements the Statement interface. func (*SetZoneConfig) StatementType() StatementType { return RowsAffected } @@ -933,6 +950,7 @@ func (n *SetClusterSetting) String() string { return AsString(n) } func (n *SetZoneConfig) String() string { return AsString(n) } func (n *SetSessionCharacteristics) String() string { return AsString(n) } func (n *SetTransaction) String() string { return AsString(n) } +func (n *SetTracing) String() string { return AsString(n) } func (n *SetVar) String() string { return AsString(n) } func (n *ShowBackup) String() string { return AsString(n) } func (n *ShowClusterSetting) String() string { return AsString(n) } diff --git a/pkg/sql/sem/tree/walk.go b/pkg/sql/sem/tree/walk.go index bedfb191a093..f80e47263a93 100644 --- a/pkg/sql/sem/tree/walk.go +++ b/pkg/sql/sem/tree/walk.go @@ -1113,6 +1113,28 @@ func (stmt *SetVar) WalkStmt(v Visitor) Statement { return ret } +// CopyNode makes a copy of this Statement without recursing in any child Statements. +func (stmt *SetTracing) CopyNode() *SetTracing { + stmtCopy := *stmt + stmtCopy.Values = append(Exprs(nil), stmt.Values...) + return &stmtCopy +} + +// WalkStmt is part of the WalkableStmt interface. +func (stmt *SetTracing) WalkStmt(v Visitor) Statement { + ret := stmt + for i, expr := range stmt.Values { + e, changed := WalkExpr(v, expr) + if changed { + if ret == stmt { + ret = stmt.CopyNode() + } + ret.Values[i] = e + } + } + return ret +} + // CopyNode makes a copy of this Statement without recursing in any child Statements. func (stmt *SetClusterSetting) CopyNode() *SetClusterSetting { stmtCopy := *stmt diff --git a/pkg/sql/session.go b/pkg/sql/session.go index 608ae037b34b..fe05677b64fa 100644 --- a/pkg/sql/session.go +++ b/pkg/sql/session.go @@ -1564,7 +1564,8 @@ func (st *SessionTracing) getRecording() ([]traceRow, error) { // are per-row. func (st *SessionTracing) StartTracing(recType tracing.RecordingType, kvTracingEnabled bool) error { if st.enabled { - return errors.Errorf("already tracing") + // We're already tracing. No-op. + return nil } // If we're inside a transaction, start recording on the txn span. @@ -1607,7 +1608,8 @@ func (st *SessionTracing) StartTracing(recType tracing.RecordingType, kvTracingE // An error is returned if tracing was not active. func (st *SessionTracing) StopTracing() error { if !st.enabled { - return errors.Errorf("not tracing") + // We're not currently tracing. No-op. + return nil } st.enabled = false diff --git a/pkg/sql/set_cluster_setting.go b/pkg/sql/set_cluster_setting.go index 6a7870b68254..d8997397f07a 100644 --- a/pkg/sql/set_cluster_setting.go +++ b/pkg/sql/set_cluster_setting.go @@ -58,11 +58,7 @@ func (p *planner) SetClusterSetting( // For DEFAULT, let the value reference be nil. That's a RESET in disguise. if _, ok := n.Value.(tree.DefaultVal); !ok { expr := n.Value - if s, ok := expr.(*tree.UnresolvedName); ok { - // Special rule for SET: because SET doesn't apply in the context - // of a table, SET ... = IDENT really means SET ... = 'IDENT'. - expr = tree.NewStrVal(tree.AsStringWithFlags(s, tree.FmtBareIdentifiers)) - } + expr = unresolvedNameToStrVal(expr) var requiredType types.T switch setting.(type) { diff --git a/pkg/sql/set_var.go b/pkg/sql/set_var.go index f4cb0af19363..9b89d2bfe07b 100644 --- a/pkg/sql/set_var.go +++ b/pkg/sql/set_var.go @@ -63,11 +63,7 @@ func (p *planner) SetVar(ctx context.Context, n *tree.SetVar) (planNode, error) if !isReset { typedValues = make([]tree.TypedExpr, len(n.Values)) for i, expr := range n.Values { - // Special rule for SET: because SET doesn't apply in the context - // of a table, SET ... = IDENT really means SET ... = 'IDENT'. - if s, ok := expr.(*tree.UnresolvedName); ok { - expr = tree.NewStrVal(tree.AsStringWithFlags(s, tree.FmtBareIdentifiers)) - } + expr = unresolvedNameToStrVal(expr) var dummyHelper tree.IndexedVarHelper typedValue, err := p.analyzeExpr( @@ -98,6 +94,15 @@ func (p *planner) SetVar(ctx context.Context, n *tree.SetVar) (planNode, error) return &setVarNode{v: v, typedValues: typedValues}, nil } +// Special rule for SET: because SET doesn't apply in the context +// of a table, SET ... = IDENT really means SET ... = 'IDENT'. +func unresolvedNameToStrVal(expr tree.Expr) tree.Expr { + if s, ok := expr.(*tree.UnresolvedName); ok { + return tree.NewStrVal(tree.AsStringWithFlags(s, tree.FmtBareIdentifiers)) + } + return expr +} + func (n *setVarNode) startExec(params runParams) error { if n.typedValues != nil { for i, v := range n.typedValues { diff --git a/pkg/sql/vars.go b/pkg/sql/vars.go index 785d8090d83c..f5609a5db069 100644 --- a/pkg/sql/vars.go +++ b/pkg/sql/vars.go @@ -582,49 +582,10 @@ var varGen = map[string]sessionVar{ } return stopTracing(m) }, - Set: func( - _ context.Context, m *sessionDataMutator, - evalCtx *extendedEvalContext, values []tree.TypedExpr, - ) error { - return enableTracing(&evalCtx.EvalContext, m, values) - }, + // Setting is done by the SetTracing statement. }, } -func enableTracing( - evalCtx *tree.EvalContext, m *sessionDataMutator, values []tree.TypedExpr, -) error { - traceKV := false - recordingType := tracing.SnowballRecording - enableMode := true - - for _, v := range values { - s, err := datumAsString(evalCtx, "trace", v) - if err != nil { - return err - } - - switch strings.ToLower(s) { - case "on": - enableMode = true - case "off": - enableMode = false - case "kv": - traceKV = true - case "local": - recordingType = tracing.SingleNodeRecording - case "cluster": - recordingType = tracing.SnowballRecording - default: - return errors.Errorf("set tracing: unknown mode %q", s) - } - } - if !enableMode { - return stopTracing(m) - } - return m.StartSessionTracing(recordingType, traceKV) -} - func stopTracing(m *sessionDataMutator) error { if err := m.StopSessionTracing(); err != nil { return errors.Wrapf(err, "error stopping tracing")