From 8fc49e56ec0297b497b04c9530e0d09e7c6d4422 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Knut=20Olav=20L=C3=B8ite?= <koloite@gmail.com>
Date: Sun, 14 Sep 2025 08:27:03 +0200
Subject: [PATCH] test: correctly return tx id for update counts

The mock Spanner server did not correctly return the transaction ID when
ExecuteStreamingSql was used in combination with an update count result.

Instead of returning the transaction ID of the transaction that was
actually started, it returned the transaction ID that was set on the
request. This would be empty for requests that used a Begin transaction
selector, which again meant that it would return an empty transaction ID.

This change also returns the UPDATE_COUNT field in the result set, which
is also not returned by real Spanner.
---
 .../cloud/spanner/MockSpannerServiceImpl.java | 37 +++++++++++--------
 1 file changed, 21 insertions(+), 16 deletions(-)

diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MockSpannerServiceImpl.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MockSpannerServiceImpl.java
index cccfaead555..a47aecdccc4 100644
--- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MockSpannerServiceImpl.java
+++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MockSpannerServiceImpl.java
@@ -1245,7 +1245,13 @@ public void executeStreamingSql(
             throw firstRes.getException();
           case UPDATE_COUNT:
             returnPartialResultSet(
-                session, 0L, !isPartitioned, responseObserver, request.getTransaction(), false);
+                session,
+                0L,
+                !isPartitioned,
+                responseObserver,
+                request.getTransaction(),
+                transactionId,
+                false);
             break;
           case RESULT_SET:
           default:
@@ -1290,7 +1296,8 @@ public void executeStreamingSql(
               res.getUpdateCount(),
               !isPartitioned,
               responseObserver,
-              request.getTransaction());
+              request.getTransaction(),
+              transactionId);
           break;
         default:
           throw new IllegalStateException("Unknown result type: " + res.getType());
@@ -1798,8 +1805,10 @@ private void returnPartialResultSet(
       Long updateCount,
       boolean exact,
       StreamObserver<PartialResultSet> responseObserver,
-      TransactionSelector transaction) {
-    returnPartialResultSet(session, updateCount, exact, responseObserver, transaction, true);
+      TransactionSelector transactionSelector,
+      ByteString transactionId) {
+    returnPartialResultSet(
+        session, updateCount, exact, responseObserver, transactionSelector, transactionId, true);
   }
 
   private void returnPartialResultSet(
@@ -1807,23 +1816,19 @@ private void returnPartialResultSet(
       Long updateCount,
       boolean exact,
       StreamObserver<PartialResultSet> responseObserver,
-      TransactionSelector transaction,
+      TransactionSelector transactionSelector,
+      ByteString transactionId,
       boolean complete) {
-    Field field =
-        Field.newBuilder()
-            .setName("UPDATE_COUNT")
-            .setType(Type.newBuilder().setCode(TypeCode.INT64).build())
-            .build();
     if (exact) {
       responseObserver.onNext(
           PartialResultSet.newBuilder()
               .setMetadata(
                   ResultSetMetadata.newBuilder()
-                      .setRowType(StructType.newBuilder().addFields(field).build())
+                      .setRowType(StructType.newBuilder().build())
                       .setTransaction(
-                          ignoreInlineBeginRequest.get()
+                          ignoreInlineBeginRequest.get() || !transactionSelector.hasBegin()
                               ? Transaction.getDefaultInstance()
-                              : Transaction.newBuilder().setId(transaction.getId()).build())
+                              : Transaction.newBuilder().setId(transactionId).build())
                       .build())
               .setStats(ResultSetStats.newBuilder().setRowCountExact(updateCount).build())
               .build());
@@ -1832,11 +1837,11 @@ private void returnPartialResultSet(
           PartialResultSet.newBuilder()
               .setMetadata(
                   ResultSetMetadata.newBuilder()
-                      .setRowType(StructType.newBuilder().addFields(field).build())
+                      .setRowType(StructType.newBuilder().build())
                       .setTransaction(
-                          ignoreInlineBeginRequest.get()
+                          ignoreInlineBeginRequest.get() || !transactionSelector.hasBegin()
                               ? Transaction.getDefaultInstance()
-                              : Transaction.newBuilder().setId(transaction.getId()).build())
+                              : Transaction.newBuilder().setId(transactionId).build())
                       .build())
               .setStats(ResultSetStats.newBuilder().setRowCountLowerBound(updateCount).build())
               .build());