All Downloads are FREE. Search and download functionalities are using the official Maven repository.

com.google.cloud.spanner.spi.v1.GapicSpannerRpc Maven / Gradle / Ivy

There is a newer version: 6.81.1
Show newest version
/*
 * 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());

      try {
        this.spannerStub =
            GrpcSpannerStub.create(
                options
                    .getSpannerStubSettings()
                    .toBuilder()
                    .setTransportChannelProvider(channelProvider)
                    .setCredentialsProvider(credentialsProvider)
                    .setStreamWatchdogProvider(watchdogProvider)
                    .setTracerFactory(options.getApiTracerFactory())
                    .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())
            .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())
                .build();
        this.instanceAdminStub = GrpcInstanceAdminStub.create(instanceAdminStubSettings);

        this.databaseAdminStubSettings =
            options
                .getDatabaseAdminStubSettings()
                .toBuilder()
                .setTransportChannelProvider(channelProvider)
                .setCredentialsProvider(credentialsProvider)
                .setStreamWatchdogProvider(watchdogProvider)
                .setTracerFactory(options.getApiTracerFactory())
                .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);
      } 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)
      throws IOException {
    final String emulatorHost = System.getenv("SPANNER_EMULATOR_HOST");
    // Only do the check if the emulator environment variable has been set to localhost.
    if (options.getChannelProvider() == null
        && emulatorHost != null
        && options.getHost() != null
        && options.getHost().startsWith("http://localhost")
        && options.getHost().endsWith(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 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)); } }