test: correctly return tx id for update counts (#4045)

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.

Author: olavloite

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,32 +1805,30 @@ 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(
       Session session,
       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());