Skip to content
Permalink
Browse files
feat: add backend query options (#90)
* feat: add backend query options

Adds support for setting QueryOptions that will be used by the backend
to execute queries.

* fix: set QueryOptions on Statement

QueryOptions should be an option on a Statement instead of a parameter
to the executeQuery method. By setting these options on a Statement, it
is possible to use it with analyzeQuery as well.

* feat: add toBuilder() method to Statement

* fix: code review comments

* fix: remove unused interface and class
  • Loading branch information
olavloite committed Mar 12, 2020
1 parent 57f5fd0 commit e96e17246bee9691171b46857806d03d1f8e19b4
Showing with 983 additions and 85 deletions.
  1. +162 −41 google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractReadContext.java
  2. +22 −15 google-cloud-spanner/src/main/java/com/google/cloud/spanner/BatchClientImpl.java
  3. +4 −0 google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionClient.java
  4. +31 −9 google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionImpl.java
  5. +6 −0 google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerImpl.java
  6. +126 −1 google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerOptions.java
  7. +37 −5 google-cloud-spanner/src/main/java/com/google/cloud/spanner/Statement.java
  8. +23 −7 google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionRunnerImpl.java
  9. +109 −0 google-cloud-spanner/src/test/java/com/google/cloud/spanner/AbstractReadContextTest.java
  10. +124 −0 google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java
  11. +5 −1 google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolTest.java
  12. +42 −0 google-cloud-spanner/src/test/java/com/google/cloud/spanner/SpannerImplTest.java
  13. +43 −0 google-cloud-spanner/src/test/java/com/google/cloud/spanner/SpannerOptionsTest.java
  14. +5 −1 google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionContextImplTest.java
  15. +1 −1 google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionManagerImplTest.java
  16. +8 −4 google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionRunnerImplTest.java
  17. +235 −0 google-cloud-spanner/src/test/java/com/google/cloud/spanner/it/ITQueryOptionsTest.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<B extends Builder<?, T>, 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<Builder, SingleReadContext> {
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<Builder, MultiUseReadOnlyTransaction> {
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:
*
* <ol>
* <li>Specific {@link QueryOptions} passed in for this query.
* <li>Any value specified in a valid environment variable when the {@link SpannerOptions}
* instance was created.
* <li>The default {@link SpannerOptions#getDefaultQueryOptions()} specified for the database
* where the query is executed.
* </ol>
*/
@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,15 +521,15 @@ ExecuteBatchDmlRequest.Builder getExecuteBatchDmlRequestBuilder(Iterable<Stateme
ResultSet executeQueryInternalWithOptions(
Statement statement,
com.google.spanner.v1.ExecuteSqlRequest.QueryMode queryMode,
Options readOptions,
Options options,
ByteString partitionToken) {
beforeReadOrQuery();
final ExecuteSqlRequest.Builder request = getExecuteSqlRequestBuilder(statement, queryMode);
if (partitionToken != null) {
request.setPartitionToken(partitionToken);
}
final int prefetchChunks =
readOptions.hasPrefetchChunks() ? readOptions.prefetchChunks() : defaultPrefetchChunks;
options.hasPrefetchChunks() ? options.prefetchChunks() : defaultPrefetchChunks;
ResumableStreamIterator stream =
new ResumableStreamIterator(MAX_BUFFERED_CHUNKS, SpannerImpl.QUERY, span) {
@Override
@@ -45,41 +45,48 @@ public class BatchClientImpl implements BatchClient {
public BatchReadOnlyTransaction batchReadOnlyTransaction(TimestampBound bound) {
SessionImpl session = sessionClient.createSession();
return new BatchReadOnlyTransactionImpl(
sessionClient.getSpanner(), session, checkNotNull(bound));
MultiUseReadOnlyTransaction.newBuilder()
.setSession(session)
.setRpc(sessionClient.getSpanner().getRpc())
.setTimestampBound(bound)
.setDefaultQueryOptions(
sessionClient.getSpanner().getDefaultQueryOptions(sessionClient.getDatabaseId()))
.setDefaultPrefetchChunks(sessionClient.getSpanner().getDefaultPrefetchChunks()),
checkNotNull(bound));
}

@Override
public BatchReadOnlyTransaction batchReadOnlyTransaction(BatchTransactionId batchTransactionId) {
SessionImpl session =
sessionClient.sessionWithId(checkNotNull(batchTransactionId).getSessionId());
return new BatchReadOnlyTransactionImpl(
sessionClient.getSpanner(), session, batchTransactionId);
MultiUseReadOnlyTransaction.newBuilder()
.setSession(session)
.setRpc(sessionClient.getSpanner().getRpc())
.setTransactionId(batchTransactionId.getTransactionId())
.setTimestamp(batchTransactionId.getTimestamp())
.setDefaultQueryOptions(
sessionClient.getSpanner().getDefaultQueryOptions(sessionClient.getDatabaseId()))
.setDefaultPrefetchChunks(sessionClient.getSpanner().getDefaultPrefetchChunks()),
batchTransactionId);
}

private static class BatchReadOnlyTransactionImpl extends MultiUseReadOnlyTransaction
implements BatchReadOnlyTransaction {
private final String sessionName;
private final Map<SpannerRpc.Option, ?> 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();
}
@@ -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,

0 comments on commit e96e172

Please sign in to comment.