Skip to content

Commit 1e81863

Browse files
olavloiteolavloite
authored
feat: add fallback to PDML mode (#1841)
DML statements that are executed in auto-commit mode can use either atomic transactions, or partitioned non-atomic transactions. The former is bound by the mutation limits in Spanner. The latter may update/delete any number of rows. The transaction type that is used to execute DML statements in auto-commit mode is determined by the connection variable `autocommit_dml_mode`. This connection variable now supports a third value. The supported values are: - TRANSACTIONAL (default): Uses atomic read/write transactions. - PARTITIONED_NON_ATOMIC: Use Partitioned DML for DML statements in auto-commit mode. Use this mode to execute DML statements that exceed the transaction mutation limit in Spanner. - TRANSACTIONAL_WITH_FALLBACK_TO_PARTITIONED_NON_ATOMIC: Execute DML statements using atomic read/write transactions. If this fails because the mutation limit on Spanner has been exceeded, the DML statement is retried using a Partitioned DML transaction.
1 parent 3cd9cd6 commit 1e81863

File tree

2 files changed

+268
-0
lines changed

2 files changed

+268
-0
lines changed

‎README.md

+9
Original file line numberDiff line numberDiff line change
@@ -118,6 +118,15 @@ these can also be supplied in a Properties instance that is passed to the
118118
- maxSessions (int): Sets the maximum number of sessions in the backing session pool. Defaults to 400.
119119
- numChannels (int): Sets the number of gRPC channels to use. Defaults to 4.
120120
- retryAbortsInternally (boolean): The JDBC driver will by default automatically retry aborted transactions internally. This is done by keeping track of all statements and the results of these during a transaction, and if the transaction is aborted by Cloud Spanner, it will replay the statements on a new transaction and compare the results with the initial attempt. Disable this option if you want to handle aborted transactions in your own application.
121+
- autocommit_dml_mode (string): Determines the transaction type that is used to execute
122+
[DML statements](https://cloud.google.com/spanner/docs/dml-tasks#using-dml) when the connection is
123+
in auto-commit mode. The following values are supported:
124+
- TRANSACTIONAL (default): Uses atomic read/write transactions.
125+
- PARTITIONED_NON_ATOMIC: Use Partitioned DML for DML statements in auto-commit mode. Use this mode
126+
to execute DML statements that exceed the transaction mutation limit in Spanner.
127+
- TRANSACTIONAL_WITH_FALLBACK_TO_PARTITIONED_NON_ATOMIC: Execute DML statements using atomic read/write
128+
transactions. If this fails because the mutation limit on Spanner has been exceeded, the DML statement
129+
is retried using a Partitioned DML transaction.
121130
- auto_batch_dml (boolean): Automatically buffer DML statements and execute them as one batch,
122131
instead of executing them on Spanner directly. The buffered DML statements are executed on Spanner
123132
in one batch when a query is executed, or when the transaction is committed. This option can for

‎src/test/java/com/google/cloud/spanner/jdbc/FallbackToPartitionedDMLMockServerTest.java

+259
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,259 @@
1+
/*
2+
* Copyright 2024 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package com.google.cloud.spanner.jdbc;
18+
19+
import static org.junit.Assert.assertEquals;
20+
import static org.junit.Assert.assertFalse;
21+
import static org.junit.Assert.assertNotNull;
22+
import static org.junit.Assert.assertThrows;
23+
import static org.junit.Assert.assertTrue;
24+
25+
import com.google.cloud.spanner.Dialect;
26+
import com.google.cloud.spanner.ErrorCode;
27+
import com.google.cloud.spanner.MockSpannerServiceImpl;
28+
import com.google.cloud.spanner.MockSpannerServiceImpl.SimulatedExecutionTime;
29+
import com.google.cloud.spanner.SpannerException;
30+
import com.google.cloud.spanner.TransactionMutationLimitExceededException;
31+
import com.google.cloud.spanner.connection.AbstractMockServerTest;
32+
import com.google.cloud.spanner.connection.AutocommitDmlMode;
33+
import com.google.cloud.spanner.connection.SpannerPool;
34+
import com.google.protobuf.Any;
35+
import com.google.rpc.Help;
36+
import com.google.rpc.Help.Link;
37+
import com.google.spanner.v1.BeginTransactionRequest;
38+
import com.google.spanner.v1.CommitRequest;
39+
import com.google.spanner.v1.ExecuteSqlRequest;
40+
import io.grpc.Metadata;
41+
import io.grpc.Status;
42+
import io.grpc.StatusRuntimeException;
43+
import java.sql.Connection;
44+
import java.sql.DriverManager;
45+
import java.sql.ResultSet;
46+
import java.sql.SQLException;
47+
import java.util.Properties;
48+
import org.junit.Test;
49+
import org.junit.runner.RunWith;
50+
import org.junit.runners.JUnit4;
51+
52+
@RunWith(JUnit4.class)
53+
public class FallbackToPartitionedDMLMockServerTest extends AbstractMockServerTest {
54+
55+
static StatusRuntimeException createTransactionMutationLimitExceededException() {
56+
Metadata.Key<byte[]> key =
57+
Metadata.Key.of("grpc-status-details-bin", Metadata.BINARY_BYTE_MARSHALLER);
58+
Help help =
59+
Help.newBuilder()
60+
.addLinks(
61+
Link.newBuilder()
62+
.setDescription("Cloud Spanner limits documentation.")
63+
.setUrl("https://cloud.google.com/spanner/docs/limits")
64+
.build())
65+
.build();
66+
com.google.rpc.Status status =
67+
com.google.rpc.Status.newBuilder().addDetails(Any.pack(help)).build();
68+
69+
Metadata trailers = new Metadata();
70+
trailers.put(key, status.toByteArray());
71+
72+
return Status.INVALID_ARGUMENT
73+
.withDescription("The transaction contains too many mutations.")
74+
.asRuntimeException(trailers);
75+
}
76+
77+
@Test
78+
public void testConnectionProperty() throws SQLException {
79+
for (AutocommitDmlMode mode : AutocommitDmlMode.values()) {
80+
Properties properties = new Properties();
81+
properties.put("autocommit_dml_mode", mode.name());
82+
try (Connection connection =
83+
DriverManager.getConnection("jdbc:" + getBaseUrl(), properties)) {
84+
assertEquals(
85+
mode, connection.unwrap(CloudSpannerJdbcConnection.class).getAutocommitDmlMode());
86+
}
87+
}
88+
}
89+
90+
@Test
91+
public void testTransactionMutationLimitExceeded_isNotRetriedByDefault() throws SQLException {
92+
mockSpanner.setExecuteSqlExecutionTime(
93+
SimulatedExecutionTime.ofException(createTransactionMutationLimitExceededException()));
94+
95+
try (Connection connection = createJdbcConnection()) {
96+
connection.setAutoCommit(true);
97+
assertEquals(
98+
AutocommitDmlMode.TRANSACTIONAL,
99+
connection.unwrap(CloudSpannerJdbcConnection.class).getAutocommitDmlMode());
100+
101+
SQLException exception =
102+
assertThrows(
103+
SQLException.class,
104+
() ->
105+
connection.createStatement().executeUpdate("update test set value=1 where true"));
106+
assertNotNull(exception.getCause());
107+
assertEquals(
108+
TransactionMutationLimitExceededException.class, exception.getCause().getClass());
109+
TransactionMutationLimitExceededException transactionMutationLimitExceededException =
110+
(TransactionMutationLimitExceededException) exception.getCause();
111+
assertEquals(0, transactionMutationLimitExceededException.getSuppressed().length);
112+
}
113+
assertEquals(1, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class));
114+
assertEquals(0, mockSpanner.countRequestsOfType(CommitRequest.class));
115+
}
116+
117+
@Test
118+
public void testTransactionMutationLimitExceeded_canBeRetriedAsPDML() throws SQLException {
119+
String sql = "update test set value=1 where true";
120+
com.google.cloud.spanner.Statement statement = com.google.cloud.spanner.Statement.of(sql);
121+
mockSpanner.setExecuteSqlExecutionTime(
122+
SimulatedExecutionTime.ofException(createTransactionMutationLimitExceededException()));
123+
mockSpanner.putStatementResult(
124+
MockSpannerServiceImpl.StatementResult.update(statement, 100000L));
125+
126+
try (Connection connection = createJdbcConnection()) {
127+
connection.setAutoCommit(true);
128+
connection
129+
.unwrap(CloudSpannerJdbcConnection.class)
130+
.setAutocommitDmlMode(
131+
AutocommitDmlMode.TRANSACTIONAL_WITH_FALLBACK_TO_PARTITIONED_NON_ATOMIC);
132+
133+
long updateCount = connection.createStatement().executeUpdate(sql);
134+
assertEquals(100000L, updateCount);
135+
}
136+
// Verify that the request is retried as Partitioned DML.
137+
assertEquals(2, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class));
138+
// The transactional request uses inline-begin.
139+
ExecuteSqlRequest transactionalRequest =
140+
mockSpanner.getRequestsOfType(ExecuteSqlRequest.class).get(0);
141+
assertTrue(transactionalRequest.getTransaction().getBegin().hasReadWrite());
142+
143+
// Partitioned DML uses an explicit BeginTransaction RPC.
144+
assertEquals(1, mockSpanner.countRequestsOfType(BeginTransactionRequest.class));
145+
BeginTransactionRequest beginRequest =
146+
mockSpanner.getRequestsOfType(BeginTransactionRequest.class).get(0);
147+
assertTrue(beginRequest.getOptions().hasPartitionedDml());
148+
ExecuteSqlRequest partitionedDmlRequest =
149+
mockSpanner.getRequestsOfType(ExecuteSqlRequest.class).get(1);
150+
assertTrue(partitionedDmlRequest.getTransaction().hasId());
151+
152+
// Partitioned DML transactions are not committed.
153+
assertEquals(0, mockSpanner.countRequestsOfType(CommitRequest.class));
154+
}
155+
156+
@Test
157+
public void testTransactionMutationLimitExceeded_retryAsPDMLFails() throws SQLException {
158+
String sql = "insert into test (id, value) select -id, value from test";
159+
com.google.cloud.spanner.Statement statement = com.google.cloud.spanner.Statement.of(sql);
160+
// The transactional update statement uses ExecuteSql(..).
161+
mockSpanner.setExecuteSqlExecutionTime(
162+
SimulatedExecutionTime.ofException(createTransactionMutationLimitExceededException()));
163+
mockSpanner.putStatementResult(
164+
MockSpannerServiceImpl.StatementResult.exception(
165+
statement,
166+
Status.INVALID_ARGUMENT
167+
.withDescription("This statement is not supported with Partitioned DML")
168+
.asRuntimeException()));
169+
170+
try (Connection connection = createJdbcConnection()) {
171+
connection.setAutoCommit(true);
172+
connection
173+
.unwrap(CloudSpannerJdbcConnection.class)
174+
.setAutocommitDmlMode(
175+
AutocommitDmlMode.TRANSACTIONAL_WITH_FALLBACK_TO_PARTITIONED_NON_ATOMIC);
176+
177+
// The connection throws TransactionMutationLimitExceededException if the retry using
178+
// partitioned DML fails. The exception from the failed retry is returned as a suppressed
179+
// exception of the TransactionMutationLimitExceededException.
180+
SQLException exception =
181+
assertThrows(SQLException.class, () -> connection.createStatement().executeUpdate(sql));
182+
assertNotNull(exception.getCause());
183+
assertEquals(
184+
TransactionMutationLimitExceededException.class, exception.getCause().getClass());
185+
TransactionMutationLimitExceededException transactionMutationLimitExceededException =
186+
(TransactionMutationLimitExceededException) exception.getCause();
187+
assertEquals(1, transactionMutationLimitExceededException.getSuppressed().length);
188+
assertEquals(
189+
SpannerException.class,
190+
transactionMutationLimitExceededException.getSuppressed()[0].getClass());
191+
SpannerException spannerException =
192+
(SpannerException) transactionMutationLimitExceededException.getSuppressed()[0];
193+
assertEquals(ErrorCode.INVALID_ARGUMENT, spannerException.getErrorCode());
194+
assertTrue(
195+
spannerException.getMessage(),
196+
spannerException
197+
.getMessage()
198+
.contains("This statement is not supported with Partitioned DML"));
199+
}
200+
// Verify that the request was retried as Partitioned DML.
201+
assertEquals(2, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class));
202+
// The transactional request uses inline-begin.
203+
ExecuteSqlRequest transactionalRequest =
204+
mockSpanner.getRequestsOfType(ExecuteSqlRequest.class).get(0);
205+
assertTrue(transactionalRequest.getTransaction().getBegin().hasReadWrite());
206+
207+
// Partitioned DML uses an explicit BeginTransaction RPC.
208+
assertEquals(1, mockSpanner.countRequestsOfType(BeginTransactionRequest.class));
209+
BeginTransactionRequest beginRequest =
210+
mockSpanner.getRequestsOfType(BeginTransactionRequest.class).get(0);
211+
assertTrue(beginRequest.getOptions().hasPartitionedDml());
212+
ExecuteSqlRequest partitionedDmlRequest =
213+
mockSpanner.getRequestsOfType(ExecuteSqlRequest.class).get(1);
214+
assertTrue(partitionedDmlRequest.getTransaction().hasId());
215+
216+
// Partitioned DML transactions are not committed.
217+
assertEquals(0, mockSpanner.countRequestsOfType(CommitRequest.class));
218+
}
219+
220+
@Test
221+
public void testSqlStatements() throws SQLException {
222+
for (Dialect dialect : Dialect.values()) {
223+
SpannerPool.closeSpannerPool();
224+
mockSpanner.putStatementResult(
225+
MockSpannerServiceImpl.StatementResult.detectDialectResult(dialect));
226+
String prefix = dialect == Dialect.POSTGRESQL ? "SPANNER." : "";
227+
228+
try (Connection connection = createJdbcConnection()) {
229+
connection.setAutoCommit(true);
230+
try (ResultSet resultSet =
231+
connection
232+
.createStatement()
233+
.executeQuery(String.format("show variable %sautocommit_dml_mode", prefix))) {
234+
assertTrue(resultSet.next());
235+
assertEquals(
236+
AutocommitDmlMode.TRANSACTIONAL.name(),
237+
resultSet.getString(String.format("%sAUTOCOMMIT_DML_MODE", prefix)));
238+
assertFalse(resultSet.next());
239+
}
240+
connection
241+
.createStatement()
242+
.execute(
243+
String.format(
244+
"set %sautocommit_dml_mode = 'transactional_with_fallback_to_partitioned_non_atomic'",
245+
prefix));
246+
try (ResultSet resultSet =
247+
connection
248+
.createStatement()
249+
.executeQuery(String.format("show variable %sautocommit_dml_mode", prefix))) {
250+
assertTrue(resultSet.next());
251+
assertEquals(
252+
AutocommitDmlMode.TRANSACTIONAL_WITH_FALLBACK_TO_PARTITIONED_NON_ATOMIC.name(),
253+
resultSet.getString(String.format("%sAUTOCOMMIT_DML_MODE", prefix)));
254+
assertFalse(resultSet.next());
255+
}
256+
}
257+
}
258+
}
259+
}

0 commit comments

Comments
 (0)