[release-24.0] vtgate: report RowsAffected for stored-procedure calls over the streaming path (#20402)

go/mysql/streaming_query_test.go

@@ -31,8 +31,9 @@ import (
 
 // TestExecuteStreamFetchOKPacket verifies that a streaming query which returns an
 // OK packet instead of a result set (e.g. a CALL of a procedure that performs DML)
-// exposes the OK-packet RowsAffected and InsertID via StreamOKResult. This mirrors
-// the buffered ExecuteFetch path, which builds the same Result from the OK packet.
+// exposes the OK-packet RowsAffected, InsertID and SessionStateChanges via
+// StreamOKResult. This mirrors the buffered ExecuteFetch path, which builds the
+// same Result from the OK packet.
 func TestExecuteStreamFetchOKPacket(t *testing.T) {
 	listener, sConn, cConn := createSocketPair(t)
 	defer func() {
@@ -41,6 +42,11 @@ func TestExecuteStreamFetchOKPacket(t *testing.T) {
 		cConn.Close()
 	}()
 
+	// Session-state data is only written and parsed when both sides negotiated
+	// session tracking.
+	sConn.Capabilities |= CapabilityClientSessionTrack
+	cConn.Capabilities |= CapabilityClientSessionTrack
+
 	wg := sync.WaitGroup{}
 	var streamErr error
 	var okRes *sqltypes.Result
@@ -53,14 +59,16 @@ func TestExecuteStreamFetchOKPacket(t *testing.T) {
 	})
 
 	// The server reads the COM_QUERY and responds with an OK packet carrying
-	// RowsAffected, InsertID and Info but no result set.
+	// RowsAffected, InsertID and session-state data but no result set.
 	data, err := sConn.readEphemeralPacket()
 	require.NoError(t, err)
 	require.EqualValues(t, ComQuery, data[0])
 	sConn.recycleReadPacket()
 	require.NoError(t, sConn.writeOKPacket(&PacketOK{
-		affectedRows: 7,
-		lastInsertID: 99,
+		affectedRows:     7,
+		lastInsertID:     99,
+		statusFlags:      ServerSessionStateChanged,
+		sessionStateData: "8bb25b46-16bd-11ea-8ffa-98af65266957:8",
 	}))
 
 	wg.Wait()
@@ -69,6 +77,8 @@ func TestExecuteStreamFetchOKPacket(t *testing.T) {
 	assert.EqualValues(t, 7, okRes.RowsAffected)
 	assert.EqualValues(t, 99, okRes.InsertID)
 	assert.True(t, okRes.InsertIDChanged)
+	assert.Equal(t, "8bb25b46-16bd-11ea-8ffa-98af65266957:8", okRes.SessionStateChanges,
+		"streaming OK packet must carry the session-state data")
 }
 
 // TestExecuteStreamFetchNoOKResultForRows verifies that a streaming query which

go/test/endtoend/vtgate/unsharded/streaming_okpacket_test.go

@@ -0,0 +1,67 @@
+/*
+Copyright 2026 The Vitess Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package unsharded
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"vitess.io/vitess/go/mysql"
+	"vitess.io/vitess/go/test/endtoend/utils"
+)
+
+// streamingConn opens a vtgate connection pinned to the OLAP workload so that
+// queries are routed through vtgate's StreamExecute path instead of the buffered
+// Execute path.
+func streamingConn(t *testing.T) *mysql.Conn {
+	t.Helper()
+	vtParams := mysql.ConnParams{
+		Host:   "localhost",
+		Port:   clusterInstance.VtgateMySQLPort,
+		Flags:  mysql.CapabilityClientMultiResults,
+		DbName: "@primary",
+	}
+	conn, err := mysql.Connect(t.Context(), &vtParams)
+	require.NoError(t, err)
+	t.Cleanup(conn.Close)
+
+	utils.Exec(t, conn, "set workload = olap")
+	return conn
+}
+
+// TestStreamingCallProcedureRowsAffected verifies that a CALL of a procedure that
+// performs DML reports the affected-row count to the client over the streaming
+// (OLAP) path, matching the buffered Execute path that TestCallProcedure pins for
+// the OLTP path.
+func TestStreamingCallProcedureRowsAffected(t *testing.T) {
+	conn := streamingConn(t)
+
+	// Clean up the inserted row so later tests see the table unchanged. The
+	// delete runs on the buffered (OLTP) path so the cleanup also works on
+	// release branches without streamed DML support.
+	t.Cleanup(func() {
+		utils.Exec(t, conn, `set workload = oltp`)
+		utils.Exec(t, conn, `delete from allDefaults`)
+	})
+
+	// sp_insert performs `insert into allDefaults () values ()`, returning an OK
+	// packet that carries the affected-row count.
+	qr := utils.Exec(t, conn, `CALL sp_insert()`)
+	assert.EqualValues(t, 1, qr.RowsAffected, "streamed CALL must report affected rows to the client")
+}

go/vt/dbconnpool/connection.go

@@ -98,7 +98,19 @@ func (dbc *DBConnection) ExecuteStreamFetch(query string, callback func(*sqltype
 	if err != nil {
 		return err
 	}
-	err = callback(&sqltypes.Result{Fields: flds})
+	firstResult := &sqltypes.Result{Fields: flds}
+	// If the query produced no result set but an OK packet (e.g. a CALL that
+	// performs DML), carry its RowsAffected/InsertID/Info/SessionStateChanges
+	// through so the streaming path reports them like the buffered ExecuteFetch
+	// path does.
+	if okRes := dbc.StreamOKResult(); okRes != nil {
+		firstResult.RowsAffected = okRes.RowsAffected
+		firstResult.InsertID = okRes.InsertID
+		firstResult.InsertIDChanged = okRes.InsertIDChanged
+		firstResult.Info = okRes.Info
+		firstResult.SessionStateChanges = okRes.SessionStateChanges
+	}
+	err = callback(firstResult)
 	if err != nil {
 		return fmt.Errorf("stream send error: %v", err)
 	}

go/vt/dbconnpool/connection_test.go

@@ -0,0 +1,63 @@
+/*
+Copyright 2026 The Vitess Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package dbconnpool
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"vitess.io/vitess/go/mysql/fakesqldb"
+	"vitess.io/vitess/go/sqltypes"
+	"vitess.io/vitess/go/vt/dbconfigs"
+)
+
+// TestExecuteStreamFetchCarriesOKPacket verifies that ExecuteStreamFetch forwards
+// the OK-packet RowsAffected and InsertID to the callback for a query that
+// produces no result set (e.g. a CALL of a procedure that performs DML), matching
+// the buffered ExecuteFetch path.
+func TestExecuteStreamFetchCarriesOKPacket(t *testing.T) {
+	db := fakesqldb.New(t)
+	t.Cleanup(db.Close)
+
+	const query = "CALL sp_insert()"
+	db.AddQuery(query, &sqltypes.Result{
+		RowsAffected:    7,
+		InsertID:        99,
+		InsertIDChanged: true,
+	})
+
+	conn, err := NewDBConnection(t.Context(), dbconfigs.New(db.ConnParams()))
+	require.NoError(t, err)
+	t.Cleanup(conn.Close)
+
+	got := &sqltypes.Result{}
+	err = conn.ExecuteStreamFetch(query, func(r *sqltypes.Result) error {
+		got.RowsAffected += r.RowsAffected
+		if r.InsertIDChanged {
+			got.InsertID = r.InsertID
+			got.InsertIDChanged = true
+		}
+		return nil
+	}, func() *sqltypes.Result { return &sqltypes.Result{} }, 4096)
+	require.NoError(t, err)
+
+	assert.EqualValues(t, 7, got.RowsAffected, "streamed OK packet must carry RowsAffected to the callback")
+	assert.EqualValues(t, 99, got.InsertID, "streamed OK packet must carry InsertID to the callback")
+	assert.True(t, got.InsertIDChanged)
+}

go/vt/vtgate/executor.go

@@ -335,13 +335,32 @@ func (e *Executor) StreamExecute(
 			srr.callback = func(qr *sqltypes.Result) error {
 				resultMu.Lock()
 				defer resultMu.Unlock()
+				// Carry over the OK-packet data (affected rows, last insert id, info,
+				// session state changes) so statements that return an OK packet (e.g.
+				// CALL of a procedure that performs DML) report it to the client,
+				// matching the buffered Execute path.
+				result.RowsAffected += qr.RowsAffected
+				if qr.InsertIDUpdated() {
+					result.InsertID = qr.InsertID
+					result.InsertIDChanged = true
+				}
+				if qr.SessionStateChanges != "" {
+					result.SessionStateChanges = qr.SessionStateChanges
+				}
+				if qr.Info != "" {
+					result.Info = qr.Info
+				}
 				// If the row has field info, send it separately.
 				// TODO(sougou): this behavior is for handling tests because
 				// the framework currently sends all results as one packet.
 				byteCount := 0
 				if len(qr.Fields) > 0 {
 					result.Fields = qr.Fields
-					if err := callback(qr.Metadata()); err != nil {
+					// Send only the fields here: any OK-packet data on qr was
+					// accumulated into result above and goes to the client with
+					// the left-over result, so repeating it in this packet would
+					// deliver it twice to clients that sum across packets.
+					if err := callback(&sqltypes.Result{Fields: qr.Fields}); err != nil {
 						return err
 					}
 					seenResults.Store(true)
@@ -372,8 +391,68 @@ func (e *Executor) StreamExecute(
 		err := vc.StreamExecutePrimitive(ctx, plan.Instructions, bindVars, true, func(qr *sqltypes.Result) error {
 			return srr.storeResultStats(plan.QueryType, qr)
 		})
+<<<<<<< HEAD
+||||||| parent of 8055e6946d (vtgate: report RowsAffected for stored-procedure calls over the streaming path (#20402))
+
+		updateLogStats := func() {
+			logStats.StmtType = plan.QueryType.String()
+			logStats.PlanType = plan.Type.String()
+			logStats.TablesUsed = plan.TablesUsed
+			executedRoot := vc.ExecutedPrimitive()
+			if executedRoot == nil {
+				executedRoot = plan.Instructions
+			}
+			logStats.RoutingIndexesUsed = engine.GetRoutingIndexes(executedRoot)
+			logStats.TabletType = vc.TabletType().String()
+			logStats.ExecuteTime = time.Since(execStart)
+			logStats.ActiveKeyspace = vc.GetKeyspace()
+
+			e.updateQueryStats(plan.QueryType.String(), plan.Type.String(), vc.TabletType().String(), int64(logStats.ShardQueries), plan.TablesUsed)
+		}
+
+=======
+
+		updateLogStats := func(err error) {
+			logStats.StmtType = plan.QueryType.String()
+			logStats.PlanType = plan.Type.String()
+			logStats.TabletType = vc.TabletType().String()
+			logStats.ExecuteTime = time.Since(execStart)
+			logStats.ActiveKeyspace = vc.GetKeyspace()
+
+			// On error, leave the tables, routing indexes and row counts unset so the
+			// per-table counters are not incremented, matching the buffered Execute path.
+			var tablesUsed []string
+			var errCount uint64
+			if err != nil {
+				logStats.Error = err
+				errCount = 1
+			} else {
+				srr.mu.Lock()
+				logStats.RowsAffected = srr.rowsAffected
+				logStats.RowsReturned = uint64(srr.rowsReturned)
+				srr.mu.Unlock()
+				logStats.TablesUsed = plan.TablesUsed
+				tablesUsed = plan.TablesUsed
+				executedRoot := vc.ExecutedPrimitive()
+				if executedRoot == nil {
+					executedRoot = plan.Instructions
+				}
+				logStats.RoutingIndexesUsed = engine.GetRoutingIndexes(executedRoot)
+			}
+
+			e.updateQueryStats(plan.QueryType.String(), plan.Type.String(), vc.TabletType().String(), int64(logStats.ShardQueries), tablesUsed)
+			plan.AddStats(1, time.Since(logStats.StartTime), logStats.ShardQueries, logStats.RowsAffected, logStats.RowsReturned, errCount)
+		}
+
+>>>>>>> 8055e6946d (vtgate: report RowsAffected for stored-procedure calls over the streaming path (#20402))
 		// Check if there was partial DML execution. If so, rollback the effect of the partially executed query.
 		if err != nil {
+			// Record query stats for a failed row-returning query before any
+			// rollback handling, matching the buffered Execute path which always
+			// records them ahead of its own rollback handling.
+			if canReturnRows(plan.QueryType) {
+				updateLogStats(err)
+			}
 			if safeSession.InTransaction() && e.rollbackOnFatalTxError(ctx, safeSession, err) {
 				return err
 			}
@@ -384,23 +463,44 @@ func (e *Executor) StreamExecute(
 		}
 
 		if !canReturnRows(plan.QueryType) {
+<<<<<<< HEAD
+||||||| parent of 8055e6946d (vtgate: report RowsAffected for stored-procedure calls over the streaming path (#20402))
+			updateLogStats()
+=======
+			updateLogStats(nil)
+>>>>>>> 8055e6946d (vtgate: report RowsAffected for stored-procedure calls over the streaming path (#20402))
 			return nil
 		}
 
-		// Send left-over rows if there is no error on execution.
-		if len(result.Rows) > 0 || !seenResults.Load() {
+		// Send left-over rows if there is no error on execution. The left-over
+		// result must also go out when it carries OK-packet data with no rows —
+		// e.g. an affected-row count that arrived after a result set was already
+		// sent — so that data is not dropped.
+		hasOKData := result.RowsAffected > 0 || result.InsertIDUpdated() ||
+			result.SessionStateChanges != "" || result.Info != ""
+		if len(result.Rows) > 0 || hasOKData || !seenResults.Load() {
 			if err := callback(result); err != nil {
+				// The query executed; only the delivery to the client failed.
+				// Record it as an error, like a mid-stream send failure that
+				// surfaces through the execution error above.
+				updateLogStats(err)
 				return err
 			}
 		}
 
 		// 5: Log and add statistics
+<<<<<<< HEAD
 		logStats.TablesUsed = plan.TablesUsed
 		logStats.TabletType = vc.TabletType().String()
 		logStats.ExecuteTime = time.Since(execStart)
 		logStats.ActiveKeyspace = vc.GetKeyspace()
 
 		e.updateQueryStats(plan.QueryType.String(), plan.Type.String(), vc.TabletType().String(), int64(logStats.ShardQueries), plan.TablesUsed)
+||||||| parent of 8055e6946d (vtgate: report RowsAffected for stored-procedure calls over the streaming path (#20402))
+		updateLogStats()
+=======
+		updateLogStats(nil)
+>>>>>>> 8055e6946d (vtgate: report RowsAffected for stored-procedure calls over the streaming path (#20402))
 
 		return err
 	}