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 f191a858a1..9fff15ed62 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 @@ -49,6 +49,7 @@ import com.google.spanner.v1.TransactionOptions; import com.google.spanner.v1.TransactionSelector; import io.opencensus.trace.Span; +import io.opencensus.trace.Tracing; import java.util.Map; import java.util.concurrent.atomic.AtomicLong; import javax.annotation.Nullable; @@ -67,6 +68,7 @@ abstract static class Builder, T extends AbstractReadCon private Span span = Tracing.getTracer().getCurrentSpan(); private int defaultPrefetchChunks = SpannerOptions.Builder.DEFAULT_PREFETCH_CHUNKS; private QueryOptions defaultQueryOptions = SpannerOptions.Builder.DEFAULT_QUERY_OPTIONS; + private ExecutorProvider executorProvider; Builder() {} @@ -100,6 +102,11 @@ B setDefaultQueryOptions(QueryOptions defaultQueryOptions) { return self(); } + B setExecutorProvider(ExecutorProvider executorProvider) { + this.executorProvider = executorProvider; + return self(); + } + abstract T build(); } @@ -391,6 +398,7 @@ void initTransaction() { this.defaultPrefetchChunks = builder.defaultPrefetchChunks; this.defaultQueryOptions = builder.defaultQueryOptions; this.span = builder.span; + this.executorProvider = builder.executorProvider; } @Override diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/BatchClientImpl.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/BatchClientImpl.java index 39827647b6..c84bef77cf 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/BatchClientImpl.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/BatchClientImpl.java @@ -52,6 +52,7 @@ public BatchReadOnlyTransaction batchReadOnlyTransaction(TimestampBound bound) { .setTimestampBound(bound) .setDefaultQueryOptions( sessionClient.getSpanner().getDefaultQueryOptions(sessionClient.getDatabaseId())) + .setExecutorProvider(sessionClient.getSpanner().getAsyncExecutorProvider()) .setDefaultPrefetchChunks(sessionClient.getSpanner().getDefaultPrefetchChunks()), checkNotNull(bound)); } @@ -68,6 +69,7 @@ public BatchReadOnlyTransaction batchReadOnlyTransaction(BatchTransactionId batc .setTimestamp(batchTransactionId.getTimestamp()) .setDefaultQueryOptions( sessionClient.getSpanner().getDefaultQueryOptions(sessionClient.getDatabaseId())) + .setExecutorProvider(sessionClient.getSpanner().getAsyncExecutorProvider()) .setDefaultPrefetchChunks(sessionClient.getSpanner().getDefaultPrefetchChunks()), batchTransactionId); } 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 1d9c99eeb5..77f8946bb1 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 @@ -25,7 +25,6 @@ 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.TransactionRunner.TransactionCallable; import com.google.cloud.spanner.TransactionRunnerImpl.TransactionContextImpl; import com.google.cloud.spanner.spi.v1.SpannerRpc; import com.google.common.collect.Lists; @@ -177,6 +176,8 @@ public ReadContext singleUse(TimestampBound bound) { .setRpc(spanner.getRpc()) .setDefaultQueryOptions(spanner.getDefaultQueryOptions(databaseId)) .setDefaultPrefetchChunks(spanner.getDefaultPrefetchChunks()) + .setSpan(currentSpan) + .setExecutorProvider(spanner.getAsyncExecutorProvider()) .build()); } @@ -194,6 +195,8 @@ public ReadOnlyTransaction singleUseReadOnlyTransaction(TimestampBound bound) { .setRpc(spanner.getRpc()) .setDefaultQueryOptions(spanner.getDefaultQueryOptions(databaseId)) .setDefaultPrefetchChunks(spanner.getDefaultPrefetchChunks()) + .setSpan(currentSpan) + .setExecutorProvider(spanner.getAsyncExecutorProvider()) .buildSingleUseReadOnlyTransaction()); } @@ -211,6 +214,8 @@ public ReadOnlyTransaction readOnlyTransaction(TimestampBound bound) { .setRpc(spanner.getRpc()) .setDefaultQueryOptions(spanner.getDefaultQueryOptions(databaseId)) .setDefaultPrefetchChunks(spanner.getDefaultPrefetchChunks()) + .setSpan(currentSpan) + .setExecutorProvider(spanner.getAsyncExecutorProvider()) .build()); } @@ -281,6 +286,7 @@ TransactionContextImpl newTransaction() { .setDefaultQueryOptions(spanner.getDefaultQueryOptions(databaseId)) .setDefaultPrefetchChunks(spanner.getDefaultPrefetchChunks()) .setSpan(currentSpan) + .setExecutorProvider(spanner.getAsyncExecutorProvider()) .build(); } 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 d161daa264..f29d36c61a 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 @@ -43,8 +43,8 @@ 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 com.google.common.util.concurrent.ThreadFactoryBuilder; +import com.google.spanner.v1.ExecuteSqlRequest.QueryOptions; import io.grpc.CallCredentials; import io.grpc.ManagedChannelBuilder; import java.io.IOException; @@ -54,11 +54,11 @@ import java.util.Map; import java.util.Map.Entry; import java.util.Set; -import javax.annotation.Nonnull; import java.util.concurrent.ScheduledThreadPoolExecutor; import java.util.concurrent.ThreadFactory; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicInteger; +import javax.annotation.Nonnull; import org.threeten.bp.Duration; /** Options for the Cloud Spanner service. */ 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 e4fae3e3a6..0f098f4f6e 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 @@ -24,7 +24,6 @@ import com.google.api.core.ApiFunction; import com.google.api.core.ApiFuture; import com.google.api.core.ApiFutures; -import com.google.api.gax.core.ExecutorProvider; import com.google.cloud.Timestamp; import com.google.cloud.spanner.Options.QueryOption; import com.google.cloud.spanner.Options.ReadOption; 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 index bfd739d553..f9cee1c488 100644 --- 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 @@ -20,6 +20,7 @@ import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; +import com.google.api.gax.core.ExecutorProvider; import com.google.cloud.spanner.spi.v1.SpannerRpc; import com.google.spanner.v1.ExecuteSqlRequest; import com.google.spanner.v1.ExecuteSqlRequest.QueryMode; @@ -80,6 +81,7 @@ public void setup() { .setSession(session) .setRpc(mock(SpannerRpc.class)) .setDefaultQueryOptions(defaultQueryOptions) + .setExecutorProvider(mock(ExecutorProvider.class)) .build(); } 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 7d7b919fc1..f6814af8dc 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 @@ -36,25 +36,19 @@ 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.common.util.concurrent.SettableFuture; -import com.google.protobuf.ListValue; +import com.google.protobuf.AbstractMessage; 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; -import com.google.spanner.v1.TypeCode; import io.grpc.Server; import io.grpc.Status; import io.grpc.StatusRuntimeException; import io.grpc.inprocess.InProcessServerBuilder; import java.io.IOException; -import java.util.List; import java.util.ArrayList; import java.util.Arrays; -import java.util.concurrent.CountDownLatch; +import java.util.List; import java.util.concurrent.ExecutionException; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; 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 db64a5be79..984851f698 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 @@ -35,7 +35,6 @@ import com.google.api.core.ApiFuture; import com.google.api.core.ApiFutures; -import com.google.api.gax.core.ExecutorProvider; import com.google.cloud.Timestamp; import com.google.cloud.spanner.MetricRegistryTestUtils.FakeMetricRegistry; import com.google.cloud.spanner.MetricRegistryTestUtils.MetricsRecord; 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 3446ba9fc8..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 @@ -19,7 +19,6 @@ import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; -import com.google.api.gax.core.ExecutorProvider; import com.google.cloud.spanner.TransactionRunnerImpl.TransactionContextImpl; import com.google.cloud.spanner.spi.v1.SpannerRpc; import com.google.protobuf.ByteString; @@ -27,7 +26,6 @@ import com.google.rpc.Status; import com.google.spanner.v1.ExecuteBatchDmlRequest; import com.google.spanner.v1.ExecuteBatchDmlResponse; -import io.opencensus.trace.Span; import java.util.Arrays; import org.junit.Test; import org.junit.runner.RunWith; 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 8bcc05aa00..b790ed93b2 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 @@ -25,7 +25,6 @@ import static org.mockito.Mockito.when; import com.google.api.core.ApiFutures; -import com.google.api.gax.core.ExecutorProvider; import com.google.cloud.grpc.GrpcTransportOptions; import com.google.cloud.grpc.GrpcTransportOptions.ExecutorFactory; import com.google.cloud.spanner.SessionClient.SessionId;