Merge #153905

153905: pgwire: show RowDescription for EXECUTE when needed r=rafiss a=rafiss

fixes #153505
Release note (bug fix): Fixed a bug where we would not show the pgwire RowDescription for EXECUTE statements that were themselves prepared using the pgwire Parse command.

Co-authored-by: Rafi Shamim <[email protected]>

Author: craig[bot]

pkg/sql/conn_executor.go

@@ -3430,8 +3430,19 @@ func (ex *connExecutor) execCopyIn(
 
 // stmtHasNoData returns true if describing a result of the input statement
 // type should return NoData.
-func stmtHasNoData(stmt tree.Statement) bool {
-	return stmt == nil || stmt.StatementReturnType() != tree.Rows
+func stmtHasNoData(stmt tree.Statement, resultColumns colinfo.ResultColumns) bool {
+	if stmt == nil {
+		return true
+	}
+	if stmt.StatementReturnType() == tree.Rows {
+		return false
+	}
+	// The statement may not always return rows (e.g. EXECUTE), but if it does,
+	// resultColumns will be non-empty.
+	if stmt.StatementReturnType() == tree.Unknown {
+		return len(resultColumns) == 0
+	}
+	return true
 }
 
 // commitPrepStmtNamespace deallocates everything in

pkg/sql/conn_executor_prepare.go

@@ -632,7 +632,7 @@ func (ex *connExecutor) execDescribe(
 			return retErr(sqlerrors.NewTransactionAbortedError("" /* customMsg */))
 		}
 		res.SetInferredTypes(ps.InferredTypes)
-		if stmtHasNoData(ast) {
+		if stmtHasNoData(ast, ps.Columns) {
 			res.SetNoDataRowDescription()
 		} else {
 			res.SetPrepStmtOutput(ctx, ps.Columns)
@@ -661,7 +661,7 @@ func (ex *connExecutor) execDescribe(
 		if isAbortedTxn && !ex.isAllowedInAbortedTxn(ast) {
 			return retErr(sqlerrors.NewTransactionAbortedError("" /* customMsg */))
 		}
-		if stmtHasNoData(ast) {
+		if stmtHasNoData(ast, portal.Stmt.Columns) {
 			res.SetNoDataRowDescription()
 		} else {
 			res.SetPortalOutput(ctx, portal.Stmt.Columns, portal.OutFormats)

pkg/sql/pgwire/testdata/pgtest/execute

@@ -0,0 +1,78 @@
+# This test verifies that the Parse/Bind/Describe/Execute protocol works with
+# prepared statements created with SQL-level PREPARE (so therefore must be
+# run with the EXECUTE statement).
+
+send
+Query {"String": "DROP TABLE IF EXISTS t0;"}
+Query {"String": "CREATE TABLE t0(c0 INT8);"}
+----
+
+until ignore=NoticeResponse
+ReadyForQuery
+ReadyForQuery
+----
+{"Type":"CommandComplete","CommandTag":"DROP TABLE"}
+{"Type":"ReadyForQuery","TxStatus":"I"}
+{"Type":"CommandComplete","CommandTag":"CREATE TABLE"}
+{"Type":"ReadyForQuery","TxStatus":"I"}
+
+# Prepare INSERT statement with a SQL-level PREPARE.
+send
+Query {"String": "PREPARE insert_query (int8) AS INSERT INTO t0 VALUES ($1);"}
+----
+
+until
+ReadyForQuery
+----
+{"Type":"CommandComplete","CommandTag":"PREPARE"}
+{"Type":"ReadyForQuery","TxStatus":"I"}
+
+# Execute INSERT statement using Parse/Bind/Describe/Execute protocol. The
+# Describe should return NoData, since INSERT doesn't return rows.
+send
+Parse {"Name": "insert_query_stmt", "Query": "EXECUTE insert_query(1)"}
+Bind {"DestinationPortal": "insert_query_portal", "PreparedStatement": "insert_query_stmt"}
+Describe {"ObjectType": "P", "Name": "insert_query_portal"}
+Execute {"Portal": "insert_query_portal"}
+Sync
+----
+
+until ignore_table_oids
+ReadyForQuery
+----
+{"Type":"ParseComplete"}
+{"Type":"BindComplete"}
+{"Type":"NoData"}
+{"Type":"CommandComplete","CommandTag":"INSERT 0 1"}
+{"Type":"ReadyForQuery","TxStatus":"I"}
+
+# Prepare SELECT statement with a SQL-level PREPARE.
+send
+Query {"String": "PREPARE select_query (int8) AS SELECT * FROM t0 WHERE c0 = $1;"}
+----
+
+until
+ReadyForQuery
+----
+{"Type":"CommandComplete","CommandTag":"PREPARE"}
+{"Type":"ReadyForQuery","TxStatus":"I"}
+
+# Execute SELECT statement using Parse/Bind/Describe/Execute protocol. The
+# Describe should return a RowDescription, since SELECT returns rows.
+send
+Parse {"Name": "select_query_stmt", "Query": "EXECUTE select_query(1)"}
+Bind {"DestinationPortal": "select_query_portal", "PreparedStatement": "select_query_stmt"}
+Describe {"ObjectType": "P", "Name": "select_query_portal"}
+Execute {"Portal": "select_query_portal"}
+Sync
+----
+
+until ignore_table_oids
+ReadyForQuery
+----
+{"Type":"ParseComplete"}
+{"Type":"BindComplete"}
+{"Type":"RowDescription","Fields":[{"Name":"c0","TableOID":0,"TableAttributeNumber":1,"DataTypeOID":20,"DataTypeSize":8,"TypeModifier":-1,"Format":0}]}
+{"Type":"DataRow","Values":[{"text":"1"}]}
+{"Type":"CommandComplete","CommandTag":"SELECT 1"}
+{"Type":"ReadyForQuery","TxStatus":"I"}