diff --git a/clirr-ignored-differences.xml b/clirr-ignored-differences.xml index 431f54069..e82c99291 100644 --- a/clirr-ignored-differences.xml +++ b/clirr-ignored-differences.xml @@ -76,4 +76,36 @@ 8001 com/google/cloud/spanner/connection/ConnectionHelper + + + + 7012 + com/google/cloud/spanner/jdbc/CloudSpannerJdbcConnection + boolean isAutoBatchDml() + + + 7012 + com/google/cloud/spanner/jdbc/CloudSpannerJdbcConnection + void setAutoBatchDml(boolean) + + + 7012 + com/google/cloud/spanner/jdbc/CloudSpannerJdbcConnection + long getAutoBatchDmlUpdateCount() + + + 7012 + com/google/cloud/spanner/jdbc/CloudSpannerJdbcConnection + void setAutoBatchDmlUpdateCount(long) + + + 7012 + com/google/cloud/spanner/jdbc/CloudSpannerJdbcConnection + boolean isAutoBatchDmlUpdateCountVerification() + + + 7012 + com/google/cloud/spanner/jdbc/CloudSpannerJdbcConnection + void setAutoBatchDmlUpdateCountVerification(boolean) + diff --git a/src/main/java/com/google/cloud/spanner/jdbc/CloudSpannerJdbcConnection.java b/src/main/java/com/google/cloud/spanner/jdbc/CloudSpannerJdbcConnection.java index 6795054d9..b8491d12c 100644 --- a/src/main/java/com/google/cloud/spanner/jdbc/CloudSpannerJdbcConnection.java +++ b/src/main/java/com/google/cloud/spanner/jdbc/CloudSpannerJdbcConnection.java @@ -406,6 +406,57 @@ default int getMaxPartitionedParallelism() throws SQLException { throw new UnsupportedOperationException(); } + /** + * Enables or disables automatic batching of DML statements. When enabled, DML statements that are + * executed on this connection will be buffered in memory instead of actually being executed. The + * buffered DML statements are flushed to Spanner when a statement that cannot be part of a DML + * batch is executed on the connection. This can be a query, a DDL statement with a THEN RETURN + * clause, or a Commit call. The update count that is returned for DML statements that are + * buffered is determined by the value that has been set with {@link + * #setAutoBatchDmlUpdateCount(long)}. The default is 1. The connection verifies that the update + * counts that were returned while buffering DML statements match the actual update counts that + * are returned by Spanner when the batch is executed. This verification can be disabled by + * calling {@link #setAutoBatchDmlUpdateCountVerification(boolean)}. + */ + default void setAutoBatchDml(boolean autoBatchDml) throws SQLException { + throw new UnsupportedOperationException(); + } + + /** Returns whether automatic DML batching is enabled on this connection. */ + default boolean isAutoBatchDml() throws SQLException { + throw new UnsupportedOperationException(); + } + + /** + * Sets the update count that is returned for DML statements that are buffered during an automatic + * DML batch. This value is only used if {@link #isAutoBatchDml()} is enabled. + */ + default void setAutoBatchDmlUpdateCount(long updateCount) throws SQLException { + throw new UnsupportedOperationException(); + } + + /** + * Returns the update count that is returned for DML statements that are buffered during an + * automatic DML batch. + */ + default long getAutoBatchDmlUpdateCount() throws SQLException { + throw new UnsupportedOperationException(); + } + + /** + * Sets whether the update count that is returned by Spanner after executing an automatic DML + * batch should be verified against the update counts that were returned during the buffering of + * those statements. + */ + default void setAutoBatchDmlUpdateCountVerification(boolean verification) throws SQLException { + throw new UnsupportedOperationException(); + } + + /** Indicates whether the update counts of automatic DML batches should be verified. */ + default boolean isAutoBatchDmlUpdateCountVerification() throws SQLException { + throw new UnsupportedOperationException(); + } + /** * @see * com.google.cloud.spanner.connection.Connection#addTransactionRetryListener(com.google.cloud.spanner.connection.TransactionRetryListener) diff --git a/src/main/java/com/google/cloud/spanner/jdbc/JdbcConnection.java b/src/main/java/com/google/cloud/spanner/jdbc/JdbcConnection.java index 26b381dd2..e603c29ac 100644 --- a/src/main/java/com/google/cloud/spanner/jdbc/JdbcConnection.java +++ b/src/main/java/com/google/cloud/spanner/jdbc/JdbcConnection.java @@ -770,6 +770,36 @@ public int getMaxPartitionedParallelism() throws SQLException { return get(Connection::getMaxPartitionedParallelism); } + @Override + public void setAutoBatchDml(boolean autoBatchDml) throws SQLException { + set(Connection::setAutoBatchDml, autoBatchDml); + } + + @Override + public boolean isAutoBatchDml() throws SQLException { + return get(Connection::isAutoBatchDml); + } + + @Override + public void setAutoBatchDmlUpdateCount(long updateCount) throws SQLException { + set(Connection::setAutoBatchDmlUpdateCount, updateCount); + } + + @Override + public long getAutoBatchDmlUpdateCount() throws SQLException { + return get(Connection::getAutoBatchDmlUpdateCount); + } + + @Override + public void setAutoBatchDmlUpdateCountVerification(boolean verification) throws SQLException { + set(Connection::setAutoBatchDmlUpdateCountVerification, verification); + } + + @Override + public boolean isAutoBatchDmlUpdateCountVerification() throws SQLException { + return get(Connection::isAutoBatchDmlUpdateCountVerification); + } + @SuppressWarnings("deprecation") private static final class JdbcToSpannerTransactionRetryListener implements com.google.cloud.spanner.connection.TransactionRetryListener { diff --git a/src/main/java/com/google/cloud/spanner/jdbc/JdbcStatement.java b/src/main/java/com/google/cloud/spanner/jdbc/JdbcStatement.java index 19b325654..5d8dfd9ee 100644 --- a/src/main/java/com/google/cloud/spanner/jdbc/JdbcStatement.java +++ b/src/main/java/com/google/cloud/spanner/jdbc/JdbcStatement.java @@ -391,7 +391,6 @@ private BatchType determineStatementBatchType(String sql) throws SQLException { * client side statement) or if the connection of this statement has an active batch. */ void checkAndSetBatchType(String sql) throws SQLException { - checkConnectionHasNoActiveBatch(); BatchType type = determineStatementBatchType(sql); if (this.currentBatchType == BatchType.NONE) { this.currentBatchType = type; @@ -401,15 +400,6 @@ void checkAndSetBatchType(String sql) throws SQLException { } } - private void checkConnectionHasNoActiveBatch() throws SQLException { - if (getConnection().getSpannerConnection().isDdlBatchActive() - || getConnection().getSpannerConnection().isDmlBatchActive()) { - throw JdbcSqlExceptionFactory.of( - "Calling addBatch() is not allowed when a DML or DDL batch has been started on the connection.", - Code.FAILED_PRECONDITION); - } - } - @Override public void addBatch(String sql) throws SQLException { checkClosed(); @@ -420,7 +410,6 @@ public void addBatch(String sql) throws SQLException { @Override public void clearBatch() throws SQLException { checkClosed(); - checkConnectionHasNoActiveBatch(); batchedStatements.clear(); this.currentBatchType = BatchType.NONE; } @@ -436,7 +425,6 @@ public long[] executeLargeBatch() throws SQLException { private long[] executeBatch(boolean large) throws SQLException { checkClosed(); - checkConnectionHasNoActiveBatch(); StatementTimeout originalTimeout = setTemporaryStatementTimeout(); try { switch (this.currentBatchType) { diff --git a/src/test/java/com/google/cloud/spanner/jdbc/AutoBatchDmlMockServerTest.java b/src/test/java/com/google/cloud/spanner/jdbc/AutoBatchDmlMockServerTest.java new file mode 100644 index 000000000..8ba1a3c0b --- /dev/null +++ b/src/test/java/com/google/cloud/spanner/jdbc/AutoBatchDmlMockServerTest.java @@ -0,0 +1,332 @@ +/* + * Copyright 2024 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.cloud.spanner.jdbc; + +import static org.junit.Assert.assertArrayEquals; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; + +import com.google.cloud.spanner.MockSpannerServiceImpl.StatementResult; +import com.google.cloud.spanner.connection.AbstractMockServerTest; +import com.google.spanner.v1.ExecuteBatchDmlRequest; +import com.google.spanner.v1.ExecuteSqlRequest; +import java.sql.Connection; +import java.sql.PreparedStatement; +import java.sql.SQLException; +import java.sql.Statement; +import java.util.concurrent.ThreadLocalRandom; +import org.junit.After; +import org.junit.BeforeClass; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +public class AutoBatchDmlMockServerTest extends AbstractMockServerTest { + private static final String NON_PARAMETERIZED_INSERT = + "insert into foo (id, value) values (1, 'One')"; + private static final String NON_PARAMETERIZED_UPDATE = "update foo set value='Zero' where id=0"; + private static final String PARAMETERIZED_INSERT = + "insert into foo (id, value) values (@p1, @p2)"; + private static final String PARAMETERIZED_UPDATE = "update foo set value=@p1 where id=@p2"; + + @BeforeClass + public static void setup() { + mockSpanner.putStatementResult( + StatementResult.update( + com.google.cloud.spanner.Statement.of(NON_PARAMETERIZED_INSERT), 1L)); + mockSpanner.putStatementResult( + StatementResult.update( + com.google.cloud.spanner.Statement.of(NON_PARAMETERIZED_UPDATE), 1L)); + + mockSpanner.putStatementResult( + StatementResult.update( + com.google.cloud.spanner.Statement.newBuilder(PARAMETERIZED_INSERT) + .bind("p1") + .to(1L) + .bind("p2") + .to("One") + .build(), + 1L)); + mockSpanner.putStatementResult( + StatementResult.update( + com.google.cloud.spanner.Statement.newBuilder(PARAMETERIZED_INSERT) + .bind("p1") + .to(2L) + .bind("p2") + .to("Two") + .build(), + 1L)); + mockSpanner.putStatementResult( + StatementResult.update( + com.google.cloud.spanner.Statement.newBuilder(PARAMETERIZED_UPDATE) + .bind("p2") + .to(1L) + .bind("p1") + .to("One") + .build(), + 1L)); + mockSpanner.putStatementResult( + StatementResult.update( + com.google.cloud.spanner.Statement.newBuilder(PARAMETERIZED_UPDATE) + .bind("p2") + .to(2L) + .bind("p1") + .to("Two") + .build(), + 1L)); + } + + @After + public void clearRequests() { + mockSpanner.clearRequests(); + } + + @Test + public void testStatementExecute() throws SQLException { + try (Connection connection = createJdbcConnection()) { + connection.setAutoCommit(false); + connection.unwrap(CloudSpannerJdbcConnection.class).setAutoBatchDml(true); + + try (Statement statement = connection.createStatement()) { + assertFalse(statement.execute(NON_PARAMETERIZED_INSERT)); + assertFalse(statement.execute(NON_PARAMETERIZED_UPDATE)); + } + connection.commit(); + } + assertEquals(1, mockSpanner.countRequestsOfType(ExecuteBatchDmlRequest.class)); + assertEquals(0, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class)); + } + + @Test + public void testStatementExecuteUpdate() throws SQLException { + try (Connection connection = createJdbcConnection()) { + connection.setAutoCommit(false); + connection.unwrap(CloudSpannerJdbcConnection.class).setAutoBatchDml(true); + + try (Statement statement = connection.createStatement()) { + assertEquals(1, statement.executeUpdate(NON_PARAMETERIZED_INSERT)); + assertEquals(1, statement.executeUpdate(NON_PARAMETERIZED_UPDATE)); + } + connection.commit(); + } + assertEquals(1, mockSpanner.countRequestsOfType(ExecuteBatchDmlRequest.class)); + assertEquals(0, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class)); + } + + @Test + public void testStatementBatch() throws SQLException { + try (Connection connection = createJdbcConnection()) { + connection.setAutoCommit(false); + connection.unwrap(CloudSpannerJdbcConnection.class).setAutoBatchDml(true); + + try (Statement statement = connection.createStatement()) { + repeat( + () -> { + statement.addBatch(NON_PARAMETERIZED_INSERT); + statement.addBatch(NON_PARAMETERIZED_UPDATE); + assertArrayEquals(new int[] {1, 1}, statement.executeBatch()); + }, + 2); + } + connection.commit(); + } + assertEquals(1, mockSpanner.countRequestsOfType(ExecuteBatchDmlRequest.class)); + assertEquals(0, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class)); + } + + @Test + public void testStatementCombination() throws SQLException { + try (Connection connection = createJdbcConnection()) { + connection.setAutoCommit(false); + connection.unwrap(CloudSpannerJdbcConnection.class).setAutoBatchDml(true); + + try (Statement statement = connection.createStatement()) { + statement.executeUpdate(NON_PARAMETERIZED_UPDATE); + repeat( + () -> { + statement.addBatch(NON_PARAMETERIZED_INSERT); + statement.addBatch(NON_PARAMETERIZED_UPDATE); + assertArrayEquals(new int[] {1, 1}, statement.executeBatch()); + }, + ThreadLocalRandom.current().nextInt(1, 5)); + repeat( + () -> { + statement.execute(NON_PARAMETERIZED_INSERT); + statement.executeUpdate(NON_PARAMETERIZED_UPDATE); + }, + ThreadLocalRandom.current().nextInt(1, 5)); + } + connection.commit(); + } + assertEquals(1, mockSpanner.countRequestsOfType(ExecuteBatchDmlRequest.class)); + assertEquals(0, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class)); + } + + @Test + public void testPreparedStatementExecute() throws SQLException { + try (Connection connection = createJdbcConnection()) { + connection.setAutoCommit(false); + connection.unwrap(CloudSpannerJdbcConnection.class).setAutoBatchDml(true); + + try (PreparedStatement statement = connection.prepareStatement(PARAMETERIZED_INSERT)) { + statement.setLong(1, 1L); + statement.setString(2, "One"); + assertFalse(statement.execute()); + statement.setLong(1, 2L); + statement.setString(2, "Two"); + assertFalse(statement.execute()); + } + try (PreparedStatement statement = connection.prepareStatement(PARAMETERIZED_UPDATE)) { + statement.setLong(2, 1L); + statement.setString(1, "One"); + assertFalse(statement.execute()); + statement.setLong(2, 2L); + statement.setString(1, "Two"); + assertFalse(statement.execute()); + } + connection.commit(); + } + assertEquals(1, mockSpanner.countRequestsOfType(ExecuteBatchDmlRequest.class)); + assertEquals(0, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class)); + } + + @Test + public void testPreparedStatementExecuteUpdate() throws SQLException { + try (Connection connection = createJdbcConnection()) { + connection.setAutoCommit(false); + connection.unwrap(CloudSpannerJdbcConnection.class).setAutoBatchDml(true); + + try (PreparedStatement statement = connection.prepareStatement(PARAMETERIZED_INSERT)) { + statement.setLong(1, 1L); + statement.setString(2, "One"); + assertEquals(1, statement.executeUpdate()); + statement.setLong(1, 2L); + statement.setString(2, "Two"); + assertEquals(1, statement.executeUpdate()); + } + try (PreparedStatement statement = connection.prepareStatement(PARAMETERIZED_UPDATE)) { + statement.setLong(2, 1L); + statement.setString(1, "One"); + assertEquals(1, statement.executeUpdate()); + statement.setLong(2, 2L); + statement.setString(1, "Two"); + assertEquals(1, statement.executeUpdate()); + } + connection.commit(); + } + assertEquals(1, mockSpanner.countRequestsOfType(ExecuteBatchDmlRequest.class)); + assertEquals(0, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class)); + } + + @Test + public void testPreparedStatementBatch() throws SQLException { + try (Connection connection = createJdbcConnection()) { + connection.setAutoCommit(false); + connection.unwrap(CloudSpannerJdbcConnection.class).setAutoBatchDml(true); + + repeat( + () -> { + try (PreparedStatement statement = connection.prepareStatement(PARAMETERIZED_INSERT)) { + statement.setLong(1, 1L); + statement.setString(2, "One"); + statement.addBatch(); + statement.setLong(1, 2L); + statement.setString(2, "Two"); + statement.addBatch(); + statement.executeBatch(); + } + try (PreparedStatement statement = connection.prepareStatement(PARAMETERIZED_UPDATE)) { + statement.setLong(2, 1L); + statement.setString(1, "One"); + statement.addBatch(); + statement.setLong(2, 2L); + statement.setString(1, "Two"); + statement.addBatch(); + statement.executeBatch(); + } + }, + 2); + connection.commit(); + } + assertEquals(1, mockSpanner.countRequestsOfType(ExecuteBatchDmlRequest.class)); + assertEquals(0, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class)); + } + + @Test + public void testPreparedStatementCombination() throws SQLException { + try (Connection connection = createJdbcConnection()) { + connection.setAutoCommit(false); + connection.unwrap(CloudSpannerJdbcConnection.class).setAutoBatchDml(true); + + try (PreparedStatement statement = connection.prepareStatement(PARAMETERIZED_INSERT)) { + statement.setLong(1, 1L); + statement.setString(2, "One"); + assertFalse(statement.execute()); + } + repeat( + () -> { + try (PreparedStatement statement = connection.prepareStatement(PARAMETERIZED_INSERT)) { + statement.setLong(1, 1L); + statement.setString(2, "One"); + statement.addBatch(); + statement.setLong(1, 2L); + statement.setString(2, "Two"); + statement.addBatch(); + assertArrayEquals(new int[] {1, 1}, statement.executeBatch()); + } + try (PreparedStatement statement = connection.prepareStatement(PARAMETERIZED_UPDATE)) { + statement.setLong(2, 1L); + statement.setString(1, "One"); + statement.addBatch(); + statement.setLong(2, 2L); + statement.setString(1, "Two"); + statement.addBatch(); + assertArrayEquals(new int[] {1, 1}, statement.executeBatch()); + } + }, + ThreadLocalRandom.current().nextInt(1, 5)); + repeat( + () -> { + try (PreparedStatement statement = connection.prepareStatement(PARAMETERIZED_INSERT)) { + statement.setLong(1, 1L); + statement.setString(2, "One"); + assertEquals(1, statement.executeUpdate()); + } + try (PreparedStatement statement = connection.prepareStatement(PARAMETERIZED_UPDATE)) { + statement.setLong(2, 2L); + statement.setString(1, "Two"); + assertFalse(statement.execute()); + } + }, + ThreadLocalRandom.current().nextInt(1, 5)); + connection.commit(); + } + assertEquals(1, mockSpanner.countRequestsOfType(ExecuteBatchDmlRequest.class)); + assertEquals(0, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class)); + } + + interface SQLRunnable { + void run() throws SQLException; + } + + static void repeat(SQLRunnable runnable, int count) throws SQLException { + for (int i = 0; i < count; i++) { + runnable.run(); + } + } +} diff --git a/src/test/java/com/google/cloud/spanner/jdbc/it/ITJdbcSimpleStatementsTest.java b/src/test/java/com/google/cloud/spanner/jdbc/it/ITJdbcSimpleStatementsTest.java index f6948a9c2..7dd7153e4 100644 --- a/src/test/java/com/google/cloud/spanner/jdbc/it/ITJdbcSimpleStatementsTest.java +++ b/src/test/java/com/google/cloud/spanner/jdbc/it/ITJdbcSimpleStatementsTest.java @@ -17,11 +17,11 @@ package com.google.cloud.spanner.jdbc.it; import static com.google.common.truth.Truth.assertThat; +import static org.junit.Assert.assertArrayEquals; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertFalse; import static org.junit.Assert.assertThrows; import static org.junit.Assert.assertTrue; -import static org.junit.Assert.fail; import static org.junit.Assume.assumeFalse; import com.google.cloud.spanner.Database; @@ -102,12 +102,6 @@ public void testSelect1PreparedStatement() throws SQLException { @Test public void testPreparedStatement() throws SQLException { - // skipping the test when dialect is POSTGRESQL because of exception below - // INVALID_ARGUMENT: io.grpc.StatusRuntimeException: INVALID_ARGUMENT: Statements with set - // operations in subqueries are not supported - assumeFalse( - "select array of structs is not supported on POSTGRESQL", - dialect.dialect == Dialect.POSTGRESQL); String sql = "select * from (select 1 as number union all select 2 union all select 3) numbers where number=?"; try (Connection connection = createConnection(env, database)) { @@ -188,15 +182,16 @@ public void testBatchedDdlStatements() throws SQLException { } @Test - public void testAddBatchWhenAlreadyInBatch() { + public void testAddBatchWhenAlreadyInBatch() throws SQLException { try (Connection connection = createConnection(env, database)) { - connection.createStatement().execute("START BATCH DML"); - connection.createStatement().addBatch("INSERT INTO Singers (SingerId) VALUES (-1)"); - fail("missing expected exception"); - } catch (SQLException e) { - assertThat(e.getMessage()) - .contains( - "Calling addBatch() is not allowed when a DML or DDL batch has been started on the connection."); + try (Statement statement = connection.createStatement()) { + statement.execute("START BATCH DML"); + statement.addBatch("INSERT INTO Singers (SingerId) VALUES (-1)"); + statement.addBatch("INSERT INTO Singers (SingerId) VALUES (-2)"); + // The returned update count for DML statements in a batch is -1. + assertArrayEquals(new int[] {-1, -1}, statement.executeBatch()); + // Note: The 'Singers' table does not actually exist, so we're not executing the batch. + } } }