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 index 3cd9917e1a8..b11e4f613ce 100644 --- 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 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 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 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 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 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 callable = transaction -> transaction.batchUpdate(stmts); + TransactionRunner runner = client.readWriteTransaction(); + long[] actualRowCounts = runner.run(callable); + + assertThat(actualRowCounts.length).isEqualTo(80); + assertThat(expectedRowCounts).isEqualTo(actualRowCounts); + } }