Skip to content
This repository was archived by the owner on Feb 24, 2026. It is now read-only.

Commit 7887f32

Browse files
kolea2kolea2
authored
feat: multiple dbs support (#1102)
This PR consists of cherry picks from the https://togithub.com/googleapis/java-datastore/tree/multi-db branch. They include: #928 #940 #942 This also enables parameterized testing for ITDatastoreTest.
1 parent c226997 commit 7887f32

39 files changed

Lines changed: 908 additions & 273 deletions

‎.kokoro/nightly/integration.cfg‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,11 @@ env_vars: {
2121
value: "java-docs-samples-testing"
2222
}
2323

24+
env_vars: {
25+
key: "DATASTORE_PROJECT_ID"
26+
value: "java-docs-samples-testing"
27+
}
28+
2429
env_vars: {
2530
key: "ENABLE_FLAKYBOT"
2631
value: "true"

‎.kokoro/nightly/java11-integration.cfg‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,11 @@ env_vars: {
2121
value: "gcloud-devel"
2222
}
2323

24+
env_vars: {
25+
key: "DATASTORE_PROJECT_ID"
26+
value: "gcloud-devel"
27+
}
28+
2429
env_vars: {
2530
key: "ENABLE_FLAKYBOT"
2631
value: "true"

‎.kokoro/presubmit/graalvm-native-17.cfg‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,4 +30,9 @@ env_vars: {
3030
env_vars: {
3131
key: "SECRET_MANAGER_KEYS"
3232
value: "java-it-service-account"
33+
}
34+
35+
env_vars: {
36+
key: "DATASTORE_PROJECT_ID"
37+
value: "gcloud-devel"
3338
}

‎.kokoro/presubmit/graalvm-native.cfg‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,3 +31,8 @@ env_vars: {
3131
key: "SECRET_MANAGER_KEYS"
3232
value: "java-it-service-account"
3333
}
34+
35+
env_vars: {
36+
key: "DATASTORE_PROJECT_ID"
37+
value: "gcloud-devel"
38+
}

‎.kokoro/presubmit/integration.cfg‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,11 @@ env_vars: {
2222
value: "gcloud-devel"
2323
}
2424

25+
env_vars: {
26+
key: "DATASTORE_PROJECT_ID"
27+
value: "gcloud-devel"
28+
}
29+
2530
env_vars: {
2631
key: "GOOGLE_APPLICATION_CREDENTIALS"
2732
value: "secret_manager/java-it-service-account"

‎datastore-v1-proto-client/src/main/java/com/google/datastore/v1/client/Datastore.java‎

Lines changed: 18 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -67,7 +67,8 @@ private DatastoreException invalidResponseException(String method, IOException e
6767
}
6868

6969
public AllocateIdsResponse allocateIds(AllocateIdsRequest request) throws DatastoreException {
70-
try (InputStream is = remoteRpc.call("allocateIds", request)) {
70+
try (InputStream is =
71+
remoteRpc.call("allocateIds", request, request.getProjectId(), request.getDatabaseId())) {
7172
return AllocateIdsResponse.parseFrom(is);
7273
} catch (IOException exception) {
7374
throw invalidResponseException("allocateIds", exception);
@@ -76,47 +77,54 @@ public AllocateIdsResponse allocateIds(AllocateIdsRequest request) throws Datast
7677

7778
public BeginTransactionResponse beginTransaction(BeginTransactionRequest request)
7879
throws DatastoreException {
79-
try (InputStream is = remoteRpc.call("beginTransaction", request)) {
80+
try (InputStream is =
81+
remoteRpc.call(
82+
"beginTransaction", request, request.getProjectId(), request.getDatabaseId())) {
8083
return BeginTransactionResponse.parseFrom(is);
8184
} catch (IOException exception) {
8285
throw invalidResponseException("beginTransaction", exception);
8386
}
8487
}
8588

8689
public CommitResponse commit(CommitRequest request) throws DatastoreException {
87-
try (InputStream is = remoteRpc.call("commit", request)) {
90+
try (InputStream is =
91+
remoteRpc.call("commit", request, request.getProjectId(), request.getDatabaseId())) {
8892
return CommitResponse.parseFrom(is);
8993
} catch (IOException exception) {
9094
throw invalidResponseException("commit", exception);
9195
}
9296
}
9397

9498
public LookupResponse lookup(LookupRequest request) throws DatastoreException {
95-
try (InputStream is = remoteRpc.call("lookup", request)) {
99+
try (InputStream is =
100+
remoteRpc.call("lookup", request, request.getProjectId(), request.getDatabaseId())) {
96101
return LookupResponse.parseFrom(is);
97102
} catch (IOException exception) {
98103
throw invalidResponseException("lookup", exception);
99104
}
100105
}
101106

102107
public ReserveIdsResponse reserveIds(ReserveIdsRequest request) throws DatastoreException {
103-
try (InputStream is = remoteRpc.call("reserveIds", request)) {
108+
try (InputStream is =
109+
remoteRpc.call("reserveIds", request, request.getProjectId(), request.getDatabaseId())) {
104110
return ReserveIdsResponse.parseFrom(is);
105111
} catch (IOException exception) {
106112
throw invalidResponseException("reserveIds", exception);
107113
}
108114
}
109115

110116
public RollbackResponse rollback(RollbackRequest request) throws DatastoreException {
111-
try (InputStream is = remoteRpc.call("rollback", request)) {
117+
try (InputStream is =
118+
remoteRpc.call("rollback", request, request.getProjectId(), request.getDatabaseId())) {
112119
return RollbackResponse.parseFrom(is);
113120
} catch (IOException exception) {
114121
throw invalidResponseException("rollback", exception);
115122
}
116123
}
117124

118125
public RunQueryResponse runQuery(RunQueryRequest request) throws DatastoreException {
119-
try (InputStream is = remoteRpc.call("runQuery", request)) {
126+
try (InputStream is =
127+
remoteRpc.call("runQuery", request, request.getProjectId(), request.getDatabaseId())) {
120128
return RunQueryResponse.parseFrom(is);
121129
} catch (IOException exception) {
122130
throw invalidResponseException("runQuery", exception);
@@ -125,7 +133,9 @@ public RunQueryResponse runQuery(RunQueryRequest request) throws DatastoreExcept
125133

126134
public RunAggregationQueryResponse runAggregationQuery(RunAggregationQueryRequest request)
127135
throws DatastoreException {
128-
try (InputStream is = remoteRpc.call("runAggregationQuery", request)) {
136+
try (InputStream is =
137+
remoteRpc.call(
138+
"runAggregationQuery", request, request.getProjectId(), request.getDatabaseId())) {
129139
return RunAggregationQueryResponse.parseFrom(is);
130140
} catch (IOException exception) {
131141
throw invalidResponseException("runAggregationQuery", exception);

‎datastore-v1-proto-client/src/main/java/com/google/datastore/v1/client/DatastoreOptions.java‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020
import com.google.api.client.auth.oauth2.Credential;
2121
import com.google.api.client.http.HttpRequestInitializer;
2222
import com.google.api.client.http.HttpTransport;
23+
import com.google.api.core.BetaApi;
2324
import java.util.Arrays;
2425
import java.util.List;
2526

@@ -40,6 +41,7 @@
4041
*/
4142
public class DatastoreOptions {
4243
private final String projectId;
44+
private final String databaseId;
4345
private final String projectEndpoint;
4446
private final String host;
4547
private final String localHost;
@@ -56,6 +58,7 @@ public class DatastoreOptions {
5658
b.projectId != null || b.projectEndpoint != null,
5759
"Either project ID or project endpoint must be provided.");
5860
this.projectId = b.projectId;
61+
this.databaseId = b.databaseId;
5962
this.projectEndpoint = b.projectEndpoint;
6063
this.host = b.host;
6164
this.localHost = b.localHost;
@@ -72,6 +75,7 @@ public static class Builder {
7275
"Can set at most one of project endpoint, host, and local host.";
7376

7477
private String projectId;
78+
private String databaseId;
7579
private String projectEndpoint;
7680
private String host;
7781
private String localHost;
@@ -83,6 +87,7 @@ public Builder() {}
8387

8488
public Builder(DatastoreOptions options) {
8589
this.projectId = options.projectId;
90+
this.databaseId = options.databaseId;
8691
this.projectEndpoint = options.projectEndpoint;
8792
this.host = options.host;
8893
this.localHost = options.localHost;
@@ -102,6 +107,13 @@ public Builder projectId(String projectId) {
102107
return this;
103108
}
104109

110+
/** Sets the database ID used to access Cloud Datastore. */
111+
@BetaApi
112+
public Builder databaseId(String databaseId) {
113+
this.databaseId = databaseId;
114+
return this;
115+
}
116+
105117
/**
106118
* Sets the host used to access Cloud Datastore. To connect to the Cloud Datastore Emulator, use
107119
* {@link #localHost} instead.
@@ -176,6 +188,10 @@ public String getProjectId() {
176188
return projectId;
177189
}
178190

191+
public String getDatabaseId() {
192+
return databaseId;
193+
}
194+
179195
public String getProjectEndpoint() {
180196
return projectEndpoint;
181197
}

‎datastore-v1-proto-client/src/main/java/com/google/datastore/v1/client/QuerySplitterImpl.java‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -221,6 +221,8 @@ private List<Key> getScatterKeys(
221221
do {
222222
RunQueryRequest.Builder scatterRequest =
223223
RunQueryRequest.newBuilder().setPartitionId(partition).setQuery(scatterPointQuery);
224+
scatterRequest.setProjectId(partition.getProjectId());
225+
scatterRequest.setDatabaseId(partition.getDatabaseId());
224226
if (readTime != null) {
225227
scatterRequest.setReadOptions(ReadOptions.newBuilder().setReadTime(readTime).build());
226228
}

‎datastore-v1-proto-client/src/main/java/com/google/datastore/v1/client/RemoteRpc.java‎

Lines changed: 16 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@
2424
import com.google.api.client.http.protobuf.ProtoHttpContent;
2525
import com.google.api.client.util.IOUtils;
2626
import com.google.common.annotations.VisibleForTesting;
27+
import com.google.common.base.Strings;
2728
import com.google.protobuf.MessageLite;
2829
import com.google.rpc.Code;
2930
import com.google.rpc.Status;
@@ -46,6 +47,8 @@ class RemoteRpc {
4647
@VisibleForTesting static final String API_FORMAT_VERSION_HEADER = "X-Goog-Api-Format-Version";
4748
private static final String API_FORMAT_VERSION = "2";
4849

50+
@VisibleForTesting static final String X_GOOG_REQUEST_PARAMS_HEADER = "x-goog-request-params";
51+
4952
private final HttpRequestFactory client;
5053
private final HttpRequestInitializer initializer;
5154
private final String url;
@@ -74,7 +77,9 @@ class RemoteRpc {
7477
*
7578
* @throws DatastoreException if the RPC fails.
7679
*/
77-
public InputStream call(String methodName, MessageLite request) throws DatastoreException {
80+
public InputStream call(
81+
String methodName, MessageLite request, String projectId, String databaseId)
82+
throws DatastoreException {
7883
logger.fine("remote datastore call " + methodName);
7984

8085
long startTime = System.currentTimeMillis();
@@ -84,7 +89,7 @@ public InputStream call(String methodName, MessageLite request) throws Datastore
8489
rpcCount.incrementAndGet();
8590
ProtoHttpContent payload = new ProtoHttpContent(request);
8691
HttpRequest httpRequest = client.buildPostRequest(resolveURL(methodName), payload);
87-
setHeaders(request, httpRequest);
92+
setHeaders(request, httpRequest, projectId, databaseId);
8893
// Don't throw an HTTPResponseException on error. It converts the response to a String and
8994
// throws away the original, whereas we need the raw bytes to parse it as a proto.
9095
httpRequest.setThrowExceptionOnExecuteError(false);
@@ -123,8 +128,16 @@ public InputStream call(String methodName, MessageLite request) throws Datastore
123128
}
124129

125130
@VisibleForTesting
126-
void setHeaders(MessageLite request, HttpRequest httpRequest) {
131+
void setHeaders(
132+
MessageLite request, HttpRequest httpRequest, String projectId, String databaseId) {
127133
httpRequest.getHeaders().put(API_FORMAT_VERSION_HEADER, API_FORMAT_VERSION);
134+
StringBuilder builder = new StringBuilder("project_id=");
135+
builder.append(projectId);
136+
if (!Strings.isNullOrEmpty(databaseId)) {
137+
builder.append("&database_id=");
138+
builder.append(databaseId);
139+
}
140+
httpRequest.getHeaders().put(X_GOOG_REQUEST_PARAMS_HEADER, builder.toString());
128141
if (enableE2EChecksum && request != null) {
129142
String checksum = EndToEndChecksumHandler.computeChecksum(request.toByteArray());
130143
if (checksum != null) {

‎datastore-v1-proto-client/src/test/java/com/google/datastore/v1/client/DatastoreClientTest.java‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -215,6 +215,18 @@ public void create_LocalHost() {
215215
.isEqualTo("http://localhost:8080/v1/projects/project-id");
216216
}
217217

218+
@Test
219+
public void setDatabaseId() {
220+
DatastoreOptions options =
221+
new DatastoreOptions.Builder()
222+
.projectId(PROJECT_ID)
223+
.databaseId("test-db")
224+
.localHost("localhost:8080")
225+
.build();
226+
assertThat(options.getProjectId()).isEqualTo(PROJECT_ID);
227+
assertThat(options.getDatabaseId()).isEqualTo("test-db");
228+
}
229+
218230
@Test
219231
public void create_LocalHostIp() {
220232
Datastore datastore =

0 commit comments

Comments
 (0)