fix: retries of updates in the Connection API ignored analyze mode (#2010)

* fix: retries of updates in the Connection API ignored analyze mode

Retries of read/write transactions in the Connection API did not respect the analyze mode
of update statements. This would cause updates to be executed using AnalyzeMode.NORMAL
during retries, regardless of what was used during the initial attempt.

Fixes #2009

* fix: remove null check and force not-null

Author: olavloite

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

@@ -444,7 +444,7 @@ private ApiFuture<ResultSetStats> internalExecuteUpdateAsync(
                                       options);
                         }
                         createAndAddRetriableUpdate(
-                            update, updateCount.getRowCountExact(), options);
+                            update, analyzeMode, updateCount.getRowCountExact(), options);
                         return updateCount;
                       } catch (AbortedException e) {
                         throw e;
@@ -712,9 +712,9 @@ private void createAndAddFailedQuery(
   }
 
   private void createAndAddRetriableUpdate(
-      ParsedStatement update, long updateCount, UpdateOption... options) {
+      ParsedStatement update, AnalyzeMode analyzeMode, long updateCount, UpdateOption... options) {
     if (retryAbortsInternally) {
-      addRetryStatement(new RetriableUpdate(this, update, updateCount, options));
+      addRetryStatement(new RetriableUpdate(this, update, analyzeMode, updateCount, options));
     }
   }
 

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

@@ -31,18 +31,19 @@
 final class RetriableUpdate implements RetriableStatement {
   private final ReadWriteTransaction transaction;
   private final ParsedStatement statement;
+  private final AnalyzeMode analyzeMode;
   private final long updateCount;
   private final UpdateOption[] options;
 
   RetriableUpdate(
       ReadWriteTransaction transaction,
       ParsedStatement statement,
+      AnalyzeMode analyzeMode,
       long updateCount,
       UpdateOption... options) {
-    Preconditions.checkNotNull(transaction);
-    Preconditions.checkNotNull(statement);
-    this.transaction = transaction;
-    this.statement = statement;
+    this.transaction = Preconditions.checkNotNull(transaction);
+    this.statement = Preconditions.checkNotNull(statement);
+    this.analyzeMode = Preconditions.checkNotNull(analyzeMode);
     this.updateCount = updateCount;
     this.options = options;
   }
@@ -54,7 +55,15 @@ public void retry(AbortedException aborted) throws AbortedException {
       transaction
           .getStatementExecutor()
           .invokeInterceptors(statement, StatementExecutionStep.RETRY_STATEMENT, transaction);
-      newCount = transaction.getReadContext().executeUpdate(statement.getStatement(), options);
+      if (analyzeMode == AnalyzeMode.NONE) {
+        newCount = transaction.getReadContext().executeUpdate(statement.getStatement(), options);
+      } else {
+        newCount =
+            transaction
+                .getReadContext()
+                .analyzeUpdate(statement.getStatement(), analyzeMode.getQueryAnalyzeMode())
+                .getRowCountExact();
+      }
     } catch (AbortedException e) {
       // Just re-throw the AbortedException and let the retry logic determine whether another try
       // should be executed or not.

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

@@ -26,6 +26,7 @@
 import com.google.cloud.Timestamp;
 import com.google.cloud.spanner.ErrorCode;
 import com.google.cloud.spanner.MockSpannerServiceImpl.StatementResult;
+import com.google.cloud.spanner.ReadContext.QueryAnalyzeMode;
 import com.google.cloud.spanner.ResultSet;
 import com.google.cloud.spanner.SpannerException;
 import com.google.cloud.spanner.Statement;
@@ -37,9 +38,11 @@
 import com.google.spanner.v1.CommitRequest;
 import com.google.spanner.v1.ExecuteBatchDmlRequest;
 import com.google.spanner.v1.ExecuteSqlRequest;
+import com.google.spanner.v1.ExecuteSqlRequest.QueryMode;
 import io.grpc.Status;
 import io.grpc.StatusRuntimeException;
 import java.util.Collections;
+import java.util.List;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
@@ -299,6 +302,41 @@ public void testRetryUsesTags() {
     assertEquals(2L, commitRequestCount);
   }
 
+  @Test
+  public void testRetryUsesAnalyzeModeForUpdate() {
+    mockSpanner.putStatementResult(
+        StatementResult.query(SELECT_COUNT_STATEMENT, SELECT_COUNT_RESULTSET_BEFORE_INSERT));
+    mockSpanner.putStatementResult(StatementResult.update(INSERT_STATEMENT, 0));
+    try (ITConnection connection = createConnection()) {
+      assertEquals(
+          0L, connection.analyzeUpdate(INSERT_STATEMENT, QueryAnalyzeMode.PLAN).getRowCountExact());
+
+      mockSpanner.abortNextStatement();
+      connection.executeQuery(SELECT_COUNT_STATEMENT);
+
+      mockSpanner.putStatementResult(StatementResult.update(INSERT_STATEMENT, 1));
+      assertEquals(1L, connection.executeUpdate(INSERT_STATEMENT));
+
+      connection.commit();
+    }
+    // 5 requests because:
+    // 1. Analyze INSERT
+    // 2. Execute SELECT COUNT(*) (Aborted)
+    // 3. Analyze INSERT (retry)
+    // 4. Execute SELECT COUNT(*) (retry)
+    // 5. Execute INSERT
+    assertEquals(5, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class));
+    List<ExecuteSqlRequest> requests = mockSpanner.getRequestsOfType(ExecuteSqlRequest.class);
+    assertEquals(QueryMode.PLAN, requests.get(0).getQueryMode());
+    assertEquals(QueryMode.NORMAL, requests.get(1).getQueryMode());
+
+    // This used NORMAL because of https://github.com/googleapis/java-spanner/issues/2009.
+    assertEquals(QueryMode.PLAN, requests.get(2).getQueryMode());
+
+    assertEquals(QueryMode.NORMAL, requests.get(3).getQueryMode());
+    assertEquals(QueryMode.NORMAL, requests.get(4).getQueryMode());
+  }
+
   ITConnection createConnection(TransactionRetryListener listener) {
     ITConnection connection =
         super.createConnection(ImmutableList.of(), ImmutableList.of(listener));