cherry pick of googleapis#928

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/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,60 @@ 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();
+
+    assertArrayEquals(expectedSplitQueryRequest.toByteArray(), mockClient.getLastBody());
+  }
+
+  @Test
+  public void getSplitsWithDatabaseId() throws Exception {
+    Datastore datastore = factory.create(options.build());
+    MockDatastoreFactory mockClient = (MockDatastoreFactory) factory;
+
+    PartitionId partition =
+        PartitionId.newBuilder().setProjectId(PROJECT_ID).setDatabaseId("test-database").build();
+
+    RunQueryResponse splitQueryResponse =
+        RunQueryResponse.newBuilder()
+            .setQuery(splitQuery)
+            .setBatch(
+                QueryResultBatch.newBuilder()
+                    .setEntityResultType(ResultType.KEY_ONLY)
+                    .setMoreResults(MoreResultsType.NO_MORE_RESULTS)
+                    .addEntityResults(makeKeyOnlyEntity(splitKey0))
+                    .addEntityResults(makeKeyOnlyEntity(splitKey1))
+                    .addEntityResults(makeKeyOnlyEntity(splitKey2))
+                    .addEntityResults(makeKeyOnlyEntity(splitKey3))
+                    .build())
+            .build();
+
+    mockClient.setNextResponse(splitQueryResponse);
+
+    List<Query> splitQueries = QuerySplitterImpl.INSTANCE.getSplits(query, partition, 3, datastore);
+
+    assertThat(splitQueries)
+        .containsExactly(
+            query
+                .toBuilder()
+                .setFilter(makeFilterWithKeyRange(propertyFilter, null, splitKey1))
+                .build(),
+            query
+                .toBuilder()
+                .setFilter(makeFilterWithKeyRange(propertyFilter, splitKey1, splitKey3))
+                .build(),
+            query
+                .toBuilder()
+                .setFilter(makeFilterWithKeyRange(propertyFilter, splitKey3, null))
+                .build());
+
+    RunQueryRequest expectedSplitQueryRequest =
+        RunQueryRequest.newBuilder()
+            .setPartitionId(partition)
+            .setProjectId(PROJECT_ID)
+            .setDatabaseId("test-database")
             .setQuery(
                 splitQuery.toBuilder().setLimit(Int32Value.newBuilder().setValue(2 * 32).build()))
             .build();
@@ -235,6 +289,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 +341,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,91 @@
+/*
+ * 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;
+
+public class ITDatastoreProtoClientTest {
+
+  private static Datastore DATASTORE;
+
+  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();
+  }
+
+  @Test
+  public void testQuerySplitterWithDefaultDb() 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();
+
+    PARTITION = PartitionId.newBuilder().setProjectId(PROJECT_ID).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());
+        });
+  }
+
+  @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();
+
+    PARTITION = PartitionId.newBuilder().setProjectId(PROJECT_ID).setDatabaseId("test-db").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());
+        });
+  }
+}