fix: allow DML THEN RETURN with retryAbortsInternally=false

Author: rahul2393

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/ReadWriteTransactionTest.java

@@ -287,6 +287,22 @@ 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);