fix: allow DML THEN RETURN with retryAbortsInternally=false

README.md

@@ -19,7 +19,7 @@ If you are using Maven with [BOM][libraries-bom], add this to your pom.xml file:
     <dependency>
       <groupId>com.google.cloud</groupId>
       <artifactId>libraries-bom</artifactId>
-      <version>26.70.0</version>
+      <version>26.71.0</version>
       <type>pom</type>
       <scope>import</scope>
     </dependency>
@@ -41,7 +41,7 @@ If you are using Maven without the BOM, add this to your dependencies:
 <dependency>
   <groupId>com.google.cloud</groupId>
   <artifactId>google-cloud-spanner</artifactId>
-  <version>6.102.0</version>
+  <version>6.102.1</version>
 </dependency>
 
 ```

google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ReadWriteTransaction.java

@@ -689,7 +689,19 @@ public ApiFuture<ResultSet> executeQueryAsync(
                 InterceptorsUsage.IGNORE_INTERCEPTORS,
                 ImmutableList.of(SpannerGrpc.getExecuteStreamingSqlMethod()));
       } else {
-        res = super.executeQueryAsync(callType, statement, analyzeMode, options);
+        // Handle both SELECT queries and DML with THEN RETURN without delegating to the base class,
+        // which rejects non-SELECT statements.
+        res =
+            executeStatementAsync(
+                callType,
+                statement,
+                () -> {
+                  checkTimedOut();
+                  checkAborted();
+                  return DirectExecuteResultSet.ofResultSet(
+                      internalExecuteQuery(statement, analyzeMode, options));
+                },
+                SpannerGrpc.getExecuteStreamingSqlMethod());
       }
       ApiFutures.addCallback(res, new StatementResultCallback<>(), MoreExecutors.directExecutor());
       return res;

google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/AbstractMockServerTest.java

@@ -110,7 +110,7 @@ public abstract class AbstractMockServerTest {
           .setMetadata(SINGLE_COL_INT64_RESULTSET_METADATA)
           .build();
   public static final com.google.spanner.v1.ResultSet UPDATE_RETURNING_RESULTSET =
-      com.google.spanner.v1.ResultSet.newBuilder()
+      ResultSet.newBuilder()
           .setStats(ResultSetStats.newBuilder().setRowCountExact(1))
           .setMetadata(
               ResultSetMetadata.newBuilder()
@@ -121,6 +121,10 @@ public abstract class AbstractMockServerTest {
                                   .setName("col")
                                   .setType(Type.newBuilder().setCodeValue(TypeCode.INT64_VALUE))
                                   .build())))
+          .addRows(
+              ListValue.newBuilder()
+                  .addValues(Value.newBuilder().setStringValue("1").build())
+                  .build())
           .build();
 
   protected static final ResultSet SELECT1_RESULTSET =

google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/AutoDmlBatchMockServerTest.java

@@ -97,6 +97,8 @@ public void testDmlWithReturningAfterDml() {
       // DML with a THEN RETURN clause cannot be batched. This therefore flushes the batch and
       // executes the INSERT ... THEN RETURN statement as a separate ExecuteSqlRequest.
       try (ResultSet resultSet = connection.executeQuery(INSERT_RETURNING_STATEMENT)) {
+        assertTrue(resultSet.next());
+        assertEquals(1L, resultSet.getLong(0));
         assertFalse(resultSet.next());
       }
 
@@ -123,6 +125,8 @@ public void testDmlWithReturningAfterDml_usingExecute() {
       StatementResult result = connection.execute(INSERT_RETURNING_STATEMENT);
       assertEquals(ResultType.RESULT_SET, result.getResultType());
       try (ResultSet resultSet = result.getResultSet()) {
+        assertTrue(resultSet.next());
+        assertEquals(1L, resultSet.getLong(0));
         assertFalse(resultSet.next());
       }
 

google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ReadWriteTransactionTest.java

@@ -287,6 +287,21 @@ public void testExecuteUpdate() {
     assertThat(get(transaction.executeUpdateAsync(CallType.SYNC, parsedStatement)), is(1L));
   }
 
+  @Test
+  public void testExecuteQueryWithDmlReturningWithoutRetry() {
+    ParsedStatement parsedStatement = mock(ParsedStatement.class);
+    when(parsedStatement.getType()).thenReturn(StatementType.UPDATE);
+    when(parsedStatement.isUpdate()).thenReturn(true);
+    when(parsedStatement.hasReturningClause()).thenReturn(true);
+    Statement statement = Statement.of("INSERT INTO TEST (ID, NAME) VALUES (1, 'x') THEN RETURN *");
+    when(parsedStatement.getStatement()).thenReturn(statement);
+
+    ReadWriteTransaction transaction = createSubject(/* commitBehavior= */ CommitBehavior.SUCCEED);
+    ResultSet rs =
+        get(transaction.executeQueryAsync(CallType.SYNC, parsedStatement, AnalyzeMode.NONE));
+    assertThat(rs, is(notNullValue()));
+  }
+
   @Test
   public void testGetCommitTimestampBeforeCommit() {
     ParsedStatement parsedStatement = mock(ParsedStatement.class);

google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/TransactionMockServerTest.java

@@ -31,6 +31,7 @@
 import com.google.cloud.spanner.SpannerException;
 import com.google.cloud.spanner.Statement;
 import com.google.cloud.spanner.connection.ITAbstractSpannerTest.ITConnection;
+import com.google.cloud.spanner.connection.StatementResult.ResultType;
 import com.google.spanner.v1.BeginTransactionRequest;
 import com.google.spanner.v1.CommitRequest;
 import com.google.spanner.v1.ExecuteBatchDmlRequest;
@@ -357,5 +358,45 @@ public long nanoTime() {
   }
 
   @Test
-  public void testTransactionTimeoutInAutoCommit() {}
+  public void testCanUseAllMethodsWithInternalRetriesDisabled() {
+    // Verify that all query/update methods work as expected when internal retries have been
+    // disabled.
+    try (Connection connection = createConnection()) {
+      connection.setAutocommit(false);
+      connection.setRetryAbortsInternally(false);
+
+      try (ResultSet result = connection.executeQuery(SELECT1_STATEMENT)) {
+        assertTrue(result.next());
+        assertEquals(1L, result.getLong(0));
+        assertFalse(result.next());
+      }
+      assertEquals(1, connection.executeUpdate(INSERT_STATEMENT));
+      try (ResultSet result = connection.executeQuery(INSERT_RETURNING_STATEMENT)) {
+        assertTrue(result.next());
+        assertEquals(1L, result.getLong(0));
+        assertFalse(result.next());
+      }
+
+      StatementResult statementResult = connection.execute(SELECT1_STATEMENT);
+      assertEquals(ResultType.RESULT_SET, statementResult.getResultType());
+      try (ResultSet result = statementResult.getResultSet()) {
+        assertTrue(result.next());
+        assertEquals(1L, result.getLong(0));
+        assertFalse(result.next());
+      }
+
+      statementResult = connection.execute(INSERT_STATEMENT);
+      assertEquals(ResultType.UPDATE_COUNT, statementResult.getResultType());
+      assertEquals(1L, statementResult.getUpdateCount().longValue());
+
+      statementResult = connection.execute(INSERT_RETURNING_STATEMENT);
+      assertEquals(ResultType.RESULT_SET, statementResult.getResultType());
+      try (ResultSet result = statementResult.getResultSet()) {
+        assertTrue(result.next());
+        assertEquals(1L, result.getLong(0));
+        assertFalse(result.next());
+      }
+      connection.commit();
+    }
+  }
 }