diff --git a/google-cloud-datastore/pom.xml b/google-cloud-datastore/pom.xml
index 59b4b30ed..a4b0434eb 100644
--- a/google-cloud-datastore/pom.xml
+++ b/google-cloud-datastore/pom.xml
@@ -106,10 +106,6 @@
com.google.oauth-client
google-oauth-client
-
- com.google.auth
- google-auth-library-oauth2-http
-
io.opencensus
opencensus-api
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreOptions.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreOptions.java
index 1eb7f5105..a166a45ae 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreOptions.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreOptions.java
@@ -17,6 +17,7 @@
package com.google.cloud.datastore;
import static com.google.cloud.datastore.Validator.validateNamespace;
+import static com.google.cloud.datastore.spi.v1.DatastoreRpc.Transport.GRPC;
import com.google.api.core.BetaApi;
import com.google.cloud.ServiceDefaults;
@@ -25,9 +26,12 @@
import com.google.cloud.TransportOptions;
import com.google.cloud.datastore.spi.DatastoreRpcFactory;
import com.google.cloud.datastore.spi.v1.DatastoreRpc;
+import com.google.cloud.datastore.spi.v1.DatastoreRpc.Transport;
import com.google.cloud.datastore.spi.v1.GrpcDatastoreRpc;
+import com.google.cloud.datastore.spi.v1.HttpDatastoreRpc;
import com.google.cloud.datastore.v1.DatastoreSettings;
import com.google.cloud.http.HttpTransportOptions;
+import com.google.common.annotations.VisibleForTesting;
import com.google.common.base.MoreObjects;
import com.google.common.collect.ImmutableSet;
import java.io.IOException;
@@ -45,6 +49,7 @@ public class DatastoreOptions extends ServiceOptions, Void> retrySettingsSetter =
- builder -> {
- builder.setRetrySettings(datastoreOptions.getRetrySettings());
- return null;
- };
+
DatastoreStubSettings datastoreStubSettings =
DatastoreStubSettings.newBuilder(clientContext)
- .applyToAllUnaryMethods(retrySettingsSetter)
+ .applyToAllUnaryMethods(retrySettingSetter(datastoreOptions))
.build();
datastoreStub = GrpcDatastoreStub.create(datastoreStubSettings);
} catch (IOException e) {
@@ -202,18 +197,4 @@ private String getResourceToken(DatastoreOptions datastoreOptions) {
}
return builder.toString();
}
-
- // This class is needed solely to get access to protected method setInternalHeaderProvider()
- private static class DatastoreSettingsBuilder extends DatastoreSettings.Builder {
-
- private DatastoreSettingsBuilder(DatastoreSettings settings) {
- super(settings);
- }
-
- @Override
- protected DatastoreSettings.Builder setInternalHeaderProvider(
- HeaderProvider internalHeaderProvider) {
- return super.setInternalHeaderProvider(internalHeaderProvider);
- }
- }
}
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/spi/v1/HttpDatastoreRpc.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/spi/v1/HttpDatastoreRpc.java
index 66bb0497b..f5a6ebb42 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/spi/v1/HttpDatastoreRpc.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/spi/v1/HttpDatastoreRpc.java
@@ -16,14 +16,17 @@
package com.google.cloud.datastore.spi.v1;
-import com.google.api.client.http.HttpRequest;
-import com.google.api.client.http.HttpRequestInitializer;
-import com.google.api.client.http.HttpTransport;
-import com.google.cloud.datastore.DatastoreException;
+import static com.google.cloud.datastore.DatastoreUtils.isLocalHost;
+import static com.google.cloud.datastore.spi.v1.RpcUtils.retrySettingSetter;
+import static java.util.concurrent.TimeUnit.SECONDS;
+
+import com.google.api.core.InternalApi;
+import com.google.api.gax.core.BackgroundResource;
+import com.google.api.gax.rpc.ClientContext;
import com.google.cloud.datastore.DatastoreOptions;
-import com.google.cloud.datastore.TraceUtil;
-import com.google.cloud.http.CensusHttpModule;
-import com.google.cloud.http.HttpTransportOptions;
+import com.google.cloud.datastore.v1.DatastoreSettings;
+import com.google.cloud.datastore.v1.stub.DatastoreStubSettings;
+import com.google.cloud.datastore.v1.stub.HttpJsonDatastoreStub;
import com.google.datastore.v1.AllocateIdsRequest;
import com.google.datastore.v1.AllocateIdsResponse;
import com.google.datastore.v1.BeginTransactionRequest;
@@ -41,184 +44,104 @@
import com.google.datastore.v1.RunQueryRequest;
import com.google.datastore.v1.RunQueryResponse;
import java.io.IOException;
-import java.net.InetAddress;
-import java.net.URL;
+@InternalApi
public class HttpDatastoreRpc implements DatastoreRpc {
- private final com.google.datastore.v1.client.Datastore client;
-
- public HttpDatastoreRpc(DatastoreOptions options) {
- HttpTransportOptions httpTransportOptions =
- (HttpTransportOptions) options.getTransportOptions();
- HttpTransport transport = httpTransportOptions.getHttpTransportFactory().create();
- com.google.datastore.v1.client.DatastoreOptions.Builder clientBuilder =
- new com.google.datastore.v1.client.DatastoreOptions.Builder()
- .projectId(options.getProjectId())
- .initializer(getHttpRequestInitializer(options, httpTransportOptions))
- .transport(transport);
- String normalizedHost = options.getHost() != null ? options.getHost().toLowerCase() : "";
- if (isLocalHost(normalizedHost)) {
- clientBuilder = clientBuilder.localHost(removeScheme(normalizedHost));
- } else if (!removeScheme(com.google.datastore.v1.client.DatastoreFactory.DEFAULT_HOST)
- .equals(removeScheme(normalizedHost))
- && !normalizedHost.isEmpty()) {
- String fullUrl = normalizedHost;
- if (fullUrl.charAt(fullUrl.length() - 1) != '/') {
- fullUrl = fullUrl + '/';
- }
- fullUrl =
- fullUrl
- + com.google.datastore.v1.client.DatastoreFactory.VERSION
- + "/projects/"
- + options.getProjectId();
- clientBuilder = clientBuilder.projectId(null).projectEndpoint(fullUrl);
- }
- client = com.google.datastore.v1.client.DatastoreFactory.get().create(clientBuilder.build());
- }
+ private final ClientContext clientContext;
+ private final HttpJsonDatastoreStub datastoreStub;
- private HttpRequestInitializer getHttpRequestInitializer(
- final DatastoreOptions options, HttpTransportOptions httpTransportOptions) {
- // Open Census initialization
- CensusHttpModule censusHttpModule =
- new CensusHttpModule(TraceUtil.getInstance().getTracer(), true);
- final HttpRequestInitializer censusHttpModuleHttpRequestInitializer =
- censusHttpModule.getHttpRequestInitializer(
- httpTransportOptions.getHttpRequestInitializer(options));
-
- final String applicationName = options.getApplicationName();
- return new HttpRequestInitializer() {
- @Override
- public void initialize(HttpRequest httpRequest) throws IOException {
- censusHttpModuleHttpRequestInitializer.initialize(httpRequest);
- httpRequest.getHeaders().setUserAgent(applicationName);
- }
- };
- }
+ private boolean closed;
- private static boolean isLocalHost(String host) {
- if (!host.isEmpty()) {
- try {
- String normalizedHost = "http://" + removeScheme(host);
- InetAddress hostAddr = InetAddress.getByName(new URL(normalizedHost).getHost());
- return hostAddr.isAnyLocalAddress() || hostAddr.isLoopbackAddress();
- } catch (Exception e) {
- // ignore
- }
- }
- return false;
- }
+ public HttpDatastoreRpc(DatastoreOptions datastoreOptions) throws IOException {
+ DatastoreSettings datastoreSettings =
+ new DatastoreSettingsBuilder(DatastoreSettings.newBuilder().build())
+ .setInternalHeaderProvider(
+ DatastoreStubSettings.defaultHttpJsonApiClientHeaderProviderBuilder().build())
+ .setTransportChannelProvider(
+ DatastoreStubSettings.defaultHttpJsonTransportProviderBuilder().build())
+ .setEndpoint(getHost(datastoreOptions))
+ .build();
- private static String removeScheme(String url) {
- if (url != null) {
- if (url.startsWith("https://")) {
- return url.substring("https://".length());
- } else if (url.startsWith("http://")) {
- return url.substring("http://".length());
- }
- }
- return url;
- }
+ clientContext = ClientContext.create(datastoreSettings);
- private static DatastoreException translate(
- com.google.datastore.v1.client.DatastoreException exception) {
- return translate(exception, true);
- }
+ DatastoreStubSettings datastoreStubSettings =
+ DatastoreStubSettings.newBuilder(clientContext)
+ .applyToAllUnaryMethods(retrySettingSetter(datastoreOptions))
+ .build();
- private static DatastoreException translate(
- com.google.datastore.v1.client.DatastoreException exception, boolean idempotent) {
- String reason = "";
- if (exception.getCode() != null) {
- reason = exception.getCode().name();
- }
- if (reason.isEmpty()) {
- if (exception.getCause() instanceof IOException) {
- return new DatastoreException((IOException) exception.getCause());
- }
- }
- return new DatastoreException(
- exception.getCode().getNumber(), exception.getMessage(), reason, idempotent, exception);
+ datastoreStub = HttpJsonDatastoreStub.create(datastoreStubSettings);
}
@Override
public AllocateIdsResponse allocateIds(AllocateIdsRequest request) {
- try {
- return client.allocateIds(request);
- } catch (com.google.datastore.v1.client.DatastoreException ex) {
- throw translate(ex);
- }
+ return this.datastoreStub.allocateIdsCallable().call(request);
}
@Override
public BeginTransactionResponse beginTransaction(BeginTransactionRequest request) {
- try {
- return client.beginTransaction(request);
- } catch (com.google.datastore.v1.client.DatastoreException ex) {
- throw translate(ex);
- }
+ return this.datastoreStub.beginTransactionCallable().call(request);
}
@Override
public CommitResponse commit(CommitRequest request) {
- try {
- return client.commit(request);
- } catch (com.google.datastore.v1.client.DatastoreException ex) {
- throw translate(ex, request.getMode() == CommitRequest.Mode.NON_TRANSACTIONAL);
- }
+ return this.datastoreStub.commitCallable().call(request);
}
@Override
public LookupResponse lookup(LookupRequest request) {
- try {
- return client.lookup(request);
- } catch (com.google.datastore.v1.client.DatastoreException ex) {
- throw translate(ex);
- }
+ return this.datastoreStub.lookupCallable().call(request);
}
@Override
public ReserveIdsResponse reserveIds(ReserveIdsRequest request) {
- try {
- return client.reserveIds(request);
- } catch (com.google.datastore.v1.client.DatastoreException ex) {
- throw translate(ex);
- }
+ return this.datastoreStub.reserveIdsCallable().call(request);
}
@Override
public RollbackResponse rollback(RollbackRequest request) {
- try {
- return client.rollback(request);
- } catch (com.google.datastore.v1.client.DatastoreException ex) {
- throw translate(ex);
- }
+ return this.datastoreStub.rollbackCallable().call(request);
}
@Override
public RunQueryResponse runQuery(RunQueryRequest request) {
- try {
- return client.runQuery(request);
- } catch (com.google.datastore.v1.client.DatastoreException ex) {
- throw translate(ex);
- }
+ return this.datastoreStub.runQueryCallable().call(request);
}
@Override
public RunAggregationQueryResponse runAggregationQuery(RunAggregationQueryRequest request) {
- try {
- return client.runAggregationQuery(request);
- } catch (com.google.datastore.v1.client.DatastoreException ex) {
- throw translate(ex);
- }
+ return this.datastoreStub.runAggregationQueryCallable().call(request);
}
@Override
public void close() throws Exception {
- throw new UnsupportedOperationException("close() is not supported");
+ if (!closed) {
+ datastoreStub.close();
+ for (BackgroundResource resource : clientContext.getBackgroundResources()) {
+ resource.close();
+ }
+ closed = true;
+ }
+ for (BackgroundResource resource : clientContext.getBackgroundResources()) {
+ resource.awaitTermination(1, SECONDS);
+ }
}
@Override
public boolean isClosed() {
- throw new UnsupportedOperationException("isClosed() is not supported");
+ return closed && datastoreStub.isShutdown();
+ }
+
+ /**
+ * Prefixing it with http scheme when host is localhost, otherwise {@link
+ * com.google.api.gax.httpjson.HttpRequestRunnable#normalizeEndpoint(String)} will prefix it with
+ * https.
+ */
+ private String getHost(DatastoreOptions options) {
+ String host = options.getHost();
+ if (isLocalHost(host) && !host.contains("://")) {
+ return "http://" + host;
+ }
+ return host;
}
}
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/spi/v1/RpcUtils.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/spi/v1/RpcUtils.java
new file mode 100644
index 000000000..dee8d6920
--- /dev/null
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/spi/v1/RpcUtils.java
@@ -0,0 +1,34 @@
+/*
+ * Copyright 2023 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.datastore.spi.v1;
+
+import com.google.api.core.ApiFunction;
+import com.google.api.core.InternalApi;
+import com.google.api.gax.rpc.UnaryCallSettings;
+import com.google.cloud.datastore.DatastoreOptions;
+
+@InternalApi
+public class RpcUtils {
+ @InternalApi
+ static ApiFunction, Void> retrySettingSetter(
+ DatastoreOptions datastoreOptions) {
+ return builder -> {
+ builder.setRetrySettings(datastoreOptions.getRetrySettings());
+ return null;
+ };
+ }
+}
diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/DatastoreOptionsTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/DatastoreOptionsTest.java
index a545580e2..6444b1a2d 100644
--- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/DatastoreOptionsTest.java
+++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/DatastoreOptionsTest.java
@@ -16,6 +16,9 @@
package com.google.cloud.datastore;
+import static com.google.cloud.datastore.spi.v1.DatastoreRpc.Transport.GRPC;
+import static com.google.cloud.datastore.spi.v1.DatastoreRpc.Transport.HTTP;
+import static com.google.common.truth.Truth.assertThat;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertNotEquals;
import static org.junit.Assert.assertNotNull;
@@ -81,6 +84,22 @@ public void testDatastore() {
assertSame(datastoreRpc, options.build().getRpc());
}
+ @Test
+ public void testTransport() {
+ // default grpc transport
+ assertThat(options.build().getTransport()).isEqualTo(GRPC);
+
+ // custom http transport
+ DatastoreOptions httpDatastoreOptions =
+ DatastoreOptions.newBuilder().setTransport(HTTP).setProjectId(PROJECT_ID).build();
+ assertThat(httpDatastoreOptions.getTransport()).isEqualTo(HTTP);
+
+ // custom grpc transport
+ DatastoreOptions grpcDatastoreOptions =
+ DatastoreOptions.newBuilder().setTransport(GRPC).setProjectId(PROJECT_ID).build();
+ assertThat(grpcDatastoreOptions.getTransport()).isEqualTo(GRPC);
+ }
+
@Test
public void testToBuilder() {
DatastoreOptions original = options.setNamespace("ns1").build();
diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/DatastoreTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/DatastoreTest.java
index 84cbd4b73..efde25e61 100644
--- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/DatastoreTest.java
+++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/DatastoreTest.java
@@ -92,6 +92,7 @@
@RunWith(JUnit4.class)
public class DatastoreTest {
+
private static final LocalDatastoreHelper helper = LocalDatastoreHelper.create(1.0, 9090);
private static DatastoreOptions options = helper.getOptions();
private static Datastore datastore;
diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/spi/v1/RpcUtilsTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/spi/v1/RpcUtilsTest.java
new file mode 100644
index 000000000..76fd00580
--- /dev/null
+++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/spi/v1/RpcUtilsTest.java
@@ -0,0 +1,57 @@
+/*
+ * Copyright 2023 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.datastore.spi.v1;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import com.google.api.gax.retrying.RetrySettings;
+import com.google.api.gax.rpc.UnaryCallSettings;
+import com.google.cloud.datastore.DatastoreOptions;
+import org.junit.Test;
+import org.threeten.bp.Duration;
+
+public class RpcUtilsTest {
+
+ @Test
+ public void testRetrySettingSetter() {
+ UnaryCallSettings.Builder