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 builder = + UnaryCallSettings.newUnaryCallSettingsBuilder(); + + // datastoreOptions with default retry settings + DatastoreOptions datastoreOptions = + DatastoreOptions.newBuilder().setProjectId("project-id").build(); + RpcUtils.retrySettingSetter(datastoreOptions).apply(builder); + assertThat(builder.getRetrySettings()).isEqualTo(datastoreOptions.getRetrySettings()); + + // datastoreOptions with custom retry settings + RetrySettings customRetrySettings = + RetrySettings.newBuilder() + .setTotalTimeout(Duration.ofMinutes(2)) + .setInitialRpcTimeout(Duration.ofSeconds(5)) + .setMaxRpcTimeout(Duration.ofSeconds(10)) + .setRetryDelayMultiplier(1.5) + .setMaxAttempts(5) + .build(); + DatastoreOptions datastoreOptionsWithCustomRetrySettings = + DatastoreOptions.newBuilder() + .setProjectId("project-id") + .setRetrySettings(customRetrySettings) + .build(); + RpcUtils.retrySettingSetter(datastoreOptionsWithCustomRetrySettings).apply(builder); + assertThat(builder.getRetrySettings()).isEqualTo(customRetrySettings); + } +}