fix: correctly retry if a failure happens in the middle

Author: olavloite

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

@@ -249,15 +249,18 @@ public ApiFuture<long[]> runBatchAsync(CallType callType) {
                   new AtomicReference<>();
               try {
                 ddlClient.runWithRetryForMissingDefaultSequenceKind(
-                    () -> {
+                    restartIndex -> {
                       OperationFuture<Void, UpdateDatabaseDdlMetadata> operation =
-                          ddlClient.executeDdl(statements, protoDescriptors);
+                          ddlClient.executeDdl(
+                              statements.subList(restartIndex, statements.size()),
+                              protoDescriptors);
                       operationReference.set(operation);
                       // Wait until the operation has finished.
                       getWithStatementTimeout(operation, RUN_BATCH_STATEMENT);
                     },
                     connectionState.getValue(DEFAULT_SEQUENCE_KIND).getValue(),
-                    dbClient.getDialect());
+                    dbClient.getDialect(),
+                    operationReference);
                 long[] updateCounts = new long[statements.size()];
                 Arrays.fill(updateCounts, 1L);
                 state = UnitOfWorkState.RAN;

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

@@ -32,6 +32,8 @@
 import java.util.Collections;
 import java.util.List;
 import java.util.concurrent.ExecutionException;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.Consumer;
 
 /**
  * Convenience class for executing Data Definition Language statements on transactions that support
@@ -136,15 +138,26 @@ static boolean isCreateDatabaseStatement(String statement) {
   }
 
   void runWithRetryForMissingDefaultSequenceKind(
-      Runnable runnable, String defaultSequenceKind, Dialect dialect) {
+      Consumer<Integer> runnable,
+      String defaultSequenceKind,
+      Dialect dialect,
+      AtomicReference<OperationFuture<Void, UpdateDatabaseDdlMetadata>> operationReference) {
     try {
-      runnable.run();
+      runnable.accept(0);
     } catch (Throwable t) {
       SpannerException spannerException = SpannerExceptionFactory.asSpannerException(t);
       if (!Strings.isNullOrEmpty(defaultSequenceKind)
           && spannerException instanceof MissingDefaultSequenceKindException) {
         setDefaultSequenceKind(defaultSequenceKind, dialect);
-        runnable.run();
+        int restartIndex = 0;
+        if (operationReference.get() != null) {
+          try {
+            UpdateDatabaseDdlMetadata metadata = operationReference.get().getMetadata().get();
+            restartIndex = metadata.getCommitTimestampsCount();
+          } catch (Throwable ignore) {
+          }
+        }
+        runnable.accept(restartIndex);
         return;
       }
       throw t;

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

@@ -64,6 +64,7 @@
 import java.util.Arrays;
 import java.util.UUID;
 import java.util.concurrent.Callable;
+import java.util.concurrent.atomic.AtomicReference;
 import javax.annotation.Nonnull;
 
 /**
@@ -387,13 +388,14 @@ public ApiFuture<Void> executeDdlAsync(CallType callType, final ParsedStatement
                 executeCreateDatabase(ddl);
               } else {
                 ddlClient.runWithRetryForMissingDefaultSequenceKind(
-                    () -> {
+                    restartIndex -> {
                       OperationFuture<?, ?> operation =
                           ddlClient.executeDdl(ddl.getSqlWithoutComments(), protoDescriptors);
                       getWithStatementTimeout(operation, ddl);
                     },
                     connectionState.getValue(DEFAULT_SEQUENCE_KIND).getValue(),
-                    dbClient.getDialect());
+                    dbClient.getDialect(),
+                    new AtomicReference<>());
               }
               state = UnitOfWorkState.COMMITTED;
               return null;

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

@@ -119,7 +119,7 @@ private DdlClient createDefaultMockDdlClient(
       when(ddlClient.executeDdl(anyList(), any())).thenReturn(operation);
       doCallRealMethod()
           .when(ddlClient)
-          .runWithRetryForMissingDefaultSequenceKind(any(), any(), any());
+          .runWithRetryForMissingDefaultSequenceKind(any(), any(), any(), any());
       return ddlClient;
     } catch (Exception e) {
       throw new RuntimeException(e);
@@ -265,7 +265,9 @@ public void testGetStateAndIsActive() {
     SpannerException exception =
         SpannerExceptionFactory.newSpannerException(ErrorCode.FAILED_PRECONDITION, "test");
     doThrow(exception).when(client).executeDdl(anyList(), isNull());
-    doCallRealMethod().when(client).runWithRetryForMissingDefaultSequenceKind(any(), any(), any());
+    doCallRealMethod()
+        .when(client)
+        .runWithRetryForMissingDefaultSequenceKind(any(), any(), any(), any());
     batch = createSubject(client);
     assertThat(batch.getState(), is(UnitOfWorkState.STARTED));
     assertThat(batch.isActive(), is(true));
@@ -479,7 +481,9 @@ public void testFailedUpdateCount() throws InterruptedException, ExecutionExcept
             new ExecutionException(
                 "ddl statement failed", Status.INVALID_ARGUMENT.asRuntimeException()));
     when(operationFuture.getMetadata()).thenReturn(metadataFuture);
-    doCallRealMethod().when(client).runWithRetryForMissingDefaultSequenceKind(any(), any(), any());
+    doCallRealMethod()
+        .when(client)
+        .runWithRetryForMissingDefaultSequenceKind(any(), any(), any(), any());
     when(client.executeDdl(argThat(isListOfStringsWithSize(2)), isNull()))
         .thenReturn(operationFuture);
     DdlBatch batch =
@@ -511,7 +515,9 @@ public void testFailedUpdateCount() throws InterruptedException, ExecutionExcept
   @Test
   public void testFailedAfterFirstStatement() throws InterruptedException, ExecutionException {
     DdlClient client = mock(DdlClient.class);
-    doCallRealMethod().when(client).runWithRetryForMissingDefaultSequenceKind(any(), any(), any());
+    doCallRealMethod()
+        .when(client)
+        .runWithRetryForMissingDefaultSequenceKind(any(), any(), any(), any());
     UpdateDatabaseDdlMetadata metadata =
         UpdateDatabaseDdlMetadata.newBuilder()
             .addCommitTimestamps(

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

@@ -314,7 +314,7 @@ private DdlClient createDefaultMockDdlClient() {
       when(ddlClient.executeDdl(anyList(), any())).thenReturn(operation);
       doCallRealMethod()
           .when(ddlClient)
-          .runWithRetryForMissingDefaultSequenceKind(any(), any(), any());
+          .runWithRetryForMissingDefaultSequenceKind(any(), any(), any(), any());
       return ddlClient;
     } catch (Exception e) {
       throw new RuntimeException(e);

google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/it/ITDdlTest.java

@@ -36,6 +36,7 @@
 import com.google.cloud.spanner.connection.ConnectionOptions;
 import com.google.cloud.spanner.connection.ITAbstractSpannerTest;
 import com.google.cloud.spanner.connection.SqlScriptVerifier;
+import com.google.cloud.spanner.testing.EmulatorSpannerHelper;
 import java.util.Arrays;
 import java.util.Collections;
 import org.junit.Before;
@@ -172,9 +173,35 @@ public void testDefaultSequenceKindInBatch() {
       connection.execute(statement2);
       SpannerBatchUpdateException exception =
           assertThrows(SpannerBatchUpdateException.class, connection::runBatch);
-      assertEquals(0, Arrays.stream(exception.getUpdateCounts()).sum());
+      long updateCount = Arrays.stream(exception.getUpdateCounts()).sum();
+      // The emulator refuses the entire batch. Spanner executes the first statement and fails on
+      // the second statement.
+      if (EmulatorSpannerHelper.isUsingEmulator()) {
+        assertEquals(0, updateCount);
+      } else {
+        assertEquals(1, updateCount);
+      }
 
       // Setting a default sequence kind on the connection should make the statement succeed.
+      connection.setDefaultSequenceKind("bit_reversed_positive");
+      connection.startBatchDdl();
+      if (updateCount == 0) {
+        connection.execute(statement1);
+      }
+      connection.execute(statement2);
+      connection.runBatch();
+    }
+  }
+
+  @Test
+  public void testDefaultSequenceKindRetriesBatchCorrectly() {
+    try (Connection connection = createConnection()) {
+      Statement statement1 =
+          Statement.of("create table testseq1 (id1 int64 primary key, value string(max))");
+      Statement statement2 =
+          Statement.of(
+              "create table testseq2 (id2 int64 auto_increment primary key, value string(max))");
+
       connection.setDefaultSequenceKind("bit_reversed_positive");
       connection.startBatchDdl();
       connection.execute(statement1);