diff --git a/google-cloud-spanner/clirr-ignored-differences.xml b/google-cloud-spanner/clirr-ignored-differences.xml index 4e03926c66e..5c103beca6f 100644 --- a/google-cloud-spanner/clirr-ignored-differences.xml +++ b/google-cloud-spanner/clirr-ignored-differences.xml @@ -828,4 +828,16 @@ com/google/cloud/spanner/SpannerOptions$SpannerEnvironment com.google.auth.oauth2.GoogleCredentials getDefaultExternalHostCredentials() + + + + 7012 + com/google/cloud/spanner/connection/Connection + void setDefaultSequenceKind(java.lang.String) + + + 7012 + com/google/cloud/spanner/connection/Connection + java.lang.String getDefaultSequenceKind() + diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/MissingDefaultSequenceKindException.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/MissingDefaultSequenceKindException.java new file mode 100644 index 00000000000..7b61dbbd881 --- /dev/null +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/MissingDefaultSequenceKindException.java @@ -0,0 +1,52 @@ +/* + * Copyright 2025 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; + +import com.google.api.gax.rpc.ApiException; +import java.util.regex.Pattern; +import javax.annotation.Nullable; + +/** + * Exception thrown by Spanner when a DDL statement failed because no default sequence kind has been + * configured for a database. + */ +public class MissingDefaultSequenceKindException extends SpannerException { + private static final long serialVersionUID = 1L; + + private static final Pattern PATTERN = + Pattern.compile( + "The sequence kind of an identity column .+ is not specified\\. Please specify the sequence kind explicitly or set the database option `default_sequence_kind`\\."); + + /** Private constructor. Use {@link SpannerExceptionFactory} to create instances. */ + MissingDefaultSequenceKindException( + DoNotConstructDirectly token, + ErrorCode errorCode, + String message, + Throwable cause, + @Nullable ApiException apiException) { + super(token, errorCode, /*retryable = */ false, message, cause, apiException); + } + + static boolean isMissingDefaultSequenceKindException(Throwable cause) { + if (cause == null + || cause.getMessage() == null + || !PATTERN.matcher(cause.getMessage()).find()) { + return false; + } + return true; + } +} diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerExceptionFactory.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerExceptionFactory.java index a3f174cda60..6476b94b144 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerExceptionFactory.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerExceptionFactory.java @@ -16,6 +16,7 @@ package com.google.cloud.spanner; +import static com.google.cloud.spanner.MissingDefaultSequenceKindException.isMissingDefaultSequenceKindException; import static com.google.cloud.spanner.TransactionMutationLimitExceededException.isTransactionMutationLimitException; import com.google.api.gax.grpc.GrpcStatusCode; @@ -336,6 +337,9 @@ static SpannerException newSpannerExceptionPreformatted( return new TransactionMutationLimitExceededException( token, code, message, cause, apiException); } + if (isMissingDefaultSequenceKindException(apiException)) { + return new MissingDefaultSequenceKindException(token, code, message, cause, apiException); + } // Fall through to the default. default: return new SpannerException( diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/Connection.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/Connection.java index eb69ae132cc..7bf4e47bd9a 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/Connection.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/Connection.java @@ -862,6 +862,18 @@ interface TransactionCallable { /** Sets how the connection should behave if a DDL statement is executed during a transaction. */ void setDdlInTransactionMode(DdlInTransactionMode ddlInTransactionMode); + /** + * Returns the default sequence kind that will be set for this database if a DDL statement is + * executed that uses auto_increment or serial. + */ + String getDefaultSequenceKind(); + + /** + * Sets the default sequence kind that will be set for this database if a DDL statement is + * executed that uses auto_increment or serial. + */ + void setDefaultSequenceKind(String defaultSequenceKind); + /** * Creates a savepoint with the given name. * diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionImpl.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionImpl.java index 5ea249ee0ac..4c0c95a91a5 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionImpl.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionImpl.java @@ -27,6 +27,7 @@ import static com.google.cloud.spanner.connection.ConnectionProperties.AUTO_PARTITION_MODE; import static com.google.cloud.spanner.connection.ConnectionProperties.DATA_BOOST_ENABLED; import static com.google.cloud.spanner.connection.ConnectionProperties.DDL_IN_TRANSACTION_MODE; +import static com.google.cloud.spanner.connection.ConnectionProperties.DEFAULT_SEQUENCE_KIND; import static com.google.cloud.spanner.connection.ConnectionProperties.DELAY_TRANSACTION_START_UNTIL_FIRST_WRITE; import static com.google.cloud.spanner.connection.ConnectionProperties.DIRECTED_READ; import static com.google.cloud.spanner.connection.ConnectionProperties.KEEP_TRANSACTION_ALIVE; @@ -761,6 +762,16 @@ public void setDdlInTransactionMode(DdlInTransactionMode ddlInTransactionMode) { setConnectionPropertyValue(DDL_IN_TRANSACTION_MODE, ddlInTransactionMode); } + @Override + public String getDefaultSequenceKind() { + return getConnectionPropertyValue(DEFAULT_SEQUENCE_KIND); + } + + @Override + public void setDefaultSequenceKind(String defaultSequenceKind) { + setConnectionPropertyValue(DEFAULT_SEQUENCE_KIND, defaultSequenceKind); + } + @Override public void setStatementTimeout(long timeout, TimeUnit unit) { Preconditions.checkArgument(timeout > 0L, "Zero or negative timeout values are not allowed"); @@ -2152,13 +2163,9 @@ UnitOfWork createNewUnitOfWork( .setDdlClient(ddlClient) .setDatabaseClient(dbClient) .setBatchClient(batchClient) - .setReadOnly(getConnectionPropertyValue(READONLY)) - .setReadOnlyStaleness(getConnectionPropertyValue(READ_ONLY_STALENESS)) - .setAutocommitDmlMode(getConnectionPropertyValue(AUTOCOMMIT_DML_MODE)) + .setConnectionState(connectionState) .setTransactionRetryListeners(transactionRetryListeners) - .setReturnCommitStats(getConnectionPropertyValue(RETURN_COMMIT_STATS)) .setExcludeTxnFromChangeStreams(excludeTxnFromChangeStreams) - .setMaxCommitDelay(getConnectionPropertyValue(MAX_COMMIT_DELAY)) .setStatementTimeout(statementTimeout) .withStatementExecutor(statementExecutor) .setSpan( @@ -2230,6 +2237,7 @@ UnitOfWork createNewUnitOfWork( .withStatementExecutor(statementExecutor) .setSpan(createSpanForUnitOfWork(DDL_BATCH)) .setProtoDescriptors(getProtoDescriptors()) + .setConnectionState(connectionState) .build(); default: } diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionOptions.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionOptions.java index deb279d8e3e..f63d8f1f10c 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionOptions.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionOptions.java @@ -242,6 +242,7 @@ public String[] getValidValues() { static final RpcPriority DEFAULT_RPC_PRIORITY = null; static final DdlInTransactionMode DEFAULT_DDL_IN_TRANSACTION_MODE = DdlInTransactionMode.ALLOW_IN_EMPTY_TRANSACTION; + static final String DEFAULT_DEFAULT_SEQUENCE_KIND = null; static final boolean DEFAULT_RETURN_COMMIT_STATS = false; static final boolean DEFAULT_LENIENT = false; static final boolean DEFAULT_ROUTE_TO_LEADER = true; @@ -324,6 +325,7 @@ public String[] getValidValues() { public static final String RPC_PRIORITY_NAME = "rpcPriority"; public static final String DDL_IN_TRANSACTION_MODE_PROPERTY_NAME = "ddlInTransactionMode"; + public static final String DEFAULT_SEQUENCE_KIND_PROPERTY_NAME = "defaultSequenceKind"; /** Dialect to use for a connection. */ static final String DIALECT_PROPERTY_NAME = "dialect"; /** Name of the 'databaseRole' connection property. */ diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionProperties.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionProperties.java index 6f2628e5a04..f65dc533570 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionProperties.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionProperties.java @@ -41,6 +41,7 @@ import static com.google.cloud.spanner.connection.ConnectionOptions.DEFAULT_DATABASE_ROLE; import static com.google.cloud.spanner.connection.ConnectionOptions.DEFAULT_DATA_BOOST_ENABLED; import static com.google.cloud.spanner.connection.ConnectionOptions.DEFAULT_DDL_IN_TRANSACTION_MODE; +import static com.google.cloud.spanner.connection.ConnectionOptions.DEFAULT_DEFAULT_SEQUENCE_KIND; import static com.google.cloud.spanner.connection.ConnectionOptions.DEFAULT_DELAY_TRANSACTION_START_UNTIL_FIRST_WRITE; import static com.google.cloud.spanner.connection.ConnectionOptions.DEFAULT_ENABLE_API_TRACING; import static com.google.cloud.spanner.connection.ConnectionOptions.DEFAULT_ENABLE_END_TO_END_TRACING; @@ -61,6 +62,7 @@ import static com.google.cloud.spanner.connection.ConnectionOptions.DEFAULT_RETURN_COMMIT_STATS; import static com.google.cloud.spanner.connection.ConnectionOptions.DEFAULT_ROUTE_TO_LEADER; import static com.google.cloud.spanner.connection.ConnectionOptions.DEFAULT_RPC_PRIORITY; +import static com.google.cloud.spanner.connection.ConnectionOptions.DEFAULT_SEQUENCE_KIND_PROPERTY_NAME; import static com.google.cloud.spanner.connection.ConnectionOptions.DEFAULT_TRACK_CONNECTION_LEAKS; import static com.google.cloud.spanner.connection.ConnectionOptions.DEFAULT_TRACK_SESSION_LEAKS; import static com.google.cloud.spanner.connection.ConnectionOptions.DEFAULT_USER_AGENT; @@ -531,6 +533,15 @@ public class ConnectionProperties { DdlInTransactionMode.values(), DdlInTransactionModeConverter.INSTANCE, Context.USER); + static final ConnectionProperty DEFAULT_SEQUENCE_KIND = + create( + DEFAULT_SEQUENCE_KIND_PROPERTY_NAME, + "The default sequence kind that should be used for the database. " + + "This property is only used when a DDL statement that requires a default " + + "sequence kind is executed on this connection.", + DEFAULT_DEFAULT_SEQUENCE_KIND, + StringValueConverter.INSTANCE, + Context.USER); static final ConnectionProperty MAX_COMMIT_DELAY = create( "maxCommitDelay", @@ -615,16 +626,10 @@ private static ConnectionProperty create( T[] validValues, ClientSideStatementValueConverter converter, Context context) { - try { - ConnectionProperty property = - ConnectionProperty.create( - name, description, defaultValue, validValues, converter, context); - CONNECTION_PROPERTIES_BUILDER.put(property.getKey(), property); - return property; - } catch (Throwable t) { - t.printStackTrace(); - } - return null; + ConnectionProperty property = + ConnectionProperty.create(name, description, defaultValue, validValues, converter, context); + CONNECTION_PROPERTIES_BUILDER.put(property.getKey(), property); + return property; } /** Parse the connection properties that can be found in the given connection URL. */ diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/DdlBatch.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/DdlBatch.java index 6ae28822473..813f5d6e45b 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/DdlBatch.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/DdlBatch.java @@ -17,6 +17,7 @@ package com.google.cloud.spanner.connection; import static com.google.cloud.spanner.connection.AbstractStatementParser.RUN_BATCH_STATEMENT; +import static com.google.cloud.spanner.connection.ConnectionProperties.DEFAULT_SEQUENCE_KIND; import com.google.api.core.ApiFuture; import com.google.api.core.ApiFutures; @@ -45,6 +46,7 @@ import java.util.Arrays; import java.util.List; import java.util.concurrent.Callable; +import java.util.concurrent.atomic.AtomicReference; import javax.annotation.Nonnull; /** @@ -61,11 +63,13 @@ class DdlBatch extends AbstractBaseUnitOfWork { private final List statements = new ArrayList<>(); private UnitOfWorkState state = UnitOfWorkState.STARTED; private final byte[] protoDescriptors; + private final ConnectionState connectionState; static class Builder extends AbstractBaseUnitOfWork.Builder { private DdlClient ddlClient; private DatabaseClient dbClient; private byte[] protoDescriptors; + private ConnectionState connectionState; private Builder() {} @@ -86,6 +90,11 @@ Builder setProtoDescriptors(byte[] protoDescriptors) { return this; } + Builder setConnectionState(ConnectionState connectionState) { + this.connectionState = connectionState; + return this; + } + @Override DdlBatch build() { Preconditions.checkState(ddlClient != null, "No DdlClient specified"); @@ -103,6 +112,7 @@ private DdlBatch(Builder builder) { this.ddlClient = builder.ddlClient; this.dbClient = builder.dbClient; this.protoDescriptors = builder.protoDescriptors; + this.connectionState = Preconditions.checkNotNull(builder.connectionState); } @Override @@ -235,17 +245,28 @@ public ApiFuture runBatchAsync(CallType callType) { Callable callable = () -> { try { - OperationFuture operation = - ddlClient.executeDdl(statements, protoDescriptors); + AtomicReference> operationReference = + new AtomicReference<>(); try { - // Wait until the operation has finished. - getWithStatementTimeout(operation, RUN_BATCH_STATEMENT); + ddlClient.runWithRetryForMissingDefaultSequenceKind( + restartIndex -> { + OperationFuture operation = + ddlClient.executeDdl( + statements.subList(restartIndex, statements.size()), + protoDescriptors); + operationReference.set(operation); + // Wait until the operation has finished. + getWithStatementTimeout(operation, RUN_BATCH_STATEMENT); + }, + connectionState.getValue(DEFAULT_SEQUENCE_KIND).getValue(), + dbClient.getDialect(), + operationReference); long[] updateCounts = new long[statements.size()]; Arrays.fill(updateCounts, 1L); state = UnitOfWorkState.RAN; return updateCounts; } catch (SpannerException e) { - long[] updateCounts = extractUpdateCounts(operation); + long[] updateCounts = extractUpdateCounts(operationReference.get()); throw SpannerExceptionFactory.newSpannerBatchUpdateException( e.getErrorCode(), e.getMessage(), updateCounts); } diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/DdlClient.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/DdlClient.java index 7bce1ab78cd..a3dc286acb4 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/DdlClient.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/DdlClient.java @@ -22,6 +22,8 @@ import com.google.cloud.spanner.DatabaseId; import com.google.cloud.spanner.Dialect; import com.google.cloud.spanner.ErrorCode; +import com.google.cloud.spanner.MissingDefaultSequenceKindException; +import com.google.cloud.spanner.SpannerException; import com.google.cloud.spanner.SpannerExceptionFactory; import com.google.common.base.Preconditions; import com.google.common.base.Strings; @@ -29,6 +31,9 @@ import com.google.spanner.admin.database.v1.UpdateDatabaseDdlMetadata; import java.util.Collections; import java.util.List; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Consumer; /** * Convenience class for executing Data Definition Language statements on transactions that support @@ -131,4 +136,46 @@ static boolean isCreateDatabaseStatement(String statement) { && tokens[0].equalsIgnoreCase("CREATE") && tokens[1].equalsIgnoreCase("DATABASE"); } + + void runWithRetryForMissingDefaultSequenceKind( + Consumer runnable, + String defaultSequenceKind, + Dialect dialect, + AtomicReference> operationReference) { + try { + runnable.accept(0); + } catch (Throwable t) { + SpannerException spannerException = SpannerExceptionFactory.asSpannerException(t); + if (!Strings.isNullOrEmpty(defaultSequenceKind) + && spannerException instanceof MissingDefaultSequenceKindException) { + setDefaultSequenceKind(defaultSequenceKind, dialect); + int restartIndex = 0; + if (operationReference.get() != null) { + try { + UpdateDatabaseDdlMetadata metadata = operationReference.get().getMetadata().get(); + restartIndex = metadata.getCommitTimestampsCount(); + } catch (Throwable ignore) { + } + } + runnable.accept(restartIndex); + return; + } + throw t; + } + } + + private void setDefaultSequenceKind(String defaultSequenceKind, Dialect dialect) { + String ddl = + dialect == Dialect.POSTGRESQL + ? "alter database \"%s\" set spanner.default_sequence_kind = '%s'" + : "alter database `%s` set options (default_sequence_kind='%s')"; + ddl = String.format(ddl, databaseName, defaultSequenceKind); + try { + executeDdl(ddl, null).get(); + } catch (ExecutionException executionException) { + throw SpannerExceptionFactory.asSpannerException(executionException.getCause()); + } catch (InterruptedException interruptedException) { + throw SpannerExceptionFactory.propagateInterrupt(interruptedException); + } + } } diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/SingleUseTransaction.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/SingleUseTransaction.java index a827f82ba36..8a5cd6a26db 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/SingleUseTransaction.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/SingleUseTransaction.java @@ -18,6 +18,13 @@ import static com.google.cloud.spanner.connection.AbstractStatementParser.COMMIT_STATEMENT; import static com.google.cloud.spanner.connection.AbstractStatementParser.RUN_BATCH_STATEMENT; +import static com.google.cloud.spanner.connection.ConnectionProperties.AUTOCOMMIT_DML_MODE; +import static com.google.cloud.spanner.connection.ConnectionProperties.DEFAULT_SEQUENCE_KIND; +import static com.google.cloud.spanner.connection.ConnectionProperties.MAX_COMMIT_DELAY; +import static com.google.cloud.spanner.connection.ConnectionProperties.READONLY; +import static com.google.cloud.spanner.connection.ConnectionProperties.READ_ONLY_STALENESS; +import static com.google.cloud.spanner.connection.ConnectionProperties.RETURN_COMMIT_STATS; +import static com.google.cloud.spanner.connection.DdlClient.isCreateDatabaseStatement; import com.google.api.core.ApiFuture; import com.google.api.core.ApiFutureCallback; @@ -43,7 +50,6 @@ import com.google.cloud.spanner.SpannerBatchUpdateException; import com.google.cloud.spanner.SpannerException; import com.google.cloud.spanner.SpannerExceptionFactory; -import com.google.cloud.spanner.TimestampBound; import com.google.cloud.spanner.TransactionMutationLimitExceededException; import com.google.cloud.spanner.TransactionRunner; import com.google.cloud.spanner.connection.AbstractStatementParser.ParsedStatement; @@ -55,10 +61,10 @@ import com.google.spanner.admin.database.v1.DatabaseAdminGrpc; import com.google.spanner.v1.SpannerGrpc; import io.opentelemetry.context.Scope; -import java.time.Duration; import java.util.Arrays; import java.util.UUID; import java.util.concurrent.Callable; +import java.util.concurrent.atomic.AtomicReference; import javax.annotation.Nonnull; /** @@ -78,15 +84,11 @@ * */ class SingleUseTransaction extends AbstractBaseUnitOfWork { - private final boolean readOnly; private final DdlClient ddlClient; private final DatabaseClient dbClient; private final BatchClient batchClient; - private final TimestampBound readOnlyStaleness; - private final AutocommitDmlMode autocommitDmlMode; - private final boolean returnCommitStats; - private final Duration maxCommitDelay; - private final boolean internalMetdataQuery; + private final ConnectionState connectionState; + private final boolean internalMetadataQuery; private final byte[] protoDescriptors; private volatile SettableApiFuture readTimestamp = null; private volatile TransactionRunner writeTransaction; @@ -97,11 +99,7 @@ static class Builder extends AbstractBaseUnitOfWork.Builder executeQueryAsync( // Do not use a read-only staleness for internal metadata queries. final ReadOnlyTransaction currentTransaction = - internalMetdataQuery + internalMetadataQuery ? dbClient.singleUseReadOnlyTransaction() - : dbClient.singleUseReadOnlyTransaction(readOnlyStaleness); + : dbClient.singleUseReadOnlyTransaction( + connectionState.getValue(READ_ONLY_STALENESS).getValue()); Callable callable = () -> { try { @@ -325,7 +300,8 @@ public ApiFuture partitionQueryAsync( Callable callable = () -> { try (BatchReadOnlyTransaction transaction = - batchClient.batchReadOnlyTransaction(readOnlyStaleness)) { + batchClient.batchReadOnlyTransaction( + connectionState.getValue(READ_ONLY_STALENESS).getValue())) { ResultSet resultSet = partitionQuery(transaction, partitionOptions, query, options); readTimestamp.set(transaction.getReadTimestamp()); state = UnitOfWorkState.COMMITTED; @@ -408,15 +384,19 @@ public ApiFuture executeDdlAsync(CallType callType, final ParsedStatement Callable callable = () -> { try { - OperationFuture operation; - if (DdlClient.isCreateDatabaseStatement(ddl.getSqlWithoutComments())) { - operation = - ddlClient.executeCreateDatabase( - ddl.getSqlWithoutComments(), dbClient.getDialect()); + if (isCreateDatabaseStatement(ddl.getSqlWithoutComments())) { + executeCreateDatabase(ddl); } else { - operation = ddlClient.executeDdl(ddl.getSqlWithoutComments(), protoDescriptors); + ddlClient.runWithRetryForMissingDefaultSequenceKind( + restartIndex -> { + OperationFuture operation = + ddlClient.executeDdl(ddl.getSqlWithoutComments(), protoDescriptors); + getWithStatementTimeout(operation, ddl); + }, + connectionState.getValue(DEFAULT_SEQUENCE_KIND).getValue(), + dbClient.getDialect(), + new AtomicReference<>()); } - getWithStatementTimeout(operation, ddl); state = UnitOfWorkState.COMMITTED; return null; } catch (Throwable t) { @@ -429,6 +409,12 @@ public ApiFuture executeDdlAsync(CallType callType, final ParsedStatement } } + private void executeCreateDatabase(ParsedStatement ddl) { + OperationFuture operation = + ddlClient.executeCreateDatabase(ddl.getSqlWithoutComments(), dbClient.getDialect()); + getWithStatementTimeout(operation, ddl); + } + @Override public ApiFuture executeUpdateAsync( CallType callType, ParsedStatement update, UpdateOption... options) { @@ -440,7 +426,7 @@ public ApiFuture executeUpdateAsync( checkAndMarkUsed(); ApiFuture res; - switch (autocommitDmlMode) { + switch (getAutocommitDmlMode()) { case TRANSACTIONAL: case TRANSACTIONAL_WITH_FALLBACK_TO_PARTITIONED_NON_ATOMIC: res = @@ -454,7 +440,7 @@ public ApiFuture executeUpdateAsync( break; default: throw SpannerExceptionFactory.newSpannerException( - ErrorCode.FAILED_PRECONDITION, "Unknown dml mode: " + autocommitDmlMode); + ErrorCode.FAILED_PRECONDITION, "Unknown dml mode: " + getAutocommitDmlMode()); } return res; } @@ -468,7 +454,7 @@ public ApiFuture analyzeUpdateAsync( ConnectionPreconditions.checkState( !isReadOnly(), "Update statements are not allowed in read-only mode"); ConnectionPreconditions.checkState( - autocommitDmlMode != AutocommitDmlMode.PARTITIONED_NON_ATOMIC, + getAutocommitDmlMode() != AutocommitDmlMode.PARTITIONED_NON_ATOMIC, "Analyzing update statements is not supported for Partitioned DML"); try (Scope ignore = span.makeCurrent()) { checkAndMarkUsed(); @@ -494,16 +480,16 @@ public ApiFuture executeBatchUpdateAsync( try (Scope ignore = span.makeCurrent()) { checkAndMarkUsed(); - switch (autocommitDmlMode) { + switch (getAutocommitDmlMode()) { case TRANSACTIONAL: return executeTransactionalBatchUpdateAsync(callType, updates, options); case PARTITIONED_NON_ATOMIC: throw SpannerExceptionFactory.newSpannerException( ErrorCode.FAILED_PRECONDITION, - "Batch updates are not allowed in " + autocommitDmlMode); + "Batch updates are not allowed in " + getAutocommitDmlMode()); default: throw SpannerExceptionFactory.newSpannerException( - ErrorCode.FAILED_PRECONDITION, "Unknown dml mode: " + autocommitDmlMode); + ErrorCode.FAILED_PRECONDITION, "Unknown dml mode: " + getAutocommitDmlMode()); } } } @@ -513,13 +499,13 @@ private TransactionRunner createWriteTransaction() { if (this.rpcPriority != null) { numOptions++; } - if (returnCommitStats) { + if (connectionState.getValue(RETURN_COMMIT_STATS).getValue()) { numOptions++; } if (excludeTxnFromChangeStreams) { numOptions++; } - if (maxCommitDelay != null) { + if (connectionState.getValue(MAX_COMMIT_DELAY).getValue() != null) { numOptions++; } if (numOptions == 0) { @@ -530,14 +516,15 @@ private TransactionRunner createWriteTransaction() { if (this.rpcPriority != null) { options[index++] = Options.priority(this.rpcPriority); } - if (returnCommitStats) { + if (connectionState.getValue(RETURN_COMMIT_STATS).getValue()) { options[index++] = Options.commitStats(); } if (excludeTxnFromChangeStreams) { options[index++] = Options.excludeTxnFromChangeStreams(); } - if (maxCommitDelay != null) { - options[index++] = Options.maxCommitDelay(maxCommitDelay); + if (connectionState.getValue(MAX_COMMIT_DELAY).getValue() != null) { + options[index++] = + Options.maxCommitDelay(connectionState.getValue(MAX_COMMIT_DELAY).getValue()); } return dbClient.readWriteTransaction(options); } diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/DdlBatchTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/DdlBatchTest.java index 93ae60891fb..4d582e09007 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/DdlBatchTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/DdlBatchTest.java @@ -30,6 +30,7 @@ import static org.mockito.Mockito.anyList; import static org.mockito.Mockito.anyString; import static org.mockito.Mockito.argThat; +import static org.mockito.Mockito.doCallRealMethod; import static org.mockito.Mockito.doThrow; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.never; @@ -59,6 +60,7 @@ import java.io.InputStream; import java.util.Arrays; import java.util.Collections; +import java.util.HashMap; import java.util.List; import java.util.concurrent.ExecutionException; import java.util.concurrent.Executors; @@ -115,6 +117,9 @@ private DdlClient createDefaultMockDdlClient( when(operation.getMetadata()).thenReturn(metadataFuture); when(ddlClient.executeDdl(anyString(), any())).thenReturn(operation); when(ddlClient.executeDdl(anyList(), any())).thenReturn(operation); + doCallRealMethod() + .when(ddlClient) + .runWithRetryForMissingDefaultSequenceKind(any(), any(), any(), any()); return ddlClient; } catch (Exception e) { throw new RuntimeException(e); @@ -135,6 +140,7 @@ private DdlBatch createSubject(DdlClient ddlClient, DatabaseClient dbClient) { .setDatabaseClient(dbClient) .withStatementExecutor(new StatementExecutor()) .setSpan(Span.getInvalid()) + .setConnectionState(new ConnectionState(new HashMap<>())) .build(); } @@ -256,9 +262,12 @@ public void testGetStateAndIsActive() { assertThat(batch.isActive(), is(false)); DdlClient client = mock(DdlClient.class); - SpannerException exception = mock(SpannerException.class); - when(exception.getErrorCode()).thenReturn(ErrorCode.FAILED_PRECONDITION); + SpannerException exception = + SpannerExceptionFactory.newSpannerException(ErrorCode.FAILED_PRECONDITION, "test"); doThrow(exception).when(client).executeDdl(anyList(), isNull()); + doCallRealMethod() + .when(client) + .runWithRetryForMissingDefaultSequenceKind(any(), any(), any(), any()); batch = createSubject(client); assertThat(batch.getState(), is(UnitOfWorkState.STARTED)); assertThat(batch.isActive(), is(true)); @@ -380,6 +389,7 @@ public void testRunBatch() { .withStatementExecutor(new StatementExecutor()) .setSpan(Span.getInvalid()) .setProtoDescriptors(null) + .setConnectionState(new ConnectionState(new HashMap<>())) .build(); batch.executeDdlAsync(CallType.SYNC, statement); batch.executeDdlAsync(CallType.SYNC, statement); @@ -406,6 +416,7 @@ public void testRunBatch() { .withStatementExecutor(new StatementExecutor()) .setSpan(Span.getInvalid()) .setProtoDescriptors(protoDescriptors) + .setConnectionState(new ConnectionState(new HashMap<>())) .build(); batch.executeDdlAsync(CallType.SYNC, statement); batch.executeDdlAsync(CallType.SYNC, statement); @@ -437,6 +448,7 @@ public void testUpdateCount() throws InterruptedException, ExecutionException { .setDdlClient(client) .setDatabaseClient(mock(DatabaseClient.class)) .setSpan(Span.getInvalid()) + .setConnectionState(new ConnectionState(new HashMap<>())) .build(); batch.executeDdlAsync( CallType.SYNC, @@ -469,6 +481,9 @@ public void testFailedUpdateCount() throws InterruptedException, ExecutionExcept new ExecutionException( "ddl statement failed", Status.INVALID_ARGUMENT.asRuntimeException())); when(operationFuture.getMetadata()).thenReturn(metadataFuture); + doCallRealMethod() + .when(client) + .runWithRetryForMissingDefaultSequenceKind(any(), any(), any(), any()); when(client.executeDdl(argThat(isListOfStringsWithSize(2)), isNull())) .thenReturn(operationFuture); DdlBatch batch = @@ -477,6 +492,7 @@ public void testFailedUpdateCount() throws InterruptedException, ExecutionExcept .setDdlClient(client) .setDatabaseClient(mock(DatabaseClient.class)) .setSpan(Span.getInvalid()) + .setConnectionState(new ConnectionState(new HashMap<>())) .build(); batch.executeDdlAsync( CallType.SYNC, @@ -499,6 +515,9 @@ public void testFailedUpdateCount() throws InterruptedException, ExecutionExcept @Test public void testFailedAfterFirstStatement() throws InterruptedException, ExecutionException { DdlClient client = mock(DdlClient.class); + doCallRealMethod() + .when(client) + .runWithRetryForMissingDefaultSequenceKind(any(), any(), any(), any()); UpdateDatabaseDdlMetadata metadata = UpdateDatabaseDdlMetadata.newBuilder() .addCommitTimestamps( @@ -521,6 +540,7 @@ public void testFailedAfterFirstStatement() throws InterruptedException, Executi .setDdlClient(client) .setDatabaseClient(mock(DatabaseClient.class)) .setSpan(Span.getInvalid()) + .setConnectionState(new ConnectionState(new HashMap<>())) .build(); batch.executeDdlAsync( CallType.SYNC, diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/DdlTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/DdlTest.java index 44a2f4d9ff7..209796b9873 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/DdlTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/DdlTest.java @@ -17,15 +17,20 @@ package com.google.cloud.spanner.connection; import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNull; import static org.junit.Assert.assertThrows; import com.google.cloud.spanner.ErrorCode; +import com.google.cloud.spanner.MissingDefaultSequenceKindException; +import com.google.cloud.spanner.SpannerBatchUpdateException; import com.google.cloud.spanner.SpannerException; import com.google.cloud.spanner.Statement; import com.google.cloud.spanner.connection.StatementResult.ResultType; import com.google.longrunning.Operation; +import com.google.protobuf.AbstractMessage; import com.google.protobuf.Any; import com.google.protobuf.Empty; +import com.google.rpc.Code; import com.google.spanner.admin.database.v1.UpdateDatabaseDdlMetadata; import com.google.spanner.admin.database.v1.UpdateDatabaseDdlRequest; import com.google.spanner.v1.CommitRequest; @@ -65,6 +70,21 @@ private void addUpdateDdlResponse() { .build()); } + private void addUpdateDdlResponse(com.google.rpc.Status error) { + mockDatabaseAdmin.addResponse( + Operation.newBuilder() + .setMetadata( + Any.pack( + UpdateDatabaseDdlMetadata.newBuilder() + .setDatabase("projects/proj/instances/inst/databases/db") + .build())) + .setName("projects/proj/instances/inst/databases/db/operations/1") + .setDone(true) + // .setResponse(Any.pack(Empty.getDefaultInstance())) + .setError(error) + .build()); + } + @Test public void testSingleAnalyzeStatement() { addUpdateDdlResponse(); @@ -230,4 +250,114 @@ public void testDdlBatchInTransaction() { } } } + + @Test + public void testMissingDefaultSequenceKindException() { + addUpdateDdlResponse( + com.google.rpc.Status.newBuilder() + .setCode(Code.INVALID_ARGUMENT_VALUE) + .setMessage( + "The sequence kind of an identity column id2 is not specified. " + + "Please specify the sequence kind explicitly or set the database option `default_sequence_kind`.") + .build()); + try (Connection connection = createConnection()) { + assertNull(connection.getDefaultSequenceKind()); + assertThrows( + MissingDefaultSequenceKindException.class, + () -> + connection.execute( + Statement.of("create table foo (id2 int64 auto_increment primary key"))); + } + // The request should not be retried. + assertEquals(1, mockDatabaseAdmin.getRequests().size()); + } + + @Test + public void testSetsDefaultSequenceKindAndRetriesStatement() { + addUpdateDdlResponse( + com.google.rpc.Status.newBuilder() + .setCode(Code.INVALID_ARGUMENT_VALUE) + .setMessage( + "The sequence kind of an identity column id2 is not specified. " + + "Please specify the sequence kind explicitly or set the database option `default_sequence_kind`.") + .build()); + // This will be the response for the 'alter database' statement. + addUpdateDdlResponse(); + // This will be the response for the 'create table' statement after the retry. + addUpdateDdlResponse(); + try (Connection connection = createConnection()) { + connection.setDefaultSequenceKind("bit_reversed_positive"); + connection.execute(Statement.of("create table foo (id2 int64 auto_increment primary key")); + } + List requests = mockDatabaseAdmin.getRequests(); + assertEquals(3, requests.size()); + assertEquals( + "create table foo (id2 int64 auto_increment primary key", + ((UpdateDatabaseDdlRequest) requests.get(0)).getStatements(0)); + assertEquals( + "alter database `db` set options (default_sequence_kind='bit_reversed_positive')", + ((UpdateDatabaseDdlRequest) requests.get(1)).getStatements(0)); + assertEquals( + "create table foo (id2 int64 auto_increment primary key", + ((UpdateDatabaseDdlRequest) requests.get(2)).getStatements(0)); + } + + @Test + public void testMissingDefaultSequenceKindExceptionInBatch() { + addUpdateDdlResponse( + com.google.rpc.Status.newBuilder() + .setCode(Code.INVALID_ARGUMENT_VALUE) + .setMessage( + "The sequence kind of an identity column id2 is not specified. " + + "Please specify the sequence kind explicitly or set the database option `default_sequence_kind`.") + .build()); + try (Connection connection = createConnection()) { + assertNull(connection.getDefaultSequenceKind()); + connection.startBatchDdl(); + connection.execute(Statement.of("create table foo (id2 int64 auto_increment primary key")); + SpannerBatchUpdateException exception = + assertThrows(SpannerBatchUpdateException.class, connection::runBatch); + } + // The request should not be retried. + assertEquals(1, mockDatabaseAdmin.getRequests().size()); + } + + @Test + public void testSetsDefaultSequenceKindAndRetriesBatch() { + addUpdateDdlResponse( + com.google.rpc.Status.newBuilder() + .setCode(Code.INVALID_ARGUMENT_VALUE) + .setMessage( + "The sequence kind of an identity column id2 is not specified. " + + "Please specify the sequence kind explicitly or set the database option `default_sequence_kind`.") + .build()); + // This will be the response for the 'alter database' statement. + addUpdateDdlResponse(); + // This will be the response for the 'create table' statements after the retry. + addUpdateDdlResponse(); + try (Connection connection = createConnection()) { + connection.setDefaultSequenceKind("bit_reversed_positive"); + connection.startBatchDdl(); + connection.execute(Statement.of("create table foo (id1 int64 auto_increment primary key")); + connection.execute(Statement.of("create table bar (id2 int64 auto_increment primary key")); + connection.runBatch(); + } + List requests = mockDatabaseAdmin.getRequests(); + assertEquals(3, requests.size()); + assertEquals( + "create table foo (id1 int64 auto_increment primary key", + ((UpdateDatabaseDdlRequest) requests.get(0)).getStatements(0)); + assertEquals( + "create table bar (id2 int64 auto_increment primary key", + ((UpdateDatabaseDdlRequest) requests.get(0)).getStatements(1)); + assertEquals( + "alter database `db` set options (default_sequence_kind='bit_reversed_positive')", + ((UpdateDatabaseDdlRequest) requests.get(1)).getStatements(0)); + assertEquals( + "create table foo (id1 int64 auto_increment primary key", + ((UpdateDatabaseDdlRequest) requests.get(0)).getStatements(0)); + assertEquals( + "create table bar (id2 int64 auto_increment primary key", + ((UpdateDatabaseDdlRequest) requests.get(0)).getStatements(1)); + } } diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/SingleUseTransactionTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/SingleUseTransactionTest.java index 0e2a322023d..5ac926d5956 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/SingleUseTransactionTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/SingleUseTransactionTest.java @@ -17,6 +17,9 @@ package com.google.cloud.spanner.connection; import static com.google.cloud.spanner.SpannerApiFutures.get; +import static com.google.cloud.spanner.connection.ConnectionProperties.AUTOCOMMIT_DML_MODE; +import static com.google.cloud.spanner.connection.ConnectionProperties.READONLY; +import static com.google.cloud.spanner.connection.ConnectionProperties.READ_ONLY_STALENESS; import static com.google.common.truth.Truth.assertThat; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertNotNull; @@ -25,6 +28,7 @@ import static org.mockito.Mockito.any; import static org.mockito.Mockito.anyList; import static org.mockito.Mockito.anyString; +import static org.mockito.Mockito.doCallRealMethod; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; @@ -55,6 +59,7 @@ import com.google.cloud.spanner.TransactionRunner; import com.google.cloud.spanner.connection.AbstractStatementParser.ParsedStatement; import com.google.cloud.spanner.connection.AbstractStatementParser.StatementType; +import com.google.cloud.spanner.connection.ConnectionProperty.Context; import com.google.cloud.spanner.connection.StatementExecutor.StatementTimeout; import com.google.cloud.spanner.connection.UnitOfWork.CallType; import com.google.common.base.Preconditions; @@ -66,6 +71,7 @@ import java.util.Arrays; import java.util.Calendar; import java.util.Collections; +import java.util.HashMap; import java.util.List; import java.util.concurrent.TimeUnit; import org.junit.Test; @@ -306,6 +312,9 @@ private DdlClient createDefaultMockDdlClient() { when(operation.get()).thenReturn(null); when(ddlClient.executeDdl(anyString(), any())).thenCallRealMethod(); when(ddlClient.executeDdl(anyList(), any())).thenReturn(operation); + doCallRealMethod() + .when(ddlClient) + .runWithRetryForMissingDefaultSequenceKind(any(), any(), any(), any()); return ddlClient; } catch (Exception e) { throw new RuntimeException(e); @@ -418,6 +427,11 @@ private SingleUseTransaction createSubject( .thenThrow( SpannerExceptionFactory.newSpannerException(ErrorCode.UNKNOWN, "invalid update")); + ConnectionState connectionState = new ConnectionState(new HashMap<>()); + connectionState.setValue(AUTOCOMMIT_DML_MODE, dmlMode, Context.STARTUP, false); + connectionState.setValue(READONLY, readOnly, Context.STARTUP, false); + connectionState.setValue(READ_ONLY_STALENESS, staleness, Context.STARTUP, false); + when(dbClient.readWriteTransaction()) .thenAnswer( new Answer() { @@ -473,9 +487,7 @@ public TransactionRunner allowNestedTransaction() { .setDatabaseClient(dbClient) .setBatchClient(mock(BatchClient.class)) .setDdlClient(ddlClient) - .setAutocommitDmlMode(dmlMode) - .setReadOnly(readOnly) - .setReadOnlyStaleness(staleness) + .setConnectionState(connectionState) .setStatementTimeout( timeout == 0L ? nullTimeout() : timeout(timeout, TimeUnit.MILLISECONDS)) .withStatementExecutor(executor) @@ -664,14 +676,18 @@ public void testExecuteQueryWithOptionsTest() { when(tx.executeQuery(Statement.of(sql), option)).thenReturn(mock(ResultSet.class)); when(client.singleUseReadOnlyTransaction(TimestampBound.strong())).thenReturn(tx); + ConnectionState connectionState = new ConnectionState(new HashMap<>()); + connectionState.setValue( + AUTOCOMMIT_DML_MODE, AutocommitDmlMode.TRANSACTIONAL, Context.STARTUP, false); + connectionState.setValue(READ_ONLY_STALENESS, TimestampBound.strong(), Context.STARTUP, false); + SingleUseTransaction transaction = SingleUseTransaction.newBuilder() .setDatabaseClient(client) .setBatchClient(mock(BatchClient.class)) .setDdlClient(mock(DdlClient.class)) - .setAutocommitDmlMode(AutocommitDmlMode.TRANSACTIONAL) + .setConnectionState(connectionState) .withStatementExecutor(executor) - .setReadOnlyStaleness(TimestampBound.strong()) .setSpan(Span.getInvalid()) .build(); assertThat( diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/it/ITDdlTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/it/ITDdlTest.java index 7a9c5aa9262..affc7ad2a18 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/it/ITDdlTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/it/ITDdlTest.java @@ -16,16 +16,31 @@ package com.google.cloud.spanner.connection.it; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertNull; import static org.junit.Assert.assertThrows; +import static org.junit.Assert.assertTrue; +import com.google.cloud.spanner.Database; import com.google.cloud.spanner.DatabaseAdminClient; import com.google.cloud.spanner.DatabaseNotFoundException; +import com.google.cloud.spanner.Dialect; +import com.google.cloud.spanner.MissingDefaultSequenceKindException; import com.google.cloud.spanner.ParallelIntegrationTest; +import com.google.cloud.spanner.ResultSet; +import com.google.cloud.spanner.SpannerBatchUpdateException; import com.google.cloud.spanner.Statement; import com.google.cloud.spanner.connection.Connection; +import com.google.cloud.spanner.connection.ConnectionOptions; import com.google.cloud.spanner.connection.ITAbstractSpannerTest; import com.google.cloud.spanner.connection.SqlScriptVerifier; +import com.google.cloud.spanner.testing.EmulatorSpannerHelper; +import java.util.Arrays; +import java.util.Collections; +import org.junit.Before; +import org.junit.BeforeClass; import org.junit.Test; import org.junit.experimental.categories.Category; import org.junit.runner.RunWith; @@ -35,6 +50,16 @@ @Category(ParallelIntegrationTest.class) @RunWith(JUnit4.class) public class ITDdlTest extends ITAbstractSpannerTest { + @BeforeClass + public static void setup() { + // This overrides the default behavior that creates a single database for the test class. This + // test needs a separate database per method. + } + + @Before + public void createTestDatabase() { + database = env.getTestHelper().createTestDatabase(); + } @Test public void testSqlScript() throws Exception { @@ -57,4 +82,131 @@ public void testCreateDatabase() { client.dropDatabase(instance, name); } } + + @Test + public void testDefaultSequenceKind() { + try (Connection connection = createConnection()) { + Statement statement = + Statement.of( + "create table test (id int64 auto_increment primary key, value string(max))"); + + // Creating a table with an auto_increment column fails if no default sequence kind has been + // set. + assertNull(connection.getDefaultSequenceKind()); + assertThrows(MissingDefaultSequenceKindException.class, () -> connection.execute(statement)); + + // Setting a default sequence kind on the connection should make the statement succeed. + connection.setDefaultSequenceKind("bit_reversed_positive"); + connection.execute(statement); + + assertEquals( + 1L, connection.executeUpdate(Statement.of("insert into test (value) values ('One')"))); + try (ResultSet resultSet = connection.executeQuery(Statement.of("select * from test"))) { + assertTrue(resultSet.next()); + assertEquals("One", resultSet.getString(1)); + assertFalse(resultSet.next()); + } + } + } + + @Test + public void testDefaultSequenceKind_PostgreSQL() throws Exception { + DatabaseAdminClient client = getTestEnv().getTestHelper().getClient().getDatabaseAdminClient(); + String instance = getTestEnv().getTestHelper().getInstanceId().getInstance(); + String name = getTestEnv().getTestHelper().getUniqueDatabaseId(); + + Database database = + client + .createDatabase( + instance, + "create database \"" + name + "\"", + Dialect.POSTGRESQL, + Collections.emptyList()) + .get(); + + StringBuilder url = extractConnectionUrl(getTestEnv().getTestHelper().getOptions(), database); + ConnectionOptions.Builder builder = ConnectionOptions.newBuilder().setUri(url.toString()); + if (hasValidKeyFile()) { + builder.setCredentialsUrl(getKeyFile()); + } + ConnectionOptions options = builder.build(); + + try (Connection connection = options.getConnection()) { + Statement statement = + Statement.of("create table test (id serial primary key, value varchar)"); + + // Creating a table with an auto_increment column fails if no default sequence kind has been + // set. + assertNull(connection.getDefaultSequenceKind()); + assertThrows(MissingDefaultSequenceKindException.class, () -> connection.execute(statement)); + + // Setting a default sequence kind on the connection should make the statement succeed. + connection.setDefaultSequenceKind("bit_reversed_positive"); + connection.execute(statement); + + assertEquals( + 1L, connection.executeUpdate(Statement.of("insert into test (value) values ('One')"))); + try (ResultSet resultSet = connection.executeQuery(Statement.of("select * from test"))) { + assertTrue(resultSet.next()); + assertEquals("One", resultSet.getString(1)); + assertFalse(resultSet.next()); + } + } finally { + client.dropDatabase(instance, name); + } + } + + @Test + public void testDefaultSequenceKindInBatch() { + try (Connection connection = createConnection()) { + Statement statement1 = + Statement.of("create table testseq1 (id1 int64 primary key, value string(max))"); + Statement statement2 = + Statement.of( + "create table testseq2 (id2 int64 auto_increment primary key, value string(max))"); + + // Creating a table with an auto_increment column fails if no default sequence kind has been + // set. + assertNull(connection.getDefaultSequenceKind()); + connection.startBatchDdl(); + connection.execute(statement1); + connection.execute(statement2); + SpannerBatchUpdateException exception = + assertThrows(SpannerBatchUpdateException.class, connection::runBatch); + long updateCount = Arrays.stream(exception.getUpdateCounts()).sum(); + // The emulator refuses the entire batch. Spanner executes the first statement and fails on + // the second statement. + if (EmulatorSpannerHelper.isUsingEmulator()) { + assertEquals(0, updateCount); + } else { + assertEquals(1, updateCount); + } + + // Setting a default sequence kind on the connection should make the statement succeed. + connection.setDefaultSequenceKind("bit_reversed_positive"); + connection.startBatchDdl(); + if (updateCount == 0) { + connection.execute(statement1); + } + connection.execute(statement2); + connection.runBatch(); + } + } + + @Test + public void testDefaultSequenceKindRetriesBatchCorrectly() { + try (Connection connection = createConnection()) { + Statement statement1 = + Statement.of("create table testseq1 (id1 int64 primary key, value string(max))"); + Statement statement2 = + Statement.of( + "create table testseq2 (id2 int64 auto_increment primary key, value string(max))"); + + connection.setDefaultSequenceKind("bit_reversed_positive"); + connection.startBatchDdl(); + connection.execute(statement1); + connection.execute(statement2); + connection.runBatch(); + } + } }