From 908dbec0dd9bbd11bf9205e4fef54912ca92a141 Mon Sep 17 00:00:00 2001
From: Radu Berinde <radu@cockroachlabs.com>
Date: Thu, 3 Jun 2021 20:41:17 -0700
Subject: [PATCH] sql: fix statement diagnostics for EXECUTE

Previously, prepared statements ran through the EXECUTE statements
could not trigger collection of statement diagnostics. This is because
we consider the fingerprint of the EXECUTE statement instead of the
one from the prepared statement.

The fix is to move up the special handling code in the executor, and
replace the AST and fingerprint before setting up the instrumentation
helper.

Release note (bug fix): queries ran through the EXECUTE statement can
now generate statement diagnostic bundles as expected.

Fixes #66048.
---
 pkg/sql/conn_executor_exec.go                 | 64 ++++++++++---------
 .../statement_diagnostics_test.go             |  9 +++
 2 files changed, 44 insertions(+), 29 deletions(-)

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.