sql: fix statement diagnostics for EXECUTE

pkg/sql/conn_executor_exec.go

@@ -358,6 +358,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:
@@ -395,6 +396,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,
@@ -531,35 +566,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)

pkg/sql/stmtdiagnostics/statement_diagnostics_test.go

@@ -88,6 +88,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.