com.google.cloud.spanner.spi.v1.GapicSpannerRpc Maven / Gradle / Ivy
Show all versions of google-cloud-spanner Show documentation
/*
* Copyright 2018 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.spi.v1;
import static com.google.cloud.spanner.SpannerExceptionFactory.newSpannerException;
import static com.google.cloud.spanner.ThreadFactoryUtil.tryCreateVirtualThreadPerTaskExecutor;
import com.google.api.core.ApiFunction;
import com.google.api.core.ApiFuture;
import com.google.api.core.InternalApi;
import com.google.api.core.NanoClock;
import com.google.api.gax.core.CredentialsProvider;
import com.google.api.gax.core.GaxProperties;
import com.google.api.gax.grpc.GaxGrpcProperties;
import com.google.api.gax.grpc.GrpcCallContext;
import com.google.api.gax.grpc.GrpcCallSettings;
import com.google.api.gax.grpc.GrpcStubCallableFactory;
import com.google.api.gax.grpc.InstantiatingGrpcChannelProvider;
import com.google.api.gax.longrunning.OperationFuture;
import com.google.api.gax.retrying.ResultRetryAlgorithm;
import com.google.api.gax.retrying.RetrySettings;
import com.google.api.gax.retrying.TimedAttemptSettings;
import com.google.api.gax.rpc.AlreadyExistsException;
import com.google.api.gax.rpc.ApiCallContext;
import com.google.api.gax.rpc.ApiClientHeaderProvider;
import com.google.api.gax.rpc.ApiException;
import com.google.api.gax.rpc.ClientContext;
import com.google.api.gax.rpc.FixedHeaderProvider;
import com.google.api.gax.rpc.HeaderProvider;
import com.google.api.gax.rpc.InstantiatingWatchdogProvider;
import com.google.api.gax.rpc.OperationCallable;
import com.google.api.gax.rpc.ResponseObserver;
import com.google.api.gax.rpc.ServerStream;
import com.google.api.gax.rpc.StatusCode;
import com.google.api.gax.rpc.StatusCode.Code;
import com.google.api.gax.rpc.StreamController;
import com.google.api.gax.rpc.TransportChannelProvider;
import com.google.api.gax.rpc.UnaryCallSettings;
import com.google.api.gax.rpc.UnaryCallable;
import com.google.api.gax.rpc.UnavailableException;
import com.google.api.gax.rpc.WatchdogProvider;
import com.google.api.pathtemplate.PathTemplate;
import com.google.cloud.RetryHelper;
import com.google.cloud.RetryHelper.RetryHelperException;
import com.google.cloud.grpc.GcpManagedChannel;
import com.google.cloud.grpc.GcpManagedChannelBuilder;
import com.google.cloud.grpc.GcpManagedChannelOptions;
import com.google.cloud.grpc.GcpManagedChannelOptions.GcpMetricsOptions;
import com.google.cloud.grpc.GrpcTransportOptions;
import com.google.cloud.spanner.AdminRequestsPerMinuteExceededException;
import com.google.cloud.spanner.BackupId;
import com.google.cloud.spanner.ErrorCode;
import com.google.cloud.spanner.Restore;
import com.google.cloud.spanner.SpannerException;
import com.google.cloud.spanner.SpannerExceptionFactory;
import com.google.cloud.spanner.SpannerOptions;
import com.google.cloud.spanner.SpannerOptions.CallContextConfigurator;
import com.google.cloud.spanner.SpannerOptions.CallCredentialsProvider;
import com.google.cloud.spanner.admin.database.v1.stub.DatabaseAdminStub;
import com.google.cloud.spanner.admin.database.v1.stub.DatabaseAdminStubSettings;
import com.google.cloud.spanner.admin.database.v1.stub.GrpcDatabaseAdminCallableFactory;
import com.google.cloud.spanner.admin.database.v1.stub.GrpcDatabaseAdminStub;
import com.google.cloud.spanner.admin.instance.v1.stub.GrpcInstanceAdminStub;
import com.google.cloud.spanner.admin.instance.v1.stub.InstanceAdminStub;
import com.google.cloud.spanner.admin.instance.v1.stub.InstanceAdminStubSettings;
import com.google.cloud.spanner.encryption.EncryptionConfigProtoMapper;
import com.google.cloud.spanner.v1.stub.GrpcSpannerStub;
import com.google.cloud.spanner.v1.stub.SpannerStub;
import com.google.cloud.spanner.v1.stub.SpannerStubSettings;
import com.google.common.annotations.VisibleForTesting;
import com.google.common.base.Function;
import com.google.common.base.MoreObjects;
import com.google.common.base.Preconditions;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableSet;
import com.google.common.io.Resources;
import com.google.common.util.concurrent.RateLimiter;
import com.google.common.util.concurrent.ThreadFactoryBuilder;
import com.google.iam.v1.GetIamPolicyRequest;
import com.google.iam.v1.GetPolicyOptions;
import com.google.iam.v1.Policy;
import com.google.iam.v1.SetIamPolicyRequest;
import com.google.iam.v1.TestIamPermissionsRequest;
import com.google.iam.v1.TestIamPermissionsResponse;
import com.google.longrunning.CancelOperationRequest;
import com.google.longrunning.GetOperationRequest;
import com.google.longrunning.Operation;
import com.google.longrunning.OperationsGrpc;
import com.google.protobuf.Empty;
import com.google.protobuf.FieldMask;
import com.google.protobuf.InvalidProtocolBufferException;
import com.google.protobuf.Message;
import com.google.protobuf.Timestamp;
import com.google.spanner.admin.database.v1.Backup;
import com.google.spanner.admin.database.v1.CopyBackupMetadata;
import com.google.spanner.admin.database.v1.CopyBackupRequest;
import com.google.spanner.admin.database.v1.CreateBackupMetadata;
import com.google.spanner.admin.database.v1.CreateBackupRequest;
import com.google.spanner.admin.database.v1.CreateDatabaseMetadata;
import com.google.spanner.admin.database.v1.CreateDatabaseRequest;
import com.google.spanner.admin.database.v1.Database;
import com.google.spanner.admin.database.v1.DatabaseAdminGrpc;
import com.google.spanner.admin.database.v1.DatabaseRole;
import com.google.spanner.admin.database.v1.DeleteBackupRequest;
import com.google.spanner.admin.database.v1.DropDatabaseRequest;
import com.google.spanner.admin.database.v1.GetBackupRequest;
import com.google.spanner.admin.database.v1.GetDatabaseDdlRequest;
import com.google.spanner.admin.database.v1.GetDatabaseDdlResponse;
import com.google.spanner.admin.database.v1.GetDatabaseRequest;
import com.google.spanner.admin.database.v1.ListBackupOperationsRequest;
import com.google.spanner.admin.database.v1.ListBackupOperationsResponse;
import com.google.spanner.admin.database.v1.ListBackupsRequest;
import com.google.spanner.admin.database.v1.ListBackupsResponse;
import com.google.spanner.admin.database.v1.ListDatabaseOperationsRequest;
import com.google.spanner.admin.database.v1.ListDatabaseOperationsResponse;
import com.google.spanner.admin.database.v1.ListDatabaseRolesRequest;
import com.google.spanner.admin.database.v1.ListDatabaseRolesResponse;
import com.google.spanner.admin.database.v1.ListDatabasesRequest;
import com.google.spanner.admin.database.v1.ListDatabasesResponse;
import com.google.spanner.admin.database.v1.RestoreDatabaseMetadata;
import com.google.spanner.admin.database.v1.RestoreDatabaseRequest;
import com.google.spanner.admin.database.v1.UpdateBackupRequest;
import com.google.spanner.admin.database.v1.UpdateDatabaseDdlMetadata;
import com.google.spanner.admin.database.v1.UpdateDatabaseDdlRequest;
import com.google.spanner.admin.database.v1.UpdateDatabaseMetadata;
import com.google.spanner.admin.database.v1.UpdateDatabaseRequest;
import com.google.spanner.admin.instance.v1.CreateInstanceConfigMetadata;
import com.google.spanner.admin.instance.v1.CreateInstanceConfigRequest;
import com.google.spanner.admin.instance.v1.CreateInstanceMetadata;
import com.google.spanner.admin.instance.v1.CreateInstanceRequest;
import com.google.spanner.admin.instance.v1.DeleteInstanceConfigRequest;
import com.google.spanner.admin.instance.v1.DeleteInstanceRequest;
import com.google.spanner.admin.instance.v1.GetInstanceConfigRequest;
import com.google.spanner.admin.instance.v1.GetInstanceRequest;
import com.google.spanner.admin.instance.v1.Instance;
import com.google.spanner.admin.instance.v1.InstanceAdminGrpc;
import com.google.spanner.admin.instance.v1.InstanceConfig;
import com.google.spanner.admin.instance.v1.ListInstanceConfigOperationsRequest;
import com.google.spanner.admin.instance.v1.ListInstanceConfigOperationsResponse;
import com.google.spanner.admin.instance.v1.ListInstanceConfigsRequest;
import com.google.spanner.admin.instance.v1.ListInstanceConfigsResponse;
import com.google.spanner.admin.instance.v1.ListInstancesRequest;
import com.google.spanner.admin.instance.v1.ListInstancesResponse;
import com.google.spanner.admin.instance.v1.UpdateInstanceConfigMetadata;
import com.google.spanner.admin.instance.v1.UpdateInstanceConfigRequest;
import com.google.spanner.admin.instance.v1.UpdateInstanceMetadata;
import com.google.spanner.admin.instance.v1.UpdateInstanceRequest;
import com.google.spanner.v1.BatchCreateSessionsRequest;
import com.google.spanner.v1.BatchWriteRequest;
import com.google.spanner.v1.BatchWriteResponse;
import com.google.spanner.v1.BeginTransactionRequest;
import com.google.spanner.v1.CommitRequest;
import com.google.spanner.v1.CommitResponse;
import com.google.spanner.v1.CreateSessionRequest;
import com.google.spanner.v1.DeleteSessionRequest;
import com.google.spanner.v1.ExecuteBatchDmlRequest;
import com.google.spanner.v1.ExecuteBatchDmlResponse;
import com.google.spanner.v1.ExecuteSqlRequest;
import com.google.spanner.v1.PartialResultSet;
import com.google.spanner.v1.PartitionQueryRequest;
import com.google.spanner.v1.PartitionReadRequest;
import com.google.spanner.v1.PartitionResponse;
import com.google.spanner.v1.ReadRequest;
import com.google.spanner.v1.ResultSet;
import com.google.spanner.v1.RollbackRequest;
import com.google.spanner.v1.Session;
import com.google.spanner.v1.SpannerGrpc;
import com.google.spanner.v1.Transaction;
import io.grpc.CallCredentials;
import io.grpc.Context;
import io.grpc.ManagedChannelBuilder;
import io.grpc.MethodDescriptor;
import io.opencensus.metrics.Metrics;
import java.io.IOException;
import java.io.UnsupportedEncodingException;
import java.net.URLDecoder;
import java.nio.charset.Charset;
import java.nio.charset.StandardCharsets;
import java.util.Comparator;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Objects;
import java.util.Optional;
import java.util.Set;
import java.util.concurrent.Callable;
import java.util.concurrent.CancellationException;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentLinkedDeque;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.stream.Collectors;
import java.util.stream.Stream;
import javax.annotation.Nullable;
import org.threeten.bp.Duration;
/** Implementation of Cloud Spanner remote calls using Gapic libraries. */
@InternalApi
public class GapicSpannerRpc implements SpannerRpc {
private static final PathTemplate PROJECT_NAME_TEMPLATE =
PathTemplate.create("projects/{project}");
private static final PathTemplate OPERATION_NAME_TEMPLATE =
PathTemplate.create("{database=projects/*/instances/*/databases/*}/operations/{operation}");
private static final int MAX_MESSAGE_SIZE = 100 * 1024 * 1024;
private static final int MAX_METADATA_SIZE = 32 * 1024; // bytes
private static final String PROPERTY_TIMEOUT_SECONDS =
"com.google.cloud.spanner.watchdogTimeoutSeconds";
private static final String PROPERTY_PERIOD_SECONDS =
"com.google.cloud.spanner.watchdogPeriodSeconds";
private static final int DEFAULT_TIMEOUT_SECONDS = 30 * 60;
private static final int DEFAULT_PERIOD_SECONDS = 10;
private static final int GRPC_KEEPALIVE_SECONDS = 2 * 60;
private static final String USER_AGENT_KEY = "user-agent";
private static final String CLIENT_LIBRARY_LANGUAGE = "spanner-java";
public static final String DEFAULT_USER_AGENT =
CLIENT_LIBRARY_LANGUAGE + "/" + GaxProperties.getLibraryVersion(GapicSpannerRpc.class);
private static final String API_FILE = "grpc-gcp-apiconfig.json";
private boolean rpcIsClosed;
private final SpannerStub spannerStub;
private final RetrySettings executeQueryRetrySettings;
private final Set executeQueryRetryableCodes;
private final RetrySettings readRetrySettings;
private final Set readRetryableCodes;
private final RetrySettings commitRetrySettings;
private final SpannerStub partitionedDmlStub;
private final RetrySettings partitionedDmlRetrySettings;
private final InstanceAdminStubSettings instanceAdminStubSettings;
private final InstanceAdminStub instanceAdminStub;
private final DatabaseAdminStubSettings databaseAdminStubSettings;
private final DatabaseAdminStub databaseAdminStub;
private final String projectId;
private final String projectName;
private final SpannerMetadataProvider metadataProvider;
private final CallCredentialsProvider callCredentialsProvider;
private final String compressorName;
private final Duration waitTimeout =
systemProperty(PROPERTY_TIMEOUT_SECONDS, DEFAULT_TIMEOUT_SECONDS);
private final Duration idleTimeout =
systemProperty(PROPERTY_TIMEOUT_SECONDS, DEFAULT_TIMEOUT_SECONDS);
private final Duration checkInterval =
systemProperty(PROPERTY_PERIOD_SECONDS, DEFAULT_PERIOD_SECONDS);
private final ScheduledExecutorService spannerWatchdog;
private final ConcurrentLinkedDeque responseObservers =
new ConcurrentLinkedDeque<>();
private final boolean throttleAdministrativeRequests;
private final RetrySettings retryAdministrativeRequestsSettings;
private static final double ADMINISTRATIVE_REQUESTS_RATE_LIMIT = 1.0D;
private static final ConcurrentMap ADMINISTRATIVE_REQUESTS_RATE_LIMITERS =
new ConcurrentHashMap<>();
private final boolean leaderAwareRoutingEnabled;
private final int numChannels;
private final boolean isGrpcGcpExtensionEnabled;
public static GapicSpannerRpc create(SpannerOptions options) {
return new GapicSpannerRpc(options);
}
public GapicSpannerRpc(final SpannerOptions options) {
this(options, true);
}
GapicSpannerRpc(final SpannerOptions options, boolean initializeStubs) {
this.projectId = options.getProjectId();
String projectNameStr = PROJECT_NAME_TEMPLATE.instantiate("project", this.projectId);
try {
// Fix use cases where projectName contains special characters.
// This would happen when projects are under an organization.
projectNameStr = URLDecoder.decode(projectNameStr, StandardCharsets.UTF_8.toString());
} catch (UnsupportedEncodingException e) { // Ignored.
}
this.projectName = projectNameStr;
this.throttleAdministrativeRequests = options.isAutoThrottleAdministrativeRequests();
if (throttleAdministrativeRequests) {
ADMINISTRATIVE_REQUESTS_RATE_LIMITERS.putIfAbsent(
projectNameStr, RateLimiter.create(ADMINISTRATIVE_REQUESTS_RATE_LIMIT));
}
this.retryAdministrativeRequestsSettings = options.getRetryAdministrativeRequestsSettings();
// create a metadataProvider which combines both internal headers and
// per-method-call extra headers for channelProvider to inject the headers
// for rpc calls
ApiClientHeaderProvider.Builder internalHeaderProviderBuilder =
ApiClientHeaderProvider.newBuilder();
ApiClientHeaderProvider internalHeaderProvider =
internalHeaderProviderBuilder
.setClientLibToken(
options.getClientLibToken(), GaxProperties.getLibraryVersion(options.getClass()))
.setTransportToken(
GaxGrpcProperties.getGrpcTokenName(), GaxGrpcProperties.getGrpcVersion())
.build();
final HeaderProvider mergedHeaderProvider =
options.getMergedHeaderProvider(internalHeaderProvider);
final HeaderProvider headerProviderWithUserAgent =
headerProviderWithUserAgentFrom(mergedHeaderProvider);
this.metadataProvider =
SpannerMetadataProvider.create(
headerProviderWithUserAgent.getHeaders(),
internalHeaderProviderBuilder.getResourceHeaderKey());
this.callCredentialsProvider = options.getCallCredentialsProvider();
this.compressorName = options.getCompressorName();
this.leaderAwareRoutingEnabled = options.isLeaderAwareRoutingEnabled();
this.numChannels = options.getNumChannels();
this.isGrpcGcpExtensionEnabled = options.isGrpcGcpExtensionEnabled();
if (initializeStubs) {
// First check if SpannerOptions provides a TransportChannelProvider. Create one
// with information gathered from SpannerOptions if none is provided
InstantiatingGrpcChannelProvider.Builder defaultChannelProviderBuilder =
InstantiatingGrpcChannelProvider.newBuilder()
.setChannelConfigurator(options.getChannelConfigurator())
.setEndpoint(options.getEndpoint())
.setMaxInboundMessageSize(MAX_MESSAGE_SIZE)
.setMaxInboundMetadataSize(MAX_METADATA_SIZE)
.setPoolSize(options.getNumChannels())
// Set a keepalive time of 120 seconds to help long running
// commit GRPC calls succeed
.setKeepAliveTime(Duration.ofSeconds(GRPC_KEEPALIVE_SECONDS))
// Then check if SpannerOptions provides an InterceptorProvider. Create a default
// SpannerInterceptorProvider if none is provided
.setInterceptorProvider(
SpannerInterceptorProvider.create(
MoreObjects.firstNonNull(
options.getInterceptorProvider(),
SpannerInterceptorProvider.createDefault(options.getOpenTelemetry())))
// This sets the response compressor (Server -> Client).
.withEncoding(compressorName))
.setHeaderProvider(headerProviderWithUserAgent)
.setAllowNonDefaultServiceAccount(true);
String directPathXdsEnv = System.getenv("GOOGLE_SPANNER_ENABLE_DIRECT_ACCESS");
boolean isAttemptDirectPathXds = Boolean.parseBoolean(directPathXdsEnv);
if (isAttemptDirectPathXds) {
defaultChannelProviderBuilder.setAttemptDirectPath(true);
defaultChannelProviderBuilder.setAttemptDirectPathXds();
}
if (options.isUseVirtualThreads()) {
ExecutorService executor =
tryCreateVirtualThreadPerTaskExecutor("spanner-virtual-grpc-executor");
if (executor != null) {
defaultChannelProviderBuilder.setExecutor(executor);
}
}
// If it is enabled in options uses the channel pool provided by the gRPC-GCP extension.
maybeEnableGrpcGcpExtension(defaultChannelProviderBuilder, options);
TransportChannelProvider channelProvider =
MoreObjects.firstNonNull(
options.getChannelProvider(), defaultChannelProviderBuilder.build());
CredentialsProvider credentialsProvider =
GrpcTransportOptions.setUpCredentialsProvider(options);
spannerWatchdog =
Executors.newSingleThreadScheduledExecutor(
new ThreadFactoryBuilder()
.setDaemon(true)
.setNameFormat("Cloud-Spanner-WatchdogProvider-%d")
.build());
WatchdogProvider watchdogProvider =
InstantiatingWatchdogProvider.create()
.withExecutor(spannerWatchdog)
.withCheckInterval(checkInterval)
.withClock(NanoClock.getDefaultClock());
final String emulatorHost = System.getenv("SPANNER_EMULATOR_HOST");
try {
this.spannerStub =
GrpcSpannerStub.create(
options
.getSpannerStubSettings()
.toBuilder()
.setTransportChannelProvider(channelProvider)
.setCredentialsProvider(credentialsProvider)
.setStreamWatchdogProvider(watchdogProvider)
.setTracerFactory(
options.getApiTracerFactory(
/* isAdminClient = */ false, isEmulatorEnabled(options, emulatorHost)))
.build());
this.readRetrySettings =
options.getSpannerStubSettings().streamingReadSettings().getRetrySettings();
this.readRetryableCodes =
options.getSpannerStubSettings().streamingReadSettings().getRetryableCodes();
this.executeQueryRetrySettings =
options.getSpannerStubSettings().executeStreamingSqlSettings().getRetrySettings();
this.executeQueryRetryableCodes =
options.getSpannerStubSettings().executeStreamingSqlSettings().getRetryableCodes();
this.commitRetrySettings =
options.getSpannerStubSettings().commitSettings().getRetrySettings();
partitionedDmlRetrySettings =
options
.getSpannerStubSettings()
.executeSqlSettings()
.getRetrySettings()
.toBuilder()
.setInitialRpcTimeout(options.getPartitionedDmlTimeout())
.setMaxRpcTimeout(options.getPartitionedDmlTimeout())
.setTotalTimeout(options.getPartitionedDmlTimeout())
.setRpcTimeoutMultiplier(1.0)
.build();
SpannerStubSettings.Builder pdmlSettings = options.getSpannerStubSettings().toBuilder();
pdmlSettings
.setTransportChannelProvider(channelProvider)
.setCredentialsProvider(credentialsProvider)
.setStreamWatchdogProvider(watchdogProvider)
.setTracerFactory(
options.getApiTracerFactory(
/* isAdminClient = */ false, isEmulatorEnabled(options, emulatorHost)))
.executeSqlSettings()
.setRetrySettings(partitionedDmlRetrySettings);
pdmlSettings.executeStreamingSqlSettings().setRetrySettings(partitionedDmlRetrySettings);
// The stream watchdog will by default only check for a timeout every 10 seconds, so if the
// timeout is less than 10 seconds, it would be ignored for the first 10 seconds unless we
// also change the StreamWatchdogCheckInterval.
if (options
.getPartitionedDmlTimeout()
.dividedBy(10L)
.compareTo(pdmlSettings.getStreamWatchdogCheckInterval())
< 0) {
pdmlSettings.setStreamWatchdogCheckInterval(
options.getPartitionedDmlTimeout().dividedBy(10));
pdmlSettings.setStreamWatchdogProvider(
pdmlSettings
.getStreamWatchdogProvider()
.withCheckInterval(pdmlSettings.getStreamWatchdogCheckInterval()));
}
this.partitionedDmlStub = GrpcSpannerStub.create(pdmlSettings.build());
this.instanceAdminStubSettings =
options
.getInstanceAdminStubSettings()
.toBuilder()
.setTransportChannelProvider(channelProvider)
.setCredentialsProvider(credentialsProvider)
.setStreamWatchdogProvider(watchdogProvider)
.setTracerFactory(
options.getApiTracerFactory(
/* isAdminClient = */ true, isEmulatorEnabled(options, emulatorHost)))
.build();
this.instanceAdminStub = GrpcInstanceAdminStub.create(instanceAdminStubSettings);
this.databaseAdminStubSettings =
options
.getDatabaseAdminStubSettings()
.toBuilder()
.setTransportChannelProvider(channelProvider)
.setCredentialsProvider(credentialsProvider)
.setStreamWatchdogProvider(watchdogProvider)
.setTracerFactory(
options.getApiTracerFactory(
/* isAdminClient = */ true, isEmulatorEnabled(options, emulatorHost)))
.build();
// Automatically retry RESOURCE_EXHAUSTED for GetOperation if auto-throttling of
// administrative requests has been set. The GetOperation RPC is called repeatedly by gax
// while polling long-running operations for their progress and can also cause these errors.
// The default behavior is not to retry these errors, and this option should normally only
// be enabled for (integration) testing.
if (options.isAutoThrottleAdministrativeRequests()) {
GrpcStubCallableFactory factory =
new GrpcDatabaseAdminCallableFactory() {
@Override
public UnaryCallable createUnaryCallable(
GrpcCallSettings grpcCallSettings,
UnaryCallSettings callSettings,
ClientContext clientContext) {
// Make GetOperation retry on RESOURCE_EXHAUSTED to prevent polling operations
// from failing with an Administrative requests limit exceeded error.
if (grpcCallSettings
.getMethodDescriptor()
.getFullMethodName()
.equals("google.longrunning.Operations/GetOperation")) {
Set codes =
ImmutableSet.builderWithExpectedSize(
callSettings.getRetryableCodes().size() + 1)
.addAll(callSettings.getRetryableCodes())
.add(StatusCode.Code.RESOURCE_EXHAUSTED)
.build();
callSettings = callSettings.toBuilder().setRetryableCodes(codes).build();
}
return super.createUnaryCallable(grpcCallSettings, callSettings, clientContext);
}
};
this.databaseAdminStub =
new GrpcDatabaseAdminStubWithCustomCallableFactory(
databaseAdminStubSettings,
ClientContext.create(databaseAdminStubSettings),
factory);
} else {
this.databaseAdminStub = GrpcDatabaseAdminStub.create(databaseAdminStubSettings);
}
// Check whether the SPANNER_EMULATOR_HOST env var has been set, and if so, if the emulator
// is actually running.
checkEmulatorConnection(options, channelProvider, credentialsProvider, emulatorHost);
} catch (Exception e) {
throw newSpannerException(e);
}
} else {
this.databaseAdminStub = null;
this.instanceAdminStub = null;
this.spannerStub = null;
this.readRetrySettings = null;
this.readRetryableCodes = null;
this.executeQueryRetrySettings = null;
this.executeQueryRetryableCodes = null;
this.commitRetrySettings =
SpannerStubSettings.newBuilder().commitSettings().getRetrySettings();
this.partitionedDmlStub = null;
this.databaseAdminStubSettings = null;
this.instanceAdminStubSettings = null;
this.spannerWatchdog = null;
this.partitionedDmlRetrySettings = null;
}
}
private static String parseGrpcGcpApiConfig() {
try {
return Resources.toString(
GapicSpannerRpc.class.getResource(API_FILE), Charset.forName("UTF8"));
} catch (IOException e) {
throw newSpannerException(e);
}
}
// Enhance metric options for gRPC-GCP extension. Adds metric registry if not specified.
private static GcpManagedChannelOptions grpcGcpOptionsWithMetrics(SpannerOptions options) {
GcpManagedChannelOptions grpcGcpOptions =
MoreObjects.firstNonNull(options.getGrpcGcpOptions(), new GcpManagedChannelOptions());
GcpMetricsOptions metricsOptions =
MoreObjects.firstNonNull(
grpcGcpOptions.getMetricsOptions(), GcpMetricsOptions.newBuilder().build());
GcpMetricsOptions.Builder metricsOptionsBuilder = GcpMetricsOptions.newBuilder(metricsOptions);
if (metricsOptions.getMetricRegistry() == null) {
metricsOptionsBuilder.withMetricRegistry(Metrics.getMetricRegistry());
}
// TODO: Add default labels with values: client_id, database, instance_id.
if (metricsOptions.getNamePrefix().equals("")) {
metricsOptionsBuilder.withNamePrefix("cloud.google.com/java/spanner/gcp-channel-pool/");
}
return GcpManagedChannelOptions.newBuilder(grpcGcpOptions)
.withMetricsOptions(metricsOptionsBuilder.build())
.build();
}
@SuppressWarnings("rawtypes")
private static void maybeEnableGrpcGcpExtension(
InstantiatingGrpcChannelProvider.Builder defaultChannelProviderBuilder,
final SpannerOptions options) {
if (!options.isGrpcGcpExtensionEnabled()) {
return;
}
final String jsonApiConfig = parseGrpcGcpApiConfig();
final GcpManagedChannelOptions grpcGcpOptions = grpcGcpOptionsWithMetrics(options);
ApiFunction apiFunction =
channelBuilder -> {
if (options.getChannelConfigurator() != null) {
channelBuilder = options.getChannelConfigurator().apply(channelBuilder);
}
return GcpManagedChannelBuilder.forDelegateBuilder(channelBuilder)
.withApiConfigJsonString(jsonApiConfig)
.withOptions(grpcGcpOptions)
.setPoolSize(options.getNumChannels());
};
// Disable the GAX channel pooling functionality by setting the GAX channel pool size to 1.
// Enable gRPC-GCP channel pool via the channel configurator.
defaultChannelProviderBuilder.setPoolSize(1).setChannelConfigurator(apiFunction);
}
private static HeaderProvider headerProviderWithUserAgentFrom(HeaderProvider headerProvider) {
final Optional> existingUserAgentEntry =
headerProvider.getHeaders().entrySet().stream()
.filter(entry -> entry.getKey().equalsIgnoreCase(USER_AGENT_KEY))
.findFirst();
final String existingUserAgentValue = existingUserAgentEntry.map(Entry::getValue).orElse(null);
final String userAgent =
Stream.of(existingUserAgentValue, DEFAULT_USER_AGENT)
.filter(Objects::nonNull)
.collect(Collectors.joining(" "));
final Map headersWithUserAgent = new HashMap<>(headerProvider.getHeaders());
existingUserAgentEntry.ifPresent(entry -> headersWithUserAgent.remove(entry.getKey()));
headersWithUserAgent.put(USER_AGENT_KEY, userAgent);
return FixedHeaderProvider.create(headersWithUserAgent);
}
private static void checkEmulatorConnection(
SpannerOptions options,
TransportChannelProvider channelProvider,
CredentialsProvider credentialsProvider,
String emulatorHost)
throws IOException {
// Only do the check if the emulator environment variable has been set to localhost.
if (isEmulatorEnabled(options, emulatorHost)) {
// Do a quick check to see if the emulator is actually running.
try {
InstanceAdminStubSettings.Builder testEmulatorSettings =
options
.getInstanceAdminStubSettings()
.toBuilder()
.setTransportChannelProvider(channelProvider)
.setCredentialsProvider(credentialsProvider);
testEmulatorSettings
.listInstanceConfigsSettings()
.setSimpleTimeoutNoRetries(Duration.ofSeconds(10L));
try (GrpcInstanceAdminStub stub =
GrpcInstanceAdminStub.create(testEmulatorSettings.build())) {
stub.listInstanceConfigsCallable()
.call(
ListInstanceConfigsRequest.newBuilder()
.setParent(String.format("projects/%s", options.getProjectId()))
.build());
}
} catch (UnavailableException e) {
throw SpannerExceptionFactory.newSpannerException(
ErrorCode.UNAVAILABLE,
String.format(
"The environment variable SPANNER_EMULATOR_HOST has been set to %s, but no running"
+ " emulator could be found at that address.\n"
+ "Did you forget to start the emulator, or to unset the environment"
+ " variable?",
emulatorHost));
}
}
}
private static boolean isEmulatorEnabled(SpannerOptions options, String emulatorHost) {
// Only do the check if the emulator environment variable has been set to localhost.
return options.getChannelProvider() == null
&& emulatorHost != null
&& options.getHost() != null
&& options.getHost().startsWith("http://localhost")
&& options.getHost().endsWith(emulatorHost);
}
private static final RetrySettings ADMIN_REQUESTS_LIMIT_EXCEEDED_RETRY_SETTINGS =
RetrySettings.newBuilder()
.setInitialRetryDelay(Duration.ofSeconds(5L))
.setRetryDelayMultiplier(2.0)
.setMaxRetryDelay(Duration.ofSeconds(60L))
.setMaxAttempts(10)
.build();
@VisibleForTesting
static final class AdminRequestsLimitExceededRetryAlgorithm
implements ResultRetryAlgorithm {
@Override
public TimedAttemptSettings createNextAttempt(
Throwable prevThrowable, T prevResponse, TimedAttemptSettings prevSettings) {
// Use default retry settings.
return null;
}
@Override
public boolean shouldRetry(Throwable prevThrowable, T prevResponse)
throws CancellationException {
return prevThrowable instanceof AdminRequestsPerMinuteExceededException;
}
}
private T runWithRetryOnAdministrativeRequestsExceeded(Callable callable) {
try {
return RetryHelper.runWithRetries(
callable,
retryAdministrativeRequestsSettings,
new AdminRequestsLimitExceededRetryAlgorithm<>(),
NanoClock.getDefaultClock());
} catch (RetryHelperException e) {
throw SpannerExceptionFactory.asSpannerException(e.getCause());
}
}
private static final class OperationFutureRetryAlgorithm
implements ResultRetryAlgorithm> {
private static final ImmutableList RETRYABLE_CODES =
ImmutableList.of(StatusCode.Code.DEADLINE_EXCEEDED, StatusCode.Code.UNAVAILABLE);
@Override
public TimedAttemptSettings createNextAttempt(
Throwable prevThrowable,
OperationFuture prevResponse,
TimedAttemptSettings prevSettings) {
// Use default retry settings.
return null;
}
@Override
public boolean shouldRetry(
Throwable prevThrowable, OperationFuture prevResponse)
throws CancellationException {
if (prevThrowable instanceof ApiException) {
ApiException e = (ApiException) prevThrowable;
return RETRYABLE_CODES.contains(e.getStatusCode().getCode());
}
if (prevResponse != null) {
try {
prevResponse.getInitialFuture().get();
} catch (ExecutionException ee) {
Throwable cause = ee.getCause();
if (cause instanceof ApiException) {
ApiException e = (ApiException) cause;
return RETRYABLE_CODES.contains(e.getStatusCode().getCode());
}
} catch (InterruptedException e) {
return false;
}
}
return false;
}
}
private final class OperationFutureCallable
implements Callable> {
final OperationCallable operationCallable;
final RequestT initialRequest;
final MethodDescriptor method;
final String instanceName;
final OperationsLister lister;
final Function getStartTimeFunction;
Timestamp initialCallTime;
boolean isRetry = false;
OperationFutureCallable(
OperationCallable operationCallable,
RequestT initialRequest,
MethodDescriptor method,
String instanceName,
OperationsLister lister,
Function getStartTimeFunction) {
this.operationCallable = operationCallable;
this.initialRequest = initialRequest;
this.method = method;
this.instanceName = instanceName;
this.lister = lister;
this.getStartTimeFunction = getStartTimeFunction;
}
@Override
public OperationFuture call() {
acquireAdministrativeRequestsRateLimiter();
return runWithRetryOnAdministrativeRequestsExceeded(
() -> {
String operationName = null;
if (isRetry) {
// Query the backend to see if the operation was actually created, and that the
// problem was caused by a network problem or other transient problem client side.
Operation operation =
mostRecentOperation(lister, getStartTimeFunction, initialCallTime);
if (operation != null) {
// Operation found, resume tracking that operation.
operationName = operation.getName();
}
} else {
initialCallTime =
Timestamp.newBuilder()
.setSeconds(
TimeUnit.SECONDS.convert(
System.currentTimeMillis(), TimeUnit.MILLISECONDS))
.build();
}
isRetry = true;
if (operationName == null) {
GrpcCallContext context = newCallContext(null, instanceName, initialRequest, method);
return operationCallable.futureCall(initialRequest, context);
} else {
return operationCallable.resumeFutureCall(operationName);
}
});
}
}
private interface OperationsLister {
Paginated listOperations(String nextPageToken);
}
private Operation mostRecentOperation(
OperationsLister lister,
Function getStartTimeFunction,
Timestamp initialCallTime) {
Operation res = null;
Timestamp currMaxStartTime = null;
String nextPageToken = null;
Paginated operations;
do {
operations = lister.listOperations(nextPageToken);
nextPageToken = operations.getNextPageToken();
for (Operation op : operations.getResults()) {
Timestamp startTime = getStartTimeFunction.apply(op);
if (res == null
|| (TimestampComparator.INSTANCE.compare(startTime, currMaxStartTime) > 0
&& TimestampComparator.INSTANCE.compare(startTime, initialCallTime) >= 0)) {
currMaxStartTime = startTime;
res = op;
}
// If the operation does not report any start time, then the operation that is not yet done
// is the one that is the most recent.
if (startTime == null && currMaxStartTime == null && !op.getDone()) {
res = op;
break;
}
}
} while (nextPageToken != null);
return res;
}
private static final class TimestampComparator implements Comparator {
private static final TimestampComparator INSTANCE = new TimestampComparator();
@Override
public int compare(Timestamp t1, Timestamp t2) {
if (t1 == null && t2 == null) {
return 0;
}
if (t1 != null && t2 == null) {
return 1;
}
if (t1 == null && t2 != null) {
return -1;
}
if (t1.getSeconds() > t2.getSeconds()
|| (t1.getSeconds() == t2.getSeconds() && t1.getNanos() > t2.getNanos())) {
return 1;
}
if (t1.getSeconds() < t2.getSeconds()
|| (t1.getSeconds() == t2.getSeconds() && t1.getNanos() < t2.getNanos())) {
return -1;
}
return 0;
}
}
private void acquireAdministrativeRequestsRateLimiter() {
if (throttleAdministrativeRequests) {
RateLimiter limiter = ADMINISTRATIVE_REQUESTS_RATE_LIMITERS.get(this.projectName);
if (limiter != null) {
limiter.acquire();
}
}
}
@Override
public Paginated listInstanceConfigs(int pageSize, @Nullable String pageToken)
throws SpannerException {
ListInstanceConfigsRequest.Builder requestBuilder =
ListInstanceConfigsRequest.newBuilder().setParent(projectName).setPageSize(pageSize);
if (pageToken != null) {
requestBuilder.setPageToken(pageToken);
}
ListInstanceConfigsRequest request = requestBuilder.build();
GrpcCallContext context =
newCallContext(
null, projectName, request, InstanceAdminGrpc.getListInstanceConfigsMethod());
ListInstanceConfigsResponse response =
get(instanceAdminStub.listInstanceConfigsCallable().futureCall(request, context));
return new Paginated<>(response.getInstanceConfigsList(), response.getNextPageToken());
}
@Override
public OperationFuture createInstanceConfig(
String parent,
String instanceConfigId,
InstanceConfig instanceConfig,
@Nullable Boolean validateOnly)
throws SpannerException {
CreateInstanceConfigRequest.Builder builder =
CreateInstanceConfigRequest.newBuilder()
.setParent(parent)
.setInstanceConfigId(instanceConfigId)
.setInstanceConfig(instanceConfig);
if (validateOnly != null) {
builder.setValidateOnly(validateOnly);
}
CreateInstanceConfigRequest request = builder.build();
GrpcCallContext context =
newCallContext(null, parent, request, InstanceAdminGrpc.getCreateInstanceConfigMethod());
return instanceAdminStub.createInstanceConfigOperationCallable().futureCall(request, context);
}
@Override
public OperationFuture updateInstanceConfig(
InstanceConfig instanceConfig, @Nullable Boolean validateOnly, FieldMask fieldMask)
throws SpannerException {
UpdateInstanceConfigRequest.Builder builder =
UpdateInstanceConfigRequest.newBuilder()
.setInstanceConfig(instanceConfig)
.setUpdateMask(fieldMask);
if (validateOnly != null) {
builder.setValidateOnly(validateOnly);
}
UpdateInstanceConfigRequest request = builder.build();
GrpcCallContext context =
newCallContext(
null,
instanceConfig.getName(),
request,
InstanceAdminGrpc.getUpdateInstanceConfigMethod());
return instanceAdminStub.updateInstanceConfigOperationCallable().futureCall(request, context);
}
@Override
public InstanceConfig getInstanceConfig(String instanceConfigName) throws SpannerException {
GetInstanceConfigRequest request =
GetInstanceConfigRequest.newBuilder().setName(instanceConfigName).build();
GrpcCallContext context =
newCallContext(null, projectName, request, InstanceAdminGrpc.getGetInstanceConfigMethod());
return get(instanceAdminStub.getInstanceConfigCallable().futureCall(request, context));
}
@Override
public void deleteInstanceConfig(
String instanceConfigName, @Nullable String etag, @Nullable Boolean validateOnly)
throws SpannerException {
DeleteInstanceConfigRequest.Builder requestBuilder =
DeleteInstanceConfigRequest.newBuilder().setName(instanceConfigName);
if (etag != null) {
requestBuilder.setEtag(etag);
}
if (validateOnly != null) {
requestBuilder.setValidateOnly(validateOnly);
}
DeleteInstanceConfigRequest request = requestBuilder.build();
GrpcCallContext context =
newCallContext(
null, instanceConfigName, request, InstanceAdminGrpc.getDeleteInstanceConfigMethod());
get(instanceAdminStub.deleteInstanceConfigCallable().futureCall(request, context));
}
@Override
public Paginated listInstanceConfigOperations(
int pageSize, @Nullable String filter, @Nullable String pageToken) {
acquireAdministrativeRequestsRateLimiter();
ListInstanceConfigOperationsRequest.Builder requestBuilder =
ListInstanceConfigOperationsRequest.newBuilder()
.setParent(projectName)
.setPageSize(pageSize);
if (filter != null) {
requestBuilder.setFilter(filter);
}
if (pageToken != null) {
requestBuilder.setPageToken(pageToken);
}
final ListInstanceConfigOperationsRequest request = requestBuilder.build();
final GrpcCallContext context =
newCallContext(
null, projectName, request, InstanceAdminGrpc.getListInstanceConfigOperationsMethod());
ListInstanceConfigOperationsResponse response =
runWithRetryOnAdministrativeRequestsExceeded(
() ->
get(
instanceAdminStub
.listInstanceConfigOperationsCallable()
.futureCall(request, context)));
return new Paginated<>(response.getOperationsList(), response.getNextPageToken());
}
@Override
public Paginated listInstances(
int pageSize, @Nullable String pageToken, @Nullable String filter) throws SpannerException {
ListInstancesRequest.Builder requestBuilder =
ListInstancesRequest.newBuilder().setParent(projectName).setPageSize(pageSize);
if (pageToken != null) {
requestBuilder.setPageToken(pageToken);
}
if (filter != null) {
requestBuilder.setFilter(filter);
}
ListInstancesRequest request = requestBuilder.build();
GrpcCallContext context =
newCallContext(null, projectName, request, InstanceAdminGrpc.getListInstancesMethod());
ListInstancesResponse response =
get(instanceAdminStub.listInstancesCallable().futureCall(request, context));
return new Paginated<>(response.getInstancesList(), response.getNextPageToken());
}
@Override
public OperationFuture createInstance(
String parent, String instanceId, Instance instance) throws SpannerException {
CreateInstanceRequest request =
CreateInstanceRequest.newBuilder()
.setParent(parent)
.setInstanceId(instanceId)
.setInstance(instance)
.build();
GrpcCallContext context =
newCallContext(null, parent, request, InstanceAdminGrpc.getCreateInstanceMethod());
return instanceAdminStub.createInstanceOperationCallable().futureCall(request, context);
}
@Override
public OperationFuture updateInstance(
Instance instance, FieldMask fieldMask) throws SpannerException {
UpdateInstanceRequest request =
UpdateInstanceRequest.newBuilder().setInstance(instance).setFieldMask(fieldMask).build();
GrpcCallContext context =
newCallContext(
null, instance.getName(), request, InstanceAdminGrpc.getUpdateInstanceMethod());
return instanceAdminStub.updateInstanceOperationCallable().futureCall(request, context);
}
@Override
public Instance getInstance(String instanceName) throws SpannerException {
GetInstanceRequest request = GetInstanceRequest.newBuilder().setName(instanceName).build();
GrpcCallContext context =
newCallContext(null, instanceName, request, InstanceAdminGrpc.getGetInstanceMethod());
return get(instanceAdminStub.getInstanceCallable().futureCall(request, context));
}
@Override
public void deleteInstance(String instanceName) throws SpannerException {
DeleteInstanceRequest request =
DeleteInstanceRequest.newBuilder().setName(instanceName).build();
GrpcCallContext context =
newCallContext(null, instanceName, request, InstanceAdminGrpc.getDeleteInstanceMethod());
get(instanceAdminStub.deleteInstanceCallable().futureCall(request, context));
}
@Override
public Paginated listBackupOperations(
String instanceName, int pageSize, @Nullable String filter, @Nullable String pageToken) {
acquireAdministrativeRequestsRateLimiter();
ListBackupOperationsRequest.Builder requestBuilder =
ListBackupOperationsRequest.newBuilder().setParent(instanceName).setPageSize(pageSize);
if (filter != null) {
requestBuilder.setFilter(filter);
}
if (pageToken != null) {
requestBuilder.setPageToken(pageToken);
}
final ListBackupOperationsRequest request = requestBuilder.build();
final GrpcCallContext context =
newCallContext(
null, instanceName, request, DatabaseAdminGrpc.getListBackupOperationsMethod());
ListBackupOperationsResponse response =
runWithRetryOnAdministrativeRequestsExceeded(
() ->
get(databaseAdminStub.listBackupOperationsCallable().futureCall(request, context)));
return new Paginated<>(response.getOperationsList(), response.getNextPageToken());
}
@Override
public Paginated listDatabaseOperations(
String instanceName, int pageSize, @Nullable String filter, @Nullable String pageToken) {
acquireAdministrativeRequestsRateLimiter();
ListDatabaseOperationsRequest.Builder requestBuilder =
ListDatabaseOperationsRequest.newBuilder().setParent(instanceName).setPageSize(pageSize);
if (filter != null) {
requestBuilder.setFilter(filter);
}
if (pageToken != null) {
requestBuilder.setPageToken(pageToken);
}
final ListDatabaseOperationsRequest request = requestBuilder.build();
final GrpcCallContext context =
newCallContext(
null, instanceName, request, DatabaseAdminGrpc.getListDatabaseOperationsMethod());
ListDatabaseOperationsResponse response =
runWithRetryOnAdministrativeRequestsExceeded(
() ->
get(
databaseAdminStub
.listDatabaseOperationsCallable()
.futureCall(request, context)));
return new Paginated<>(response.getOperationsList(), response.getNextPageToken());
}
@Override
public Paginated listDatabaseRoles(
String databaseName, int pageSize, @Nullable String pageToken) {
acquireAdministrativeRequestsRateLimiter();
ListDatabaseRolesRequest.Builder requestBuilder =
ListDatabaseRolesRequest.newBuilder().setParent(databaseName).setPageSize(pageSize);
if (pageToken != null) {
requestBuilder.setPageToken(pageToken);
}
final ListDatabaseRolesRequest request = requestBuilder.build();
final GrpcCallContext context =
newCallContext(null, databaseName, request, DatabaseAdminGrpc.getListDatabaseRolesMethod());
ListDatabaseRolesResponse response =
runWithRetryOnAdministrativeRequestsExceeded(
() -> get(databaseAdminStub.listDatabaseRolesCallable().futureCall(request, context)));
return new Paginated<>(response.getDatabaseRolesList(), response.getNextPageToken());
}
@Override
public Paginated listBackups(
String instanceName, int pageSize, @Nullable String filter, @Nullable String pageToken)
throws SpannerException {
acquireAdministrativeRequestsRateLimiter();
ListBackupsRequest.Builder requestBuilder =
ListBackupsRequest.newBuilder().setParent(instanceName).setPageSize(pageSize);
if (filter != null) {
requestBuilder.setFilter(filter);
}
if (pageToken != null) {
requestBuilder.setPageToken(pageToken);
}
final ListBackupsRequest request = requestBuilder.build();
final GrpcCallContext context =
newCallContext(null, instanceName, request, DatabaseAdminGrpc.getListBackupsMethod());
ListBackupsResponse response =
runWithRetryOnAdministrativeRequestsExceeded(
() -> get(databaseAdminStub.listBackupsCallable().futureCall(request, context)));
return new Paginated<>(response.getBackupsList(), response.getNextPageToken());
}
@Override
public Paginated listDatabases(
String instanceName, int pageSize, @Nullable String pageToken) throws SpannerException {
acquireAdministrativeRequestsRateLimiter();
ListDatabasesRequest.Builder requestBuilder =
ListDatabasesRequest.newBuilder().setParent(instanceName).setPageSize(pageSize);
if (pageToken != null) {
requestBuilder.setPageToken(pageToken);
}
final ListDatabasesRequest request = requestBuilder.build();
final GrpcCallContext context =
newCallContext(null, instanceName, request, DatabaseAdminGrpc.getListDatabasesMethod());
ListDatabasesResponse response =
runWithRetryOnAdministrativeRequestsExceeded(
() -> get(databaseAdminStub.listDatabasesCallable().futureCall(request, context)));
return new Paginated<>(response.getDatabasesList(), response.getNextPageToken());
}
@Override
public OperationFuture createDatabase(
final String instanceName,
String createDatabaseStatement,
Iterable additionalStatements,
com.google.cloud.spanner.Database databaseInfo)
throws SpannerException {
final String databaseId = databaseInfo.getId().getDatabase();
CreateDatabaseRequest.Builder requestBuilder =
CreateDatabaseRequest.newBuilder()
.setParent(instanceName)
.setCreateStatement(createDatabaseStatement)
.addAllExtraStatements(additionalStatements);
if (databaseInfo.getEncryptionConfig() != null) {
requestBuilder.setEncryptionConfig(
EncryptionConfigProtoMapper.encryptionConfig(databaseInfo.getEncryptionConfig()));
}
if (databaseInfo.getDialect() != null) {
requestBuilder.setDatabaseDialect(databaseInfo.getDialect().toProto());
}
if (databaseInfo.getProtoDescriptors() != null) {
requestBuilder.setProtoDescriptors(databaseInfo.getProtoDescriptors());
}
final CreateDatabaseRequest request = requestBuilder.build();
OperationFutureCallable callable =
new OperationFutureCallable<>(
databaseAdminStub.createDatabaseOperationCallable(),
request,
DatabaseAdminGrpc.getCreateDatabaseMethod(),
instanceName,
nextPageToken ->
listDatabaseOperations(
instanceName,
0,
String.format(
"(metadata.@type:type.googleapis.com/%s) AND (name:%s/operations/)",
CreateDatabaseMetadata.getDescriptor().getFullName(),
String.format("%s/databases/%s", instanceName, databaseId)),
nextPageToken),
input -> {
if (input.getDone() && input.hasResponse()) {
try {
Timestamp createTime = input.getResponse().unpack(Database.class).getCreateTime();
if (Timestamp.getDefaultInstance().equals(createTime)) {
// Create time was not returned by the server (proto objects never return
// null, instead they return the default instance). Return null from this
// method to indicate that there is no known create time.
return null;
}
} catch (InvalidProtocolBufferException e) {
return null;
}
}
return null;
});
return RetryHelper.runWithRetries(
callable,
databaseAdminStubSettings
.createDatabaseOperationSettings()
.getInitialCallSettings()
.getRetrySettings(),
new OperationFutureRetryAlgorithm<>(),
NanoClock.getDefaultClock());
}
/**
* If the update database ddl operation returns an ALREADY_EXISTS error, meaning the operation id
* used is already in flight, this method will simply resume the original operation. The returned
* future will be completed when the original operation finishes.
*
* This mechanism is necessary, because the update database ddl can be retried. If a retryable
* failure occurs, the backend has already started processing the update database ddl operation
* with the given id and the library issues a retry, an ALREADY_EXISTS error will be returned. If
* we were to bubble this error up, it would be confusing for the caller, who used originally
* called the method with a new operation id.
*/
@Override
public OperationFuture updateDatabaseDdl(
com.google.cloud.spanner.Database databaseInfo,
final Iterable updateDatabaseStatements,
@Nullable final String updateId)
throws SpannerException {
acquireAdministrativeRequestsRateLimiter();
Preconditions.checkNotNull(databaseInfo.getId());
UpdateDatabaseDdlRequest.Builder requestBuilder =
UpdateDatabaseDdlRequest.newBuilder()
.setDatabase(databaseInfo.getId().getName())
.addAllStatements(updateDatabaseStatements)
.setOperationId(MoreObjects.firstNonNull(updateId, ""));
if (databaseInfo.getProtoDescriptors() != null) {
requestBuilder.setProtoDescriptors(databaseInfo.getProtoDescriptors());
}
final UpdateDatabaseDdlRequest request = requestBuilder.build();
final GrpcCallContext context =
newCallContext(
null,
databaseInfo.getId().getName(),
request,
DatabaseAdminGrpc.getUpdateDatabaseDdlMethod());
final OperationCallable callable =
databaseAdminStub.updateDatabaseDdlOperationCallable();
return runWithRetryOnAdministrativeRequestsExceeded(
() -> {
OperationFuture operationFuture =
callable.futureCall(request, context);
try {
operationFuture.getInitialFuture().get();
} catch (InterruptedException e) {
throw newSpannerException(e);
} catch (ExecutionException e) {
Throwable t = e.getCause();
SpannerException se = SpannerExceptionFactory.asSpannerException(t);
if (se instanceof AdminRequestsPerMinuteExceededException) {
// Propagate this to trigger a retry.
throw se;
}
if (t instanceof AlreadyExistsException) {
String operationName =
OPERATION_NAME_TEMPLATE.instantiate(
"database", databaseInfo.getId().getName(), "operation", updateId);
return callable.resumeFutureCall(operationName, context);
}
}
return operationFuture;
});
}
@Override
public void dropDatabase(String databaseName) throws SpannerException {
acquireAdministrativeRequestsRateLimiter();
final DropDatabaseRequest request =
DropDatabaseRequest.newBuilder().setDatabase(databaseName).build();
final GrpcCallContext context =
newCallContext(null, databaseName, request, DatabaseAdminGrpc.getDropDatabaseMethod());
runWithRetryOnAdministrativeRequestsExceeded(
() -> {
get(databaseAdminStub.dropDatabaseCallable().futureCall(request, context));
return null;
});
}
@Override
public Database getDatabase(String databaseName) throws SpannerException {
acquireAdministrativeRequestsRateLimiter();
final GetDatabaseRequest request =
GetDatabaseRequest.newBuilder().setName(databaseName).build();
final GrpcCallContext context =
newCallContext(null, databaseName, request, DatabaseAdminGrpc.getGetDatabaseMethod());
return runWithRetryOnAdministrativeRequestsExceeded(
() -> get(databaseAdminStub.getDatabaseCallable().futureCall(request, context)));
}
@Override
public OperationFuture updateDatabase(
Database database, FieldMask updateMask) throws SpannerException {
UpdateDatabaseRequest request =
UpdateDatabaseRequest.newBuilder().setDatabase(database).setUpdateMask(updateMask).build();
GrpcCallContext context =
newCallContext(
null, database.getName(), request, DatabaseAdminGrpc.getUpdateDatabaseMethod());
return databaseAdminStub.updateDatabaseOperationCallable().futureCall(request, context);
}
@Override
public GetDatabaseDdlResponse getDatabaseDdl(String databaseName) throws SpannerException {
acquireAdministrativeRequestsRateLimiter();
final GetDatabaseDdlRequest request =
GetDatabaseDdlRequest.newBuilder().setDatabase(databaseName).build();
final GrpcCallContext context =
newCallContext(null, databaseName, request, DatabaseAdminGrpc.getGetDatabaseDdlMethod());
return runWithRetryOnAdministrativeRequestsExceeded(
() -> get(databaseAdminStub.getDatabaseDdlCallable().futureCall(request, context)));
}
@Override
public OperationFuture createBackup(
final com.google.cloud.spanner.Backup backupInfo) throws SpannerException {
final String instanceName = backupInfo.getInstanceId().getName();
final String databaseName = backupInfo.getDatabase().getName();
final String backupId = backupInfo.getId().getBackup();
final Backup.Builder backupBuilder =
com.google.spanner.admin.database.v1.Backup.newBuilder()
.setDatabase(databaseName)
.setExpireTime(backupInfo.getExpireTime().toProto());
if (backupInfo.getVersionTime() != null) {
backupBuilder.setVersionTime(backupInfo.getVersionTime().toProto());
}
final Backup backup = backupBuilder.build();
final CreateBackupRequest.Builder requestBuilder =
CreateBackupRequest.newBuilder()
.setParent(instanceName)
.setBackupId(backupId)
.setBackup(backup);
if (backupInfo.getEncryptionConfig() != null) {
requestBuilder.setEncryptionConfig(
EncryptionConfigProtoMapper.createBackupEncryptionConfig(
backupInfo.getEncryptionConfig()));
}
final CreateBackupRequest request = requestBuilder.build();
final OperationFutureCallable callable =
new OperationFutureCallable<>(
databaseAdminStub.createBackupOperationCallable(),
request,
DatabaseAdminGrpc.getCreateBackupMethod(),
instanceName,
nextPageToken ->
listBackupOperations(
instanceName,
0,
String.format(
"(metadata.@type:type.googleapis.com/%s) AND (metadata.name:%s)",
CreateBackupMetadata.getDescriptor().getFullName(),
String.format("%s/backups/%s", instanceName, backupId)),
nextPageToken),
input -> {
try {
return input
.getMetadata()
.unpack(CreateBackupMetadata.class)
.getProgress()
.getStartTime();
} catch (InvalidProtocolBufferException e) {
return null;
}
});
return RetryHelper.runWithRetries(
callable,
databaseAdminStubSettings
.createBackupOperationSettings()
.getInitialCallSettings()
.getRetrySettings(),
new OperationFutureRetryAlgorithm<>(),
NanoClock.getDefaultClock());
}
@Override
public OperationFuture copyBackup(
BackupId sourceBackupId, final com.google.cloud.spanner.Backup destinationBackup)
throws SpannerException {
Preconditions.checkNotNull(sourceBackupId);
Preconditions.checkNotNull(destinationBackup);
final String instanceName = destinationBackup.getInstanceId().getName();
final String backupId = destinationBackup.getId().getBackup();
final CopyBackupRequest.Builder requestBuilder =
CopyBackupRequest.newBuilder()
.setParent(instanceName)
.setBackupId(backupId)
.setSourceBackup(sourceBackupId.getName())
.setExpireTime(destinationBackup.getExpireTime().toProto());
if (destinationBackup.getEncryptionConfig() != null) {
requestBuilder.setEncryptionConfig(
EncryptionConfigProtoMapper.copyBackupEncryptionConfig(
destinationBackup.getEncryptionConfig()));
}
final CopyBackupRequest request = requestBuilder.build();
final OperationFutureCallable callable =
new OperationFutureCallable<>(
databaseAdminStub.copyBackupOperationCallable(),
request,
// calling copy backup method of dbClientImpl
DatabaseAdminGrpc.getCopyBackupMethod(),
instanceName,
nextPageToken ->
listBackupOperations(
instanceName,
0,
String.format(
"(metadata.@type:type.googleapis.com/%s) AND (metadata.name:%s)",
CopyBackupMetadata.getDescriptor().getFullName(),
String.format("%s/backups/%s", instanceName, backupId)),
nextPageToken),
input -> {
try {
return input
.getMetadata()
.unpack(CopyBackupMetadata.class)
.getProgress()
.getStartTime();
} catch (InvalidProtocolBufferException e) {
return null;
}
});
return RetryHelper.runWithRetries(
callable,
databaseAdminStubSettings
.copyBackupOperationSettings()
.getInitialCallSettings()
.getRetrySettings(),
new OperationFutureRetryAlgorithm<>(),
NanoClock.getDefaultClock());
}
@Override
public OperationFuture restoreDatabase(final Restore restore) {
final String databaseInstanceName = restore.getDestination().getInstanceId().getName();
final String databaseId = restore.getDestination().getDatabase();
final RestoreDatabaseRequest.Builder requestBuilder =
RestoreDatabaseRequest.newBuilder()
.setParent(databaseInstanceName)
.setDatabaseId(databaseId)
.setBackup(restore.getSource().getName());
if (restore.getEncryptionConfig() != null) {
requestBuilder.setEncryptionConfig(
EncryptionConfigProtoMapper.restoreDatabaseEncryptionConfig(
restore.getEncryptionConfig()));
}
final OperationFutureCallable
callable =
new OperationFutureCallable<>(
databaseAdminStub.restoreDatabaseOperationCallable(),
requestBuilder.build(),
DatabaseAdminGrpc.getRestoreDatabaseMethod(),
databaseInstanceName,
nextPageToken ->
listDatabaseOperations(
databaseInstanceName,
0,
String.format(
"(metadata.@type:type.googleapis.com/%s) AND (metadata.name:%s)",
RestoreDatabaseMetadata.getDescriptor().getFullName(),
String.format("%s/databases/%s", databaseInstanceName, databaseId)),
nextPageToken),
input -> {
try {
return input
.getMetadata()
.unpack(RestoreDatabaseMetadata.class)
.getProgress()
.getStartTime();
} catch (InvalidProtocolBufferException e) {
return null;
}
});
return RetryHelper.runWithRetries(
callable,
databaseAdminStubSettings
.restoreDatabaseOperationSettings()
.getInitialCallSettings()
.getRetrySettings(),
new OperationFutureRetryAlgorithm<>(),
NanoClock.getDefaultClock());
}
@Override
public Backup updateBackup(Backup backup, FieldMask updateMask) {
acquireAdministrativeRequestsRateLimiter();
final UpdateBackupRequest request =
UpdateBackupRequest.newBuilder().setBackup(backup).setUpdateMask(updateMask).build();
final GrpcCallContext context =
newCallContext(null, backup.getName(), request, DatabaseAdminGrpc.getUpdateBackupMethod());
return runWithRetryOnAdministrativeRequestsExceeded(
() -> databaseAdminStub.updateBackupCallable().call(request, context));
}
@Override
public void deleteBackup(String backupName) {
acquireAdministrativeRequestsRateLimiter();
final DeleteBackupRequest request =
DeleteBackupRequest.newBuilder().setName(backupName).build();
final GrpcCallContext context =
newCallContext(null, backupName, request, DatabaseAdminGrpc.getDeleteBackupMethod());
runWithRetryOnAdministrativeRequestsExceeded(
() -> {
databaseAdminStub.deleteBackupCallable().call(request, context);
return null;
});
}
@Override
public Backup getBackup(String backupName) throws SpannerException {
acquireAdministrativeRequestsRateLimiter();
final GetBackupRequest request = GetBackupRequest.newBuilder().setName(backupName).build();
final GrpcCallContext context =
newCallContext(null, backupName, request, DatabaseAdminGrpc.getGetBackupMethod());
return runWithRetryOnAdministrativeRequestsExceeded(
() -> get(databaseAdminStub.getBackupCallable().futureCall(request, context)));
}
@Override
public Operation getOperation(String name) throws SpannerException {
acquireAdministrativeRequestsRateLimiter();
final GetOperationRequest request = GetOperationRequest.newBuilder().setName(name).build();
final GrpcCallContext context =
newCallContext(null, name, request, OperationsGrpc.getGetOperationMethod());
return runWithRetryOnAdministrativeRequestsExceeded(
() ->
get(
databaseAdminStub
.getOperationsStub()
.getOperationCallable()
.futureCall(request, context)));
}
@Override
public void cancelOperation(String name) throws SpannerException {
acquireAdministrativeRequestsRateLimiter();
final CancelOperationRequest request =
CancelOperationRequest.newBuilder().setName(name).build();
final GrpcCallContext context =
newCallContext(null, name, request, OperationsGrpc.getCancelOperationMethod());
runWithRetryOnAdministrativeRequestsExceeded(
() -> {
get(
databaseAdminStub
.getOperationsStub()
.cancelOperationCallable()
.futureCall(request, context));
return null;
});
}
@Override
public List batchCreateSessions(
String databaseName,
int sessionCount,
@Nullable String databaseRole,
@Nullable Map labels,
@Nullable Map options)
throws SpannerException {
BatchCreateSessionsRequest.Builder requestBuilder =
BatchCreateSessionsRequest.newBuilder()
.setDatabase(databaseName)
.setSessionCount(sessionCount);
Session.Builder sessionBuilder = Session.newBuilder();
if (labels != null && !labels.isEmpty()) {
sessionBuilder.putAllLabels(labels);
}
if (databaseRole != null && !databaseRole.isEmpty()) {
sessionBuilder.setCreatorRole(databaseRole);
}
requestBuilder.setSessionTemplate(sessionBuilder);
BatchCreateSessionsRequest request = requestBuilder.build();
GrpcCallContext context =
newCallContext(
options, databaseName, request, SpannerGrpc.getBatchCreateSessionsMethod(), true);
return get(spannerStub.batchCreateSessionsCallable().futureCall(request, context))
.getSessionList();
}
@Override
public Session createSession(
String databaseName,
@Nullable String databaseRole,
@Nullable Map labels,
@Nullable Map options)
throws SpannerException {
// By default sessions are not multiplexed
return createSession(databaseName, databaseRole, labels, options, false);
}
@Override
public Session createSession(
String databaseName,
@Nullable String databaseRole,
@Nullable Map labels,
@Nullable Map options,
boolean isMultiplexed)
throws SpannerException {
CreateSessionRequest.Builder requestBuilder =
CreateSessionRequest.newBuilder().setDatabase(databaseName);
Session.Builder sessionBuilder = Session.newBuilder();
if (labels != null && !labels.isEmpty()) {
sessionBuilder.putAllLabels(labels);
}
if (databaseRole != null && !databaseRole.isEmpty()) {
sessionBuilder.setCreatorRole(databaseRole);
}
sessionBuilder.setMultiplexed(isMultiplexed);
requestBuilder.setSession(sessionBuilder);
CreateSessionRequest request = requestBuilder.build();
GrpcCallContext context =
newCallContext(options, databaseName, request, SpannerGrpc.getCreateSessionMethod(), true);
return get(spannerStub.createSessionCallable().futureCall(request, context));
}
@Override
public void deleteSession(String sessionName, @Nullable Map options)
throws SpannerException {
get(asyncDeleteSession(sessionName, options));
}
@Override
public ApiFuture asyncDeleteSession(String sessionName, @Nullable Map options) {
DeleteSessionRequest request = DeleteSessionRequest.newBuilder().setName(sessionName).build();
GrpcCallContext context =
newCallContext(options, sessionName, request, SpannerGrpc.getDeleteSessionMethod());
return spannerStub.deleteSessionCallable().futureCall(request, context);
}
@Override
public RetrySettings getReadRetrySettings() {
return readRetrySettings;
}
@Override
public Set getReadRetryableCodes() {
return readRetryableCodes;
}
@Override
public StreamingCall read(
ReadRequest request,
ResultStreamConsumer consumer,
@Nullable Map options,
boolean routeToLeader) {
GrpcCallContext context =
newCallContext(
options, request.getSession(), request, SpannerGrpc.getReadMethod(), routeToLeader);
SpannerResponseObserver responseObserver = new SpannerResponseObserver(consumer);
spannerStub.streamingReadCallable().call(request, responseObserver, context);
return new GrpcStreamingCall(context, responseObserver.getController());
}
@Override
public RetrySettings getExecuteQueryRetrySettings() {
return executeQueryRetrySettings;
}
@Override
public Set getExecuteQueryRetryableCodes() {
return executeQueryRetryableCodes;
}
@Override
public ResultSet executeQuery(
ExecuteSqlRequest request, @Nullable Map options, boolean routeToLeader) {
return get(executeQueryAsync(request, options, routeToLeader));
}
@Override
public ApiFuture executeQueryAsync(
ExecuteSqlRequest request, @Nullable Map options, boolean routeToLeader) {
GrpcCallContext context =
newCallContext(
options,
request.getSession(),
request,
SpannerGrpc.getExecuteSqlMethod(),
routeToLeader);
return spannerStub.executeSqlCallable().futureCall(request, context);
}
@Override
public ResultSet executePartitionedDml(
ExecuteSqlRequest request, @Nullable Map options) {
GrpcCallContext context =
newCallContext(
options, request.getSession(), request, SpannerGrpc.getExecuteSqlMethod(), true);
return get(partitionedDmlStub.executeSqlCallable().futureCall(request, context));
}
@Override
public RetrySettings getPartitionedDmlRetrySettings() {
return partitionedDmlRetrySettings;
}
@Override
public ServerStream executeStreamingPartitionedDml(
ExecuteSqlRequest request, Map options, Duration timeout) {
GrpcCallContext context =
newCallContext(
options,
request.getSession(),
request,
SpannerGrpc.getExecuteStreamingSqlMethod(),
true);
// Override any timeout settings that might have been set on the call context.
context = context.withTimeout(timeout).withStreamWaitTimeout(timeout);
return partitionedDmlStub.executeStreamingSqlCallable().call(request, context);
}
@Override
public ServerStream batchWriteAtLeastOnce(
BatchWriteRequest request, @Nullable Map options) {
GrpcCallContext context =
newCallContext(options, request.getSession(), request, SpannerGrpc.getBatchWriteMethod());
return spannerStub.batchWriteCallable().call(request, context);
}
@Override
public StreamingCall executeQuery(
ExecuteSqlRequest request,
ResultStreamConsumer consumer,
@Nullable Map options,
boolean routeToLeader) {
GrpcCallContext context =
newCallContext(
options,
request.getSession(),
request,
SpannerGrpc.getExecuteStreamingSqlMethod(),
routeToLeader);
SpannerResponseObserver responseObserver = new SpannerResponseObserver(consumer);
spannerStub.executeStreamingSqlCallable().call(request, responseObserver, context);
return new GrpcStreamingCall(context, responseObserver.getController());
}
@Override
public ExecuteBatchDmlResponse executeBatchDml(
ExecuteBatchDmlRequest request, @Nullable Map options) {
return get(executeBatchDmlAsync(request, options));
}
@Override
public ApiFuture executeBatchDmlAsync(
ExecuteBatchDmlRequest request, @Nullable Map options) {
GrpcCallContext context =
newCallContext(
options, request.getSession(), request, SpannerGrpc.getExecuteBatchDmlMethod(), true);
return spannerStub.executeBatchDmlCallable().futureCall(request, context);
}
@Override
public ApiFuture beginTransactionAsync(
BeginTransactionRequest request, @Nullable Map options, boolean routeToLeader) {
GrpcCallContext context =
newCallContext(
options,
request.getSession(),
request,
SpannerGrpc.getBeginTransactionMethod(),
routeToLeader);
return spannerStub.beginTransactionCallable().futureCall(request, context);
}
@Override
public Transaction beginTransaction(
BeginTransactionRequest request, @Nullable Map options, boolean routeToLeader)
throws SpannerException {
return get(beginTransactionAsync(request, options, routeToLeader));
}
@Override
public ApiFuture commitAsync(
CommitRequest request, @Nullable Map options) {
GrpcCallContext context =
newCallContext(options, request.getSession(), request, SpannerGrpc.getCommitMethod(), true);
return spannerStub.commitCallable().futureCall(request, context);
}
@Override
public CommitResponse commit(CommitRequest commitRequest, @Nullable Map options)
throws SpannerException {
return get(commitAsync(commitRequest, options));
}
@Override
public RetrySettings getCommitRetrySettings() {
return commitRetrySettings;
}
@Override
public ApiFuture rollbackAsync(RollbackRequest request, @Nullable Map options) {
GrpcCallContext context =
newCallContext(
options, request.getSession(), request, SpannerGrpc.getRollbackMethod(), true);
return spannerStub.rollbackCallable().futureCall(request, context);
}
@Override
public void rollback(RollbackRequest request, @Nullable Map options)
throws SpannerException {
get(rollbackAsync(request, options));
}
@Override
public PartitionResponse partitionQuery(
PartitionQueryRequest request, @Nullable Map options) throws SpannerException {
GrpcCallContext context =
newCallContext(
options, request.getSession(), request, SpannerGrpc.getPartitionQueryMethod(), true);
return get(spannerStub.partitionQueryCallable().futureCall(request, context));
}
@Override
public PartitionResponse partitionRead(
PartitionReadRequest request, @Nullable Map options) throws SpannerException {
GrpcCallContext context =
newCallContext(
options, request.getSession(), request, SpannerGrpc.getPartitionReadMethod(), true);
return get(spannerStub.partitionReadCallable().futureCall(request, context));
}
@Override
public Policy getDatabaseAdminIAMPolicy(String resource, @Nullable GetPolicyOptions options) {
acquireAdministrativeRequestsRateLimiter();
GetIamPolicyRequest.Builder builder = GetIamPolicyRequest.newBuilder().setResource(resource);
if (options != null) {
builder.setOptions(options);
}
final GetIamPolicyRequest request = builder.build();
final GrpcCallContext context =
newCallContext(null, resource, request, DatabaseAdminGrpc.getGetIamPolicyMethod());
return runWithRetryOnAdministrativeRequestsExceeded(
() -> get(databaseAdminStub.getIamPolicyCallable().futureCall(request, context)));
}
@Override
public Policy setDatabaseAdminIAMPolicy(String resource, Policy policy) {
acquireAdministrativeRequestsRateLimiter();
final SetIamPolicyRequest request =
SetIamPolicyRequest.newBuilder().setResource(resource).setPolicy(policy).build();
final GrpcCallContext context =
newCallContext(null, resource, request, DatabaseAdminGrpc.getSetIamPolicyMethod());
return runWithRetryOnAdministrativeRequestsExceeded(
() -> get(databaseAdminStub.setIamPolicyCallable().futureCall(request, context)));
}
@Override
public TestIamPermissionsResponse testDatabaseAdminIAMPermissions(
String resource, Iterable permissions) {
acquireAdministrativeRequestsRateLimiter();
final TestIamPermissionsRequest request =
TestIamPermissionsRequest.newBuilder()
.setResource(resource)
.addAllPermissions(permissions)
.build();
final GrpcCallContext context =
newCallContext(null, resource, request, DatabaseAdminGrpc.getTestIamPermissionsMethod());
return runWithRetryOnAdministrativeRequestsExceeded(
() -> get(databaseAdminStub.testIamPermissionsCallable().futureCall(request, context)));
}
@Override
public Policy getInstanceAdminIAMPolicy(String resource) {
acquireAdministrativeRequestsRateLimiter();
final GetIamPolicyRequest request =
GetIamPolicyRequest.newBuilder().setResource(resource).build();
final GrpcCallContext context =
newCallContext(null, resource, request, InstanceAdminGrpc.getGetIamPolicyMethod());
return runWithRetryOnAdministrativeRequestsExceeded(
() -> get(instanceAdminStub.getIamPolicyCallable().futureCall(request, context)));
}
@Override
public Policy setInstanceAdminIAMPolicy(String resource, Policy policy) {
acquireAdministrativeRequestsRateLimiter();
final SetIamPolicyRequest request =
SetIamPolicyRequest.newBuilder().setResource(resource).setPolicy(policy).build();
final GrpcCallContext context =
newCallContext(null, resource, request, InstanceAdminGrpc.getSetIamPolicyMethod());
return runWithRetryOnAdministrativeRequestsExceeded(
() -> get(instanceAdminStub.setIamPolicyCallable().futureCall(request, context)));
}
@Override
public TestIamPermissionsResponse testInstanceAdminIAMPermissions(
String resource, Iterable permissions) {
acquireAdministrativeRequestsRateLimiter();
final TestIamPermissionsRequest request =
TestIamPermissionsRequest.newBuilder()
.setResource(resource)
.addAllPermissions(permissions)
.build();
final GrpcCallContext context =
newCallContext(null, resource, request, InstanceAdminGrpc.getTestIamPermissionsMethod());
return runWithRetryOnAdministrativeRequestsExceeded(
() -> get(instanceAdminStub.testIamPermissionsCallable().futureCall(request, context)));
}
/** Gets the result of an async RPC call, handling any exceptions encountered. */
private static T get(final Future future) throws SpannerException {
final Context context = Context.current();
try {
return future.get();
} catch (InterruptedException e) {
// We are the sole consumer of the future, so cancel it.
future.cancel(true);
throw SpannerExceptionFactory.propagateInterrupt(e);
} catch (Exception e) {
throw newSpannerException(context, e);
}
}
// Before removing this method, please verify with a code owner that it is not used
// in any internal testing infrastructure.
@VisibleForTesting
@Deprecated
GrpcCallContext newCallContext(@Nullable Map options, String resource) {
return newCallContext(options, resource, null, null);
}
@VisibleForTesting
GrpcCallContext newCallContext(
@Nullable Map options,
String resource,
ReqT request,
MethodDescriptor method) {
return newCallContext(options, resource, request, method, false);
}
@VisibleForTesting
GrpcCallContext newCallContext(
@Nullable Map options,
String resource,
ReqT request,
MethodDescriptor method,
boolean routeToLeader) {
GrpcCallContext context = GrpcCallContext.createDefault();
if (options != null) {
if (this.isGrpcGcpExtensionEnabled) {
// Set channel affinity in gRPC-GCP.
// Compute bounded channel hint to prevent gRPC-GCP affinity map from getting unbounded.
int boundedChannelHint = Option.CHANNEL_HINT.getLong(options).intValue() % this.numChannels;
context =
context.withCallOptions(
context
.getCallOptions()
.withOption(
GcpManagedChannel.AFFINITY_KEY, String.valueOf(boundedChannelHint)));
} else {
// Set channel affinity in GAX.
context = context.withChannelAffinity(Option.CHANNEL_HINT.getLong(options).intValue());
}
}
if (compressorName != null) {
// This sets the compressor for Client -> Server.
context = context.withCallOptions(context.getCallOptions().withCompression(compressorName));
}
context = context.withExtraHeaders(metadataProvider.newExtraHeaders(resource, projectName));
if (routeToLeader && leaderAwareRoutingEnabled) {
context = context.withExtraHeaders(metadataProvider.newRouteToLeaderHeader());
}
if (callCredentialsProvider != null) {
CallCredentials callCredentials = callCredentialsProvider.getCallCredentials();
if (callCredentials != null) {
context =
context.withCallOptions(context.getCallOptions().withCallCredentials(callCredentials));
}
}
context = context.withStreamWaitTimeout(waitTimeout).withStreamIdleTimeout(idleTimeout);
CallContextConfigurator configurator = SpannerOptions.CALL_CONTEXT_CONFIGURATOR_KEY.get();
ApiCallContext apiCallContextFromContext = null;
if (configurator != null) {
apiCallContextFromContext = configurator.configure(context, request, method);
}
return (GrpcCallContext) context.merge(apiCallContextFromContext);
}
void registerResponseObserver(SpannerResponseObserver responseObserver) {
responseObservers.add(responseObserver);
}
void unregisterResponseObserver(SpannerResponseObserver responseObserver) {
responseObservers.remove(responseObserver);
}
void closeResponseObservers() {
responseObservers.forEach(SpannerResponseObserver::close);
responseObservers.clear();
}
@InternalApi
@VisibleForTesting
public int getNumActiveResponseObservers() {
return responseObservers.size();
}
@Override
public void shutdown() {
this.rpcIsClosed = true;
closeResponseObservers();
if (this.spannerStub != null) {
this.spannerStub.close();
this.partitionedDmlStub.close();
this.instanceAdminStub.close();
this.databaseAdminStub.close();
this.spannerWatchdog.shutdown();
try {
this.spannerStub.awaitTermination(10L, TimeUnit.SECONDS);
this.partitionedDmlStub.awaitTermination(10L, TimeUnit.SECONDS);
this.instanceAdminStub.awaitTermination(10L, TimeUnit.SECONDS);
this.databaseAdminStub.awaitTermination(10L, TimeUnit.SECONDS);
this.spannerWatchdog.awaitTermination(10L, TimeUnit.SECONDS);
} catch (InterruptedException e) {
throw SpannerExceptionFactory.propagateInterrupt(e);
}
}
}
public void shutdownNow() {
this.rpcIsClosed = true;
closeResponseObservers();
this.spannerStub.close();
this.partitionedDmlStub.close();
this.instanceAdminStub.close();
this.databaseAdminStub.close();
this.spannerWatchdog.shutdown();
this.spannerStub.shutdownNow();
this.partitionedDmlStub.shutdownNow();
this.instanceAdminStub.shutdownNow();
this.databaseAdminStub.shutdownNow();
this.spannerWatchdog.shutdownNow();
}
@Override
public boolean isClosed() {
return rpcIsClosed;
}
@Override
public DatabaseAdminStubSettings getDatabaseAdminStubSettings() {
return databaseAdminStubSettings;
}
@Override
public InstanceAdminStubSettings getInstanceAdminStubSettings() {
return instanceAdminStubSettings;
}
private static final class GrpcStreamingCall implements StreamingCall {
private final ApiCallContext callContext;
private final StreamController controller;
GrpcStreamingCall(ApiCallContext callContext, StreamController controller) {
this.callContext = callContext;
this.controller = controller;
}
@Override
public ApiCallContext getCallContext() {
return callContext;
}
@Override
public void request(int numMessages) {
controller.request(numMessages);
}
@Override
public void cancel(@Nullable String message) {
controller.cancel();
}
}
/**
* A {@code ResponseObserver} that exposes the {@code StreamController} and delegates callbacks to
* the {@link ResultStreamConsumer}.
*/
private class SpannerResponseObserver implements ResponseObserver {
private StreamController controller;
private final ResultStreamConsumer consumer;
public SpannerResponseObserver(ResultStreamConsumer consumer) {
this.consumer = consumer;
}
void close() {
if (this.controller != null) {
this.controller.cancel();
}
}
@Override
public void onStart(StreamController controller) {
// Disable the auto flow control to allow client library
// set the number of messages it prefers to request
controller.disableAutoInboundFlowControl();
this.controller = controller;
if (this.consumer.cancelQueryWhenClientIsClosed()) {
registerResponseObserver(this);
}
}
@Override
public void onResponse(PartialResultSet response) {
consumer.onPartialResultSet(response);
}
@Override
public void onError(Throwable t) {
// Unregister the response observer when the query has completed with an error.
if (this.consumer.cancelQueryWhenClientIsClosed()) {
unregisterResponseObserver(this);
}
consumer.onError(newSpannerException(t));
}
@Override
public void onComplete() {
// Unregister the response observer when the query has completed normally.
if (this.consumer.cancelQueryWhenClientIsClosed()) {
unregisterResponseObserver(this);
}
consumer.onCompleted();
}
StreamController getController() {
return Preconditions.checkNotNull(this.controller);
}
}
private static Duration systemProperty(String name, int defaultValue) {
String stringValue = System.getProperty(name, "");
return Duration.ofSeconds(stringValue.isEmpty() ? defaultValue : Integer.parseInt(stringValue));
}
}