pgwire: fix JDBC driver compatibility for procedures with COMMIT statements

Previously, when using the extended protocol
(Parse/Bind/Describe/Execute/Sync) to call a PL/pgSQL procedure
containing COMMIT statements, CockroachDB would send extra
RowDescription messages after the COMMIT, causing the JDBC driver to
throw NoSuchElementException due to unexpected message sequences.

This change fixes the message flow to match JDBC driver expectations
when procedures execute COMMIT statements internally. The fix ensures
that the proper sequence of PostgreSQL wire protocol messages is sent,
preventing the driver from entering an inconsistent state.

Fixes #158771

Release Notes (Bug Fix): Fixed compatibility issue with JDBC driver when
calling PL/pgSQL procedures containing COMMIT statements via prepared
statements. The driver would previously throw NoSuchElementException due
to unexpected PostgreSQL wire protocol message sequences.

Author: ZhouXing19

pkg/sql/conn_executor.go

@@ -1712,6 +1712,11 @@ type connExecutor struct {
 			// will be used to synthesize an ExecStmt command to avoid attempting to
 			// execute the portal twice.
 			resumeStmt statements.Statement[tree.Statement]
+
+			// callRowDescSent tracks whether RowDescription has already been
+			// sent for a CALL statement to prevent duplicate RowDescription
+			// messages when stored procedures contain COMMIT/ROLLBACK.
+			callRowDescSent bool
 		}
 
 		// shouldExecuteOnTxnRestart indicates that ex.onTxnRestart will be
@@ -2801,6 +2806,7 @@ func (ex *connExecutor) execCmd() (retErr error) {
 		ex.extraTxnState.storedProcTxnState.resumeProc = nil
 		ex.extraTxnState.storedProcTxnState.resumeStmt = statements.Statement[tree.Statement]{}
 		ex.extraTxnState.storedProcTxnState.txnModes = nil
+		ex.extraTxnState.storedProcTxnState.callRowDescSent = false
 	}
 
 	if err := ex.updateTxnRewindPosMaybe(ctx, cmd, pos, advInfo); err != nil {
@@ -3502,6 +3508,11 @@ func stmtHasNoData(stmt tree.Statement, resultColumns colinfo.ResultColumns) boo
 		return true
 	}
 	if stmt.StatementReturnType() == tree.Rows {
+		// If the procedure doesn't contain output parameters, write a NoData
+		// message.
+		if stmt.StatementTag() == tree.CallStmtTag {
+			return len(resultColumns) == 0
+		}
 		return false
 	}
 	// The statement may not always return rows (e.g. EXECUTE), but if it does,
@@ -4384,8 +4395,17 @@ func (ex *connExecutor) initStatementResult(
 	// ANALYZE), then the columns will be set later.
 	if ex.planner.instrumentation.outputMode == unmodifiedOutput &&
 		ast.StatementReturnType() == tree.Rows {
+		// Only write RowDescription message if the procedure has output parameters.
+		skipRowDescription := false
+		if ast.StatementTag() == tree.CallStmtTag {
+			if len(cols) == 0 || ex.extraTxnState.storedProcTxnState.callRowDescSent {
+				skipRowDescription = true
+			} else {
+				ex.extraTxnState.storedProcTxnState.callRowDescSent = true
+			}
+		}
 		// Note that this call is necessary even if cols is nil.
-		res.SetColumns(ctx, cols)
+		res.SetColumns(ctx, cols, skipRowDescription)
 	}
 	return nil
 }

pkg/sql/conn_executor_exec.go

@@ -3781,7 +3781,7 @@ func (ex *connExecutor) runObserverStatement(
 func (ex *connExecutor) runShowSyntax(
 	ctx context.Context, stmt string, res RestrictedCommandResult,
 ) error {
-	res.SetColumns(ctx, colinfo.ShowSyntaxColumns)
+	res.SetColumns(ctx, colinfo.ShowSyntaxColumns, false /* skipRowDescription */)
 	var commErr error
 	parser.RunShowSyntax(ctx, stmt,
 		func(ctx context.Context, field, msg string) {
@@ -3800,7 +3800,7 @@ func (ex *connExecutor) runShowSyntax(
 func (ex *connExecutor) runShowTransactionState(
 	ctx context.Context, res RestrictedCommandResult,
 ) error {
-	res.SetColumns(ctx, colinfo.ResultColumns{{Name: "TRANSACTION STATUS", Typ: types.String}})
+	res.SetColumns(ctx, colinfo.ResultColumns{{Name: "TRANSACTION STATUS", Typ: types.String}}, false /* skipRowDescription */)
 
 	state := fmt.Sprintf("%s", ex.machine.CurState())
 	return res.AddRow(ctx, tree.Datums{tree.NewDString(state)})
@@ -3863,7 +3863,7 @@ func (ex *connExecutor) runShowTransferState(
 	for i := 0; i < len(colNames); i++ {
 		cols[i] = colinfo.ResultColumn{Name: colNames[i], Typ: types.String}
 	}
-	res.SetColumns(ctx, cols)
+	res.SetColumns(ctx, cols, false /* skipRowDescription */)
 
 	var sessionState, sessionRevivalToken tree.Datum
 	var row tree.Datums
@@ -3897,7 +3897,7 @@ func (ex *connExecutor) runShowTransferState(
 func (ex *connExecutor) runShowCompletions(
 	ctx context.Context, n *tree.ShowCompletions, res RestrictedCommandResult,
 ) error {
-	res.SetColumns(ctx, colinfo.ShowCompletionsColumns)
+	res.SetColumns(ctx, colinfo.ShowCompletionsColumns, false /* skipRowDescription */)
 	log.Dev.Warningf(ctx, "COMPLETION GENERATOR FOR: %+v", *n)
 	sd := ex.planner.SessionData()
 	override := sessiondata.InternalExecutorOverride{
@@ -3964,7 +3964,7 @@ func (ex *connExecutor) runShowLastQueryStatistics(
 	for i, n := range stmt.Columns {
 		resColumns[i] = colinfo.ResultColumn{Name: string(n), Typ: types.String}
 	}
-	res.SetColumns(ctx, resColumns)
+	res.SetColumns(ctx, resColumns, false /* skipRowDescription */)
 
 	phaseTimes := ex.statsCollector.PreviousPhaseTimes()
 

pkg/sql/conn_executor_savepoints.go

@@ -410,7 +410,7 @@ func (ex *connExecutor) runShowSavepointState(
 	res.SetColumns(ctx, colinfo.ResultColumns{
 		{Name: "savepoint_name", Typ: types.String},
 		{Name: "is_initial_savepoint", Typ: types.Bool},
-	})
+	}, false /* skipRowDescription */)
 
 	for _, entry := range ex.extraTxnState.savepoints {
 		if err := res.AddRow(ctx, tree.Datums{

pkg/sql/conn_executor_show_commit_timestamp.go

@@ -141,6 +141,6 @@ func (ex *connExecutor) execShowCommitTimestampInNoTxnState(
 func writeShowCommitTimestampRow(
 	ctx context.Context, res RestrictedCommandResult, ts hlc.Timestamp,
 ) error {
-	res.SetColumns(ctx, colinfo.ShowCommitTimestampColumns)
+	res.SetColumns(ctx, colinfo.ShowCommitTimestampColumns, false /* skipRowDescription */)
 	return res.AddRow(ctx, tree.Datums{eval.TimestampToDecimalDatum(ts)})
 }

pkg/sql/conn_io.go

@@ -867,7 +867,7 @@ type RestrictedCommandResult interface {
 	// can be nil.
 	//
 	// This needs to be called (once) before AddRow.
-	SetColumns(context.Context, colinfo.ResultColumns)
+	SetColumns(ctx context.Context, cols colinfo.ResultColumns, skipRowDescription bool)
 
 	// ResetStmtType allows a client to change the statement type of the current
 	// result, from the original one set when the result was created trough
@@ -948,7 +948,7 @@ type DescribeResult interface {
 	SetPrepStmtOutput(context.Context, colinfo.ResultColumns)
 	// SetPortalOutput tells the client about the results schema and formatting of
 	// a portal.
-	SetPortalOutput(context.Context, colinfo.ResultColumns, []pgwirebase.FormatCode)
+	SetPortalOutput(ctx context.Context, cols colinfo.ResultColumns, fmtCode []pgwirebase.FormatCode)
 }
 
 // ParseResult represents the result of a Parse command.
@@ -1114,7 +1114,9 @@ func (r *streamingCommandResult) RevokePortalPausability() error {
 }
 
 // SetColumns is part of the RestrictedCommandResult interface.
-func (r *streamingCommandResult) SetColumns(ctx context.Context, cols colinfo.ResultColumns) {
+func (r *streamingCommandResult) SetColumns(
+	ctx context.Context, cols colinfo.ResultColumns, skipRowDescription bool,
+) {
 	// The interface allows for cols to be nil, yet the iterator result expects
 	// non-nil value to indicate that it was the column metadata.
 	if cols == nil {

pkg/sql/explain_bundle.go

@@ -60,7 +60,7 @@ func setExplainBundleResult(
 	warnings []string,
 ) error {
 	res.ResetStmtType(&tree.ExplainAnalyze{})
-	res.SetColumns(ctx, colinfo.ExplainPlanColumns)
+	res.SetColumns(ctx, colinfo.ExplainPlanColumns, false /* skipRowDescription */)
 
 	var text []string
 	if bundle.collectionErr != nil {

pkg/sql/instrumentation.go

@@ -976,7 +976,7 @@ func (ih *instrumentationHelper) setExplainAnalyzeResult(
 	trace tracingpb.Recording,
 ) (commErr error) {
 	res.ResetStmtType(&tree.ExplainAnalyze{})
-	res.SetColumns(ctx, colinfo.ExplainPlanColumns)
+	res.SetColumns(ctx, colinfo.ExplainPlanColumns, false /* skipRowDescription */)
 
 	if res.Err() != nil {
 		// Can't add rows if there was an error.

pkg/sql/isession/command_result.go

@@ -73,7 +73,9 @@ func (i *internalCommandResult) SendNotice(
 	return nil
 }
 
-func (i *internalCommandResult) SetColumns(ctx context.Context, cols colinfo.ResultColumns) {
+func (i *internalCommandResult) SetColumns(
+	ctx context.Context, cols colinfo.ResultColumns, skipRowDescription bool,
+) {
 	// We don't need this because the datums include type information.
 }
 
@@ -166,7 +168,7 @@ func (i *internalCommandResult) SendCopyOut(
 }
 
 func (i *internalCommandResult) SetPortalOutput(
-	ctx context.Context, cols colinfo.ResultColumns, formatCodes []pgwirebase.FormatCode,
+	ctx context.Context, cols colinfo.ResultColumns, fmtCode []pgwirebase.FormatCode,
 ) {
 	i.SetError(errors.AssertionFailedf("SetPortalOutput is not supported by the internal session"))
 }

pkg/sql/opt/optbuilder/plpgsql.go

@@ -1272,7 +1272,7 @@ func (b *plpgsqlBuilder) buildPLpgSQLStatements(stmts []ast.Statement, s *scope)
 					))
 				}
 			}
-			b.checkDuplicateTargets(target, "CALL")
+			b.checkDuplicateTargets(target, tree.CallStmtTag)
 			if len(target) == 0 {
 				// When there is no INTO target, build the nested procedure call into a
 				// body statement that is only executed for its side effects.

pkg/sql/pgwire/command_result.go

@@ -343,10 +343,12 @@ func (r *commandResult) SendNotice(
 }
 
 // SetColumns is part of the sql.RestrictedCommandResult interface.
-func (r *commandResult) SetColumns(ctx context.Context, cols colinfo.ResultColumns) {
+func (r *commandResult) SetColumns(
+	ctx context.Context, cols colinfo.ResultColumns, skipRowDescription bool,
+) {
 	r.assertNotReleased()
 	r.conn.writerState.fi.registerCmd(r.pos)
-	if r.descOpt == sql.NeedRowDesc {
+	if r.descOpt == sql.NeedRowDesc && !skipRowDescription {
 		_ /* err */ = r.conn.writeRowDescription(ctx, cols, r.formatCodes, &r.conn.writerState.buf)
 	}
 	r.types = make([]*types.T, len(cols))