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.

Author: kolea2

.kokoro/nightly/integration.cfg

@@ -21,6 +21,11 @@ env_vars: {
  value: "java-docs-samples-testing"
 }
 
+env_vars: {
+ key: "DATASTORE_PROJECT_ID"
+ value: "java-docs-samples-testing"
+}
+
 env_vars: {
  key: "ENABLE_FLAKYBOT"
  value: "true"

.kokoro/nightly/java11-integration.cfg

@@ -21,6 +21,11 @@ env_vars: {
  value: "gcloud-devel"
 }
 
+env_vars: {
+ key: "DATASTORE_PROJECT_ID"
+ value: "gcloud-devel"
+}
+
 env_vars: {
  key: "ENABLE_FLAKYBOT"
  value: "true"

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

@@ -30,4 +30,9 @@ env_vars: {
 env_vars: {
  key: "SECRET_MANAGER_KEYS"
  value: "java-it-service-account"
+}
+
+env_vars: {
+ key: "DATASTORE_PROJECT_ID"
+ value: "gcloud-devel"
 }

.kokoro/presubmit/graalvm-native.cfg

@@ -31,3 +31,8 @@ env_vars: {
  key: "SECRET_MANAGER_KEYS"
  value: "java-it-service-account"
 }
+
+env_vars: {
+ key: "DATASTORE_PROJECT_ID"
+ value: "gcloud-devel"
+}

.kokoro/presubmit/integration.cfg

@@ -22,6 +22,11 @@ env_vars: {
  value: "gcloud-devel"
 }
 
+env_vars: {
+ key: "DATASTORE_PROJECT_ID"
+ value: "gcloud-devel"
+}
+
 env_vars: {
  key: "GOOGLE_APPLICATION_CREDENTIALS"
  value: "secret_manager/java-it-service-account"

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

@@ -67,7 +67,8 @@ private DatastoreException invalidResponseException(String method, IOException e
  }
 
  public AllocateIdsResponse allocateIds(AllocateIdsRequest request) throws DatastoreException {
- try (InputStream is = remoteRpc.call("allocateIds", request)) {
+ try (InputStream is =
+ remoteRpc.call("allocateIds", request, request.getProjectId(), request.getDatabaseId())) {
  return AllocateIdsResponse.parseFrom(is);
  } catch (IOException exception) {
  throw invalidResponseException("allocateIds", exception);
@@ -76,47 +77,54 @@ public AllocateIdsResponse allocateIds(AllocateIdsRequest request) throws Datast
 
  public BeginTransactionResponse beginTransaction(BeginTransactionRequest request)
  throws DatastoreException {
- try (InputStream is = remoteRpc.call("beginTransaction", request)) {
+ try (InputStream is =
+ remoteRpc.call(
+ "beginTransaction", request, request.getProjectId(), request.getDatabaseId())) {
  return BeginTransactionResponse.parseFrom(is);
  } catch (IOException exception) {
  throw invalidResponseException("beginTransaction", exception);
  }
  }
 
  public CommitResponse commit(CommitRequest request) throws DatastoreException {
- try (InputStream is = remoteRpc.call("commit", request)) {
+ try (InputStream is =
+ remoteRpc.call("commit", request, request.getProjectId(), request.getDatabaseId())) {
  return CommitResponse.parseFrom(is);
  } catch (IOException exception) {
  throw invalidResponseException("commit", exception);
  }
  }
 
  public LookupResponse lookup(LookupRequest request) throws DatastoreException {
- try (InputStream is = remoteRpc.call("lookup", request)) {
+ try (InputStream is =
+ remoteRpc.call("lookup", request, request.getProjectId(), request.getDatabaseId())) {
  return LookupResponse.parseFrom(is);
  } catch (IOException exception) {
  throw invalidResponseException("lookup", exception);
  }
  }
 
  public ReserveIdsResponse reserveIds(ReserveIdsRequest request) throws DatastoreException {
- try (InputStream is = remoteRpc.call("reserveIds", request)) {
+ try (InputStream is =
+ remoteRpc.call("reserveIds", request, request.getProjectId(), request.getDatabaseId())) {
  return ReserveIdsResponse.parseFrom(is);
  } catch (IOException exception) {
  throw invalidResponseException("reserveIds", exception);
  }
  }
 
  public RollbackResponse rollback(RollbackRequest request) throws DatastoreException {
- try (InputStream is = remoteRpc.call("rollback", request)) {
+ try (InputStream is =
+ remoteRpc.call("rollback", request, request.getProjectId(), request.getDatabaseId())) {
  return RollbackResponse.parseFrom(is);
  } catch (IOException exception) {
  throw invalidResponseException("rollback", exception);
  }
  }
 
  public RunQueryResponse runQuery(RunQueryRequest request) throws DatastoreException {
- try (InputStream is = remoteRpc.call("runQuery", request)) {
+ try (InputStream is =
+ remoteRpc.call("runQuery", request, request.getProjectId(), request.getDatabaseId())) {
  return RunQueryResponse.parseFrom(is);
  } catch (IOException exception) {
  throw invalidResponseException("runQuery", exception);
@@ -125,7 +133,9 @@ public RunQueryResponse runQuery(RunQueryRequest request) throws DatastoreExcept
 
  public RunAggregationQueryResponse runAggregationQuery(RunAggregationQueryRequest request)
  throws DatastoreException {
- try (InputStream is = remoteRpc.call("runAggregationQuery", request)) {
+ try (InputStream is =
+ remoteRpc.call(
+ "runAggregationQuery", request, request.getProjectId(), request.getDatabaseId())) {
  return RunAggregationQueryResponse.parseFrom(is);
  } catch (IOException exception) {
  throw invalidResponseException("runAggregationQuery", exception);

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

@@ -20,6 +20,7 @@
 import com.google.api.client.auth.oauth2.Credential;
 import com.google.api.client.http.HttpRequestInitializer;
 import com.google.api.client.http.HttpTransport;
+import com.google.api.core.BetaApi;
 import java.util.Arrays;
 import java.util.List;
 
@@ -40,6 +41,7 @@
  */
 public class DatastoreOptions {
  private final String projectId;
+ private final String databaseId;
  private final String projectEndpoint;
  private final String host;
  private final String localHost;
@@ -56,6 +58,7 @@ public class DatastoreOptions {
  b.projectId != null || b.projectEndpoint != null,
  "Either project ID or project endpoint must be provided.");
  this.projectId = b.projectId;
+ this.databaseId = b.databaseId;
  this.projectEndpoint = b.projectEndpoint;
  this.host = b.host;
  this.localHost = b.localHost;
@@ -72,6 +75,7 @@ public static class Builder {
  "Can set at most one of project endpoint, host, and local host.";
 
  private String projectId;
+ private String databaseId;
  private String projectEndpoint;
  private String host;
  private String localHost;
@@ -83,6 +87,7 @@ public Builder() {}
 
  public Builder(DatastoreOptions options) {
  this.projectId = options.projectId;
+ this.databaseId = options.databaseId;
  this.projectEndpoint = options.projectEndpoint;
  this.host = options.host;
  this.localHost = options.localHost;
@@ -102,6 +107,13 @@ public Builder projectId(String projectId) {
  return this;
  }
 
+ /** Sets the database ID used to access Cloud Datastore. */
+ @BetaApi
+ public Builder databaseId(String databaseId) {
+ this.databaseId = databaseId;
+ return this;
+ }
+
  /**
  * Sets the host used to access Cloud Datastore. To connect to the Cloud Datastore Emulator, use
  * {@link #localHost} instead.
@@ -176,6 +188,10 @@ public String getProjectId() {
  return projectId;
  }
 
+ public String getDatabaseId() {
+ return databaseId;
+ }
+
  public String getProjectEndpoint() {
  return projectEndpoint;
  }

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

@@ -221,6 +221,8 @@ private List<Key> getScatterKeys(
  do {
  RunQueryRequest.Builder scatterRequest =
  RunQueryRequest.newBuilder().setPartitionId(partition).setQuery(scatterPointQuery);
+ scatterRequest.setProjectId(partition.getProjectId());
+ scatterRequest.setDatabaseId(partition.getDatabaseId());
  if (readTime != null) {
  scatterRequest.setReadOptions(ReadOptions.newBuilder().setReadTime(readTime).build());
  }

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

@@ -24,6 +24,7 @@
 import com.google.api.client.http.protobuf.ProtoHttpContent;
 import com.google.api.client.util.IOUtils;
 import com.google.common.annotations.VisibleForTesting;
+import com.google.common.base.Strings;
 import com.google.protobuf.MessageLite;
 import com.google.rpc.Code;
 import com.google.rpc.Status;
@@ -46,6 +47,8 @@ class RemoteRpc {
  @VisibleForTesting static final String API_FORMAT_VERSION_HEADER = "X-Goog-Api-Format-Version";
  private static final String API_FORMAT_VERSION = "2";
 
+ @VisibleForTesting static final String X_GOOG_REQUEST_PARAMS_HEADER = "x-goog-request-params";
+
  private final HttpRequestFactory client;
  private final HttpRequestInitializer initializer;
  private final String url;
@@ -74,7 +77,9 @@ class RemoteRpc {
  *
  * @throws DatastoreException if the RPC fails.
  */
- public InputStream call(String methodName, MessageLite request) throws DatastoreException {
+ public InputStream call(
+ String methodName, MessageLite request, String projectId, String databaseId)
+ throws DatastoreException {
  logger.fine("remote datastore call " + methodName);
 
  long startTime = System.currentTimeMillis();
@@ -84,7 +89,7 @@ public InputStream call(String methodName, MessageLite request) throws Datastore
  rpcCount.incrementAndGet();
  ProtoHttpContent payload = new ProtoHttpContent(request);
  HttpRequest httpRequest = client.buildPostRequest(resolveURL(methodName), payload);
- setHeaders(request, httpRequest);
+ setHeaders(request, httpRequest, projectId, databaseId);
  // Don't throw an HTTPResponseException on error. It converts the response to a String and
  // throws away the original, whereas we need the raw bytes to parse it as a proto.
  httpRequest.setThrowExceptionOnExecuteError(false);
@@ -123,8 +128,16 @@ public InputStream call(String methodName, MessageLite request) throws Datastore
  }
 
  @VisibleForTesting
- void setHeaders(MessageLite request, HttpRequest httpRequest) {
+ void setHeaders(
+ MessageLite request, HttpRequest httpRequest, String projectId, String databaseId) {
  httpRequest.getHeaders().put(API_FORMAT_VERSION_HEADER, API_FORMAT_VERSION);
+ StringBuilder builder = new StringBuilder("project_id=");
+ builder.append(projectId);
+ if (!Strings.isNullOrEmpty(databaseId)) {
+ builder.append("&database_id=");
+ builder.append(databaseId);
+ }
+ httpRequest.getHeaders().put(X_GOOG_REQUEST_PARAMS_HEADER, builder.toString());
  if (enableE2EChecksum && request != null) {
  String checksum = EndToEndChecksumHandler.computeChecksum(request.toByteArray());
  if (checksum != null) {

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

@@ -215,6 +215,18 @@ public void create_LocalHost() {
  .isEqualTo("http://localhost:8080/v1/projects/project-id");
  }
 
+ @Test
+ public void setDatabaseId() {
+ DatastoreOptions options =
+ new DatastoreOptions.Builder()
+ .projectId(PROJECT_ID)
+ .databaseId("test-db")
+ .localHost("localhost:8080")
+ .build();
+ assertThat(options.getProjectId()).isEqualTo(PROJECT_ID);
+ assertThat(options.getDatabaseId()).isEqualTo("test-db");
+ }
+
  @Test
  public void create_LocalHostIp() {
  Datastore datastore =