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/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"

README.md

@@ -49,20 +49,20 @@ If you are using Maven without BOM, add this to your dependencies:
 If you are using Gradle 5.x or later, add this to your dependencies:
 
 ```Groovy
-implementation platform('com.google.cloud:libraries-bom:26.1.4')
+implementation platform('com.google.cloud:libraries-bom:26.2.0')
 
 implementation 'com.google.cloud:google-cloud-datastore'
 ```
 If you are using Gradle without BOM, add this to your dependencies:
 
 ```Groovy
-implementation 'com.google.cloud:google-cloud-datastore:2.12.5'
+implementation 'com.google.cloud:google-cloud-datastore:2.13.0'
 ```
 
 If you are using SBT, add this to your dependencies:
 
 ```Scala
-libraryDependencies += "com.google.cloud" % "google-cloud-datastore" % "2.12.5"
+libraryDependencies += "com.google.cloud" % "google-cloud-datastore" % "2.13.0"
 ```
 
 ## Authentication

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/main/resources/META-INF/native-image/reflect-config.json

@@ -29,5 +29,55 @@
  "allPublicMethods":true,
  "allDeclaredConstructors" : true,
  "allPublicConstructors" : true
- }
+ },
+ [
+ {
+ "name":"com.google.api.client.auth.oauth2.TokenRequest",
+ "allDeclaredFields":true,
+ "queryAllDeclaredMethods":true
+ },
+ {
+ "name":"com.google.api.client.auth.oauth2.TokenResponse",
+ "allDeclaredFields":true,
+ "queryAllDeclaredMethods":true,
+ "methods":[
+ {"name":"<init>","parameterTypes":[] },
+ {"name":"setAccessToken","parameterTypes":["java.lang.String"] },
+ {"name":"setExpiresInSeconds","parameterTypes":["java.lang.Long"] },
+ {"name":"setTokenType","parameterTypes":["java.lang.String"] }
+ ]
+ },
+ {
+ "name":"com.google.api.client.http.GenericUrl",
+ "allDeclaredFields":true
+ },
+ {
+ "name":"com.google.api.client.json.GenericJson",
+ "allDeclaredFields":true,
+ "methods":[{"name":"<init>","parameterTypes":[] }]
+ },
+ {
+ "name":"com.google.api.client.json.webtoken.JsonWebSignature$Header",
+ "allDeclaredFields":true,
+ "queryAllDeclaredMethods":true
+ },
+ {
+ "name":"com.google.api.client.json.webtoken.JsonWebToken$Header",
+ "allDeclaredFields":true,
+ "queryAllDeclaredMethods":true
+ },
+ {
+ "name":"com.google.api.client.json.webtoken.JsonWebToken$Payload",
+ "allDeclaredFields":true,
+ "queryAllDeclaredMethods":true
+ },
+ {
+ "name":"com.google.api.client.util.GenericData",
+ "allDeclaredFields":true
+ },
+ {
+ "name":"com.google.protobuf.ExtensionRegistry",
+ "methods":[{"name":"getEmptyRegistry","parameterTypes":[] }]
+ }
+ ]
 ]

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());
+ });
+ }
+}