feat: allow setting channel and credential providers in DatastoreOptions

README.md

@@ -50,20 +50,20 @@ If you are using Maven without the BOM, add this to your dependencies:
 If you are using Gradle 5.x or later, add this to your dependencies:
 
 ```Groovy
-implementation platform('com.google.cloud:libraries-bom:26.29.0')
+implementation platform('com.google.cloud:libraries-bom:26.31.0')
 
 implementation 'com.google.cloud:google-cloud-datastore'
 ```
 If you are using Gradle without BOM, add this to your dependencies:
 
 ```Groovy
-implementation 'com.google.cloud:google-cloud-datastore:2.18.0'
+implementation 'com.google.cloud:google-cloud-datastore:2.18.2'
 ```
 
 If you are using SBT, add this to your dependencies:
 
 ```Scala
-libraryDependencies += "com.google.cloud" % "google-cloud-datastore" % "2.18.0"
+libraryDependencies += "com.google.cloud" % "google-cloud-datastore" % "2.18.2"
 ```
 <!-- {x-version-update-end} -->
 
@@ -380,7 +380,7 @@ Java is a registered trademark of Oracle and/or its affiliates.
 [kokoro-badge-link-5]: http://storage.googleapis.com/cloud-devrel-public/java/badges/java-datastore/java11.html
 [stability-image]: https://img.shields.io/badge/stability-stable-green
 [maven-version-image]: https://img.shields.io/maven-central/v/com.google.cloud/google-cloud-datastore.svg
-[maven-version-link]: https://central.sonatype.com/artifact/com.google.cloud/google-cloud-datastore/2.18.0
+[maven-version-link]: https://central.sonatype.com/artifact/com.google.cloud/google-cloud-datastore/2.18.2
 [authentication]: https://github.com/googleapis/google-cloud-java#authentication
 [auth-scopes]: https://developers.google.com/identity/protocols/oauth2/scopes
 [predefined-iam-roles]: https://cloud.google.com/iam/docs/understanding-roles#predefined_roles

google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreOptions.java

@@ -18,6 +18,9 @@
 
 import static com.google.cloud.datastore.Validator.validateNamespace;
 
+import com.google.api.gax.core.CredentialsProvider;
+import com.google.api.gax.grpc.InstantiatingGrpcChannelProvider;
+import com.google.api.gax.rpc.TransportChannelProvider;
 import com.google.cloud.ServiceDefaults;
 import com.google.cloud.ServiceOptions;
 import com.google.cloud.ServiceRpc;
@@ -46,6 +49,9 @@ public class DatastoreOptions extends ServiceOptions<Datastore, DatastoreOptions
   public static final String PROJECT_ID_ENV_VAR = "DATASTORE_PROJECT_ID";
   public static final String LOCAL_HOST_ENV_VAR = "DATASTORE_EMULATOR_HOST";
 
+  private final transient TransportChannelProvider channelProvider;
+  private final transient CredentialsProvider credentialsProvider;
+
   private final String namespace;
   private final String databaseId;
 
@@ -69,6 +75,10 @@ public ServiceRpc create(DatastoreOptions options) {
         if (options.getTransportOptions() instanceof GrpcTransportOptions) {
           return new GrpcDatastoreRpc(options);
         } else if (options.getTransportOptions() instanceof HttpTransportOptions) {
+          // todo see if we can remove this check
+          if (DatastoreUtils.isEmulator(options)) {
+            throw new IllegalArgumentException("Only GRPC channels are allowed for emulator.");
+          }
           return new HttpDatastoreRpc(options);
         } else {
           throw new IllegalArgumentException(
@@ -84,20 +94,50 @@ public static class Builder extends ServiceOptions.Builder<Datastore, DatastoreO
 
     private String namespace;
     private String databaseId;
+    private TransportChannelProvider channelProvider = null;
+    private CredentialsProvider credentialsProvider = null;
 
     private Builder() {}
 
     private Builder(DatastoreOptions options) {
       super(options);
       namespace = options.namespace;
       databaseId = options.databaseId;
+      this.channelProvider = options.channelProvider;
+      this.credentialsProvider = options.credentialsProvider;
     }
 
     @Override
     public Builder setTransportOptions(TransportOptions transportOptions) {
       return super.setTransportOptions(transportOptions);
     }
 
+    /**
+     * Sets the {@link TransportChannelProvider} to use with this Datastore client.
+     *
+     * @param channelProvider A InstantiatingGrpcChannelProvider object that defines the transport
+     *     provider for this client.
+     */
+    public Builder setChannelProvider(TransportChannelProvider channelProvider) {
+      if (!(channelProvider instanceof InstantiatingGrpcChannelProvider)) {
+        throw new IllegalArgumentException(
+            "Only GRPC channels are allowed for " + API_SHORT_NAME + ".");
+      }
+      this.channelProvider = channelProvider;
+      return this;
+    }
+
+    /**
+     * Sets the {@link CredentialsProvider} to use with this Datastore client.
+     *
+     * @param credentialsProvider A CredentialsProvider object that defines the credential provider
+     *     for this client.
+     */
+    public Builder setCredentialsProvider(CredentialsProvider credentialsProvider) {
+      this.credentialsProvider = credentialsProvider;
+      return this;
+    }
+
     @Override
     public DatastoreOptions build() {
       return new DatastoreOptions(this);
@@ -119,6 +159,32 @@ private DatastoreOptions(Builder builder) {
     super(DatastoreFactory.class, DatastoreRpcFactory.class, builder, new DatastoreDefaults());
     namespace = MoreObjects.firstNonNull(builder.namespace, defaultNamespace());
     databaseId = MoreObjects.firstNonNull(builder.databaseId, DEFAULT_DATABASE_ID);
+
+    // todo see if we can update this
+    if (getTransportOptions() instanceof HttpTransportOptions
+        && (builder.channelProvider != null || builder.credentialsProvider != null)) {
+      throw new IllegalArgumentException(
+          "Only gRPC transport allows setting of channel provider or credentials provider");
+    } else {
+      this.channelProvider =
+          builder.channelProvider != null
+              ? builder.channelProvider
+              : GrpcTransportOptions.setUpChannelProvider(
+                  DatastoreSettings.defaultGrpcTransportProviderBuilder(), this);
+
+      this.credentialsProvider =
+          builder.credentialsProvider != null
+              ? builder.credentialsProvider
+              : GrpcTransportOptions.setUpCredentialsProvider(this);
+    }
+  }
+
+  public CredentialsProvider getCredentialsProvider() {
+    return credentialsProvider;
+  }
+
+  public TransportChannelProvider getTransportChannelProvider() {
+    return channelProvider;
   }
 
   @Override
@@ -134,6 +200,7 @@ protected String getDefaultProject() {
   }
 
   private static class DatastoreDefaults implements ServiceDefaults<Datastore, DatastoreOptions> {
+    private final TransportOptions TRANSPORT_OPTIONS = getDefaultTransportOptionsBuilder().build();
 
     @Override
     public DatastoreFactory getDefaultServiceFactory() {
@@ -147,7 +214,11 @@ public DatastoreRpcFactory getDefaultRpcFactory() {
 
     @Override
     public TransportOptions getDefaultTransportOptions() {
-      return getDefaultGrpcTransportOptions();
+      return TRANSPORT_OPTIONS;
+    }
+
+    public static GrpcTransportOptions.Builder getDefaultTransportOptionsBuilder() {
+      return GrpcTransportOptions.newBuilder();
     }
   }
 

google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreUtils.java

@@ -17,13 +17,19 @@
 package com.google.cloud.datastore;
 
 import com.google.api.core.InternalApi;
+import com.google.cloud.NoCredentials;
 import com.google.common.base.Strings;
 import java.net.InetAddress;
 import java.net.URL;
 
 @InternalApi
 public class DatastoreUtils {
 
+  public static boolean isEmulator(DatastoreOptions datastoreOptions) {
+    return isLocalHost(datastoreOptions.getHost())
+        || NoCredentials.getInstance().equals(datastoreOptions.getCredentials());
+  }
+
   public static boolean isLocalHost(String host) {
     if (Strings.isNullOrEmpty(host)) {
       return false;

google-cloud-datastore/src/main/java/com/google/cloud/datastore/spi/v1/GrpcDatastoreRpc.java

@@ -16,7 +16,7 @@
 
 package com.google.cloud.datastore.spi.v1;
 
-import static com.google.cloud.datastore.DatastoreUtils.isLocalHost;
+import static com.google.cloud.datastore.DatastoreUtils.isEmulator;
 import static com.google.cloud.datastore.DatastoreUtils.removeScheme;
 import static com.google.cloud.datastore.spi.v1.RpcUtils.retrySettingSetter;
 import static java.util.concurrent.TimeUnit.SECONDS;
@@ -30,14 +30,12 @@
 import com.google.api.gax.rpc.HeaderProvider;
 import com.google.api.gax.rpc.NoHeaderProvider;
 import com.google.api.gax.rpc.TransportChannel;
-import com.google.cloud.NoCredentials;
 import com.google.cloud.ServiceOptions;
 import com.google.cloud.datastore.DatastoreException;
 import com.google.cloud.datastore.DatastoreOptions;
 import com.google.cloud.datastore.v1.DatastoreSettings;
 import com.google.cloud.datastore.v1.stub.DatastoreStubSettings;
 import com.google.cloud.datastore.v1.stub.GrpcDatastoreStub;
-import com.google.cloud.grpc.GrpcTransportOptions;
 import com.google.common.base.Strings;
 import com.google.datastore.v1.AllocateIdsRequest;
 import com.google.datastore.v1.AllocateIdsResponse;
@@ -69,7 +67,6 @@ public class GrpcDatastoreRpc implements DatastoreRpc {
   private boolean closed;
 
   public GrpcDatastoreRpc(DatastoreOptions datastoreOptions) throws IOException {
-
     try {
       clientContext =
           isEmulator(datastoreOptions)
@@ -146,11 +143,6 @@ public boolean isClosed() {
     return closed && datastoreStub.isShutdown();
   }
 
-  private boolean isEmulator(DatastoreOptions datastoreOptions) {
-    return isLocalHost(datastoreOptions.getHost())
-        || NoCredentials.getInstance().equals(datastoreOptions.getCredentials());
-  }
-
   private ClientContext getClientContextForEmulator(DatastoreOptions datastoreOptions)
       throws IOException {
     ManagedChannel managedChannel =
@@ -177,11 +169,8 @@ private ClientContext getClientContext(DatastoreOptions datastoreOptions) throws
 
     DatastoreSettingsBuilder settingsBuilder =
         new DatastoreSettingsBuilder(DatastoreSettings.newBuilder().build());
-    settingsBuilder.setCredentialsProvider(
-        GrpcTransportOptions.setUpCredentialsProvider(datastoreOptions));
-    settingsBuilder.setTransportChannelProvider(
-        GrpcTransportOptions.setUpChannelProvider(
-            DatastoreSettings.defaultGrpcTransportProviderBuilder(), datastoreOptions));
+    settingsBuilder.setCredentialsProvider(datastoreOptions.getCredentialsProvider());
+    settingsBuilder.setTransportChannelProvider(datastoreOptions.getTransportChannelProvider());
     settingsBuilder.setInternalHeaderProvider(internalHeaderProvider);
     settingsBuilder.setHeaderProvider(
         datastoreOptions.getMergedHeaderProvider(new NoHeaderProvider()));

google-cloud-datastore/src/test/java/com/google/cloud/datastore/DatastoreOptionsTest.java

@@ -22,11 +22,17 @@
 import static org.junit.Assert.assertSame;
 import static org.junit.Assert.assertTrue;
 
+import com.google.api.gax.core.NoCredentialsProvider;
+import com.google.api.gax.grpc.ChannelPoolSettings;
+import com.google.api.gax.grpc.InstantiatingGrpcChannelProvider;
+import com.google.cloud.NoCredentials;
 import com.google.cloud.datastore.spi.DatastoreRpcFactory;
 import com.google.cloud.datastore.spi.v1.DatastoreRpc;
+import com.google.cloud.datastore.v1.DatastoreSettings;
 import com.google.cloud.grpc.GrpcTransportOptions;
 import com.google.cloud.http.HttpTransportOptions;
 import org.easymock.EasyMock;
+import org.junit.Assert;
 import org.junit.Before;
 import org.junit.Test;
 
@@ -38,6 +44,16 @@ public class DatastoreOptionsTest {
   private DatastoreRpcFactory datastoreRpcFactory;
   private DatastoreRpc datastoreRpc;
   private DatastoreOptions.Builder options;
+  private DatastoreOptions.Builder options2;
+  private final InstantiatingGrpcChannelProvider channelProvider =
+      DatastoreSettings.defaultGrpcTransportProviderBuilder()
+          .setChannelPoolSettings(
+              ChannelPoolSettings.builder()
+                  .setInitialChannelCount(10)
+                  .setMaxChannelCount(20)
+                  .build())
+          .build();
+  private final NoCredentialsProvider noCredentialsProvider = NoCredentialsProvider.create();
 
   @Before
   public void setUp() {
@@ -48,6 +64,16 @@ public void setUp() {
             .setServiceRpcFactory(datastoreRpcFactory)
             .setProjectId(PROJECT_ID)
             .setDatabaseId(DATABASE_ID)
+            .setCredentials(NoCredentials.getInstance())
+            .setHost("http://localhost:" + PORT);
+
+    options2 =
+        DatastoreOptions.newBuilder()
+            .setServiceRpcFactory(datastoreRpcFactory)
+            .setProjectId(PROJECT_ID)
+            .setDatabaseId(DATABASE_ID)
+            .setChannelProvider(channelProvider)
+            .setCredentialsProvider(noCredentialsProvider)
             .setHost("http://localhost:" + PORT);
     EasyMock.expect(datastoreRpcFactory.create(EasyMock.anyObject(DatastoreOptions.class)))
         .andReturn(datastoreRpc)
@@ -81,6 +107,27 @@ public void testDatastore() {
     assertSame(datastoreRpc, options.build().getRpc());
   }
 
+  @Test
+  public void testCustomChannelAndCredentials() {
+    DatastoreOptions datastoreOptions = options2.build();
+    assertEquals(datastoreOptions.getTransportChannelProvider(), channelProvider);
+    assertEquals(datastoreOptions.getCredentialsProvider(), noCredentialsProvider);
+  }
+
+  @Test
+  public void testInvalidConfigForHttp() {
+    DatastoreOptions.Builder options =
+        DatastoreOptions.newBuilder()
+            .setServiceRpcFactory(datastoreRpcFactory)
+            .setProjectId(PROJECT_ID)
+            .setDatabaseId(DATABASE_ID)
+            .setTransportOptions(HttpTransportOptions.newBuilder().build())
+            .setChannelProvider(channelProvider)
+            .setCredentialsProvider(noCredentialsProvider)
+            .setHost("http://localhost:" + PORT);
+    Assert.assertThrows(IllegalArgumentException.class, options::build);
+  }
+
   @Test
   public void testTransport() {
     // default grpc transport