diff --git a/pkg/sql/conn_executor_exec.go b/pkg/sql/conn_executor_exec.go index 9fef9d4ba6cd..4dbe13444c0a 100644 --- a/pkg/sql/conn_executor_exec.go +++ b/pkg/sql/conn_executor_exec.go @@ -351,6 +351,7 @@ func (ex *connExecutor) execStmtInOpenState( p.noticeSender = res ih := &p.instrumentation + // Special top-level handling for EXPLAIN ANALYZE. if e, ok := ast.(*tree.ExplainAnalyze); ok { switch e.Mode { case tree.ExplainDebug: @@ -388,6 +389,40 @@ func (ex *connExecutor) execStmtInOpenState( stmt.ExpectedTypes = nil } + // Special top-level handling for EXECUTE. This must happen after the handling + // for EXPLAIN ANALYZE (in order to support EXPLAIN ANALYZE EXECUTE) but + // before setting up the instrumentation helper. + if e, ok := ast.(*tree.Execute); ok { + // Replace the `EXECUTE foo` statement with the prepared statement, and + // continue execution. + name := e.Name.String() + ps, ok := ex.extraTxnState.prepStmtsNamespace.prepStmts[name] + if !ok { + err := pgerror.Newf( + pgcode.InvalidSQLStatementName, + "prepared statement %q does not exist", name, + ) + return makeErrEvent(err) + } + var err error + pinfo, err = fillInPlaceholders(ctx, ps, name, e.Params, ex.sessionData.SearchPath) + if err != nil { + return makeErrEvent(err) + } + + // TODO(radu): what about .SQL, .NumAnnotations, .NumPlaceholders? + stmt.Statement = ps.Statement + stmt.Prepared = ps + stmt.ExpectedTypes = ps.Columns + stmt.AnonymizedStr = ps.AnonymizedStr + res.ResetStmtType(ps.AST) + + if e.DiscardRows { + ih.SetDiscardRows() + } + ast = stmt.Statement.AST + } + var needFinish bool ctx, needFinish = ih.Setup( ctx, ex.server.cfg, ex.appStats, p, ex.stmtDiagnosticsRecorder, @@ -523,35 +558,6 @@ func (ex *connExecutor) execStmtInOpenState( return makeErrEvent(err) } return nil, nil, nil - - case *tree.Execute: - // Replace the `EXECUTE foo` statement with the prepared statement, and - // continue execution below. - name := s.Name.String() - ps, ok := ex.extraTxnState.prepStmtsNamespace.prepStmts[name] - if !ok { - err := pgerror.Newf( - pgcode.InvalidSQLStatementName, - "prepared statement %q does not exist", name, - ) - return makeErrEvent(err) - } - var err error - pinfo, err = fillInPlaceholders(ctx, ps, name, s.Params, ex.sessionData.SearchPath) - if err != nil { - return makeErrEvent(err) - } - - stmt.Statement = ps.Statement - ast = stmt.AST - stmt.Prepared = ps - stmt.ExpectedTypes = ps.Columns - stmt.AnonymizedStr = ps.AnonymizedStr - res.ResetStmtType(ps.AST) - - if s.DiscardRows { - ih.SetDiscardRows() - } } p.semaCtx.Annotations = tree.MakeAnnotations(stmt.NumAnnotations) diff --git a/pkg/sql/stmtdiagnostics/statement_diagnostics_test.go b/pkg/sql/stmtdiagnostics/statement_diagnostics_test.go index 1443a253bc5c..fd303050b43a 100644 --- a/pkg/sql/stmtdiagnostics/statement_diagnostics_test.go +++ b/pkg/sql/stmtdiagnostics/statement_diagnostics_test.go @@ -86,6 +86,15 @@ func TestDiagnosticsRequest(t *testing.T) { _, err = db.Exec("INSERT INTO test VALUES (2)") require.NoError(t, err) checkCompleted(id1) + + // Verify that EXECUTE triggers diagnostics collection (#66048). + id4, err := registry.InsertRequestInternal(ctx, "SELECT x + $1 FROM test") + require.NoError(t, err) + _, err = db.Exec("PREPARE stmt AS SELECT x + $1 FROM test") + require.NoError(t, err) + _, err = db.Exec("EXECUTE stmt(1)") + require.NoError(t, err) + checkCompleted(id4) } // Test that a different node can service a diagnostics request.