feat: multiple dbs support

.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/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/DatastoreOptions.java

@@ -40,6 +40,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 +57,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 +74,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 +86,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 +106,12 @@ public Builder projectId(String projectId) {
       return this;
     }
 
+    /** Sets the database ID used to access Cloud Datastore. */
+    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 +186,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/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 =

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

@@ -193,6 +193,7 @@ public void getSplits() throws Exception {
     RunQueryRequest expectedSplitQueryRequest =
         RunQueryRequest.newBuilder()
             .setPartitionId(PARTITION)
+            .setProjectId(PROJECT_ID)
             .setQuery(
                 splitQuery.toBuilder().setLimit(Int32Value.newBuilder().setValue(2 * 32).build()))
             .build();
@@ -235,6 +236,7 @@ public void notEnoughSplits() throws Exception {
     RunQueryRequest expectedSplitQueryRequest =
         RunQueryRequest.newBuilder()
             .setPartitionId(PARTITION)
+            .setProjectId(PROJECT_ID)
             .setQuery(
                 splitQuery.toBuilder().setLimit(Int32Value.newBuilder().setValue(99 * 32).build()))
             .build();
@@ -286,6 +288,7 @@ public void getSplits_withReadTime() throws Exception {
     RunQueryRequest expectedSplitQueryRequest =
         RunQueryRequest.newBuilder()
             .setPartitionId(PARTITION)
+            .setProjectId(PROJECT_ID)
             .setQuery(
                 splitQuery.toBuilder().setLimit(Int32Value.newBuilder().setValue(2 * 32).build()))
             .setReadOptions(ReadOptions.newBuilder().setReadTime(readTime))

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

@@ -0,0 +1,82 @@
+/*
+ * Copyright 2022 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.datastore.v1.client.it;
+
+import static com.google.datastore.v1.client.DatastoreHelper.makeFilter;
+import static com.google.datastore.v1.client.DatastoreHelper.makeValue;
+
+import com.google.datastore.v1.Filter;
+import com.google.datastore.v1.KindExpression;
+import com.google.datastore.v1.PartitionId;
+import com.google.datastore.v1.PropertyFilter;
+import com.google.datastore.v1.Query;
+import com.google.datastore.v1.client.Datastore;
+import com.google.datastore.v1.client.DatastoreException;
+import com.google.datastore.v1.client.DatastoreHelper;
+import java.io.IOException;
+import java.security.GeneralSecurityException;
+import java.util.List;
+import org.junit.Assert;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.Parameterized;
+
+@RunWith(Parameterized.class)
+public class ITDatastoreProtoClientTest {
+
+  private static Datastore DATASTORE;
+  private final String databaseId;
+
+  @Parameterized.Parameters(name = "database id: {0}")
+  public static Object[] data() {
+    return new Object[] {"", "test-db"};
+  }
+
+  public ITDatastoreProtoClientTest(String databaseId) {
+    this.databaseId = databaseId;
+  }
+
+  private static PartitionId PARTITION;
+
+  private static final String KIND = "test-kind";
+  private static final String PROJECT_ID = System.getenv(DatastoreHelper.PROJECT_ID_ENV_VAR);
+
+  @Before
+  public void setUp() throws GeneralSecurityException, IOException {
+    DATASTORE = DatastoreHelper.getDatastoreFromEnv();
+    PARTITION = PartitionId.newBuilder().setProjectId(PROJECT_ID).setDatabaseId(databaseId).build();
+  }
+
+  @Test
+  public void testQuerySplitterWithDb() throws DatastoreException {
+    Filter propertyFilter =
+        makeFilter("foo", PropertyFilter.Operator.EQUAL, makeValue("value")).build();
+    Query query =
+        Query.newBuilder()
+            .addKind(KindExpression.newBuilder().setName(KIND).build())
+            .setFilter(propertyFilter)
+            .build();
+
+    List<Query> splits =
+        DatastoreHelper.getQuerySplitter().getSplits(query, PARTITION, 2, DATASTORE);
+    splits.forEach(
+        split -> {
+          Assert.assertEquals(KIND, split.getKind(0).getName());
+          Assert.assertEquals(propertyFilter, split.getFilter());
+        });
+  }
+}

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

@@ -16,9 +16,9 @@
 
 package com.google.cloud.datastore;
 
-import static com.google.cloud.datastore.Validator.validateDatabase;
 import static com.google.cloud.datastore.Validator.validateKind;
 import static com.google.cloud.datastore.Validator.validateNamespace;
+import static com.google.cloud.datastore.Validator.validateProjectId;
 
 import com.google.common.base.MoreObjects;
 import com.google.common.base.Preconditions;
@@ -35,6 +35,7 @@ public abstract class BaseKey implements Serializable {
 
   private final String projectId;
   private final String namespace;
+  private final String databaseId;
   private final ImmutableList<PathElement> path;
 
   /**
@@ -46,13 +47,14 @@ public abstract static class Builder<B extends Builder<B>> {
 
     String projectId = "";
     String namespace = "";
+    String databaseId = "";
     String kind;
     final List<PathElement> ancestors;
 
     private static final int MAX_PATH = 100;
 
     Builder(String projectId) {
-      this.projectId = validateDatabase(projectId);
+      this.projectId = validateProjectId(projectId);
       ancestors = new LinkedList<>();
     }
 
@@ -64,6 +66,7 @@ public abstract static class Builder<B extends Builder<B>> {
     Builder(BaseKey copyFrom) {
       projectId = copyFrom.getProjectId();
       namespace = copyFrom.getNamespace();
+      databaseId = copyFrom.getDatabaseId();
       ancestors = new LinkedList<>(copyFrom.getAncestors());
       kind = copyFrom.getKind();
     }
@@ -102,7 +105,7 @@ public B setKind(String kind) {
 
     /** Sets the project ID of the key. */
     public B setProjectId(String projectId) {
-      this.projectId = validateDatabase(projectId);
+      this.projectId = validateProjectId(projectId);
       return self();
     }
 
@@ -112,13 +115,24 @@ public B setNamespace(String namespace) {
       return self();
     }
 
+    /** Sets the database id of the key. */
+    public B setDatabaseId(String databaseId) {
+      this.databaseId = databaseId;
+      return self();
+    }
+
     protected abstract BaseKey build();
   }
 
   BaseKey(String projectId, String namespace, ImmutableList<PathElement> path) {
+    this(projectId, namespace, "", path);
+  }
+
+  BaseKey(String projectId, String namespace, String databaseId, ImmutableList<PathElement> path) {
     Preconditions.checkArgument(!path.isEmpty(), "Path must not be empty");
     this.projectId = projectId;
     this.namespace = namespace;
+    this.databaseId = databaseId;
     this.path = path;
   }
 
@@ -132,6 +146,10 @@ public String getNamespace() {
     return namespace;
   }
 
+  public String getDatabaseId() {
+    return databaseId;
+  }
+
   /** Returns an immutable list with the key's ancestors. */
   public List<PathElement> getAncestors() {
     return getPath().subList(0, getPath().size() - 1);
@@ -158,13 +176,14 @@ public String toString() {
     return MoreObjects.toStringHelper(this)
         .add("projectId", projectId)
         .add("namespace", namespace)
+        .add("databaseId", databaseId)
         .add("path", path)
         .toString();
   }
 
   @Override
   public int hashCode() {
-    return Objects.hash(getProjectId(), getNamespace(), getPath());
+    return Objects.hash(getProjectId(), getNamespace(), getDatabaseId(), getPath());
   }
 
   @Override
@@ -178,6 +197,7 @@ public boolean equals(Object obj) {
     BaseKey other = (BaseKey) obj;
     return Objects.equals(getProjectId(), other.getProjectId())
         && Objects.equals(getNamespace(), other.getNamespace())
+        && Objects.equals(getDatabaseId(), other.getDatabaseId())
         && Objects.equals(getPath(), other.getPath());
   }
 
@@ -186,6 +206,7 @@ com.google.datastore.v1.Key toPb() {
     com.google.datastore.v1.PartitionId.Builder partitionIdPb =
         com.google.datastore.v1.PartitionId.newBuilder();
     partitionIdPb.setProjectId(projectId);
+    partitionIdPb.setDatabaseId(databaseId);
     partitionIdPb.setNamespaceId(namespace);
     keyPb.setPartitionId(partitionIdPb.build());
     for (PathElement pathEntry : path) {

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

@@ -59,6 +59,8 @@ public Batch.Response submit() {
         com.google.datastore.v1.CommitRequest.newBuilder();
     requestPb.setMode(com.google.datastore.v1.CommitRequest.Mode.NON_TRANSACTIONAL);
     requestPb.addAllMutations(mutationsPb);
+    requestPb.setProjectId(datastore.getOptions().getProjectId());
+    requestPb.setDatabaseId(datastore.getOptions().getDatabaseId());
     com.google.datastore.v1.CommitResponse responsePb = datastore.commit(requestPb.build());
     deactivate();
     return new ResponseImpl(responsePb, toAddAutoId().size());

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

@@ -54,7 +54,7 @@ static Entity put(DatastoreWriter writer, FullEntity<?> entity) {
   }
 
   static KeyFactory newKeyFactory(DatastoreOptions options) {
-    return new KeyFactory(options.getProjectId(), options.getNamespace());
+    return new KeyFactory(options.getProjectId(), options.getNamespace(), options.getDatabaseId());
   }
 
   /**