feat: automatically set default sequence kind in JDBC and PGAdapter (… · googleapis/java-spanner@e8abf33 (original) (raw)
`@@ -17,6 +17,7 @@
`
17
17
`package com.google.cloud.spanner.connection;
`
18
18
``
19
19
`import static com.google.cloud.spanner.connection.AbstractStatementParser.RUN_BATCH_STATEMENT;
`
``
20
`+
import static com.google.cloud.spanner.connection.ConnectionProperties.DEFAULT_SEQUENCE_KIND;
`
20
21
``
21
22
`import com.google.api.core.ApiFuture;
`
22
23
`import com.google.api.core.ApiFutures;
`
45
46
`import java.util.Arrays;
`
46
47
`import java.util.List;
`
47
48
`import java.util.concurrent.Callable;
`
``
49
`+
import java.util.concurrent.atomic.AtomicReference;
`
48
50
`import javax.annotation.Nonnull;
`
49
51
``
50
52
`/**
`
`@@ -61,11 +63,13 @@ class DdlBatch extends AbstractBaseUnitOfWork {
`
61
63
`private final List statements = new ArrayList<>();
`
62
64
`private UnitOfWorkState state = UnitOfWorkState.STARTED;
`
63
65
`private final byte[] protoDescriptors;
`
``
66
`+
private final ConnectionState connectionState;
`
64
67
``
65
68
`static class Builder extends AbstractBaseUnitOfWork.Builder<Builder, DdlBatch> {
`
66
69
`private DdlClient ddlClient;
`
67
70
`private DatabaseClient dbClient;
`
68
71
`private byte[] protoDescriptors;
`
``
72
`+
private ConnectionState connectionState;
`
69
73
``
70
74
`private Builder() {}
`
71
75
``
`@@ -86,6 +90,11 @@ Builder setProtoDescriptors(byte[] protoDescriptors) {
`
86
90
`return this;
`
87
91
` }
`
88
92
``
``
93
`+
Builder setConnectionState(ConnectionState connectionState) {
`
``
94
`+
this.connectionState = connectionState;
`
``
95
`+
return this;
`
``
96
`+
}
`
``
97
+
89
98
`@Override
`
90
99
`DdlBatch build() {
`
91
100
`Preconditions.checkState(ddlClient != null, "No DdlClient specified");
`
`@@ -103,6 +112,7 @@ private DdlBatch(Builder builder) {
`
103
112
`this.ddlClient = builder.ddlClient;
`
104
113
`this.dbClient = builder.dbClient;
`
105
114
`this.protoDescriptors = builder.protoDescriptors;
`
``
115
`+
this.connectionState = Preconditions.checkNotNull(builder.connectionState);
`
106
116
` }
`
107
117
``
108
118
`@Override
`
`@@ -235,17 +245,28 @@ public ApiFuture<long[]> runBatchAsync(CallType callType) {
`
235
245
`Callable<long[]> callable =
`
236
246
` () -> {
`
237
247
`try {
`
238
``
`-
OperationFuture<Void, UpdateDatabaseDdlMetadata> operation =
`
239
``
`-
ddlClient.executeDdl(statements, protoDescriptors);
`
``
248
`+
AtomicReference<OperationFuture<Void, UpdateDatabaseDdlMetadata>> operationReference =
`
``
249
`+
new AtomicReference<>();
`
240
250
`try {
`
241
``
`-
// Wait until the operation has finished.
`
242
``
`-
getWithStatementTimeout(operation, RUN_BATCH_STATEMENT);
`
``
251
`+
ddlClient.runWithRetryForMissingDefaultSequenceKind(
`
``
252
`+
restartIndex -> {
`
``
253
`+
OperationFuture<Void, UpdateDatabaseDdlMetadata> operation =
`
``
254
`+
ddlClient.executeDdl(
`
``
255
`+
statements.subList(restartIndex, statements.size()),
`
``
256
`+
protoDescriptors);
`
``
257
`+
operationReference.set(operation);
`
``
258
`+
// Wait until the operation has finished.
`
``
259
`+
getWithStatementTimeout(operation, RUN_BATCH_STATEMENT);
`
``
260
`+
},
`
``
261
`+
connectionState.getValue(DEFAULT_SEQUENCE_KIND).getValue(),
`
``
262
`+
dbClient.getDialect(),
`
``
263
`+
operationReference);
`
243
264
`long[] updateCounts = new long[statements.size()];
`
244
265
`Arrays.fill(updateCounts, 1L);
`
245
266
`state = UnitOfWorkState.RAN;
`
246
267
`return updateCounts;
`
247
268
` } catch (SpannerException e) {
`
248
``
`-
long[] updateCounts = extractUpdateCounts(operation);
`
``
269
`+
long[] updateCounts = extractUpdateCounts(operationReference.get());
`
249
270
`throw SpannerExceptionFactory.newSpannerBatchUpdateException(
`
250
271
`e.getErrorCode(), e.getMessage(), updateCounts);
`
251
272
` }
`