Skip to content

Commit 644cbdb

Browse files
olavloiteolavloite
committed
refactor: include the option in autocommit_dml_mode
1 parent ff8e0df commit 644cbdb

File tree

15 files changed

+1003
-2445
lines changed

15 files changed

+1003
-2445
lines changed

‎google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/AutocommitDmlMode.java‎

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,8 +18,18 @@
1818

1919
/** Enum used to define the behavior of DML statements in autocommit mode */
2020
public enum AutocommitDmlMode {
21+
/** TRANSACTIONAL: DML statements use a standard atomic transaction. */
2122
TRANSACTIONAL,
22-
PARTITIONED_NON_ATOMIC;
23+
/** PARTITIONED_NON_ATOMIC: DML statements use a Partitioned DML transaction. */
24+
PARTITIONED_NON_ATOMIC,
25+
/**
26+
* TRANSACTIONAL_WITH_FALLBACK_TO_PARTITIONED_NON_ATOMIC: DML statements are first executed with a
27+
* standard atomic transaction. If that fails due to the mutation limit being exceeded, the
28+
* statement will automatically be retried using a Partitioned DML transaction. These statements
29+
* are not guaranteed to be atomic. The corresponding {@link TransactionRetryListener} methods
30+
* will be invoked when a DML statement falls back to Partitioned DML.
31+
*/
32+
TRANSACTIONAL_WITH_FALLBACK_TO_PARTITIONED_NON_ATOMIC;
2333

2434
private final String statementString;
2535

‎google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/Connection.java‎

Lines changed: 0 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -524,14 +524,6 @@ default byte[] getProtoDescriptors() {
524524
*/
525525
AutocommitDmlMode getAutocommitDmlMode();
526526

527-
default void setFallbackToPartitionedDml(boolean fallbackToPartitionedDml) {
528-
throw new UnsupportedOperationException();
529-
}
530-
531-
default boolean isFallbackToPartitionedDml() {
532-
throw new UnsupportedOperationException();
533-
}
534-
535527
/**
536528
* Sets the staleness to use for the current read-only transaction. This method may only be called
537529
* when the transaction mode of the current transaction is {@link

‎google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionImpl.java‎

Lines changed: 0 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,6 @@
2929
import static com.google.cloud.spanner.connection.ConnectionProperties.DDL_IN_TRANSACTION_MODE;
3030
import static com.google.cloud.spanner.connection.ConnectionProperties.DELAY_TRANSACTION_START_UNTIL_FIRST_WRITE;
3131
import static com.google.cloud.spanner.connection.ConnectionProperties.DIRECTED_READ;
32-
import static com.google.cloud.spanner.connection.ConnectionProperties.FALLBACK_TO_PARTITIONED_DML;
3332
import static com.google.cloud.spanner.connection.ConnectionProperties.KEEP_TRANSACTION_ALIVE;
3433
import static com.google.cloud.spanner.connection.ConnectionProperties.MAX_COMMIT_DELAY;
3534
import static com.google.cloud.spanner.connection.ConnectionProperties.MAX_PARTITIONED_PARALLELISM;
@@ -664,18 +663,6 @@ public AutocommitDmlMode getAutocommitDmlMode() {
664663
return getConnectionPropertyValue(AUTOCOMMIT_DML_MODE);
665664
}
666665

667-
@Override
668-
public void setFallbackToPartitionedDml(boolean fallbackToPartitionedDml) {
669-
ConnectionPreconditions.checkState(!isClosed(), CLOSED_ERROR_MSG);
670-
setConnectionPropertyValue(FALLBACK_TO_PARTITIONED_DML, fallbackToPartitionedDml);
671-
}
672-
673-
@Override
674-
public boolean isFallbackToPartitionedDml() {
675-
ConnectionPreconditions.checkState(!isClosed(), CLOSED_ERROR_MSG);
676-
return getConnectionPropertyValue(FALLBACK_TO_PARTITIONED_DML);
677-
}
678-
679666
@Override
680667
public void setReadOnlyStaleness(TimestampBound staleness) {
681668
Preconditions.checkNotNull(staleness);
@@ -2138,7 +2125,6 @@ UnitOfWork createNewUnitOfWork(
21382125
.setReadOnly(getConnectionPropertyValue(READONLY))
21392126
.setReadOnlyStaleness(getConnectionPropertyValue(READ_ONLY_STALENESS))
21402127
.setAutocommitDmlMode(getConnectionPropertyValue(AUTOCOMMIT_DML_MODE))
2141-
.setRetryDmlAsPartitionedDml(getConnectionPropertyValue(FALLBACK_TO_PARTITIONED_DML))
21422128
.setTransactionRetryListeners(transactionRetryListeners)
21432129
.setReturnCommitStats(getConnectionPropertyValue(RETURN_COMMIT_STATS))
21442130
.setExcludeTxnFromChangeStreams(excludeTxnFromChangeStreams)

‎google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionProperties.java‎

Lines changed: 0 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -361,15 +361,6 @@ class ConnectionProperties {
361361
AutocommitDmlMode.TRANSACTIONAL,
362362
AutocommitDmlModeConverter.INSTANCE,
363363
Context.USER);
364-
static final ConnectionProperty<Boolean> FALLBACK_TO_PARTITIONED_DML =
365-
create(
366-
"fallback_to_partitioned_dml",
367-
"Automatically retry DML statements as Partitioned DML if the atomic DML "
368-
+ "statement failed due to exceeding the Spanner transaction mutation limit. "
369-
+ "This option only affects DML statements that are executed in auto-commit mode.",
370-
false,
371-
BooleanConverter.INSTANCE,
372-
Context.USER);
373364
static final ConnectionProperty<Boolean> RETRY_ABORTS_INTERNALLY =
374365
create(
375366
// TODO: Add support for synonyms for connection properties.

‎google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionStatementExecutor.java‎

Lines changed: 0 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -187,8 +187,4 @@ StatementResult statementSetPgSessionCharacteristicsTransactionMode(
187187
StatementResult statementSetAutoBatchDmlUpdateCountVerification(Boolean verification);
188188

189189
StatementResult statementShowAutoBatchDmlUpdateCountVerification();
190-
191-
StatementResult statementSetFallbackToPartitionedDml(Boolean fallbackToPartitionedDml);
192-
193-
StatementResult statementShowFallbackToPartitionedDml();
194190
}

‎google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionStatementExecutorImpl.java‎

Lines changed: 0 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -34,7 +34,6 @@
3434
import static com.google.cloud.spanner.connection.StatementResult.ClientSideStatementType.SET_DELAY_TRANSACTION_START_UNTIL_FIRST_WRITE;
3535
import static com.google.cloud.spanner.connection.StatementResult.ClientSideStatementType.SET_DIRECTED_READ;
3636
import static com.google.cloud.spanner.connection.StatementResult.ClientSideStatementType.SET_EXCLUDE_TXN_FROM_CHANGE_STREAMS;
37-
import static com.google.cloud.spanner.connection.StatementResult.ClientSideStatementType.SET_FALLBACK_TO_PARTITIONED_DML;
3837
import static com.google.cloud.spanner.connection.StatementResult.ClientSideStatementType.SET_KEEP_TRANSACTION_ALIVE;
3938
import static com.google.cloud.spanner.connection.StatementResult.ClientSideStatementType.SET_MAX_COMMIT_DELAY;
4039
import static com.google.cloud.spanner.connection.StatementResult.ClientSideStatementType.SET_MAX_PARTITIONED_PARALLELISM;
@@ -65,7 +64,6 @@
6564
import static com.google.cloud.spanner.connection.StatementResult.ClientSideStatementType.SHOW_DELAY_TRANSACTION_START_UNTIL_FIRST_WRITE;
6665
import static com.google.cloud.spanner.connection.StatementResult.ClientSideStatementType.SHOW_DIRECTED_READ;
6766
import static com.google.cloud.spanner.connection.StatementResult.ClientSideStatementType.SHOW_EXCLUDE_TXN_FROM_CHANGE_STREAMS;
68-
import static com.google.cloud.spanner.connection.StatementResult.ClientSideStatementType.SHOW_FALLBACK_TO_PARTITIONED_DML;
6967
import static com.google.cloud.spanner.connection.StatementResult.ClientSideStatementType.SHOW_KEEP_TRANSACTION_ALIVE;
7068
import static com.google.cloud.spanner.connection.StatementResult.ClientSideStatementType.SHOW_MAX_COMMIT_DELAY;
7169
import static com.google.cloud.spanner.connection.StatementResult.ClientSideStatementType.SHOW_MAX_PARTITIONED_PARALLELISM;
@@ -732,20 +730,6 @@ public StatementResult statementShowAutoBatchDmlUpdateCountVerification() {
732730
SHOW_AUTO_BATCH_DML_UPDATE_COUNT_VERIFICATION);
733731
}
734732

735-
@Override
736-
public StatementResult statementSetFallbackToPartitionedDml(Boolean fallbackToPartitionedDml) {
737-
getConnection().setFallbackToPartitionedDml(fallbackToPartitionedDml);
738-
return noResult(SET_FALLBACK_TO_PARTITIONED_DML);
739-
}
740-
741-
@Override
742-
public StatementResult statementShowFallbackToPartitionedDml() {
743-
return resultSet(
744-
String.format("%sFALLBACK_TO_PARTITIONED_DML", getNamespace(connection.getDialect())),
745-
getConnection().isFallbackToPartitionedDml(),
746-
SHOW_FALLBACK_TO_PARTITIONED_DML);
747-
}
748-
749733
private String processQueryPlan(PlanNode planNode) {
750734
StringBuilder planNodeDescription = new StringBuilder(" : { ");
751735
com.google.protobuf.Struct metadata = planNode.getMetadata();

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

Lines changed: 7 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -83,7 +83,6 @@ class SingleUseTransaction extends AbstractBaseUnitOfWork {
8383
private final BatchClient batchClient;
8484
private final TimestampBound readOnlyStaleness;
8585
private final AutocommitDmlMode autocommitDmlMode;
86-
private final boolean retryDmlAsPartitionedDml;
8786
private final boolean returnCommitStats;
8887
private final Duration maxCommitDelay;
8988
private final boolean internalMetdataQuery;
@@ -100,7 +99,6 @@ static class Builder extends AbstractBaseUnitOfWork.Builder<Builder, SingleUseTr
10099
private boolean readOnly;
101100
private TimestampBound readOnlyStaleness;
102101
private AutocommitDmlMode autocommitDmlMode;
103-
private boolean retryDmlAsPartitionedDml;
104102
private boolean returnCommitStats;
105103
private Duration maxCommitDelay;
106104
private boolean internalMetadataQuery;
@@ -142,11 +140,6 @@ Builder setAutocommitDmlMode(AutocommitDmlMode dmlMode) {
142140
return this;
143141
}
144142

145-
Builder setRetryDmlAsPartitionedDml(boolean retryDmlAsPartitionedDml) {
146-
this.retryDmlAsPartitionedDml = retryDmlAsPartitionedDml;
147-
return this;
148-
}
149-
150143
Builder setReturnCommitStats(boolean returnCommitStats) {
151144
this.returnCommitStats = returnCommitStats;
152145
return this;
@@ -190,7 +183,6 @@ private SingleUseTransaction(Builder builder) {
190183
this.readOnly = builder.readOnly;
191184
this.readOnlyStaleness = builder.readOnlyStaleness;
192185
this.autocommitDmlMode = builder.autocommitDmlMode;
193-
this.retryDmlAsPartitionedDml = builder.retryDmlAsPartitionedDml;
194186
this.returnCommitStats = builder.returnCommitStats;
195187
this.maxCommitDelay = builder.maxCommitDelay;
196188
this.internalMetdataQuery = builder.internalMetadataQuery;
@@ -228,6 +220,11 @@ public boolean supportsDirectedReads(ParsedStatement parsedStatement) {
228220
return parsedStatement.isQuery();
229221
}
230222

223+
private boolean isRetryDmlAsPartitionedDml() {
224+
return this.autocommitDmlMode
225+
== AutocommitDmlMode.TRANSACTIONAL_WITH_FALLBACK_TO_PARTITIONED_NON_ATOMIC;
226+
}
227+
231228
private void checkAndMarkUsed() {
232229
Preconditions.checkState(!used, "This single-use transaction has already been used");
233230
used = true;
@@ -443,6 +440,7 @@ public ApiFuture<Long> executeUpdateAsync(
443440
ApiFuture<Long> res;
444441
switch (autocommitDmlMode) {
445442
case TRANSACTIONAL:
443+
case TRANSACTIONAL_WITH_FALLBACK_TO_PARTITIONED_NON_ATOMIC:
446444
res =
447445
ApiFutures.transform(
448446
executeTransactionalUpdateAsync(callType, update, AnalyzeMode.NONE, options),
@@ -578,7 +576,7 @@ private ApiFuture<Tuple<Long, ResultSet>> executeTransactionalUpdateAsync(
578576
ImmutableList.of(SpannerGrpc.getExecuteSqlMethod(), SpannerGrpc.getCommitMethod()));
579577
// Retry as Partitioned DML if the statement fails due to exceeding the mutation limit if that
580578
// option has been enabled.
581-
if (this.retryDmlAsPartitionedDml) {
579+
if (isRetryDmlAsPartitionedDml()) {
582580
return addRetryUpdateAsPartitionedDmlCallback(transactionalResult, callType, update, options);
583581
}
584582
return transactionalResult;

‎google-cloud-spanner/src/main/resources/com/google/cloud/spanner/connection/ClientSideStatements.json‎

Lines changed: 3 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -221,15 +221,6 @@
221221
"method": "statementShowAutoBatchDmlUpdateCountVerification",
222222
"exampleStatements": ["show variable auto_batch_dml_update_count_verification"]
223223
},
224-
{
225-
"name": "SHOW VARIABLE FALLBACK_TO_PARTITIONED_DML",
226-
"executorName": "ClientSideStatementNoParamExecutor",
227-
"resultType": "RESULT_SET",
228-
"statementType": "SHOW_FALLBACK_TO_PARTITIONED_DML",
229-
"regex": "(?is)\\A\\s*show\\s+variable\\s+fallback_to_partitioned_dml\\s*\\z",
230-
"method": "statementShowFallbackToPartitionedDml",
231-
"exampleStatements": ["show variable fallback_to_partitioned_dml"]
232-
},
233224
{
234225
"name": "PARTITION <sql>",
235226
"executorName": "ClientSideStatementPartitionExecutor",
@@ -380,17 +371,17 @@
380371
}
381372
},
382373
{
383-
"name": "SET AUTOCOMMIT_DML_MODE = 'PARTITIONED_NON_ATOMIC'|'TRANSACTIONAL'",
374+
"name": "SET AUTOCOMMIT_DML_MODE = 'PARTITIONED_NON_ATOMIC'|'TRANSACTIONAL'|'TRANSACTIONAL_WITH_FALLBACK_TO_PARTITIONED_NON_ATOMIC'",
384375
"executorName": "ClientSideStatementSetExecutor",
385376
"resultType": "NO_RESULT",
386377
"statementType": "SET_AUTOCOMMIT_DML_MODE",
387378
"regex": "(?is)\\A\\s*set\\s+autocommit_dml_mode\\s*(?:=)\\s*(.*)\\z",
388379
"method": "statementSetAutocommitDmlMode",
389-
"exampleStatements": ["set autocommit_dml_mode='PARTITIONED_NON_ATOMIC'", "set autocommit_dml_mode='TRANSACTIONAL'"],
380+
"exampleStatements": ["set autocommit_dml_mode='PARTITIONED_NON_ATOMIC'", "set autocommit_dml_mode='TRANSACTIONAL'", "set autocommit_dml_mode='TRANSACTIONAL_WITH_FALLBACK_TO_PARTITIONED_NON_ATOMIC'"],
390381
"setStatement": {
391382
"propertyName": "AUTOCOMMIT_DML_MODE",
392383
"separator": "=",
393-
"allowedValues": "'(PARTITIONED_NON_ATOMIC|TRANSACTIONAL)'",
384+
"allowedValues": "'(PARTITIONED_NON_ATOMIC|TRANSACTIONAL|TRANSACTIONAL_WITH_FALLBACK_TO_PARTITIONED_NON_ATOMIC)'",
394385
"converterName": "ClientSideStatementValueConverters$AutocommitDmlModeConverter"
395386
}
396387
},
@@ -710,21 +701,6 @@
710701
"converterName": "ClientSideStatementValueConverters$BooleanConverter"
711702
}
712703
},
713-
{
714-
"name": "SET FALLBACK_TO_PARTITIONED_DML = TRUE|FALSE",
715-
"executorName": "ClientSideStatementSetExecutor",
716-
"resultType": "NO_RESULT",
717-
"statementType": "SET_FALLBACK_TO_PARTITIONED_DML",
718-
"regex": "(?is)\\A\\s*set\\s+fallback_to_partitioned_dml\\s*(?:=)\\s*(.*)\\z",
719-
"method": "statementSetFallbackToPartitionedDml",
720-
"exampleStatements": ["set fallback_to_partitioned_dml = true", "set fallback_to_partitioned_dml = false"],
721-
"setStatement": {
722-
"propertyName": "FALLBACK_TO_PARTITIONED_DML",
723-
"separator": "=",
724-
"allowedValues": "(TRUE|FALSE)",
725-
"converterName": "ClientSideStatementValueConverters$BooleanConverter"
726-
}
727-
},
728704
{
729705
"name": "SHOW VARIABLE DATA_BOOST_ENABLED",
730706
"executorName": "ClientSideStatementNoParamExecutor",

‎google-cloud-spanner/src/main/resources/com/google/cloud/spanner/connection/PG_ClientSideStatements.json‎

Lines changed: 5 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -221,15 +221,6 @@
221221
"method": "statementShowAutoBatchDmlUpdateCountVerification",
222222
"exampleStatements": ["show spanner.auto_batch_dml_update_count_verification","show variable spanner.auto_batch_dml_update_count_verification"]
223223
},
224-
{
225-
"name": "SHOW [VARIABLE] SPANNER.FALLBACK_TO_PARTITIONED_DML",
226-
"executorName": "ClientSideStatementNoParamExecutor",
227-
"resultType": "RESULT_SET",
228-
"statementType": "SHOW_FALLBACK_TO_PARTITIONED_DML",
229-
"regex": "(?is)\\A\\s*show\\s+(?:variable\\s+)?spanner\\.fallback_to_partitioned_dml\\s*\\z",
230-
"method": "statementShowFallbackToPartitionedDml",
231-
"exampleStatements": ["show spanner.fallback_to_partitioned_dml", "show variable spanner.fallback_to_partitioned_dml"]
232-
},
233224
{
234225
"name": "SHOW [VARIABLE] TRANSACTION ISOLATION LEVEL",
235226
"executorName": "ClientSideStatementNoParamExecutor",
@@ -429,7 +420,7 @@
429420
}
430421
},
431422
{
432-
"name": "SET SPANNER.AUTOCOMMIT_DML_MODE =|TO 'PARTITIONED_NON_ATOMIC'|'TRANSACTIONAL'",
423+
"name": "SET SPANNER.AUTOCOMMIT_DML_MODE =|TO 'PARTITIONED_NON_ATOMIC'|'TRANSACTIONAL'|'TRANSACTIONAL_WITH_FALLBACK_TO_PARTITIONED_NON_ATOMIC'",
433424
"executorName": "ClientSideStatementSetExecutor",
434425
"resultType": "NO_RESULT",
435426
"statementType": "SET_AUTOCOMMIT_DML_MODE",
@@ -438,13 +429,15 @@
438429
"exampleStatements": [
439430
"set spanner.autocommit_dml_mode='PARTITIONED_NON_ATOMIC'",
440431
"set spanner.autocommit_dml_mode='TRANSACTIONAL'",
432+
"set spanner.autocommit_dml_mode='TRANSACTIONAL_WITH_FALLBACK_TO_PARTITIONED_NON_ATOMIC'",
441433
"set spanner.autocommit_dml_mode to 'PARTITIONED_NON_ATOMIC'",
442-
"set spanner.autocommit_dml_mode to 'TRANSACTIONAL'"
434+
"set spanner.autocommit_dml_mode to 'TRANSACTIONAL'",
435+
"set spanner.autocommit_dml_mode to 'TRANSACTIONAL_WITH_FALLBACK_TO_PARTITIONED_NON_ATOMIC'"
443436
],
444437
"setStatement": {
445438
"propertyName": "SPANNER.AUTOCOMMIT_DML_MODE",
446439
"separator": "(?:=|\\s+TO\\s+)",
447-
"allowedValues": "'(PARTITIONED_NON_ATOMIC|TRANSACTIONAL)'",
440+
"allowedValues": "'(PARTITIONED_NON_ATOMIC|TRANSACTIONAL|TRANSACTIONAL_WITH_FALLBACK_TO_PARTITIONED_NON_ATOMIC)'",
448441
"converterName": "ClientSideStatementValueConverters$AutocommitDmlModeConverter"
449442
}
450443
},
@@ -864,21 +857,6 @@
864857
"converterName": "ClientSideStatementValueConverters$PgBooleanConverter"
865858
}
866859
},
867-
{
868-
"name": "SET SPANNER.FALLBACK_TO_PARTITIONED_DML = TRUE|FALSE",
869-
"executorName": "ClientSideStatementSetExecutor",
870-
"resultType": "NO_RESULT",
871-
"statementType": "SET_FALLBACK_TO_PARTITIONED_DML",
872-
"regex": "(?is)\\A\\s*set\\s+spanner\\.fallback_to_partitioned_dml(?:\\s*=\\s*|\\s+to\\s+)(.*)\\z",
873-
"method": "statementSetFallbackToPartitionedDml",
874-
"exampleStatements": ["set spanner.fallback_to_partitioned_dml = true", "set spanner.fallback_to_partitioned_dml = false", "set spanner.fallback_to_partitioned_dml to true", "set spanner.fallback_to_partitioned_dml to false", "set spanner.fallback_to_partitioned_dml to off"],
875-
"setStatement": {
876-
"propertyName": "SPANNER.FALLBACK_TO_PARTITIONED_DML",
877-
"separator": "(?:=|\\s+TO\\s+)",
878-
"allowedValues": "(.+)",
879-
"converterName": "ClientSideStatementValueConverters$PgBooleanConverter"
880-
}
881-
},
882860
{
883861
"name": "SET SPANNER.AUTO_BATCH_DML_UPDATE_COUNT = <bigint>",
884862
"executorName": "ClientSideStatementSetExecutor",

‎google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/RetryDmlAsPartitionedDmlMockServerTest.java‎

Lines changed: 18 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -74,7 +74,7 @@ public void testTransactionMutationLimitExceeded_isNotRetriedByDefault() {
7474

7575
try (Connection connection = createConnection()) {
7676
connection.setAutocommit(true);
77-
assertFalse(connection.isFallbackToPartitionedDml());
77+
assertEquals(AutocommitDmlMode.TRANSACTIONAL, connection.getAutocommitDmlMode());
7878

7979
TransactionMutationLimitExceededException exception =
8080
assertThrows(
@@ -96,7 +96,8 @@ public void testTransactionMutationLimitExceeded_canBeRetriedAsPDML() {
9696

9797
try (Connection connection = createConnection()) {
9898
connection.setAutocommit(true);
99-
connection.setFallbackToPartitionedDml(true);
99+
connection.setAutocommitDmlMode(
100+
AutocommitDmlMode.TRANSACTIONAL_WITH_FALLBACK_TO_PARTITIONED_NON_ATOMIC);
100101

101102
long updateCount = connection.executeUpdate(statement);
102103
assertEquals(100000L, updateCount);
@@ -136,7 +137,8 @@ public void testTransactionMutationLimitExceeded_retryAsPDMLFails() {
136137

137138
try (Connection connection = createConnection()) {
138139
connection.setAutocommit(true);
139-
connection.setFallbackToPartitionedDml(true);
140+
connection.setAutocommitDmlMode(
141+
AutocommitDmlMode.TRANSACTIONAL_WITH_FALLBACK_TO_PARTITIONED_NON_ATOMIC);
140142

141143
// The connection throws TransactionMutationLimitExceededException if the retry using
142144
// partitioned DML fails. The exception from the failed retry is returned as a suppressed
@@ -184,22 +186,28 @@ public void testSqlStatements() {
184186
String prefix = dialect == Dialect.POSTGRESQL ? "SPANNER." : "";
185187

186188
try (Connection connection = createConnection()) {
189+
connection.setAutocommit(true);
187190
try (ResultSet resultSet =
188191
connection.executeQuery(
189-
Statement.of(
190-
String.format("show variable %sfallback_to_partitioned_dml", prefix)))) {
192+
Statement.of(String.format("show variable %sautocommit_dml_mode", prefix)))) {
191193
assertTrue(resultSet.next());
192-
assertFalse(resultSet.getBoolean(String.format("%sFALLBACK_TO_PARTITIONED_DML", prefix)));
194+
assertEquals(
195+
AutocommitDmlMode.TRANSACTIONAL.name(),
196+
resultSet.getString(String.format("%sAUTOCOMMIT_DML_MODE", prefix)));
193197
assertFalse(resultSet.next());
194198
}
195199
connection.execute(
196-
Statement.of(String.format("set %sfallback_to_partitioned_dml = true", prefix)));
200+
Statement.of(
201+
String.format(
202+
"set %sautocommit_dml_mode = 'transactional_with_fallback_to_partitioned_non_atomic'",
203+
prefix)));
197204
try (ResultSet resultSet =
198205
connection.executeQuery(
199-
Statement.of(
200-
String.format("show variable %sfallback_to_partitioned_dml", prefix)))) {
206+
Statement.of(String.format("show variable %sautocommit_dml_mode", prefix)))) {
201207
assertTrue(resultSet.next());
202-
assertTrue(resultSet.getBoolean(String.format("%sFALLBACK_TO_PARTITIONED_DML", prefix)));
208+
assertEquals(
209+
AutocommitDmlMode.TRANSACTIONAL_WITH_FALLBACK_TO_PARTITIONED_NON_ATOMIC.name(),
210+
resultSet.getString(String.format("%sAUTOCOMMIT_DML_MODE", prefix)));
203211
assertFalse(resultSet.next());
204212
}
205213
}

0 commit comments

Comments
 (0)