Skip to content

Commit a7c9beb

Browse files
refactor: Swap usage of HttpDatastoreRpc with GrpcDatastoreRpc (#1240)
* Create basic structure of GrpcDatastoreRpc and using it in DatastoreOptions * applying unary settings to all the unary methods * Configuring header provider for GrpcDatastoreRpc * fixing emulator tests to be able to run successfully with grpc now * ignoring one more test which will be fixed in actionable error implementation * Making HttpDatastoreRpc completely unused * Making GrpcDatastoreRpc implement AutoCloseable * 🦉 Updates from OwlBot post-processor See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md * incorporating feedbacks * pinging emulator after each test for debugging * Revert "pinging emulator after each test for debugging" This reverts commit 60ee45d. * Reapply "pinging emulator after each test for debugging" This reverts commit d42e3b9. * more debugging * Constant ping to avoid flaky behaviour of /shutdown endpoint * fixing test * checking if emulator is running before sending a shutdown command * fix lint * implement helper method for localhost * fix header lint * moving emulator health check to src/test * fix lint * adding no extra headers * minor cleanup * using mutlipleAttemptsRule in DatastoreTest * Revert "adding no extra headers" This reverts commit 9b43798. * using classRule --------- Co-authored-by: Owl Bot <gcf-owl-bot[bot]@users.noreply.github.com>
1 parent 0f79a32 commit a7c9beb

File tree

9 files changed

+339
-6
lines changed

9 files changed

+339
-6
lines changed

google-cloud-datastore/pom.xml

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,10 @@
3030
<groupId>com.google.cloud</groupId>
3131
<artifactId>google-cloud-core-http</artifactId>
3232
</dependency>
33+
<dependency>
34+
<groupId>com.google.cloud</groupId>
35+
<artifactId>google-cloud-core-grpc</artifactId>
36+
</dependency>
3337
<dependency>
3438
<groupId>com.google.api.grpc</groupId>
3539
<artifactId>proto-google-cloud-datastore-v1</artifactId>

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

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -25,10 +25,12 @@
2525
import com.google.cloud.TransportOptions;
2626
import com.google.cloud.datastore.spi.DatastoreRpcFactory;
2727
import com.google.cloud.datastore.spi.v1.DatastoreRpc;
28-
import com.google.cloud.datastore.spi.v1.HttpDatastoreRpc;
28+
import com.google.cloud.datastore.spi.v1.GrpcDatastoreRpc;
29+
import com.google.cloud.datastore.v1.DatastoreSettings;
2930
import com.google.cloud.http.HttpTransportOptions;
3031
import com.google.common.base.MoreObjects;
3132
import com.google.common.collect.ImmutableSet;
33+
import java.io.IOException;
3234
import java.lang.reflect.Method;
3335
import java.util.Objects;
3436
import java.util.Set;
@@ -60,7 +62,11 @@ public static class DefaultDatastoreRpcFactory implements DatastoreRpcFactory {
6062

6163
@Override
6264
public ServiceRpc create(DatastoreOptions options) {
63-
return new HttpDatastoreRpc(options);
65+
try {
66+
return new GrpcDatastoreRpc(options);
67+
} catch (IOException e) {
68+
throw new RuntimeException(e);
69+
}
6470
}
6571
}
6672

@@ -116,7 +122,7 @@ protected String getDefaultHost() {
116122
System.getProperty(
117123
com.google.datastore.v1.client.DatastoreHelper.LOCAL_HOST_ENV_VAR,
118124
System.getenv(com.google.datastore.v1.client.DatastoreHelper.LOCAL_HOST_ENV_VAR));
119-
return host != null ? host : com.google.datastore.v1.client.DatastoreFactory.DEFAULT_HOST;
125+
return host != null ? host : DatastoreSettings.getDefaultEndpoint();
120126
}
121127

122128
@Override
Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,50 @@
1+
/*
2+
* Copyright 2023 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package com.google.cloud.datastore;
18+
19+
import com.google.api.core.InternalApi;
20+
import com.google.common.base.Strings;
21+
import java.net.InetAddress;
22+
import java.net.URL;
23+
24+
@InternalApi
25+
public class DatastoreUtils {
26+
27+
public static boolean isLocalHost(String host) {
28+
if (Strings.isNullOrEmpty(host)) {
29+
return false;
30+
}
31+
try {
32+
String normalizedHost = "http://" + removeScheme(host);
33+
InetAddress hostAddr = InetAddress.getByName(new URL(normalizedHost).getHost());
34+
return hostAddr.isAnyLocalAddress() || hostAddr.isLoopbackAddress();
35+
} catch (Exception e) {
36+
throw new RuntimeException(e);
37+
}
38+
}
39+
40+
public static String removeScheme(String url) {
41+
if (url != null) {
42+
if (url.startsWith("https://")) {
43+
return url.substring("https://".length());
44+
} else if (url.startsWith("http://")) {
45+
return url.substring("http://".length());
46+
}
47+
}
48+
return url;
49+
}
50+
}

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -16,14 +16,14 @@
1616

1717
package com.google.cloud.datastore;
1818

19-
import com.google.cloud.datastore.spi.v1.HttpDatastoreRpc;
19+
import com.google.cloud.datastore.spi.v1.DatastoreRpc;
2020
import io.opencensus.trace.EndSpanOptions;
2121
import io.opencensus.trace.Span;
2222
import io.opencensus.trace.Tracer;
2323
import io.opencensus.trace.Tracing;
2424

2525
/**
26-
* Helper class for tracing utility. It is used for instrumenting {@link HttpDatastoreRpc} with
26+
* Helper class for tracing utility. It is used for instrumenting {@link DatastoreRpc} with
2727
* OpenCensus APIs.
2828
*
2929
* <p>TraceUtil instances are created by the {@link TraceUtil#getInstance()} method.
Lines changed: 215 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,215 @@
1+
/*
2+
* Copyright 2023 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package com.google.cloud.datastore.spi.v1;
18+
19+
import static com.google.cloud.datastore.DatastoreUtils.isLocalHost;
20+
import static com.google.cloud.datastore.DatastoreUtils.removeScheme;
21+
import static java.util.concurrent.TimeUnit.SECONDS;
22+
23+
import com.google.api.core.ApiFunction;
24+
import com.google.api.core.InternalApi;
25+
import com.google.api.gax.core.BackgroundResource;
26+
import com.google.api.gax.core.GaxProperties;
27+
import com.google.api.gax.grpc.GrpcCallContext;
28+
import com.google.api.gax.grpc.GrpcTransportChannel;
29+
import com.google.api.gax.rpc.ClientContext;
30+
import com.google.api.gax.rpc.HeaderProvider;
31+
import com.google.api.gax.rpc.NoHeaderProvider;
32+
import com.google.api.gax.rpc.TransportChannel;
33+
import com.google.api.gax.rpc.UnaryCallSettings;
34+
import com.google.cloud.NoCredentials;
35+
import com.google.cloud.ServiceOptions;
36+
import com.google.cloud.datastore.DatastoreException;
37+
import com.google.cloud.datastore.DatastoreOptions;
38+
import com.google.cloud.datastore.v1.DatastoreSettings;
39+
import com.google.cloud.datastore.v1.stub.DatastoreStubSettings;
40+
import com.google.cloud.datastore.v1.stub.GrpcDatastoreStub;
41+
import com.google.cloud.grpc.GrpcTransportOptions;
42+
import com.google.common.base.Strings;
43+
import com.google.datastore.v1.AllocateIdsRequest;
44+
import com.google.datastore.v1.AllocateIdsResponse;
45+
import com.google.datastore.v1.BeginTransactionRequest;
46+
import com.google.datastore.v1.BeginTransactionResponse;
47+
import com.google.datastore.v1.CommitRequest;
48+
import com.google.datastore.v1.CommitResponse;
49+
import com.google.datastore.v1.LookupRequest;
50+
import com.google.datastore.v1.LookupResponse;
51+
import com.google.datastore.v1.ReserveIdsRequest;
52+
import com.google.datastore.v1.ReserveIdsResponse;
53+
import com.google.datastore.v1.RollbackRequest;
54+
import com.google.datastore.v1.RollbackResponse;
55+
import com.google.datastore.v1.RunAggregationQueryRequest;
56+
import com.google.datastore.v1.RunAggregationQueryResponse;
57+
import com.google.datastore.v1.RunQueryRequest;
58+
import com.google.datastore.v1.RunQueryResponse;
59+
import io.grpc.CallOptions;
60+
import io.grpc.ManagedChannel;
61+
import io.grpc.ManagedChannelBuilder;
62+
import java.io.IOException;
63+
import java.util.Collections;
64+
65+
@InternalApi
66+
public class GrpcDatastoreRpc implements AutoCloseable, DatastoreRpc {
67+
68+
private final GrpcDatastoreStub datastoreStub;
69+
private final ClientContext clientContext;
70+
private boolean closed;
71+
72+
public GrpcDatastoreRpc(DatastoreOptions datastoreOptions) throws IOException {
73+
74+
try {
75+
clientContext =
76+
isEmulator(datastoreOptions)
77+
? getClientContextForEmulator(datastoreOptions)
78+
: getClientContext(datastoreOptions);
79+
ApiFunction<UnaryCallSettings.Builder<?, ?>, Void> retrySettingsSetter =
80+
builder -> {
81+
builder.setRetrySettings(datastoreOptions.getRetrySettings());
82+
return null;
83+
};
84+
DatastoreStubSettings datastoreStubSettings =
85+
DatastoreStubSettings.newBuilder(clientContext)
86+
.applyToAllUnaryMethods(retrySettingsSetter)
87+
.build();
88+
datastoreStub = GrpcDatastoreStub.create(datastoreStubSettings);
89+
} catch (IOException e) {
90+
throw new IOException(e);
91+
}
92+
}
93+
94+
@Override
95+
public void close() throws Exception {
96+
if (!closed) {
97+
datastoreStub.close();
98+
for (BackgroundResource resource : clientContext.getBackgroundResources()) {
99+
resource.close();
100+
}
101+
closed = true;
102+
}
103+
for (BackgroundResource resource : clientContext.getBackgroundResources()) {
104+
resource.awaitTermination(1, SECONDS);
105+
}
106+
}
107+
108+
@Override
109+
public AllocateIdsResponse allocateIds(AllocateIdsRequest request) {
110+
return datastoreStub.allocateIdsCallable().call(request);
111+
}
112+
113+
@Override
114+
public BeginTransactionResponse beginTransaction(BeginTransactionRequest request)
115+
throws DatastoreException {
116+
return datastoreStub.beginTransactionCallable().call(request);
117+
}
118+
119+
@Override
120+
public CommitResponse commit(CommitRequest request) {
121+
return datastoreStub.commitCallable().call(request);
122+
}
123+
124+
@Override
125+
public LookupResponse lookup(LookupRequest request) {
126+
return datastoreStub.lookupCallable().call(request);
127+
}
128+
129+
@Override
130+
public ReserveIdsResponse reserveIds(ReserveIdsRequest request) {
131+
return datastoreStub.reserveIdsCallable().call(request);
132+
}
133+
134+
@Override
135+
public RollbackResponse rollback(RollbackRequest request) {
136+
return datastoreStub.rollbackCallable().call(request);
137+
}
138+
139+
@Override
140+
public RunQueryResponse runQuery(RunQueryRequest request) {
141+
return datastoreStub.runQueryCallable().call(request);
142+
}
143+
144+
@Override
145+
public RunAggregationQueryResponse runAggregationQuery(RunAggregationQueryRequest request) {
146+
return datastoreStub.runAggregationQueryCallable().call(request);
147+
}
148+
149+
private boolean isEmulator(DatastoreOptions datastoreOptions) {
150+
return isLocalHost(datastoreOptions.getHost())
151+
|| NoCredentials.getInstance().equals(datastoreOptions.getCredentials());
152+
}
153+
154+
private ClientContext getClientContextForEmulator(DatastoreOptions datastoreOptions)
155+
throws IOException {
156+
ManagedChannel managedChannel =
157+
ManagedChannelBuilder.forTarget(removeScheme(datastoreOptions.getHost()))
158+
.usePlaintext()
159+
.build();
160+
TransportChannel transportChannel = GrpcTransportChannel.create(managedChannel);
161+
return ClientContext.newBuilder()
162+
.setCredentials(null)
163+
.setTransportChannel(transportChannel)
164+
.setDefaultCallContext(GrpcCallContext.of(managedChannel, CallOptions.DEFAULT))
165+
.setBackgroundResources(Collections.singletonList(transportChannel))
166+
.build();
167+
}
168+
169+
private ClientContext getClientContext(DatastoreOptions datastoreOptions) throws IOException {
170+
HeaderProvider internalHeaderProvider =
171+
DatastoreSettings.defaultApiClientHeaderProviderBuilder()
172+
.setClientLibToken(
173+
ServiceOptions.getGoogApiClientLibName(),
174+
GaxProperties.getLibraryVersion(datastoreOptions.getClass()))
175+
.setResourceToken(getResourceToken(datastoreOptions))
176+
.build();
177+
178+
DatastoreSettingsBuilder settingsBuilder =
179+
new DatastoreSettingsBuilder(DatastoreSettings.newBuilder().build());
180+
settingsBuilder.setCredentialsProvider(
181+
GrpcTransportOptions.setUpCredentialsProvider(datastoreOptions));
182+
settingsBuilder.setTransportChannelProvider(
183+
GrpcTransportOptions.setUpChannelProvider(
184+
DatastoreSettings.defaultGrpcTransportProviderBuilder(), datastoreOptions));
185+
settingsBuilder.setInternalHeaderProvider(internalHeaderProvider);
186+
settingsBuilder.setHeaderProvider(
187+
datastoreOptions.getMergedHeaderProvider(new NoHeaderProvider()));
188+
ClientContext clientContext = ClientContext.create(settingsBuilder.build());
189+
return clientContext;
190+
}
191+
192+
private String getResourceToken(DatastoreOptions datastoreOptions) {
193+
StringBuilder builder = new StringBuilder("project_id=");
194+
builder.append(datastoreOptions.getProjectId());
195+
if (!Strings.isNullOrEmpty(datastoreOptions.getDatabaseId())) {
196+
builder.append("&database_id=");
197+
builder.append(datastoreOptions.getDatabaseId());
198+
}
199+
return builder.toString();
200+
}
201+
202+
// This class is needed solely to get access to protected method setInternalHeaderProvider()
203+
private static class DatastoreSettingsBuilder extends DatastoreSettings.Builder {
204+
205+
private DatastoreSettingsBuilder(DatastoreSettings settings) {
206+
super(settings);
207+
}
208+
209+
@Override
210+
protected DatastoreSettings.Builder setInternalHeaderProvider(
211+
HeaderProvider internalHeaderProvider) {
212+
return super.setInternalHeaderProvider(internalHeaderProvider);
213+
}
214+
}
215+
}

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

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,7 @@
3838
import com.google.cloud.datastore.Query.ResultType;
3939
import com.google.cloud.datastore.StructuredQuery.OrderBy;
4040
import com.google.cloud.datastore.StructuredQuery.PropertyFilter;
41+
import com.google.cloud.datastore.it.MultipleAttemptsRule;
4142
import com.google.cloud.datastore.spi.DatastoreRpcFactory;
4243
import com.google.cloud.datastore.spi.v1.DatastoreRpc;
4344
import com.google.cloud.datastore.testing.LocalDatastoreHelper;
@@ -84,13 +85,19 @@
8485
import org.junit.Assert;
8586
import org.junit.Before;
8687
import org.junit.BeforeClass;
88+
import org.junit.ClassRule;
89+
import org.junit.Ignore;
8790
import org.junit.Test;
8891
import org.junit.runner.RunWith;
8992
import org.junit.runners.JUnit4;
9093
import org.threeten.bp.Duration;
9194

9295
@RunWith(JUnit4.class)
9396
public class DatastoreTest {
97+
private static final int NUMBER_OF_ATTEMPTS = 5;
98+
99+
@ClassRule
100+
public static MultipleAttemptsRule rr = new MultipleAttemptsRule(NUMBER_OF_ATTEMPTS, 10);
94101

95102
private static LocalDatastoreHelper helper = LocalDatastoreHelper.create(1.0);
96103
private static final DatastoreOptions options = helper.getOptions();
@@ -231,6 +238,8 @@ public void testNewTransactionCommit() {
231238
verifyNotUsable(transaction);
232239
}
233240

241+
// TODO(gapic_upgrade): Remove the @ignore annotation
242+
@Ignore("This should be fixed with actionable error implementation")
234243
@Test
235244
public void testTransactionWithRead() {
236245
Transaction transaction = datastore.newTransaction();
@@ -252,6 +261,8 @@ public void testTransactionWithRead() {
252261
}
253262
}
254263

264+
// TODO(gapic_upgrade): Remove the @ignore annotation
265+
@Ignore("This should be fixed with actionable error implementation")
255266
@Test
256267
public void testTransactionWithQuery() {
257268
Query<Entity> query =
@@ -648,6 +659,7 @@ private List<RunQueryResponse> buildResponsesForQueryPagination() {
648659
List<RunQueryResponse> responses = new ArrayList<>();
649660
RecordQuery<Key> query = Query.newKeyQueryBuilder().build();
650661
RunQueryRequest.Builder requestPb = RunQueryRequest.newBuilder();
662+
requestPb.setProjectId(PROJECT_ID);
651663
query.populatePb(requestPb);
652664
QueryResultBatch queryResultBatchPb =
653665
RunQueryResponse.newBuilder()
@@ -757,6 +769,7 @@ private List<RunQueryResponse> buildResponsesForQueryPaginationWithLimit() {
757769
List<RunQueryResponse> responses = new ArrayList<>();
758770
RecordQuery<Entity> query = Query.newEntityQueryBuilder().build();
759771
RunQueryRequest.Builder requestPb = RunQueryRequest.newBuilder();
772+
requestPb.setProjectId(PROJECT_ID);
760773
query.populatePb(requestPb);
761774
QueryResultBatch queryResultBatchPb =
762775
RunQueryResponse.newBuilder()

0 commit comments

Comments
 (0)