test: add integration tests for large batch dml updates. (#2618)

* test: add integration tests for large batch dml updates. * chore: incorporate review comments. * fix: lint issues. * 🦉 Updates from OwlBot post-processor See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md --------- Co-authored-by: Owl Bot <gcf-owl-bot[bot]@users.noreply.github.com>
googleapis · Sep 19, 2023 · 84cd62f · 84cd62f
1 parent 6fe132a
commit 84cd62f
Showing 1 changed file with 63 additions and 0 deletions.
diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/it/ITBatchDmlTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/it/ITBatchDmlTest.java
@@ -31,7 +31,9 @@
 import com.google.cloud.spanner.TransactionRunner.TransactionCallable;
 import com.google.spanner.admin.database.v1.UpdateDatabaseDdlMetadata;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collections;
+import java.util.LinkedList;
 import java.util.List;
 import org.junit.After;
 import org.junit.Assert;
@@ -189,4 +191,65 @@ public void errorBatchDmlAlreadyExist() {
       }
     }
   }
+
+  @Test
+  public void largeBatchDml() {
+    List<Statement> stmts = new LinkedList<>();
+    for (int i = 0; i < 40; i++) {
+      stmts.add(Statement.of("INSERT INTO T (k, v) VALUES ('boo" + i + "', " + i + ");"));
+    }
+
+    for (int i = 0; i < 40; i++) {
+      stmts.add(Statement.of("DELETE FROM T WHERE T.K = 'boo" + i + "';"));
+    }
+    long[] expectedRowCounts = new long[stmts.size()];
+    Arrays.fill(expectedRowCounts, 1L);
+
+    final TransactionCallable<long[]> callable = transaction -> transaction.batchUpdate(stmts);
+    TransactionRunner runner = client.readWriteTransaction();
+    long[] actualRowCounts = runner.run(callable);
+    assertThat(actualRowCounts.length).isEqualTo(80);
+    assertThat(expectedRowCounts).isEqualTo(actualRowCounts);
+  }
+
+  @Test
+  public void largeBatchDml_withParameterisedStatements() {
+    List<Statement> stmts = new LinkedList<>();
+    String insertQuery = "INSERT INTO T(k, v) VALUES(@key, @val)";
+    for (int i = 0; i < 80; i++) {
+      stmts.add(
+          Statement.newBuilder(insertQuery)
+              .bind("key")
+              .to("'boo" + i + "'")
+              .bind("val")
+              .to(i)
+              .build());
+    }
+    long[] expectedRowCounts = new long[stmts.size()];
+    Arrays.fill(expectedRowCounts, 1L);
+
+    final TransactionCallable<long[]> callable = transaction -> transaction.batchUpdate(stmts);
+    TransactionRunner runner = client.readWriteTransaction();
+    long[] actualRowCounts = runner.run(callable);
+
+    assertThat(actualRowCounts.length).isEqualTo(80);
+    assertThat(expectedRowCounts).isEqualTo(actualRowCounts);
+  }
+
+  @Test
+  public void largeBatchDml_withNonParameterisedStatements() {
+    List<Statement> stmts = new LinkedList<>();
+    for (int i = 0; i < 80; i++) {
+      stmts.add(Statement.of("INSERT INTO T (k, v) VALUES ('boo" + i + "', " + i + ");"));
+    }
+    long[] expectedRowCounts = new long[stmts.size()];
+    Arrays.fill(expectedRowCounts, 1L);
+
+    final TransactionCallable<long[]> callable = transaction -> transaction.batchUpdate(stmts);
+    TransactionRunner runner = client.readWriteTransaction();
+    long[] actualRowCounts = runner.run(callable);
+
+    assertThat(actualRowCounts.length).isEqualTo(80);
+    assertThat(expectedRowCounts).isEqualTo(actualRowCounts);
+  }
 }