From d3cbcfd3a51a79fbfa60d0674236346fb4436dea 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 23f20c18db70..b2d20d77cef9 100644
--- a/pkg/sql/conn_executor_exec.go
+++ b/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)
diff --git a/pkg/sql/stmtdiagnostics/statement_diagnostics_test.go b/pkg/sql/stmtdiagnostics/statement_diagnostics_test.go
index 381ae220597a..65a865e78566 100644
--- a/pkg/sql/stmtdiagnostics/statement_diagnostics_test.go
+++ b/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.