diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractReadContext.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractReadContext.java index 1ef2cf03f1..685e9a1e31 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractReadContext.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractReadContext.java @@ -30,11 +30,13 @@ import com.google.cloud.spanner.Options.ReadOption; import com.google.cloud.spanner.SessionImpl.SessionTransaction; import com.google.cloud.spanner.spi.v1.SpannerRpc; +import com.google.common.annotations.VisibleForTesting; import com.google.protobuf.ByteString; import com.google.spanner.v1.BeginTransactionRequest; import com.google.spanner.v1.ExecuteBatchDmlRequest; import com.google.spanner.v1.ExecuteSqlRequest; import com.google.spanner.v1.ExecuteSqlRequest.QueryMode; +import com.google.spanner.v1.ExecuteSqlRequest.QueryOptions; import com.google.spanner.v1.PartialResultSet; import com.google.spanner.v1.ReadRequest; import com.google.spanner.v1.Transaction; @@ -53,20 +55,86 @@ */ abstract class AbstractReadContext implements ReadContext, AbstractResultSet.Listener, SessionTransaction { + + abstract static class Builder, T extends AbstractReadContext> { + private SessionImpl session; + private SpannerRpc rpc; + private Span span = Tracing.getTracer().getCurrentSpan(); + private int defaultPrefetchChunks = SpannerOptions.Builder.DEFAULT_PREFETCH_CHUNKS; + private QueryOptions defaultQueryOptions = SpannerOptions.Builder.DEFAULT_QUERY_OPTIONS; + + Builder() {} + + @SuppressWarnings("unchecked") + B self() { + return (B) this; + } + + B setSession(SessionImpl session) { + this.session = session; + return self(); + } + + B setRpc(SpannerRpc rpc) { + this.rpc = rpc; + return self(); + } + + B setSpan(Span span) { + this.span = span; + return self(); + } + + B setDefaultPrefetchChunks(int defaultPrefetchChunks) { + this.defaultPrefetchChunks = defaultPrefetchChunks; + return self(); + } + + B setDefaultQueryOptions(QueryOptions defaultQueryOptions) { + this.defaultQueryOptions = defaultQueryOptions; + return self(); + } + + abstract T build(); + } + /** * A {@code ReadContext} for standalone reads. This can only be used for a single operation, since * each standalone read may see a different timestamp of Cloud Spanner data. */ static class SingleReadContext extends AbstractReadContext { + static class Builder extends AbstractReadContext.Builder { + private TimestampBound bound; + + private Builder() {} + + Builder setTimestampBound(TimestampBound bound) { + this.bound = bound; + return self(); + } + + @Override + SingleReadContext build() { + return new SingleReadContext(this); + } + + SingleUseReadOnlyTransaction buildSingleUseReadOnlyTransaction() { + return new SingleUseReadOnlyTransaction(this); + } + } + + static Builder newBuilder() { + return new Builder(); + } + final TimestampBound bound; @GuardedBy("lock") private boolean used; - SingleReadContext( - SessionImpl session, TimestampBound bound, SpannerRpc rpc, int defaultPrefetchChunks) { - super(session, rpc, defaultPrefetchChunks); - this.bound = bound; + private SingleReadContext(Builder builder) { + super(builder); + this.bound = builder.bound; } @GuardedBy("lock") @@ -99,9 +167,8 @@ static class SingleUseReadOnlyTransaction extends SingleReadContext @GuardedBy("lock") private Timestamp timestamp; - SingleUseReadOnlyTransaction( - SessionImpl session, TimestampBound bound, SpannerRpc rpc, int defaultPrefetchChunks) { - super(session, bound, rpc, defaultPrefetchChunks); + private SingleUseReadOnlyTransaction(SingleReadContext.Builder builder) { + super(builder); } @Override @@ -139,6 +206,38 @@ public void onTransactionMetadata(Transaction transaction) { static class MultiUseReadOnlyTransaction extends AbstractReadContext implements ReadOnlyTransaction { + static class Builder extends AbstractReadContext.Builder { + private TimestampBound bound; + private Timestamp timestamp; + private ByteString transactionId; + + private Builder() {} + + Builder setTimestampBound(TimestampBound bound) { + this.bound = bound; + return this; + } + + Builder setTimestamp(Timestamp timestamp) { + this.timestamp = timestamp; + return this; + } + + Builder setTransactionId(ByteString transactionId) { + this.transactionId = transactionId; + return this; + } + + @Override + MultiUseReadOnlyTransaction build() { + return new MultiUseReadOnlyTransaction(this); + } + } + + static Builder newBuilder() { + return new Builder(); + } + private TimestampBound bound; private final Object txnLock = new Object(); @@ -148,27 +247,24 @@ static class MultiUseReadOnlyTransaction extends AbstractReadContext @GuardedBy("txnLock") private ByteString transactionId; - MultiUseReadOnlyTransaction( - SessionImpl session, TimestampBound bound, SpannerRpc rpc, int defaultPrefetchChunks) { - super(session, rpc, defaultPrefetchChunks); + MultiUseReadOnlyTransaction(Builder builder) { + super(builder); checkArgument( - bound.getMode() != TimestampBound.Mode.MAX_STALENESS - && bound.getMode() != TimestampBound.Mode.MIN_READ_TIMESTAMP, - "Bounded staleness mode %s is not supported for multi-use read-only transactions." - + " Create a single-use read or read-only transaction instead.", - bound.getMode()); - this.bound = bound; - } - - MultiUseReadOnlyTransaction( - SessionImpl session, - ByteString transactionId, - Timestamp timestamp, - SpannerRpc rpc, - int defaultPrefetchChunks) { - super(session, rpc, defaultPrefetchChunks); - this.transactionId = transactionId; - this.timestamp = timestamp; + !(builder.bound != null && builder.transactionId != null) + && !(builder.bound == null && builder.transactionId == null), + "Either TimestampBound or TransactionId must be specified"); + if (builder.bound != null) { + checkArgument( + builder.bound.getMode() != TimestampBound.Mode.MAX_STALENESS + && builder.bound.getMode() != TimestampBound.Mode.MIN_READ_TIMESTAMP, + "Bounded staleness mode %s is not supported for multi-use read-only transactions." + + " Create a single-use read or read-only transaction instead.", + builder.bound.getMode()); + this.bound = builder.bound; + } else { + this.timestamp = builder.timestamp; + this.transactionId = builder.transactionId; + } } @Override @@ -256,6 +352,7 @@ void initTransaction() { final SpannerRpc rpc; final Span span; private final int defaultPrefetchChunks; + private final QueryOptions defaultQueryOptions; @GuardedBy("lock") private boolean isValid = true; @@ -271,16 +368,12 @@ void initTransaction() { // much more frequently. private static final int MAX_BUFFERED_CHUNKS = 512; - AbstractReadContext(SessionImpl session, SpannerRpc rpc, int defaultPrefetchChunks) { - this(session, rpc, defaultPrefetchChunks, Tracing.getTracer().getCurrentSpan()); - } - - private AbstractReadContext( - SessionImpl session, SpannerRpc rpc, int defaultPrefetchChunks, Span span) { - this.session = session; - this.rpc = rpc; - this.defaultPrefetchChunks = defaultPrefetchChunks; - this.span = span; + AbstractReadContext(Builder builder) { + this.session = builder.session; + this.rpc = builder.rpc; + this.defaultPrefetchChunks = builder.defaultPrefetchChunks; + this.defaultQueryOptions = builder.defaultQueryOptions; + this.span = builder.span; } long getSeqNo() { @@ -341,9 +434,36 @@ private ResultSet executeQueryInternal( Statement statement, com.google.spanner.v1.ExecuteSqlRequest.QueryMode queryMode, QueryOption... options) { - Options readOptions = Options.fromQueryOptions(options); + Options queryOptions = Options.fromQueryOptions(options); return executeQueryInternalWithOptions( - statement, queryMode, readOptions, null /*partitionToken*/); + statement, queryMode, queryOptions, null /*partitionToken*/); + } + + /** + * Determines the {@link QueryOptions} to use for a query. This is determined using the following + * precedence: + * + *
    + *
  1. Specific {@link QueryOptions} passed in for this query. + *
  2. Any value specified in a valid environment variable when the {@link SpannerOptions} + * instance was created. + *
  3. The default {@link SpannerOptions#getDefaultQueryOptions()} specified for the database + * where the query is executed. + *
+ */ + @VisibleForTesting + QueryOptions buildQueryOptions(QueryOptions requestOptions) { + // Shortcut for the most common return value. + if (defaultQueryOptions.equals(QueryOptions.getDefaultInstance()) && requestOptions == null) { + return QueryOptions.getDefaultInstance(); + } + // Create a builder based on the default query options. + QueryOptions.Builder builder = defaultQueryOptions.toBuilder(); + // Then overwrite with specific options for this query. + if (requestOptions != null) { + builder.mergeFrom(requestOptions); + } + return builder.build(); } ExecuteSqlRequest.Builder getExecuteSqlRequestBuilder(Statement statement, QueryMode queryMode) { @@ -365,6 +485,7 @@ ExecuteSqlRequest.Builder getExecuteSqlRequestBuilder(Statement statement, Query builder.setTransaction(selector); } builder.setSeqno(getSeqNo()); + builder.setQueryOptions(buildQueryOptions(statement.getQueryOptions())); return builder; } @@ -400,7 +521,7 @@ ExecuteBatchDmlRequest.Builder getExecuteBatchDmlRequestBuilder(Iterable options; - BatchReadOnlyTransactionImpl(SpannerImpl spanner, SessionImpl session, TimestampBound bound) { - super( - checkNotNull(session), - checkNotNull(bound), - checkNotNull(spanner).getOptions().getSpannerRpcV1(), - spanner.getOptions().getPrefetchChunks()); + BatchReadOnlyTransactionImpl( + MultiUseReadOnlyTransaction.Builder builder, TimestampBound bound) { + super(builder.setTimestampBound(bound)); this.sessionName = session.getName(); this.options = session.getOptions(); initTransaction(); } BatchReadOnlyTransactionImpl( - SpannerImpl spanner, SessionImpl session, BatchTransactionId batchTransactionId) { - super( - checkNotNull(session), - checkNotNull(batchTransactionId).getTransactionId(), - batchTransactionId.getTimestamp(), - checkNotNull(spanner).getOptions().getSpannerRpcV1(), - spanner.getOptions().getPrefetchChunks()); + MultiUseReadOnlyTransaction.Builder builder, BatchTransactionId batchTransactionId) { + super(builder.setTransactionId(batchTransactionId.getTransactionId())); this.sessionName = session.getName(); this.options = session.getOptions(); } diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionClient.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionClient.java index 7cacaef497..d77a95f48b 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionClient.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionClient.java @@ -193,6 +193,10 @@ SpannerImpl getSpanner() { return spanner; } + DatabaseId getDatabaseId() { + return db; + } + /** Create a single session. */ SessionImpl createSession() { // The sessionChannelCounter could overflow, but that will just flip it to Integer.MIN_VALUE, diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionImpl.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionImpl.java index 50da01f31f..015e1862d6 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionImpl.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionImpl.java @@ -24,6 +24,7 @@ import com.google.cloud.spanner.AbstractReadContext.MultiUseReadOnlyTransaction; import com.google.cloud.spanner.AbstractReadContext.SingleReadContext; import com.google.cloud.spanner.AbstractReadContext.SingleUseReadOnlyTransaction; +import com.google.cloud.spanner.SessionClient.SessionId; import com.google.cloud.spanner.TransactionRunnerImpl.TransactionContextImpl; import com.google.cloud.spanner.spi.v1.SpannerRpc; import com.google.common.collect.Lists; @@ -79,6 +80,7 @@ static interface SessionTransaction { private final SpannerImpl spanner; private final String name; + private final DatabaseId databaseId; private SessionTransaction activeTransaction; private ByteString readyTransactionId; private final Map options; @@ -87,6 +89,7 @@ static interface SessionTransaction { this.spanner = spanner; this.options = options; this.name = checkNotNull(name); + this.databaseId = SessionId.of(name).getDatabaseId(); } @Override @@ -160,7 +163,13 @@ public ReadContext singleUse() { @Override public ReadContext singleUse(TimestampBound bound) { return setActive( - new SingleReadContext(this, bound, spanner.getRpc(), spanner.getDefaultPrefetchChunks())); + SingleReadContext.newBuilder() + .setSession(this) + .setTimestampBound(bound) + .setRpc(spanner.getRpc()) + .setDefaultQueryOptions(spanner.getDefaultQueryOptions(databaseId)) + .setDefaultPrefetchChunks(spanner.getDefaultPrefetchChunks()) + .build()); } @Override @@ -171,8 +180,13 @@ public ReadOnlyTransaction singleUseReadOnlyTransaction() { @Override public ReadOnlyTransaction singleUseReadOnlyTransaction(TimestampBound bound) { return setActive( - new SingleUseReadOnlyTransaction( - this, bound, spanner.getRpc(), spanner.getDefaultPrefetchChunks())); + SingleUseReadOnlyTransaction.newBuilder() + .setSession(this) + .setTimestampBound(bound) + .setRpc(spanner.getRpc()) + .setDefaultQueryOptions(spanner.getDefaultQueryOptions(databaseId)) + .setDefaultPrefetchChunks(spanner.getDefaultPrefetchChunks()) + .buildSingleUseReadOnlyTransaction()); } @Override @@ -183,8 +197,13 @@ public ReadOnlyTransaction readOnlyTransaction() { @Override public ReadOnlyTransaction readOnlyTransaction(TimestampBound bound) { return setActive( - new MultiUseReadOnlyTransaction( - this, bound, spanner.getRpc(), spanner.getDefaultPrefetchChunks())); + MultiUseReadOnlyTransaction.newBuilder() + .setSession(this) + .setTimestampBound(bound) + .setRpc(spanner.getRpc()) + .setDefaultQueryOptions(spanner.getDefaultQueryOptions(databaseId)) + .setDefaultPrefetchChunks(spanner.getDefaultPrefetchChunks()) + .build()); } @Override @@ -240,10 +259,13 @@ ByteString beginTransaction() { } TransactionContextImpl newTransaction() { - TransactionContextImpl txn = - new TransactionContextImpl( - this, readyTransactionId, spanner.getRpc(), spanner.getDefaultPrefetchChunks()); - return txn; + return TransactionContextImpl.newBuilder() + .setSession(this) + .setTransactionId(readyTransactionId) + .setRpc(spanner.getRpc()) + .setDefaultQueryOptions(spanner.getDefaultQueryOptions(databaseId)) + .setDefaultPrefetchChunks(spanner.getDefaultPrefetchChunks()) + .build(); } T setActive(@Nullable T ctx) { diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerImpl.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerImpl.java index bac7e753a1..08089c89e4 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerImpl.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerImpl.java @@ -31,6 +31,7 @@ import com.google.common.collect.ImmutableList; import com.google.common.util.concurrent.Futures; import com.google.common.util.concurrent.ListenableFuture; +import com.google.spanner.v1.ExecuteSqlRequest.QueryOptions; import io.opencensus.metrics.LabelValue; import io.opencensus.trace.Tracer; import io.opencensus.trace.Tracing; @@ -101,6 +102,11 @@ int getDefaultPrefetchChunks() { return getOptions().getPrefetchChunks(); } + /** Returns the default query options that should be used for the specified database. */ + QueryOptions getDefaultQueryOptions(DatabaseId databaseId) { + return getOptions().getDefaultQueryOptions(databaseId); + } + SessionImpl sessionWithId(String name) { Preconditions.checkArgument(!Strings.isNullOrEmpty(name), "name is null or empty"); SessionId id = SessionId.of(name); diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerOptions.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerOptions.java index a6aa502909..0ef5516d27 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerOptions.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerOptions.java @@ -25,6 +25,7 @@ import com.google.cloud.ServiceRpc; import com.google.cloud.TransportOptions; import com.google.cloud.grpc.GrpcTransportOptions; +import com.google.cloud.spanner.Options.QueryOption; import com.google.cloud.spanner.admin.database.v1.DatabaseAdminSettings; import com.google.cloud.spanner.admin.database.v1.stub.DatabaseAdminStubSettings; import com.google.cloud.spanner.admin.instance.v1.InstanceAdminSettings; @@ -34,21 +35,27 @@ import com.google.cloud.spanner.spi.v1.SpannerRpc; import com.google.cloud.spanner.v1.SpannerSettings; import com.google.cloud.spanner.v1.stub.SpannerStubSettings; +import com.google.common.base.MoreObjects; import com.google.common.base.Preconditions; import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableSet; +import com.google.spanner.v1.ExecuteSqlRequest.QueryOptions; import io.grpc.CallCredentials; import io.grpc.ManagedChannelBuilder; import java.io.IOException; import java.net.MalformedURLException; import java.net.URL; +import java.util.HashMap; import java.util.Map; +import java.util.Map.Entry; import java.util.Set; +import javax.annotation.Nonnull; import org.threeten.bp.Duration; /** Options for the Cloud Spanner service. */ public class SpannerOptions extends ServiceOptions { private static final long serialVersionUID = 2789571558532701170L; + private static SpannerEnvironment environment = SpannerEnvironmentImpl.INSTANCE; private static final String JDBC_API_CLIENT_LIB_TOKEN = "sp-jdbc"; private static final String HIBERNATE_API_CLIENT_LIB_TOKEN = "sp-hib"; @@ -73,6 +80,20 @@ public class SpannerOptions extends ServiceOptions { private final InstanceAdminStubSettings instanceAdminStubSettings; private final DatabaseAdminStubSettings databaseAdminStubSettings; private final Duration partitionedDmlTimeout; + /** + * These are the default {@link QueryOptions} defined by the user on this {@link SpannerOptions}. + */ + private final Map defaultQueryOptions; + /** These are the default {@link QueryOptions} defined in environment variables on this system. */ + private final QueryOptions envQueryOptions; + /** + * These are the merged query options of the {@link QueryOptions} set on this {@link + * SpannerOptions} and the {@link QueryOptions} in the environment variables. Options specified in + * environment variables take precedence above options specified in the {@link SpannerOptions} + * instance. + */ + private final Map mergedQueryOptions; + private final CallCredentialsProvider callCredentialsProvider; /** @@ -130,13 +151,55 @@ private SpannerOptions(Builder builder) { throw SpannerExceptionFactory.newSpannerException(e); } partitionedDmlTimeout = builder.partitionedDmlTimeout; + defaultQueryOptions = builder.defaultQueryOptions; + envQueryOptions = builder.getEnvironmentQueryOptions(); + if (envQueryOptions.equals(QueryOptions.getDefaultInstance())) { + this.mergedQueryOptions = ImmutableMap.copyOf(builder.defaultQueryOptions); + } else { + // Merge all specific database options with the environment options. + Map merged = new HashMap<>(builder.defaultQueryOptions); + for (Entry entry : builder.defaultQueryOptions.entrySet()) { + merged.put(entry.getKey(), entry.getValue().toBuilder().mergeFrom(envQueryOptions).build()); + } + this.mergedQueryOptions = ImmutableMap.copyOf(merged); + } callCredentialsProvider = builder.callCredentialsProvider; } + /** + * The environment to read configuration values from. The default implementation uses environment + * variables. + */ + public static interface SpannerEnvironment { + /** + * The optimizer version to use. Must return an empty string to indicate that no value has been + * set. + */ + @Nonnull + String getOptimizerVersion(); + } + + /** + * Default implementation of {@link SpannerEnvironment}. Reads all configuration from environment + * variables. + */ + private static class SpannerEnvironmentImpl implements SpannerEnvironment { + private static final SpannerEnvironmentImpl INSTANCE = new SpannerEnvironmentImpl(); + private static final String SPANNER_OPTIMIZER_VERSION_ENV_VAR = "SPANNER_OPTIMIZER_VERSION"; + + private SpannerEnvironmentImpl() {} + + @Override + public String getOptimizerVersion() { + return MoreObjects.firstNonNull(System.getenv(SPANNER_OPTIMIZER_VERSION_ENV_VAR), ""); + } + } + /** Builder for {@link SpannerOptions} instances. */ public static class Builder extends ServiceOptions.Builder { - private static final int DEFAULT_PREFETCH_CHUNKS = 4; + static final int DEFAULT_PREFETCH_CHUNKS = 4; + static final QueryOptions DEFAULT_QUERY_OPTIONS = QueryOptions.getDefaultInstance(); private final ImmutableSet allowedClientLibTokens = ImmutableSet.of( ServiceOptions.getGoogApiClientLibName(), @@ -162,6 +225,7 @@ public static class Builder private DatabaseAdminStubSettings.Builder databaseAdminStubSettingsBuilder = DatabaseAdminStubSettings.newBuilder(); private Duration partitionedDmlTimeout = Duration.ofHours(2L); + private Map defaultQueryOptions = new HashMap<>(); private CallCredentialsProvider callCredentialsProvider; private String emulatorHost = System.getenv("SPANNER_EMULATOR_HOST"); @@ -177,6 +241,7 @@ private Builder() {} this.instanceAdminStubSettingsBuilder = options.instanceAdminStubSettings.toBuilder(); this.databaseAdminStubSettingsBuilder = options.databaseAdminStubSettings.toBuilder(); this.partitionedDmlTimeout = options.partitionedDmlTimeout; + this.defaultQueryOptions = options.defaultQueryOptions; this.callCredentialsProvider = options.callCredentialsProvider; this.channelProvider = options.channelProvider; this.channelConfigurator = options.channelConfigurator; @@ -369,6 +434,37 @@ public Builder setPartitionedDmlTimeout(Duration timeout) { return this; } + /** + * Sets the default {@link QueryOptions} that will be used for all queries on the specified + * database. Query options can also be specified on a per-query basis and as environment + * variables. The precedence of these settings are: + * + *
    + *
  1. Query options for a specific query + *
  2. Environment variables + *
  3. These default query options + *
+ * + * Each {@link QueryOption} value that is used for a query is determined individually based on + * the above precedence. If for example a value for {@link QueryOptions#getOptimizerVersion()} + * is specified in an environment variable and a value for {@link + * QueryOptions#getOptimizerStatisticsPackage()} is specified for a specific query, both values + * will be used for the specific query. Environment variables are only read during the + * initialization of a {@link SpannerOptions} instance. Changing an environment variable after + * initializing a {@link SpannerOptions} instance will not have any effect on that instance. + */ + public Builder setDefaultQueryOptions(DatabaseId database, QueryOptions defaultQueryOptions) { + this.defaultQueryOptions.put(database, defaultQueryOptions); + return this; + } + + /** Gets the {@link QueryOptions} specified in the {@link SpannerEnvironment}. */ + QueryOptions getEnvironmentQueryOptions() { + return QueryOptions.newBuilder() + .setOptimizerVersion(environment.getOptimizerVersion()) + .build(); + } + /** * Sets a {@link CallCredentialsProvider} that can deliver {@link CallCredentials} to use on a * per-gRPC basis. Any credentials returned by this {@link CallCredentialsProvider} will have @@ -436,6 +532,22 @@ public static Builder newBuilder() { return new Builder(); } + /** + * Sets the environment to use to read configuration. The default will read configuration from + * environment variables. + */ + public static void useEnvironment(SpannerEnvironment environment) { + SpannerOptions.environment = environment; + } + + /** + * Sets the environment to use to read configuration to the default environment. This will read + * configuration from environment variables. + */ + public static void useDefaultEnvironment() { + SpannerOptions.environment = SpannerEnvironmentImpl.INSTANCE; + } + public TransportChannelProvider getChannelProvider() { return channelProvider; } @@ -481,6 +593,19 @@ public CallCredentialsProvider getCallCredentialsProvider() { return callCredentialsProvider; } + /** Returns the default query options to use for the specific database. */ + public QueryOptions getDefaultQueryOptions(DatabaseId databaseId) { + // Use the specific query options for the database if any have been specified. These have + // already been merged with the query options specified in the environment variables. + QueryOptions options = this.mergedQueryOptions.get(databaseId); + if (options == null) { + // Use the generic environment query options. These are initialized as a default instance of + // query options and appended with any options specified in the environment variables. + options = this.envQueryOptions; + } + return options; + } + public int getPrefetchChunks() { return prefetchChunks; } diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/Statement.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/Statement.java index fd7cfddf01..401aef9299 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/Statement.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/Statement.java @@ -22,6 +22,7 @@ import com.google.cloud.spanner.ReadContext.QueryAnalyzeMode; import com.google.common.base.Preconditions; import com.google.common.collect.ImmutableMap; +import com.google.spanner.v1.ExecuteSqlRequest.QueryOptions; import java.io.Serializable; import java.util.HashMap; import java.util.Map; @@ -57,29 +58,46 @@ public final class Statement implements Serializable { private final ImmutableMap parameters; private final String sql; + private final QueryOptions queryOptions; - private Statement(String sql, ImmutableMap parameters) { + private Statement(String sql, ImmutableMap parameters, QueryOptions queryOptions) { this.sql = sql; this.parameters = parameters; + this.queryOptions = queryOptions; } /** Builder for {@code Statement}. */ public static final class Builder { - final Map parameters = new HashMap<>(); + final Map parameters; private final StringBuilder sqlBuffer; private String currentBinding; private final ValueBinder binder = new Binder(); + private QueryOptions queryOptions; private Builder(String sql) { + parameters = new HashMap<>(); sqlBuffer = new StringBuilder(sql); } + private Builder(Statement statement) { + sqlBuffer = new StringBuilder(statement.sql); + parameters = new HashMap<>(statement.parameters); + queryOptions = + statement.queryOptions == null ? null : statement.queryOptions.toBuilder().build(); + } + /** Appends {@code sqlFragment} to the statement. */ public Builder append(String sqlFragment) { sqlBuffer.append(checkNotNull(sqlFragment)); return this; } + /** Sets the {@link QueryOptions} to use when executing this {@link Statement}. */ + public Builder withQueryOptions(QueryOptions queryOptions) { + this.queryOptions = queryOptions; + return this; + } + /** Returns a binder to bind the value of the query parameter {@code parameter}. */ public ValueBinder bind(String parameter) { checkState( @@ -94,7 +112,7 @@ public ValueBinder bind(String parameter) { public Statement build() { checkState( currentBinding == null, "Binding for parameter '%s' is incomplete.", currentBinding); - return new Statement(sqlBuffer.toString(), ImmutableMap.copyOf(parameters)); + return new Statement(sqlBuffer.toString(), ImmutableMap.copyOf(parameters), queryOptions); } private class Binder extends ValueBinder { @@ -151,11 +169,20 @@ public String getSql() { return sql; } + /** Returns the {@link QueryOptions} that will be used with this {@link Statement}. */ + public QueryOptions getQueryOptions() { + return queryOptions; + } + /** Returns the parameters bound to this {@code Statement}. */ public Map getParameters() { return parameters; } + public Builder toBuilder() { + return new Builder(this); + } + @Override public String toString() { return toString(new StringBuilder()).toString(); @@ -171,12 +198,14 @@ public boolean equals(Object o) { } Statement that = (Statement) o; - return Objects.equals(sql, that.sql) && Objects.equals(parameters, that.parameters); + return Objects.equals(sql, that.sql) + && Objects.equals(parameters, that.parameters) + && Objects.equals(queryOptions, that.queryOptions); } @Override public int hashCode() { - return Objects.hash(sql, parameters); + return Objects.hash(sql, parameters, queryOptions); } StringBuilder toString(StringBuilder b) { @@ -193,6 +222,9 @@ StringBuilder toString(StringBuilder b) { } b.append("}"); } + if (queryOptions != null) { + b.append(",queryOptions=").append(queryOptions.toString()); + } return b; } } diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionRunnerImpl.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionRunnerImpl.java index 9c2772da4f..cfa8b73c4a 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionRunnerImpl.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionRunnerImpl.java @@ -56,6 +56,26 @@ class TransactionRunnerImpl implements SessionTransaction, TransactionRunner { @VisibleForTesting static class TransactionContextImpl extends AbstractReadContext implements TransactionContext { + static class Builder extends AbstractReadContext.Builder { + private ByteString transactionId; + + private Builder() {} + + Builder setTransactionId(ByteString transactionId) { + this.transactionId = transactionId; + return self(); + } + + @Override + TransactionContextImpl build() { + return new TransactionContextImpl(this); + } + } + + static Builder newBuilder() { + return new Builder(); + } + @GuardedBy("lock") private List mutations = new ArrayList<>(); @@ -69,13 +89,9 @@ static class TransactionContextImpl extends AbstractReadContext implements Trans private ByteString transactionId; private Timestamp commitTimestamp; - TransactionContextImpl( - SessionImpl session, - @Nullable ByteString transactionId, - SpannerRpc rpc, - int defaultPrefetchChunks) { - super(session, rpc, defaultPrefetchChunks); - this.transactionId = transactionId; + private TransactionContextImpl(Builder builder) { + super(builder); + this.transactionId = builder.transactionId; } void ensureTxn() { diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AbstractReadContextTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AbstractReadContextTest.java new file mode 100644 index 0000000000..bfd739d553 --- /dev/null +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AbstractReadContextTest.java @@ -0,0 +1,109 @@ +/* + * Copyright 2020 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 static com.google.common.truth.Truth.assertThat; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +import com.google.cloud.spanner.spi.v1.SpannerRpc; +import com.google.spanner.v1.ExecuteSqlRequest; +import com.google.spanner.v1.ExecuteSqlRequest.QueryMode; +import com.google.spanner.v1.ExecuteSqlRequest.QueryOptions; +import com.google.spanner.v1.TransactionSelector; +import java.util.ArrayList; +import java.util.Collection; +import java.util.List; +import org.junit.Before; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.Parameterized; +import org.junit.runners.Parameterized.Parameter; +import org.junit.runners.Parameterized.Parameters; + +@RunWith(Parameterized.class) +public class AbstractReadContextTest { + @Parameter(0) + public QueryOptions defaultQueryOptions; + + @Parameters(name = "SpannerOptions.DefaultQueryOptions = {0}") + public static Collection parameters() { + List params = new ArrayList<>(); + params.add(new Object[] {QueryOptions.getDefaultInstance()}); + params.add( + new Object[] {QueryOptions.newBuilder().setOptimizerVersion("some-version").build()}); + return params; + } + + class TestReadContextBuilder + extends AbstractReadContext.Builder { + @Override + TestReadContext build() { + return new TestReadContext(this); + } + } + + private final class TestReadContext extends AbstractReadContext { + TestReadContext(TestReadContextBuilder builder) { + super(builder); + } + + @Override + TransactionSelector getTransactionSelector() { + return TransactionSelector.getDefaultInstance(); + } + } + + private TestReadContext context; + + @Before + public void setup() { + SessionImpl session = mock(SessionImpl.class); + when(session.getName()).thenReturn("session-1"); + TestReadContextBuilder builder = new TestReadContextBuilder(); + context = + builder + .setSession(session) + .setRpc(mock(SpannerRpc.class)) + .setDefaultQueryOptions(defaultQueryOptions) + .build(); + } + + @Test + public void executeSqlRequestBuilderWithoutQueryOptions() { + ExecuteSqlRequest request = + context + .getExecuteSqlRequestBuilder(Statement.of("SELECT FOO FROM BAR"), QueryMode.NORMAL) + .build(); + assertThat(request.getSql()).isEqualTo("SELECT FOO FROM BAR"); + assertThat(request.getQueryOptions()).isEqualTo(defaultQueryOptions); + } + + @Test + public void executeSqlRequestBuilderWithQueryOptions() { + ExecuteSqlRequest request = + context + .getExecuteSqlRequestBuilder( + Statement.newBuilder("SELECT FOO FROM BAR") + .withQueryOptions(QueryOptions.newBuilder().setOptimizerVersion("2.0").build()) + .build(), + QueryMode.NORMAL) + .build(); + assertThat(request.getSql()).isEqualTo("SELECT FOO FROM BAR"); + assertThat(request.getQueryOptions().getOptimizerVersion()).isEqualTo("2.0"); + } +} diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java index 43be8db1c0..f665d66ada 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java @@ -27,9 +27,14 @@ import com.google.cloud.NoCredentials; import com.google.cloud.spanner.MockSpannerServiceImpl.SimulatedExecutionTime; import com.google.cloud.spanner.MockSpannerServiceImpl.StatementResult; +import com.google.cloud.spanner.ReadContext.QueryAnalyzeMode; import com.google.cloud.spanner.TransactionRunner.TransactionCallable; import com.google.common.base.Stopwatch; +import com.google.protobuf.AbstractMessage; import com.google.protobuf.ListValue; +import com.google.spanner.v1.ExecuteSqlRequest; +import com.google.spanner.v1.ExecuteSqlRequest.QueryMode; +import com.google.spanner.v1.ExecuteSqlRequest.QueryOptions; import com.google.spanner.v1.ResultSetMetadata; import com.google.spanner.v1.StructType; import com.google.spanner.v1.StructType.Field; @@ -39,6 +44,7 @@ import io.grpc.StatusRuntimeException; import io.grpc.inprocess.InProcessServerBuilder; import java.io.IOException; +import java.util.List; import java.util.concurrent.ScheduledThreadPoolExecutor; import java.util.concurrent.TimeUnit; import org.junit.After; @@ -701,4 +707,122 @@ public Long run(TransactionContext transaction) throws Exception { assertThat(client1.pool.getNumberOfSessionsInPool(), is(equalTo(minSessions))); assertThat(client2.pool.getNumberOfSessionsInPool(), is(equalTo(minSessions))); } + + @Test + public void testBackendQueryOptions() { + // Use a Spanner instance with MinSession=0 and WriteFraction=0.0 to prevent background requests + // from the session pool interfering with the test case. + try (Spanner spanner = + SpannerOptions.newBuilder() + .setProjectId("[PROJECT]") + .setChannelProvider(channelProvider) + .setCredentials(NoCredentials.getInstance()) + .setSessionPoolOption( + SessionPoolOptions.newBuilder() + .setMinSessions(0) + .setWriteSessionsFraction(0.0f) + .build()) + .build() + .getService()) { + DatabaseClient client = + spanner.getDatabaseClient(DatabaseId.of("[PROJECT]", "[INSTANCE]", "[DATABASE")); + try (ResultSet rs = + client + .singleUse() + .executeQuery( + Statement.newBuilder(SELECT1.getSql()) + .withQueryOptions(QueryOptions.newBuilder().setOptimizerVersion("1").build()) + .build())) { + // Just iterate over the results to execute the query. + while (rs.next()) {} + } + // Check that the last query was executed using a custom optimizer version. + List requests = mockSpanner.getRequests(); + assertThat(requests).isNotEmpty(); + assertThat(requests.get(requests.size() - 1)).isInstanceOf(ExecuteSqlRequest.class); + ExecuteSqlRequest request = (ExecuteSqlRequest) requests.get(requests.size() - 1); + assertThat(request.getQueryOptions()).isNotNull(); + assertThat(request.getQueryOptions().getOptimizerVersion()).isEqualTo("1"); + } + } + + @Test + public void testBackendQueryOptionsWithAnalyzeQuery() { + // Use a Spanner instance with MinSession=0 and WriteFraction=0.0 to prevent background requests + // from the session pool interfering with the test case. + try (Spanner spanner = + SpannerOptions.newBuilder() + .setProjectId("[PROJECT]") + .setChannelProvider(channelProvider) + .setCredentials(NoCredentials.getInstance()) + .setSessionPoolOption( + SessionPoolOptions.newBuilder() + .setMinSessions(0) + .setWriteSessionsFraction(0.0f) + .build()) + .build() + .getService()) { + DatabaseClient client = + spanner.getDatabaseClient(DatabaseId.of("[PROJECT]", "[INSTANCE]", "[DATABASE")); + try (ReadOnlyTransaction tx = client.readOnlyTransaction()) { + try (ResultSet rs = + tx.analyzeQuery( + Statement.newBuilder(SELECT1.getSql()) + .withQueryOptions(QueryOptions.newBuilder().setOptimizerVersion("1").build()) + .build(), + QueryAnalyzeMode.PROFILE)) { + // Just iterate over the results to execute the query. + while (rs.next()) {} + } + } + // Check that the last query was executed using a custom optimizer version. + List requests = mockSpanner.getRequests(); + assertThat(requests).isNotEmpty(); + assertThat(requests.get(requests.size() - 1)).isInstanceOf(ExecuteSqlRequest.class); + ExecuteSqlRequest request = (ExecuteSqlRequest) requests.get(requests.size() - 1); + assertThat(request.getQueryOptions()).isNotNull(); + assertThat(request.getQueryOptions().getOptimizerVersion()).isEqualTo("1"); + assertThat(request.getQueryMode()).isEqualTo(QueryMode.PROFILE); + } + } + + @Test + public void testBackendPartitionQueryOptions() { + // Use a Spanner instance with MinSession=0 and WriteFraction=0.0 to prevent background requests + // from the session pool interfering with the test case. + try (Spanner spanner = + SpannerOptions.newBuilder() + .setProjectId("[PROJECT]") + .setChannelProvider(channelProvider) + .setCredentials(NoCredentials.getInstance()) + .setSessionPoolOption( + SessionPoolOptions.newBuilder() + .setMinSessions(0) + .setWriteSessionsFraction(0.0f) + .build()) + .build() + .getService()) { + BatchClient client = + spanner.getBatchClient(DatabaseId.of("[PROJECT]", "[INSTANCE]", "[DATABASE")); + BatchReadOnlyTransaction transaction = + client.batchReadOnlyTransaction(TimestampBound.strong()); + List partitions = + transaction.partitionQuery( + PartitionOptions.newBuilder().setMaxPartitions(10L).build(), + Statement.newBuilder(SELECT1.getSql()) + .withQueryOptions(QueryOptions.newBuilder().setOptimizerVersion("1").build()) + .build()); + try (ResultSet rs = transaction.execute(partitions.get(0))) { + // Just iterate over the results to execute the query. + while (rs.next()) {} + } + // Check that the last query was executed using a custom optimizer version. + List requests = mockSpanner.getRequests(); + assertThat(requests).isNotEmpty(); + assertThat(requests.get(requests.size() - 1)).isInstanceOf(ExecuteSqlRequest.class); + ExecuteSqlRequest request = (ExecuteSqlRequest) requests.get(requests.size() - 1); + assertThat(request.getQueryOptions()).isNotNull(); + assertThat(request.getQueryOptions().getOptimizerVersion()).isEqualTo("1"); + } + } } diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolTest.java index 559d0e1a05..2047b6c858 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolTest.java @@ -1196,7 +1196,11 @@ public void testSessionNotFoundReadWriteTransaction() { ByteString preparedTransactionId = hasPreparedTransaction ? ByteString.copyFromUtf8("test-txn") : null; final TransactionContextImpl closedTransactionContext = - new TransactionContextImpl(closedSession, preparedTransactionId, rpc, 10); + TransactionContextImpl.newBuilder() + .setSession(closedSession) + .setTransactionId(preparedTransactionId) + .setRpc(rpc) + .build(); when(closedSession.asyncClose()) .thenReturn(ApiFutures.immediateFuture(Empty.getDefaultInstance())); when(closedSession.newTransaction()).thenReturn(closedTransactionContext); diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SpannerImplTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SpannerImplTest.java index 23325e902f..f8b4d67c16 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SpannerImplTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SpannerImplTest.java @@ -28,6 +28,7 @@ import com.google.cloud.ServiceRpc; import com.google.cloud.grpc.GrpcTransportOptions; import com.google.cloud.spanner.spi.v1.SpannerRpc; +import com.google.spanner.v1.ExecuteSqlRequest.QueryOptions; import java.util.Collections; import java.util.HashMap; import java.util.Map; @@ -88,6 +89,47 @@ public void getDbclientAgainGivesSame() { assertThat(databaseClient1).isSameInstanceAs(databaseClient); } + @Test + public void queryOptions() { + QueryOptions queryOptions = QueryOptions.newBuilder().setOptimizerVersion("2").build(); + QueryOptions defaultOptions = QueryOptions.getDefaultInstance(); + DatabaseId db = DatabaseId.of("p", "i", "d"); + DatabaseId otherDb = DatabaseId.of("p", "i", "other"); + + // Create a SpannerOptions with and without default query options. + SpannerOptions optionsWithQueryOptions = + new SpannerOptions.Builder(SpannerOptions.getDefaultInstance()) { + @Override + QueryOptions getEnvironmentQueryOptions() { + // Override and return default instance to prevent environment variables from + // interfering with the test case. + return QueryOptions.getDefaultInstance(); + } + }.setDefaultQueryOptions(db, queryOptions).build(); + SpannerOptions optionsWithoutQueryOptions = + new SpannerOptions.Builder(SpannerOptions.getDefaultInstance()) { + @Override + QueryOptions getEnvironmentQueryOptions() { + // Override and return default instance to prevent environment variables from + // interfering with the test case. + return QueryOptions.getDefaultInstance(); + } + }.build(); + + try (SpannerImpl implWithQueryOptions = new SpannerImpl(rpc, optionsWithQueryOptions); + SpannerImpl implWithoutQueryOptions = new SpannerImpl(rpc, optionsWithoutQueryOptions)) { + + // Default query options are on a per-database basis, so we should only get the custom options + // for 'db' and not for 'otherDb'. + assertThat(implWithQueryOptions.getDefaultQueryOptions(db)).isEqualTo(queryOptions); + assertThat(implWithQueryOptions.getDefaultQueryOptions(otherDb)).isEqualTo(defaultOptions); + + // The other Spanner instance should return default options for both databases. + assertThat(implWithoutQueryOptions.getDefaultQueryOptions(db)).isEqualTo(defaultOptions); + assertThat(implWithoutQueryOptions.getDefaultQueryOptions(otherDb)).isEqualTo(defaultOptions); + } + } + @Test public void getDbclientAfterCloseThrows() { SpannerImpl imp = new SpannerImpl(rpc, spannerOptions); diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SpannerOptionsTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SpannerOptionsTest.java index 214773b9fd..8e7f946b3b 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SpannerOptionsTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SpannerOptionsTest.java @@ -29,6 +29,7 @@ import com.google.cloud.spanner.admin.database.v1.stub.DatabaseAdminStubSettings; import com.google.cloud.spanner.admin.instance.v1.stub.InstanceAdminStubSettings; import com.google.cloud.spanner.v1.stub.SpannerStubSettings; +import com.google.spanner.v1.ExecuteSqlRequest.QueryOptions; import java.util.Arrays; import java.util.HashMap; import java.util.List; @@ -427,4 +428,46 @@ public void testSetEmulatorHostWithProtocol() { assertThat(options.getHost()).isEqualTo("http://localhost:1234"); assertThat(options.getEndpoint()).isEqualTo("localhost:1234"); } + + @Test + public void testDefaultQueryOptions() { + SpannerOptions.useEnvironment( + new SpannerOptions.SpannerEnvironment() { + @Override + public String getOptimizerVersion() { + return ""; + } + }); + SpannerOptions options = + SpannerOptions.newBuilder() + .setDefaultQueryOptions( + DatabaseId.of("p", "i", "d"), + QueryOptions.newBuilder().setOptimizerVersion("1").build()) + .build(); + assertThat(options.getDefaultQueryOptions(DatabaseId.of("p", "i", "d"))) + .isEqualTo(QueryOptions.newBuilder().setOptimizerVersion("1").build()); + assertThat(options.getDefaultQueryOptions(DatabaseId.of("p", "i", "o"))) + .isEqualTo(QueryOptions.getDefaultInstance()); + + // Now simulate that the user has set an environment variable for the query optimizer version. + SpannerOptions.useEnvironment( + new SpannerOptions.SpannerEnvironment() { + @Override + public String getOptimizerVersion() { + return "2"; + } + }); + // Create options with '1' as the default query optimizer version. This should be overridden by + // the environment variable. + options = + SpannerOptions.newBuilder() + .setDefaultQueryOptions( + DatabaseId.of("p", "i", "d"), + QueryOptions.newBuilder().setOptimizerVersion("1").build()) + .build(); + assertThat(options.getDefaultQueryOptions(DatabaseId.of("p", "i", "d"))) + .isEqualTo(QueryOptions.newBuilder().setOptimizerVersion("2").build()); + assertThat(options.getDefaultQueryOptions(DatabaseId.of("p", "i", "o"))) + .isEqualTo(QueryOptions.newBuilder().setOptimizerVersion("2").build()); + } } diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionContextImplTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionContextImplTest.java index 5eb25f33e1..077b660576 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionContextImplTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionContextImplTest.java @@ -59,7 +59,11 @@ private void batchDml(int status) { when(rpc.executeBatchDml(Mockito.any(ExecuteBatchDmlRequest.class), Mockito.anyMap())) .thenReturn(response); try (TransactionContextImpl impl = - new TransactionContextImpl(session, ByteString.copyFromUtf8("test"), rpc, 10)) { + TransactionContextImpl.newBuilder() + .setSession(session) + .setRpc(rpc) + .setTransactionId(ByteString.copyFromUtf8("test")) + .build()) { impl.batchUpdate(Arrays.asList(statement)); } } diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionManagerImplTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionManagerImplTest.java index 4b064a047d..cc522f3f45 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionManagerImplTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionManagerImplTest.java @@ -211,7 +211,7 @@ public List answer(InvocationOnMock invocation) throws Throwable { return Arrays.asList( com.google.spanner.v1.Session.newBuilder() - .setName((String) invocation.getArguments()[0]) + .setName((String) invocation.getArguments()[0] + "/sessions/1") .setCreateTime( com.google.protobuf.Timestamp.newBuilder() .setSeconds(System.currentTimeMillis() * 1000)) diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionRunnerImplTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionRunnerImplTest.java index e46a74b96b..3dcd523c13 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionRunnerImplTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionRunnerImplTest.java @@ -27,6 +27,7 @@ import com.google.api.core.ApiFutures; import com.google.cloud.grpc.GrpcTransportOptions; import com.google.cloud.grpc.GrpcTransportOptions.ExecutorFactory; +import com.google.cloud.spanner.SessionClient.SessionId; import com.google.cloud.spanner.TransactionRunner.TransactionCallable; import com.google.cloud.spanner.TransactionRunnerImpl.TransactionContextImpl; import com.google.cloud.spanner.spi.v1.SpannerRpc; @@ -120,7 +121,7 @@ public List answer(InvocationOnMock invocation) throws Throwable { return Arrays.asList( com.google.spanner.v1.Session.newBuilder() - .setName((String) invocation.getArguments()[0]) + .setName((String) invocation.getArguments()[0] + "/sessions/1") .setCreateTime( Timestamp.newBuilder().setSeconds(System.currentTimeMillis() * 1000)) .build()); @@ -266,12 +267,15 @@ public void batchDmlFailedPrecondition() { private long[] batchDmlException(int status) { Preconditions.checkArgument(status != Code.OK_VALUE); TransactionContextImpl transaction = - new TransactionContextImpl( - session, ByteString.copyFromUtf8(UUID.randomUUID().toString()), rpc, 10); + TransactionContextImpl.newBuilder() + .setSession(session) + .setTransactionId(ByteString.copyFromUtf8(UUID.randomUUID().toString())) + .setRpc(rpc) + .build(); when(session.newTransaction()).thenReturn(transaction); when(session.beginTransaction()) .thenReturn(ByteString.copyFromUtf8(UUID.randomUUID().toString())); - when(session.getName()).thenReturn("test"); + when(session.getName()).thenReturn(SessionId.of("p", "i", "d", "test").getName()); TransactionRunnerImpl runner = new TransactionRunnerImpl(session, rpc, 10); ExecuteBatchDmlResponse response1 = ExecuteBatchDmlResponse.newBuilder() diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/it/ITQueryOptionsTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/it/ITQueryOptionsTest.java new file mode 100644 index 0000000000..89d789e6f0 --- /dev/null +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/it/ITQueryOptionsTest.java @@ -0,0 +1,235 @@ +/* + * Copyright 2020 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.it; + +import static com.google.common.truth.Truth.assertThat; +import static org.junit.Assert.fail; + +import com.google.cloud.spanner.Database; +import com.google.cloud.spanner.DatabaseClient; +import com.google.cloud.spanner.ErrorCode; +import com.google.cloud.spanner.IntegrationTest; +import com.google.cloud.spanner.IntegrationTestEnv; +import com.google.cloud.spanner.ResultSet; +import com.google.cloud.spanner.Spanner; +import com.google.cloud.spanner.SpannerException; +import com.google.cloud.spanner.Statement; +import com.google.cloud.spanner.TransactionContext; +import com.google.cloud.spanner.TransactionRunner.TransactionCallable; +import com.google.spanner.v1.ExecuteSqlRequest.QueryOptions; +import org.junit.BeforeClass; +import org.junit.ClassRule; +import org.junit.Rule; +import org.junit.Test; +import org.junit.experimental.categories.Category; +import org.junit.rules.ExpectedException; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@Category(IntegrationTest.class) +@RunWith(JUnit4.class) +public class ITQueryOptionsTest { + @ClassRule public static IntegrationTestEnv env = new IntegrationTestEnv(); + private static Database db; + @Rule public ExpectedException expectedException = ExpectedException.none(); + private static DatabaseClient client; + + @BeforeClass + public static void setUpDatabase() { + // Empty database. + db = + env.getTestHelper() + .createTestDatabase("CREATE TABLE TEST (ID INT64, NAME STRING(100)) PRIMARY KEY (ID)"); + client = env.getTestHelper().getDatabaseClient(db); + } + + @Test + public void executeQuery() { + // Version '1' should work. + try (ResultSet rs = + client + .singleUse() + .executeQuery( + Statement.newBuilder("SELECT 1") + .withQueryOptions(QueryOptions.newBuilder().setOptimizerVersion("1").build()) + .build())) { + while (rs.next()) { + assertThat(rs.getLong(0)).isEqualTo(1L); + } + } + // Version 'latest' should also work. + try (ResultSet rs = + client + .singleUse() + .executeQuery( + Statement.newBuilder("SELECT 1") + .withQueryOptions( + QueryOptions.newBuilder().setOptimizerVersion("latest").build()) + .build())) { + while (rs.next()) { + assertThat(rs.getLong(0)).isEqualTo(1L); + } + } + // Version '100000' should not work. + try (ResultSet rs = + client + .singleUse() + .executeQuery( + Statement.newBuilder("SELECT 1") + .withQueryOptions( + QueryOptions.newBuilder().setOptimizerVersion("100000").build()) + .build())) { + while (rs.next()) { + fail("should not get any results"); + } + } catch (SpannerException e) { + assertThat(e.getErrorCode()).isEqualTo(ErrorCode.INVALID_ARGUMENT); + assertThat(e.getMessage()).contains("Query optimizer version: 100000 is not supported"); + } + } + + @Test + public void executeUpdate() { + // Query optimizer version is ignored for DML statements by the backend, but setting it does not + // cause an error. + assertThat( + client + .readWriteTransaction() + .run( + new TransactionCallable() { + @Override + public Long run(TransactionContext transaction) throws Exception { + return transaction.executeUpdate( + Statement.newBuilder("INSERT INTO TEST (ID, NAME) VALUES (@id, @name)") + .bind("id") + .to(1L) + .bind("name") + .to("One") + .withQueryOptions( + QueryOptions.newBuilder().setOptimizerVersion("1").build()) + .build()); + } + })) + .isEqualTo(1L); + + // Version 'latest' should also work. + assertThat( + client + .readWriteTransaction() + .run( + new TransactionCallable() { + @Override + public Long run(TransactionContext transaction) throws Exception { + return transaction.executeUpdate( + Statement.newBuilder("INSERT INTO TEST (ID, NAME) VALUES (@id, @name)") + .bind("id") + .to(2L) + .bind("name") + .to("Two") + .withQueryOptions( + QueryOptions.newBuilder().setOptimizerVersion("latest").build()) + .build()); + } + })) + .isEqualTo(1L); + + // Version '100000' is an invalid value, but is ignored by the backend. + assertThat( + client + .readWriteTransaction() + .run( + new TransactionCallable() { + @Override + public Long run(TransactionContext transaction) throws Exception { + return transaction.executeUpdate( + Statement.newBuilder("INSERT INTO TEST (ID, NAME) VALUES (@id, @name)") + .bind("id") + .to(3L) + .bind("name") + .to("Three") + .withQueryOptions( + QueryOptions.newBuilder().setOptimizerVersion("10000").build()) + .build()); + } + })) + .isEqualTo(1L); + + // Verify that query options are ignored with Partitioned DML as well, and that all the above + // DML INSERT statements succeeded. + assertThat( + client.executePartitionedUpdate( + Statement.newBuilder("UPDATE TEST SET NAME='updated' WHERE 1=1") + .withQueryOptions(QueryOptions.newBuilder().setOptimizerVersion("1").build()) + .build())) + .isEqualTo(3L); + } + + @Test + public void spannerOptions() { + // Version '1' should work. + try (Spanner spanner = + env.getTestHelper() + .getOptions() + .toBuilder() + .setDefaultQueryOptions( + db.getId(), QueryOptions.newBuilder().setOptimizerVersion("1").build()) + .build() + .getService()) { + DatabaseClient client = spanner.getDatabaseClient(db.getId()); + try (ResultSet rs = client.singleUse().executeQuery(Statement.of("SELECT 1"))) { + while (rs.next()) { + assertThat(rs.getLong(0)).isEqualTo(1L); + } + } + } + // Version 'latest' should also work. + try (Spanner spanner = + env.getTestHelper() + .getOptions() + .toBuilder() + .setDefaultQueryOptions( + db.getId(), QueryOptions.newBuilder().setOptimizerVersion("latest").build()) + .build() + .getService()) { + DatabaseClient client = spanner.getDatabaseClient(db.getId()); + try (ResultSet rs = client.singleUse().executeQuery(Statement.of("SELECT 1"))) { + while (rs.next()) { + assertThat(rs.getLong(0)).isEqualTo(1L); + } + } + } + // Version '100000' should not work. + try (Spanner spanner = + env.getTestHelper() + .getOptions() + .toBuilder() + .setDefaultQueryOptions( + db.getId(), QueryOptions.newBuilder().setOptimizerVersion("100000").build()) + .build() + .getService()) { + DatabaseClient client = spanner.getDatabaseClient(db.getId()); + try (ResultSet rs = client.singleUse().executeQuery(Statement.of("SELECT 1"))) { + while (rs.next()) { + fail("should not get any results"); + } + } catch (SpannerException e) { + assertThat(e.getErrorCode()).isEqualTo(ErrorCode.INVALID_ARGUMENT); + assertThat(e.getMessage()).contains("Query optimizer version: 100000 is not supported"); + } + } + } +}