Repository: googleapis/java-datastore
Pull request #928 (2022-11-30 14:50:29)

feat: multiple dbs support


Patch:
diff --git a/.kokoro/nightly/integration.cfg b/.kokoro/nightly/integration.cfg
index a2907a25..41898301 100644
--- a/.kokoro/nightly/integration.cfg
+++ b/.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"
diff --git a/.kokoro/nightly/java11-integration.cfg b/.kokoro/nightly/java11-integration.cfg
index 58049cc3..5929bb20 100644
--- a/.kokoro/nightly/java11-integration.cfg
+++ b/.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"
diff --git a/.kokoro/presubmit/graalvm-native-17.cfg b/.kokoro/presubmit/graalvm-native-17.cfg
index a3f7fb9d..29e8b1c6 100644
--- a/.kokoro/presubmit/graalvm-native-17.cfg
+++ b/.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"
 }
\ No newline at end of file
diff --git a/.kokoro/presubmit/graalvm-native.cfg b/.kokoro/presubmit/graalvm-native.cfg
index 4c7225ec..10b0b97c 100644
--- a/.kokoro/presubmit/graalvm-native.cfg
+++ b/.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"
+}
\ No newline at end of file
diff --git a/.kokoro/presubmit/integration.cfg b/.kokoro/presubmit/integration.cfg
index dded67a9..fd5cd668 100644
--- a/.kokoro/presubmit/integration.cfg
+++ b/.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"
diff --git a/README.md b/README.md
index d6e6c482..c1419ff3 100644
--- a/README.md
+++ b/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.5')
+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
diff --git a/datastore-v1-proto-client/src/main/java/com/google/datastore/v1/client/DatastoreOptions.java b/datastore-v1-proto-client/src/main/java/com/google/datastore/v1/client/DatastoreOptions.java
index 8da86dc4..28b3cf4e 100644
--- a/datastore-v1-proto-client/src/main/java/com/google/datastore/v1/client/DatastoreOptions.java
+++ b/datastore-v1-proto-client/src/main/java/com/google/datastore/v1/client/DatastoreOptions.java
@@ -40,6 +40,7 @@ import java.util.List;
  */
 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 class DatastoreOptions {
         "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 class DatastoreOptions {
 
     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 class DatastoreOptions {
       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 class DatastoreOptions {
     return projectId;
   }
 
+  public String getDatabaseId() {
+    return databaseId;
+  }
+
   public String getProjectEndpoint() {
     return projectEndpoint;
   }
diff --git a/datastore-v1-proto-client/src/main/java/com/google/datastore/v1/client/QuerySplitterImpl.java b/datastore-v1-proto-client/src/main/java/com/google/datastore/v1/client/QuerySplitterImpl.java
index 6143bdd5..8952111e 100644
--- a/datastore-v1-proto-client/src/main/java/com/google/datastore/v1/client/QuerySplitterImpl.java
+++ b/datastore-v1-proto-client/src/main/java/com/google/datastore/v1/client/QuerySplitterImpl.java
@@ -221,6 +221,8 @@ final class QuerySplitterImpl implements QuerySplitter {
     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());
       }
diff --git a/datastore-v1-proto-client/src/main/resources/META-INF/native-image/reflect-config.json b/datastore-v1-proto-client/src/main/resources/META-INF/native-image/reflect-config.json
index 17876ff4..07c6eaa3 100644
--- a/datastore-v1-proto-client/src/main/resources/META-INF/native-image/reflect-config.json
+++ b/datastore-v1-proto-client/src/main/resources/META-INF/native-image/reflect-config.json
@@ -29,5 +29,53 @@
     "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":[] }]
   }
 ]
\ No newline at end of file
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/BaseKey.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/BaseKey.java
index 553a12bb..5547979d 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/BaseKey.java
+++ b/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 class BaseKey implements Serializable {
 
     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 class BaseKey implements Serializable {
     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 abstract class BaseKey implements Serializable {
 
     /** 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 abstract class BaseKey implements Serializable {
       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 abstract class BaseKey implements Serializable {
     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 abstract class BaseKey implements Serializable {
     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 abstract class BaseKey implements Serializable {
     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 @@ public abstract class BaseKey implements Serializable {
     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) {
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/BatchImpl.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/BatchImpl.java
index 5fc0bcd8..f1647cac 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/BatchImpl.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/BatchImpl.java
@@ -59,6 +59,8 @@ class BatchImpl extends BaseDatastoreBatchWriter implements Batch {
         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());
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreHelper.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreHelper.java
index 985b76a7..33e90a66 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreHelper.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreHelper.java
@@ -54,7 +54,7 @@ class DatastoreHelper {
   }
 
   static KeyFactory newKeyFactory(DatastoreOptions options) {
-    return new KeyFactory(options.getProjectId(), options.getNamespace());
+    return new KeyFactory(options.getProjectId(), options.getNamespace(), options.getDatabaseId());
   }
 
   /**
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java
index 4f6533ec..5a3c6687 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java
@@ -214,12 +214,7 @@ final class DatastoreImpl extends BaseService<DatastoreOptions> implements Datas
     Span span = traceUtil.startSpan(TraceUtil.SPAN_NAME_RUNQUERY);
     try (Scope scope = traceUtil.getTracer().withSpan(span)) {
       return RetryHelper.runWithRetries(
-          new Callable<com.google.datastore.v1.RunQueryResponse>() {
-            @Override
-            public com.google.datastore.v1.RunQueryResponse call() throws DatastoreException {
-              return datastoreRpc.runQuery(requestPb);
-            }
-          },
+          () -> datastoreRpc.runQuery(requestPb),
           retrySettings,
           requestPb.getReadOptions().getTransaction().isEmpty()
               ? EXCEPTION_HANDLER
@@ -259,6 +254,8 @@ final class DatastoreImpl extends BaseService<DatastoreOptions> implements Datas
     for (IncompleteKey key : keys) {
       requestPb.addKeys(trimNameOrId(key).toPb());
     }
+    requestPb.setProjectId(getOptions().getProjectId());
+    requestPb.setDatabaseId(getOptions().getDatabaseId());
     com.google.datastore.v1.AllocateIdsResponse responsePb = allocateIds(requestPb.build());
     ImmutableList.Builder<Key> keyList = ImmutableList.builder();
     for (com.google.datastore.v1.Key keyPb : responsePb.getKeysList()) {
@@ -389,6 +386,8 @@ final class DatastoreImpl extends BaseService<DatastoreOptions> implements Datas
     for (Key k : Sets.newLinkedHashSet(Arrays.asList(keys))) {
       requestPb.addKeys(k.toPb());
     }
+    requestPb.setProjectId(getOptions().getProjectId());
+    requestPb.setDatabaseId(getOptions().getDatabaseId());
     return new ResultsIterator(requestPb);
   }
 
@@ -454,6 +453,8 @@ final class DatastoreImpl extends BaseService<DatastoreOptions> implements Datas
     for (Key key : keys) {
       requestPb.addKeys(key.toPb());
     }
+    requestPb.setProjectId(getOptions().getProjectId());
+    requestPb.setDatabaseId(getOptions().getDatabaseId());
     com.google.datastore.v1.ReserveIdsResponse responsePb = reserveIds(requestPb.build());
     ImmutableList.Builder<Key> keyList = ImmutableList.builder();
     if (responsePb.isInitialized()) {
@@ -568,6 +569,8 @@ final class DatastoreImpl extends BaseService<DatastoreOptions> implements Datas
     com.google.datastore.v1.CommitRequest.Builder requestPb =
         com.google.datastore.v1.CommitRequest.newBuilder();
     requestPb.setMode(com.google.datastore.v1.CommitRequest.Mode.NON_TRANSACTIONAL);
+    requestPb.setProjectId(getOptions().getProjectId());
+    requestPb.setDatabaseId(getOptions().getDatabaseId());
     requestPb.addAllMutations(mutationsPb);
     return commit(requestPb.build());
   }
@@ -577,12 +580,7 @@ final class DatastoreImpl extends BaseService<DatastoreOptions> implements Datas
     Span span = traceUtil.startSpan(TraceUtil.SPAN_NAME_COMMIT);
     try (Scope scope = traceUtil.getTracer().withSpan(span)) {
       return RetryHelper.runWithRetries(
-          new Callable<com.google.datastore.v1.CommitResponse>() {
-            @Override
-            public com.google.datastore.v1.CommitResponse call() throws DatastoreException {
-              return datastoreRpc.commit(requestPb);
-            }
-          },
+          () -> datastoreRpc.commit(requestPb),
           retrySettings,
           requestPb.getTransaction().isEmpty()
               ? EXCEPTION_HANDLER
@@ -628,6 +626,8 @@ final class DatastoreImpl extends BaseService<DatastoreOptions> implements Datas
     com.google.datastore.v1.RollbackRequest.Builder requestPb =
         com.google.datastore.v1.RollbackRequest.newBuilder();
     requestPb.setTransaction(transaction);
+    requestPb.setProjectId(getOptions().getProjectId());
+    requestPb.setDatabaseId(getOptions().getDatabaseId());
     rollback(requestPb.build());
   }
 
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreOptions.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreOptions.java
index f754866a..6d8fdd91 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreOptions.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreOptions.java
@@ -38,8 +38,10 @@ public class DatastoreOptions extends ServiceOptions<Datastore, DatastoreOptions
   private static final String API_SHORT_NAME = "Datastore";
   private static final String DATASTORE_SCOPE = "https://www.googleapis.com/auth/datastore";
   private static final Set<String> SCOPES = ImmutableSet.of(DATASTORE_SCOPE);
+  private static final String DEFAULT_DATABASE_ID = "";
 
   private final String namespace;
+  private final String databaseId;
 
   public static class DefaultDatastoreFactory implements DatastoreFactory {
 
@@ -64,12 +66,14 @@ public class DatastoreOptions extends ServiceOptions<Datastore, DatastoreOptions
   public static class Builder extends ServiceOptions.Builder<Datastore, DatastoreOptions, Builder> {
 
     private String namespace;
+    private String databaseId;
 
     private Builder() {}
 
     private Builder(DatastoreOptions options) {
       super(options);
       namespace = options.namespace;
+      databaseId = options.databaseId;
     }
 
     @Override
@@ -91,11 +95,17 @@ public class DatastoreOptions extends ServiceOptions<Datastore, DatastoreOptions
       this.namespace = validateNamespace(namespace);
       return this;
     }
+
+    public Builder setDatabaseId(String databaseId) {
+      this.databaseId = databaseId;
+      return this;
+    }
   }
 
   private DatastoreOptions(Builder builder) {
     super(DatastoreFactory.class, DatastoreRpcFactory.class, builder, new DatastoreDefaults());
-    namespace = builder.namespace != null ? builder.namespace : defaultNamespace();
+    namespace = MoreObjects.firstNonNull(builder.namespace, defaultNamespace());
+    databaseId = MoreObjects.firstNonNull(builder.databaseId, DEFAULT_DATABASE_ID);
   }
 
   @Override
@@ -143,6 +153,10 @@ public class DatastoreOptions extends ServiceOptions<Datastore, DatastoreOptions
     return namespace;
   }
 
+  public String getDatabaseId() {
+    return this.databaseId;
+  }
+
   /** Returns a default {@code DatastoreOptions} instance. */
   public static DatastoreOptions getDefaultInstance() {
     return newBuilder().build();
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/IncompleteKey.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/IncompleteKey.java
index afabadba..db9973cb 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/IncompleteKey.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/IncompleteKey.java
@@ -31,6 +31,11 @@ public class IncompleteKey extends BaseKey {
       super(projectId, kind);
     }
 
+    private Builder(String projectId, String kind, String databaseId) {
+      super(projectId, kind);
+      this.databaseId = databaseId;
+    }
+
     private Builder(IncompleteKey copyFrom) {
       super(copyFrom);
     }
@@ -39,7 +44,7 @@ public class IncompleteKey extends BaseKey {
     public IncompleteKey build() {
       ImmutableList<PathElement> path =
           ImmutableList.<PathElement>builder().addAll(ancestors).add(PathElement.of(kind)).build();
-      return new IncompleteKey(projectId, namespace, path);
+      return new IncompleteKey(projectId, namespace, databaseId, path);
     }
   }
 
@@ -47,13 +52,20 @@ public class IncompleteKey extends BaseKey {
     super(projectId, namespace, path);
   }
 
+  IncompleteKey(
+      String projectId, String namespace, String databaseId, ImmutableList<PathElement> path) {
+    super(projectId, namespace, databaseId, path);
+  }
+
   static IncompleteKey fromPb(com.google.datastore.v1.Key keyPb) {
     String projectId = "";
     String namespace = "";
+    String databaseId = "";
     if (keyPb.hasPartitionId()) {
       com.google.datastore.v1.PartitionId partitionIdPb = keyPb.getPartitionId();
       projectId = partitionIdPb.getProjectId();
       namespace = partitionIdPb.getNamespaceId();
+      databaseId = partitionIdPb.getDatabaseId();
     }
     List<com.google.datastore.v1.Key.PathElement> pathElementsPb = keyPb.getPathList();
     Preconditions.checkArgument(!pathElementsPb.isEmpty(), "Path must not be empty");
@@ -64,9 +76,9 @@ public class IncompleteKey extends BaseKey {
     ImmutableList<PathElement> path = pathBuilder.build();
     PathElement leaf = path.get(path.size() - 1);
     if (leaf.getNameOrId() != null) {
-      return new Key(projectId, namespace, path);
+      return new Key(projectId, namespace, databaseId, path);
     }
-    return new IncompleteKey(projectId, namespace, path);
+    return new IncompleteKey(projectId, namespace, databaseId, path);
   }
 
   /** Returns the key's parent. */
@@ -94,12 +106,16 @@ public class IncompleteKey extends BaseKey {
     return new Builder(projectId, kind);
   }
 
+  public static Builder newBuilderWithDatabaseId(String projectId, String kind, String databaseId) {
+    return new Builder(projectId, kind, databaseId);
+  }
+
   public static Builder newBuilder(IncompleteKey copyFrom) {
     return new Builder(copyFrom);
   }
 
   public static Builder newBuilder(Key parent, String kind) {
-    return newBuilder(parent.getProjectId(), kind)
+    return newBuilderWithDatabaseId(parent.getProjectId(), kind, parent.getDatabaseId())
         .setNamespace(parent.getNamespace())
         .addAncestors(parent.getPath());
   }
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Key.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Key.java
index 785d3592..9e851d0c 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Key.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Key.java
@@ -51,6 +51,18 @@ public final class Key extends IncompleteKey {
       this.id = id;
     }
 
+    private Builder(String projectId, String kind, long id, String databaseId) {
+      super(projectId, kind);
+      this.id = id;
+      this.databaseId = databaseId;
+    }
+
+    private Builder(String projectId, String kind, String name, String databaseId) {
+      super(projectId, kind);
+      this.name = name;
+      this.databaseId = databaseId;
+    }
+
     private Builder(IncompleteKey copyFrom, String name) {
       super(copyFrom);
       this.name = name;
@@ -93,7 +105,7 @@ public final class Key extends IncompleteKey {
       } else {
         pathBuilder.add(PathElement.of(kind, id));
       }
-      return new Key(projectId, namespace, pathBuilder.build());
+      return new Key(projectId, namespace, databaseId, pathBuilder.build());
     }
   }
 
@@ -102,6 +114,11 @@ public final class Key extends IncompleteKey {
     Preconditions.checkArgument(getNameOrId() != null);
   }
 
+  Key(String projectId, String namespace, String databaseId, ImmutableList<PathElement> path) {
+    super(projectId, namespace, databaseId, path);
+    Preconditions.checkArgument(getNameOrId() != null);
+  }
+
   public boolean hasId() {
     return getLeaf().hasId();
   }
@@ -162,10 +179,18 @@ public final class Key extends IncompleteKey {
     return new Builder(projectId, kind, name);
   }
 
+  public static Builder newBuilder(String projectId, String kind, String name, String databaseId) {
+    return new Builder(projectId, kind, name, databaseId);
+  }
+
   public static Builder newBuilder(String projectId, String kind, long id) {
     return new Builder(projectId, kind, id);
   }
 
+  public static Builder newBuilder(String projectId, String kind, long id, String databaseId) {
+    return new Builder(projectId, kind, id, databaseId);
+  }
+
   public static Builder newBuilder(Key copyFrom) {
     return new Builder(copyFrom);
   }
@@ -179,13 +204,13 @@ public final class Key extends IncompleteKey {
   }
 
   public static Builder newBuilder(Key parent, String kind, String name) {
-    Builder builder = newBuilder(parent.getProjectId(), kind, name);
+    Builder builder = newBuilder(parent.getProjectId(), kind, name, parent.getDatabaseId());
     addParentToBuilder(parent, builder);
     return builder;
   }
 
   public static Builder newBuilder(Key parent, String kind, long id) {
-    Builder builder = newBuilder(parent.getProjectId(), kind, id);
+    Builder builder = newBuilder(parent.getProjectId(), kind, id, parent.getDatabaseId());
     addParentToBuilder(parent, builder);
     return builder;
   }
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/KeyFactory.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/KeyFactory.java
index 7c6b0734..72b21442 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/KeyFactory.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/KeyFactory.java
@@ -24,8 +24,9 @@ import com.google.common.collect.ImmutableList;
  */
 public final class KeyFactory extends BaseKey.Builder<KeyFactory> {
 
-  private final String pi;
-  private final String ns;
+  private final String initialProjectId;
+  private final String initialNamespace;
+  private final String initialDatabaseId;
 
   public KeyFactory(String projectId) {
     this(projectId, "");
@@ -34,14 +35,24 @@ public final class KeyFactory extends BaseKey.Builder<KeyFactory> {
   public KeyFactory(String projectId, String namespace) {
     super(projectId);
     setNamespace(namespace);
-    this.pi = projectId;
-    this.ns = namespace;
+    this.initialProjectId = projectId;
+    this.initialNamespace = namespace;
+    this.initialDatabaseId = "";
+  }
+
+  public KeyFactory(String projectId, String namespace, String databaseId) {
+    super(projectId);
+    setNamespace(namespace);
+    setDatabaseId(databaseId);
+    this.initialProjectId = projectId;
+    this.initialNamespace = namespace;
+    this.initialDatabaseId = databaseId;
   }
 
   public IncompleteKey newKey() {
     ImmutableList<PathElement> path =
         ImmutableList.<PathElement>builder().addAll(ancestors).add(PathElement.of(kind)).build();
-    return new IncompleteKey(projectId, namespace, path);
+    return new IncompleteKey(projectId, namespace, databaseId, path);
   }
 
   public Key newKey(String name) {
@@ -50,7 +61,7 @@ public final class KeyFactory extends BaseKey.Builder<KeyFactory> {
             .addAll(ancestors)
             .add(PathElement.of(kind, name))
             .build();
-    return new Key(projectId, namespace, path);
+    return new Key(projectId, namespace, databaseId, path);
   }
 
   public Key newKey(long id) {
@@ -59,7 +70,7 @@ public final class KeyFactory extends BaseKey.Builder<KeyFactory> {
             .addAll(ancestors)
             .add(PathElement.of(kind, id))
             .build();
-    return new Key(projectId, namespace, path);
+    return new Key(projectId, namespace, databaseId, path);
   }
 
   /**
@@ -68,8 +79,9 @@ public final class KeyFactory extends BaseKey.Builder<KeyFactory> {
    * @return {@code this} for chaining
    */
   public KeyFactory reset() {
-    setProjectId(pi);
-    setNamespace(ns);
+    setProjectId(initialProjectId);
+    setNamespace(initialNamespace);
+    setDatabaseId(initialDatabaseId);
     kind = null;
     ancestors.clear();
     return this;
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/QueryResultsImpl.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/QueryResultsImpl.java
index 6170c0b8..8fc731ac 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/QueryResultsImpl.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/QueryResultsImpl.java
@@ -53,6 +53,7 @@ class QueryResultsImpl<T> extends AbstractIterator<T> implements QueryResults<T>
     com.google.datastore.v1.PartitionId.Builder pbBuilder =
         com.google.datastore.v1.PartitionId.newBuilder();
     pbBuilder.setProjectId(datastore.getOptions().getProjectId());
+    pbBuilder.setDatabaseId(datastore.getOptions().getDatabaseId());
     if (namespace != null) {
       pbBuilder.setNamespaceId(namespace);
     } else if (datastore.getOptions().getNamespace() != null) {
@@ -72,6 +73,8 @@ class QueryResultsImpl<T> extends AbstractIterator<T> implements QueryResults<T>
         com.google.datastore.v1.RunQueryRequest.newBuilder();
     readOptionsPb.ifPresent(requestPb::setReadOptions);
     requestPb.setPartitionId(partitionIdPb);
+    requestPb.setProjectId(datastore.getOptions().getProjectId());
+    requestPb.setDatabaseId(datastore.getOptions().getDatabaseId());
     query.populatePb(requestPb);
     runQueryResponsePb = datastore.runQuery(requestPb.build());
     mostRecentQueryPb = requestPb.getQuery();
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/TransactionImpl.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/TransactionImpl.java
index fc6c5e94..3b5e5e4e 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/TransactionImpl.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/TransactionImpl.java
@@ -67,6 +67,9 @@ final class TransactionImpl extends BaseDatastoreBatchWriter implements Transact
     com.google.datastore.v1.BeginTransactionRequest.Builder requestPb =
         com.google.datastore.v1.BeginTransactionRequest.newBuilder();
 
+    requestPb.setProjectId(this.datastore.getOptions().getProjectId());
+    requestPb.setDatabaseId(this.datastore.getOptions().getDatabaseId());
+
     if (options != null) {
       requestPb.setTransactionOptions(options);
     }
@@ -116,6 +119,8 @@ final class TransactionImpl extends BaseDatastoreBatchWriter implements Transact
     requestPb.setMode(com.google.datastore.v1.CommitRequest.Mode.TRANSACTIONAL);
     requestPb.setTransaction(transactionId);
     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());
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Validator.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Validator.java
index 4a6d99d5..f1c19d15 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Validator.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Validator.java
@@ -35,7 +35,7 @@ final class Validator {
     // utility class
   }
 
-  static String validateDatabase(String projectId) {
+  static String validateProjectId(String projectId) {
     checkArgument(!Strings.isNullOrEmpty(projectId), "projectId can't be empty or null");
     checkArgument(
         PROJECT_ID_PATTERN.matcher(projectId).matches(),
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/request/AggregationQueryRequestProtoPreparer.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/request/AggregationQueryRequestProtoPreparer.java
index b5da8d9f..36e7de77 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/request/AggregationQueryRequestProtoPreparer.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/request/AggregationQueryRequestProtoPreparer.java
@@ -59,7 +59,8 @@ public class AggregationQueryRequestProtoPreparer
     RunAggregationQueryRequest.Builder aggregationQueryRequestBuilder =
         RunAggregationQueryRequest.newBuilder()
             .setPartitionId(partitionId)
-            .setProjectId(datastoreOptions.getProjectId());
+            .setProjectId(datastoreOptions.getProjectId())
+            .setDatabaseId(datastoreOptions.getDatabaseId());
 
     if (aggregationQuery.getMode() == GQL) {
       aggregationQueryRequestBuilder.setGqlQuery(buildGqlQuery(aggregationQuery));
@@ -91,7 +92,9 @@ public class AggregationQueryRequestProtoPreparer
 
   private PartitionId getPartitionId(AggregationQuery aggregationQuery) {
     PartitionId.Builder builder =
-        PartitionId.newBuilder().setProjectId(datastoreOptions.getProjectId());
+        PartitionId.newBuilder()
+            .setProjectId(datastoreOptions.getProjectId())
+            .setDatabaseId(datastoreOptions.getDatabaseId());
     if (aggregationQuery.getNamespace() != null) {
       builder.setNamespaceId(aggregationQuery.getNamespace());
     }
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/spi/v1/HttpDatastoreRpc.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/spi/v1/HttpDatastoreRpc.java
index fd3cdc65..cfbbaa7d 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/spi/v1/HttpDatastoreRpc.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/spi/v1/HttpDatastoreRpc.java
@@ -55,6 +55,7 @@ public class HttpDatastoreRpc implements DatastoreRpc {
     com.google.datastore.v1.client.DatastoreOptions.Builder clientBuilder =
         new com.google.datastore.v1.client.DatastoreOptions.Builder()
             .projectId(options.getProjectId())
+            .databaseId(options.getDatabaseId())
             .initializer(getHttpRequestInitializer(options, httpTransportOptions))
             .transport(transport);
     String normalizedHost = options.getHost() != null ? options.getHost().toLowerCase() : "";
diff --git a/owlbot.py b/owlbot.py
index 784f35a4..f96943ca 100644
--- a/owlbot.py
+++ b/owlbot.py
@@ -45,4 +45,10 @@ for library in s.get_staging_dirs():
     s.move(library)
 
 s.remove_staging_dirs()
-java.common_templates()
+java.common_templates(excludes=[
+    '.kokoro/presubmit/integration.cfg',
+    '.kokoro/presubmit/graalvm-native.cfg',
+    '.kokoro/presubmit/graalvm-native-17.cfg',
+    '.kokoro/nightly/integration.cfg',
+    '.kokoro/nightly/java11-integration.cfg',
+])


Test patch:
diff --git a/datastore-v1-proto-client/src/test/java/com/google/datastore/v1/client/DatastoreClientTest.java b/datastore-v1-proto-client/src/test/java/com/google/datastore/v1/client/DatastoreClientTest.java
index 16a6303b..3eb7a66d 100644
--- a/datastore-v1-proto-client/src/test/java/com/google/datastore/v1/client/DatastoreClientTest.java
+++ b/datastore-v1-proto-client/src/test/java/com/google/datastore/v1/client/DatastoreClientTest.java
@@ -215,6 +215,18 @@ public class DatastoreClientTest {
         .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 =
diff --git a/datastore-v1-proto-client/src/test/java/com/google/datastore/v1/client/QuerySplitterTest.java b/datastore-v1-proto-client/src/test/java/com/google/datastore/v1/client/QuerySplitterTest.java
index e8694372..b064e137 100644
--- a/datastore-v1-proto-client/src/test/java/com/google/datastore/v1/client/QuerySplitterTest.java
+++ b/datastore-v1-proto-client/src/test/java/com/google/datastore/v1/client/QuerySplitterTest.java
@@ -193,6 +193,60 @@ public class QuerySplitterTest {
     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 class QuerySplitterTest {
     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 class QuerySplitterTest {
     RunQueryRequest expectedSplitQueryRequest =
         RunQueryRequest.newBuilder()
             .setPartitionId(PARTITION)
+            .setProjectId(PROJECT_ID)
             .setQuery(
                 splitQuery.toBuilder().setLimit(Int32Value.newBuilder().setValue(2 * 32).build()))
             .setReadOptions(ReadOptions.newBuilder().setReadTime(readTime))
diff --git a/datastore-v1-proto-client/src/test/java/com/google/datastore/v1/client/it/ITDatastoreProtoClientTest.java b/datastore-v1-proto-client/src/test/java/com/google/datastore/v1/client/it/ITDatastoreProtoClientTest.java
new file mode 100644
index 00000000..f5cf1ea9
--- /dev/null
+++ b/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());
+        });
+  }
+}
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/testing/LocalDatastoreHelper.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/testing/LocalDatastoreHelper.java
index e586f7e5..1ba02d52 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/testing/LocalDatastoreHelper.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/testing/LocalDatastoreHelper.java
@@ -50,6 +50,7 @@ public class LocalDatastoreHelper extends BaseEmulatorHelper<DatastoreOptions> {
   private final double consistency;
   private final Path gcdPath;
   private boolean storeOnDisk;
+  private String databaseId;
 
   // Gcloud emulator settings
   private static final String GCLOUD_CMD_TEXT = "gcloud beta emulators datastore start";
@@ -84,6 +85,13 @@ public class LocalDatastoreHelper extends BaseEmulatorHelper<DatastoreOptions> {
     }
   }
 
+  private final DatastoreOptions.Builder optionsBuilder =
+      DatastoreOptions.newBuilder()
+          .setProjectId(getProjectId())
+          .setHost(DEFAULT_HOST + ":" + getPort())
+          .setCredentials(NoCredentials.getInstance())
+          .setRetrySettings(ServiceOptions.getNoRetrySettings());
+
   /** A builder for {@code LocalDatastoreHelper} objects. */
   public static class Builder {
     private double consistency;
@@ -177,29 +185,32 @@ public class LocalDatastoreHelper extends BaseEmulatorHelper<DatastoreOptions> {
     return LOGGER;
   }
 
-  private DatastoreOptions.Builder optionsBuilder() {
-    return DatastoreOptions.newBuilder()
-        .setProjectId(getProjectId())
-        .setHost(DEFAULT_HOST + ":" + Integer.toString(getPort()))
-        .setCredentials(NoCredentials.getInstance())
-        .setRetrySettings(ServiceOptions.getNoRetrySettings());
-  }
-
   /**
    * Returns a {@link DatastoreOptions} instance that sets the host to use the Datastore emulator on
    * localhost.
    */
   @Override
   public DatastoreOptions getOptions() {
-    return optionsBuilder().build();
+    return optionsBuilder.build();
   }
 
   /**
    * Returns a {@link DatastoreOptions} instance that sets the host to use the Datastore emulator on
    * localhost. The default namespace is set to {@code namespace}.
+   *
+   * @deprecated use setNamespace and then build() instead
    */
+  @Deprecated
   public DatastoreOptions getOptions(String namespace) {
-    return optionsBuilder().setNamespace(namespace).build();
+    return optionsBuilder.setNamespace(namespace).build();
+  }
+
+  public DatastoreOptions.Builder setNamespace(String namespace) {
+    return optionsBuilder.setNamespace(namespace);
+  }
+
+  public DatastoreOptions.Builder setDatabaseId(String databaseId) {
+    return optionsBuilder.setDatabaseId(databaseId);
   }
 
   /** Returns the consistency setting for the local Datastore emulator. */
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/testing/RemoteDatastoreHelper.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/testing/RemoteDatastoreHelper.java
index 803400ea..596ce96d 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/testing/RemoteDatastoreHelper.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/testing/RemoteDatastoreHelper.java
@@ -74,11 +74,17 @@ public class RemoteDatastoreHelper {
 
   /** Creates a {@code RemoteStorageHelper} object. */
   public static RemoteDatastoreHelper create() {
+    return create("");
+  }
+
+  /** Creates a {@code RemoteStorageHelper} object. */
+  public static RemoteDatastoreHelper create(String databaseId) {
     HttpTransportOptions transportOptions = DatastoreOptions.getDefaultHttpTransportOptions();
     transportOptions =
         transportOptions.toBuilder().setConnectTimeout(60000).setReadTimeout(60000).build();
     DatastoreOptions datastoreOption =
         DatastoreOptions.newBuilder()
+            .setDatabaseId(databaseId)
             .setNamespace(UUID.randomUUID().toString())
             .setRetrySettings(retrySettings())
             .setTransportOptions(transportOptions)
diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/BaseKeyTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/BaseKeyTest.java
index ffa54970..08b65bdc 100644
--- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/BaseKeyTest.java
+++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/BaseKeyTest.java
@@ -36,12 +36,16 @@ public class BaseKeyTest {
       super(projectId, kind);
     }
 
+    Builder(BaseKey copyFrom) {
+      super(copyFrom);
+    }
+
     @Override
     protected BaseKey build() {
       ImmutableList.Builder<PathElement> path = ImmutableList.builder();
       path.addAll(ancestors);
       path.add(PathElement.of(kind));
-      return new BaseKey(projectId, namespace, path.build()) {
+      return new BaseKey(projectId, namespace, databaseId, path.build()) {
 
         @Override
         protected BaseKey getParent() {
@@ -58,6 +62,17 @@ public class BaseKeyTest {
     assertEquals("ds1", key.getProjectId());
     key = builder.setProjectId("ds2").build();
     assertEquals("ds2", key.getProjectId());
+    assertEquals("", key.getDatabaseId());
+  }
+
+  @Test
+  public void testDatabaseId() {
+    Builder builder = new Builder("ds1", "k").setDatabaseId("test-db");
+    BaseKey key = builder.build();
+    assertEquals("ds1", key.getProjectId());
+    key = builder.setProjectId("ds2").build();
+    assertEquals("ds2", key.getProjectId());
+    assertEquals("test-db", key.getDatabaseId());
   }
 
   @Test(expected = IllegalArgumentException.class)
@@ -120,4 +135,17 @@ public class BaseKeyTest {
     key = builder.addAncestor(path.get(1)).build();
     assertEquals(path, key.getAncestors());
   }
+
+  @Test
+  public void testCopyFrom() {
+    Builder copyFrom = new Builder("test-project", "kind").setDatabaseId("test-db");
+    Builder builder = new Builder(copyFrom.build());
+    BaseKey baseKey = builder.build();
+    assertEquals("test-project", baseKey.getProjectId());
+    assertEquals("test-db", baseKey.getDatabaseId());
+    assertEquals("kind", baseKey.getKind());
+    assertEquals("", baseKey.getNamespace());
+    assertEquals(new ArrayList<>(), baseKey.getAncestors());
+    assertEquals(PathElement.of("kind"), baseKey.getLeaf());
+  }
 }
diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/DatastoreHelperTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/DatastoreHelperTest.java
index 576ad74c..010a892a 100644
--- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/DatastoreHelperTest.java
+++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/DatastoreHelperTest.java
@@ -37,6 +37,7 @@ public class DatastoreHelperTest {
     DatastoreOptions options = createMock(DatastoreOptions.class);
     expect(options.getProjectId()).andReturn("ds1").once();
     expect(options.getNamespace()).andReturn("ns1").once();
+    expect(options.getDatabaseId()).andReturn("test-db").once();
     replay(options);
     KeyFactory keyFactory = DatastoreHelper.newKeyFactory(options);
     Key key = keyFactory.setKind("k").newKey("bla");
@@ -44,6 +45,7 @@ public class DatastoreHelperTest {
     assertEquals("ns1", key.getNamespace());
     assertEquals("k", key.getKind());
     assertEquals("bla", key.getName());
+    assertEquals("test-db", key.getDatabaseId());
     verify(options);
   }
 
diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/DatastoreOptionsTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/DatastoreOptionsTest.java
index aabdc7af..19dd3f0c 100644
--- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/DatastoreOptionsTest.java
+++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/DatastoreOptionsTest.java
@@ -32,11 +32,13 @@ import org.junit.Test;
 public class DatastoreOptionsTest {
 
   private static final String PROJECT_ID = "project-id";
+  private static final String DATABASE_ID = "database-id";
   private static final int PORT = 8080;
   private DatastoreRpcFactory datastoreRpcFactory;
   private DatastoreRpc datastoreRpc;
   private DatastoreOptions.Builder options;
 
+  // todo parameterize
   @Before
   public void setUp() {
     datastoreRpcFactory = EasyMock.createMock(DatastoreRpcFactory.class);
@@ -45,6 +47,7 @@ public class DatastoreOptionsTest {
         DatastoreOptions.newBuilder()
             .setServiceRpcFactory(datastoreRpcFactory)
             .setProjectId(PROJECT_ID)
+            .setDatabaseId(DATABASE_ID)
             .setHost("http://localhost:" + PORT);
     EasyMock.expect(datastoreRpcFactory.create(EasyMock.anyObject(DatastoreOptions.class)))
         .andReturn(datastoreRpc)
@@ -57,6 +60,11 @@ public class DatastoreOptionsTest {
     assertEquals(PROJECT_ID, options.build().getProjectId());
   }
 
+  @Test
+  public void testDatabaseId() {
+    assertEquals(DATABASE_ID, options.build().getDatabaseId());
+  }
+
   @Test
   public void testHost() {
     assertEquals("http://localhost:" + PORT, options.build().getHost());
diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/DatastoreTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/DatastoreTest.java
index 7dc625ba..d0f00d79 100644
--- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/DatastoreTest.java
+++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/DatastoreTest.java
@@ -812,6 +812,7 @@ public class DatastoreTest {
         RunQueryRequest.newBuilder()
             .setReadOptions(readOption)
             .setGqlQuery(query)
+            .setProjectId(PROJECT_ID)
             .setPartitionId(PartitionId.newBuilder().setProjectId(PROJECT_ID).build());
     EasyMock.expect(rpcMock.runQuery(expectedRequest.build()))
         .andReturn(RunQueryResponse.newBuilder().build());
@@ -832,6 +833,7 @@ public class DatastoreTest {
         RunQueryRequest.newBuilder()
             .setReadOptions(readOption)
             .setGqlQuery(query)
+            .setProjectId(PROJECT_ID)
             .setPartitionId(PartitionId.newBuilder().setProjectId(PROJECT_ID).build());
     EasyMock.expect(rpcMock.runQuery(expectedRequest.build()))
         .andReturn(RunQueryResponse.newBuilder().build());
@@ -902,7 +904,7 @@ public class DatastoreTest {
   @Test
   public void testReserveIds() {
     ReserveIdsRequest reserveIdsRequest =
-        ReserveIdsRequest.newBuilder().addKeys(KEY1.toPb()).build();
+        ReserveIdsRequest.newBuilder().setProjectId(PROJECT_ID).addKeys(KEY1.toPb()).build();
     EasyMock.expect(rpcMock.reserveIds(reserveIdsRequest))
         .andReturn(ReserveIdsResponse.newBuilder().build())
         .times(1);
@@ -962,7 +964,11 @@ public class DatastoreTest {
                     .build())
             .build();
     LookupRequest lookupRequest =
-        LookupRequest.newBuilder().setReadOptions(readOption).addKeys(key).build();
+        LookupRequest.newBuilder()
+            .setProjectId(PROJECT_ID)
+            .setReadOptions(readOption)
+            .addKeys(key)
+            .build();
     EasyMock.expect(rpcMock.lookup(lookupRequest))
         .andReturn(LookupResponse.newBuilder().build())
         .times(3);
@@ -988,7 +994,11 @@ public class DatastoreTest {
                     .build())
             .build();
     LookupRequest lookupRequest =
-        LookupRequest.newBuilder().setReadOptions(readOption).addKeys(key).build();
+        LookupRequest.newBuilder()
+            .setProjectId(PROJECT_ID)
+            .setReadOptions(readOption)
+            .addKeys(key)
+            .build();
     EasyMock.expect(rpcMock.lookup(lookupRequest))
         .andReturn(LookupResponse.newBuilder().build())
         .times(3);
@@ -1069,14 +1079,17 @@ public class DatastoreTest {
     keysPb.add(KEY4.toPb());
     keysPb.add(KEY5.toPb());
     List<LookupRequest> lookupRequests = new ArrayList<>();
-    lookupRequests.add(LookupRequest.newBuilder().addAllKeys(keysPb).build());
+    lookupRequests.add(
+        LookupRequest.newBuilder().setProjectId(PROJECT_ID).addAllKeys(keysPb).build());
     lookupRequests.add(
         LookupRequest.newBuilder()
+            .setProjectId(PROJECT_ID)
             .addKeys(keysPb.get(1))
             .addKeys(keysPb.get(2))
             .addKeys(keysPb.get(4))
             .build());
-    lookupRequests.add(LookupRequest.newBuilder().addKeys(keysPb.get(4)).build());
+    lookupRequests.add(
+        LookupRequest.newBuilder().setProjectId(PROJECT_ID).addKeys(keysPb.get(4)).build());
     Entity entity4 = Entity.newBuilder(KEY4).set("value", StringValue.of("value")).build();
     Entity entity5 = Entity.newBuilder(KEY5).set("value", "value").build();
     List<LookupResponse> lookupResponses = new ArrayList<>();
@@ -1201,7 +1214,8 @@ public class DatastoreTest {
 
   @Test
   public void testRetryableException() {
-    LookupRequest requestPb = LookupRequest.newBuilder().addKeys(KEY1.toPb()).build();
+    LookupRequest requestPb =
+        LookupRequest.newBuilder().setProjectId(PROJECT_ID).addKeys(KEY1.toPb()).build();
     LookupResponse responsePb =
         LookupResponse.newBuilder()
             .addFound(EntityResult.newBuilder().setEntity(ENTITY1.toPb()))
@@ -1221,6 +1235,7 @@ public class DatastoreTest {
     ByteString txnBytes = ByteString.copyFromUtf8("txn1");
     LookupRequest requestPb =
         LookupRequest.newBuilder()
+            .setProjectId(PROJECT_ID)
             .addKeys(KEY1.toPb())
             .setReadOptions(ReadOptions.newBuilder().setTransaction(txnBytes).build())
             .build();
@@ -1246,6 +1261,7 @@ public class DatastoreTest {
     ByteString txnBytes = ByteString.copyFromUtf8("txn1");
     LookupRequest requestPb =
         LookupRequest.newBuilder()
+            .setProjectId(PROJECT_ID)
             .addKeys(KEY1.toPb())
             .setReadOptions(ReadOptions.newBuilder().setTransaction(txnBytes).build())
             .build();
@@ -1268,7 +1284,8 @@ public class DatastoreTest {
 
   @Test
   public void testNonRetryableException() {
-    LookupRequest requestPb = LookupRequest.newBuilder().addKeys(KEY1.toPb()).build();
+    LookupRequest requestPb =
+        LookupRequest.newBuilder().setProjectId(PROJECT_ID).addKeys(KEY1.toPb()).build();
     EasyMock.expect(rpcMock.lookup(requestPb))
         .andThrow(
             new DatastoreException(DatastoreException.UNKNOWN_CODE, "denied", "PERMISSION_DENIED"))
@@ -1286,7 +1303,8 @@ public class DatastoreTest {
 
   @Test
   public void testRuntimeException() {
-    LookupRequest requestPb = LookupRequest.newBuilder().addKeys(KEY1.toPb()).build();
+    LookupRequest requestPb =
+        LookupRequest.newBuilder().setProjectId(PROJECT_ID).addKeys(KEY1.toPb()).build();
     String exceptionMessage = "Artificial runtime exception";
     EasyMock.expect(rpcMock.lookup(requestPb)).andThrow(new RuntimeException(exceptionMessage));
     EasyMock.replay(rpcFactoryMock, rpcMock);
@@ -1345,6 +1363,26 @@ public class DatastoreTest {
     datastore.delete(entity1.getKey(), entity2.getKey(), entity3.getKey());
   }
 
+  @Test
+  public void testDatabaseIdKeyFactory() {
+    KeyFactory keyFactory = datastore.newKeyFactory().setKind(KIND1);
+
+    Key key1 = keyFactory.newKey("key1");
+    checkKeyProperties(key1);
+
+    Key key2 = keyFactory.newKey(123);
+    checkKeyProperties(key2);
+
+    IncompleteKey incompleteKey = keyFactory.newKey();
+    checkKeyProperties(incompleteKey);
+  }
+
+  private void checkKeyProperties(BaseKey key) {
+    assertEquals(options.getDatabaseId(), key.getDatabaseId());
+    assertEquals(options.getProjectId(), key.getProjectId());
+    assertEquals(options.getNamespace(), key.getNamespace());
+  }
+
   private RunAggregationQueryResponse placeholderAggregationQueryResponse() {
     Map<String, com.google.datastore.v1.Value> result1 =
         new HashMap<>(ImmutableMap.of("total_count", intValue(209)));
diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/IncompleteKeyTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/IncompleteKeyTest.java
index c3396a28..a1f8956b 100644
--- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/IncompleteKeyTest.java
+++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/IncompleteKeyTest.java
@@ -25,8 +25,7 @@ import org.junit.Test;
 
 public class IncompleteKeyTest {
 
-  private static IncompleteKey pk1, pk2;
-  private static IncompleteKey deprecatedPk1, deprecatedPk2;
+  private static IncompleteKey pk1, pk2, pk4;
   private static Key parent1;
 
   @Before
@@ -34,23 +33,32 @@ public class IncompleteKeyTest {
     pk1 = IncompleteKey.newBuilder("ds", "kind1").build();
     parent1 = Key.newBuilder("ds", "kind2", 10).setNamespace("ns").build();
     pk2 = IncompleteKey.newBuilder(parent1, "kind3").build();
+    pk4 = IncompleteKey.newBuilderWithDatabaseId("ds", "kind3", "test-db").build();
   }
 
   @Test
   public void testBuilders() {
     assertEquals("ds", pk1.getProjectId());
+    assertEquals("", pk1.getDatabaseId());
     assertEquals("kind1", pk1.getKind());
     assertTrue(pk1.getAncestors().isEmpty());
 
     assertEquals("ds", pk2.getProjectId());
+    assertEquals("", pk2.getDatabaseId());
     assertEquals("kind3", pk2.getKind());
     assertEquals(parent1.getPath(), pk2.getAncestors());
 
     assertEquals(pk2, IncompleteKey.newBuilder(pk2).build());
     IncompleteKey pk3 = IncompleteKey.newBuilder(pk2).setKind("kind4").build();
     assertEquals("ds", pk3.getProjectId());
+    assertEquals("", pk3.getDatabaseId());
     assertEquals("kind4", pk3.getKind());
     assertEquals(parent1.getPath(), pk3.getAncestors());
+
+    assertEquals("ds", pk4.getProjectId());
+    assertEquals("test-db", pk4.getDatabaseId());
+    assertEquals("kind3", pk4.getKind());
+    assertTrue(pk4.getAncestors().isEmpty());
   }
 
   @Test
diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/KeyFactoryTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/KeyFactoryTest.java
index 26c922e5..37017c28 100644
--- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/KeyFactoryTest.java
+++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/KeyFactoryTest.java
@@ -26,13 +26,13 @@ import org.junit.Test;
 public class KeyFactoryTest {
 
   private static final String PROJECT_ID = "projectid";
+  private static final String DATABASE_ID = "database-id";
 
   private KeyFactory keyFactory;
-  private KeyFactory deprecatedKeyFactory;
 
   @Before
   public void setUp() {
-    keyFactory = new KeyFactory(PROJECT_ID).setKind("k");
+    keyFactory = new KeyFactory(PROJECT_ID).setDatabaseId(DATABASE_ID).setKind("k");
   }
 
   @Test
@@ -40,11 +40,13 @@ public class KeyFactoryTest {
     IncompleteKey key =
         keyFactory
             .setProjectId("ds1")
+            .setDatabaseId("db")
             .setNamespace("ns1")
             .addAncestor(PathElement.of("p", 1))
             .build();
     assertEquals("k", key.getKind());
     assertEquals("ds1", key.getProjectId());
+    assertEquals("db", key.getDatabaseId());
     assertEquals("ns1", key.getNamespace());
     assertEquals(1, key.getAncestors().size());
 
@@ -58,12 +60,59 @@ public class KeyFactoryTest {
     key = keyFactory.newKey();
     assertEquals("k1", key.getKind());
     assertEquals(PROJECT_ID, key.getProjectId());
+    assertEquals("", key.getDatabaseId());
     assertTrue(key.getNamespace().isEmpty());
     assertTrue(key.getAncestors().isEmpty());
 
     keyFactory = new KeyFactory(PROJECT_ID, "ns1").setKind("k");
     key = keyFactory.newKey();
     assertEquals(PROJECT_ID, key.getProjectId());
+    assertEquals("", key.getDatabaseId());
+    assertEquals("ns1", key.getNamespace());
+    key = keyFactory.setProjectId("bla1").setNamespace("bla2").build();
+    assertEquals("bla1", key.getProjectId());
+    assertEquals("bla2", key.getNamespace());
+    keyFactory.reset().setKind("kind");
+    key = keyFactory.newKey();
+    assertEquals(PROJECT_ID, key.getProjectId());
+    assertEquals("ns1", key.getNamespace());
+    assertEquals("kind", key.getKind());
+  }
+
+  @Test
+  public void testCreatedWithDbId() {
+    KeyFactory keyFactory = new KeyFactory(PROJECT_ID, "namespace", DATABASE_ID).setKind("k");
+    IncompleteKey key =
+        keyFactory
+            .setProjectId("ds1")
+            .setDatabaseId("db")
+            .setNamespace("ns1")
+            .addAncestor(PathElement.of("p", 1))
+            .build();
+    assertEquals("k", key.getKind());
+    assertEquals("ds1", key.getProjectId());
+    assertEquals("db", key.getDatabaseId());
+    assertEquals("ns1", key.getNamespace());
+    assertEquals(1, key.getAncestors().size());
+
+    keyFactory.reset();
+    try {
+      keyFactory.newKey(1);
+    } catch (NullPointerException ex) {
+      assertEquals("kind must not be null", ex.getMessage());
+    }
+    keyFactory.setKind("k1");
+    key = keyFactory.newKey();
+    assertEquals("k1", key.getKind());
+    assertEquals(PROJECT_ID, key.getProjectId());
+    assertEquals(DATABASE_ID, key.getDatabaseId());
+    assertEquals("namespace", key.getNamespace());
+    assertTrue(key.getAncestors().isEmpty());
+
+    keyFactory = new KeyFactory(PROJECT_ID, "ns1").setKind("k");
+    key = keyFactory.newKey();
+    assertEquals(PROJECT_ID, key.getProjectId());
+    assertEquals("", key.getDatabaseId());
     assertEquals("ns1", key.getNamespace());
     key = keyFactory.setProjectId("bla1").setNamespace("bla2").build();
     assertEquals("bla1", key.getProjectId());
@@ -115,6 +164,7 @@ public class KeyFactoryTest {
   private void verifyIncompleteKey(IncompleteKey key, String namespace, PathElement... ancestors) {
     assertEquals("k", key.getKind());
     assertEquals(PROJECT_ID, key.getProjectId());
+    assertEquals(DATABASE_ID, key.getDatabaseId());
     assertEquals(namespace, key.getNamespace());
     assertEquals(ancestors.length, key.getAncestors().size());
     Iterator<PathElement> iter = key.getAncestors().iterator();
diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/KeyTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/KeyTest.java
index d6e3b4d8..cb3d13fd 100644
--- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/KeyTest.java
+++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/KeyTest.java
@@ -76,4 +76,14 @@ public class KeyTest {
     Key copy = Key.fromUrlSafe(urlSafe);
     assertEquals(key, copy);
   }
+
+  @Test
+  public void testDatabaseId() {
+    Key.Builder builder = Key.newBuilder("project-id", "kind", "name", "database-id");
+    Key key = builder.build();
+    assertEquals("database-id", key.getDatabaseId());
+    assertEquals("project-id", key.getProjectId());
+    assertEquals("name", key.getName());
+    assertEquals("kind", key.getKind());
+  }
 }
diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/request/AggregationQueryRequestProtoPreparerTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/request/AggregationQueryRequestProtoPreparerTest.java
index 6301ebef..689d9fa3 100644
--- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/request/AggregationQueryRequestProtoPreparerTest.java
+++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/request/AggregationQueryRequestProtoPreparerTest.java
@@ -49,8 +49,13 @@ public class AggregationQueryRequestProtoPreparerTest {
   private static final String KIND = "Task";
   private static final String NAMESPACE = "ns";
   private static final String PROJECT_ID = "project-id";
+  private static final String DATABASE_ID = "database-id";
   private static final DatastoreOptions DATASTORE_OPTIONS =
-      DatastoreOptions.newBuilder().setProjectId(PROJECT_ID).setNamespace(NAMESPACE).build();
+      DatastoreOptions.newBuilder()
+          .setProjectId(PROJECT_ID)
+          .setDatabaseId(DATABASE_ID)
+          .setNamespace(NAMESPACE)
+          .build();
   private static final EntityQuery COMPLETED_TASK_STRUCTURED_QUERY =
       Query.newEntityQueryBuilder()
           .setNamespace(NAMESPACE)
@@ -89,6 +94,7 @@ public class AggregationQueryRequestProtoPreparerTest {
         protoPreparer.prepare(QueryAndReadOptions.create(AGGREGATION_OVER_STRUCTURED_QUERY));
 
     assertThat(runAggregationQueryRequest.getProjectId()).isEqualTo(PROJECT_ID);
+    assertThat(runAggregationQueryRequest.getDatabaseId()).isEqualTo(DATABASE_ID);
 
     assertThat(runAggregationQueryRequest.getPartitionId().getProjectId()).isEqualTo(PROJECT_ID);
     assertThat(runAggregationQueryRequest.getPartitionId().getNamespaceId()).isEqualTo(NAMESPACE);
@@ -111,8 +117,10 @@ public class AggregationQueryRequestProtoPreparerTest {
         protoPreparer.prepare(QueryAndReadOptions.create(AGGREGATION_OVER_GQL_QUERY));
 
     assertThat(runAggregationQueryRequest.getProjectId()).isEqualTo(PROJECT_ID);
+    assertThat(runAggregationQueryRequest.getDatabaseId()).isEqualTo(DATABASE_ID);
 
     assertThat(runAggregationQueryRequest.getPartitionId().getProjectId()).isEqualTo(PROJECT_ID);
+    assertThat(runAggregationQueryRequest.getPartitionId().getDatabaseId()).isEqualTo(DATABASE_ID);
     assertThat(runAggregationQueryRequest.getPartitionId().getNamespaceId()).isEqualTo(NAMESPACE);
 
     com.google.datastore.v1.GqlQuery gqlQueryProto = runAggregationQueryRequest.getGqlQuery();
diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITDatastoreTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITDatastoreTest.java
index b8c3bb4b..0f36ed06 100644
--- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITDatastoreTest.java
+++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITDatastoreTest.java
@@ -63,7 +63,6 @@ import com.google.cloud.datastore.StructuredQuery.OrderBy;
 import com.google.cloud.datastore.StructuredQuery.PropertyFilter;
 import com.google.cloud.datastore.TimestampValue;
 import com.google.cloud.datastore.Transaction;
-import com.google.cloud.datastore.Value;
 import com.google.cloud.datastore.ValueType;
 import com.google.cloud.datastore.testing.RemoteDatastoreHelper;
 import com.google.common.base.Preconditions;
@@ -88,13 +87,24 @@ import org.junit.Rule;
 import org.junit.Test;
 import org.junit.rules.Timeout;
 
+// todo: combine this with ITDatastoreTestWithNamedDb once we resolve issues with parameterized
+// graalvm tests
 public class ITDatastoreTest {
 
   private static final RemoteDatastoreHelper HELPER = RemoteDatastoreHelper.create();
-  private static final DatastoreOptions OPTIONS = HELPER.getOptions();
-  private static final Datastore DATASTORE = OPTIONS.getService();
-  private static final String PROJECT_ID = OPTIONS.getProjectId();
-  private static final String NAMESPACE = OPTIONS.getNamespace();
+  private static final DatastoreOptions OPTIONS_1 = HELPER.getOptions();
+  private static final Datastore DATASTORE_1 = OPTIONS_1.getService();
+
+  private static final String CUSTOM_DB_ID = "test-db";
+  private static final RemoteDatastoreHelper HELPER2 = RemoteDatastoreHelper.create(CUSTOM_DB_ID);
+  private static final DatastoreOptions OPTIONS_2 = HELPER2.getOptions();
+  private static final Datastore DATASTORE_2 = OPTIONS_2.getService();
+
+  private final DatastoreOptions options = OPTIONS_1;
+  private final Datastore datastore = DATASTORE_1;
+
+  private static String PROJECT_ID;
+  private static String NAMESPACE;
   private static final String KIND1 = "kind1";
   private static final String KIND2 = "kind2";
   private static final String KIND3 = "kind3";
@@ -102,70 +112,27 @@ public class ITDatastoreTest {
   private static final StringValue STR_VALUE = StringValue.of("str");
   private static final BooleanValue BOOL_VALUE =
       BooleanValue.newBuilder(false).setExcludeFromIndexes(true).build();
-  private static final Key ROOT_KEY =
-      Key.newBuilder(PROJECT_ID, "rootkey", "default").setNamespace(NAMESPACE).build();
-  private static final IncompleteKey INCOMPLETE_KEY1 =
-      IncompleteKey.newBuilder(ROOT_KEY, KIND1).setNamespace(NAMESPACE).build();
-  private static final IncompleteKey INCOMPLETE_KEY2 =
-      IncompleteKey.newBuilder(PROJECT_ID, KIND2).setNamespace(NAMESPACE).build();
-  private static final Key KEY1 = Key.newBuilder(INCOMPLETE_KEY1, "name").build();
-  private static final Key KEY2 = Key.newBuilder(KEY1, KIND2, 1).build();
-  private static final Key KEY3 =
-      Key.newBuilder(KEY2).setName("bla").setNamespace(NAMESPACE).build();
-  private static final Key KEY4 =
-      Key.newBuilder(KEY2).setName("newName1").setNamespace(NAMESPACE).build();
-  private static final Key KEY5 =
-      Key.newBuilder(KEY2).setName("newName2").setNamespace(NAMESPACE).build();
-  private static final KeyValue KEY_VALUE = KeyValue.of(KEY1);
+  private static final ListValue EMPTY_LIST_VALUE = ListValue.of(Collections.emptyList());
   private static final ListValue LIST_VALUE1 =
       ListValue.newBuilder().addValue(NULL_VALUE).addValue(STR_VALUE, BOOL_VALUE).build();
-  private static final ListValue LIST_VALUE2 = ListValue.of(Collections.singletonList(KEY_VALUE));
-  private static final ListValue EMPTY_LIST_VALUE = ListValue.of(Collections.<Value<?>>emptyList());
   private static final TimestampValue TIMESTAMP_VALUE = new TimestampValue(Timestamp.now());
   private static final LatLngValue LAT_LNG_VALUE =
       new LatLngValue(LatLng.of(37.422035, -122.084124));
-  private static final FullEntity<IncompleteKey> PARTIAL_ENTITY1 =
-      FullEntity.newBuilder(INCOMPLETE_KEY2)
-          .set("str", STR_VALUE)
-          .set("bool", BOOL_VALUE)
-          .set("list", LIST_VALUE1)
-          .build();
-  private static final FullEntity<IncompleteKey> PARTIAL_ENTITY2 =
-      FullEntity.newBuilder(PARTIAL_ENTITY1)
-          .remove("str")
-          .set("bool", true)
-          .set("list", LIST_VALUE1.get())
-          .build();
-  private static final FullEntity<IncompleteKey> PARTIAL_ENTITY3 =
-      FullEntity.newBuilder(PARTIAL_ENTITY1)
-          .setKey(IncompleteKey.newBuilder(PROJECT_ID, KIND3).build())
-          .build();
-  private static final Entity ENTITY1 =
-      Entity.newBuilder(KEY1)
-          .set("str", STR_VALUE)
-          .set("date", TIMESTAMP_VALUE)
-          .set("latLng", LAT_LNG_VALUE)
-          .set("bool", BOOL_VALUE)
-          .set("partial1", EntityValue.of(PARTIAL_ENTITY1))
-          .set("list", LIST_VALUE2)
-          .set("emptyList", EMPTY_LIST_VALUE)
-          .build();
-  private static final Entity ENTITY2 =
-      Entity.newBuilder(ENTITY1)
-          .setKey(KEY2)
-          .remove("str")
-          .set("name", "Dan")
-          .setNull("null")
-          .set("age", 20)
-          .build();
-  private static final Entity ENTITY3 =
-      Entity.newBuilder(ENTITY1)
-          .setKey(KEY3)
-          .remove("str")
-          .set("null", NULL_VALUE)
-          .set("partial1", PARTIAL_ENTITY2)
-          .set("partial2", ENTITY2)
-          .build();
+
+  private static Key ROOT_KEY;
+  private static IncompleteKey INCOMPLETE_KEY1;
+  private static ListValue LIST_VALUE2;
+  private static Key KEY1;
+  private static Key KEY2;
+  private static Key KEY3;
+  private static Key KEY4;
+  private static Key KEY5;
+  private static FullEntity<IncompleteKey> PARTIAL_ENTITY1;
+  private static FullEntity<IncompleteKey> PARTIAL_ENTITY2;
+  private static FullEntity<IncompleteKey> PARTIAL_ENTITY3;
+  private static Entity ENTITY1;
+  private static Entity ENTITY2;
+  private static Entity ENTITY3;
 
   @Rule public Timeout globalTimeout = Timeout.seconds(100);
 
@@ -178,30 +145,99 @@ public class ITDatastoreTest {
 
   @Before
   public void setUp() {
-    DATASTORE.put(ENTITY1, ENTITY2);
+    PROJECT_ID = this.options.getProjectId();
+    NAMESPACE = this.options.getNamespace();
+
+    ROOT_KEY =
+        Key.newBuilder(PROJECT_ID, "rootkey", "default", options.getDatabaseId())
+            .setNamespace(NAMESPACE)
+            .build();
+    INCOMPLETE_KEY1 = IncompleteKey.newBuilder(ROOT_KEY, KIND1).setNamespace(NAMESPACE).build();
+
+    IncompleteKey INCOMPLETE_KEY2 =
+        IncompleteKey.newBuilder(PROJECT_ID, KIND2)
+            .setDatabaseId(options.getDatabaseId())
+            .setNamespace(NAMESPACE)
+            .build();
+
+    KEY1 = Key.newBuilder(INCOMPLETE_KEY1, "name").build();
+    KEY2 = Key.newBuilder(KEY1, KIND2, 1).build();
+    KEY3 = Key.newBuilder(KEY2).setName("bla").setNamespace(NAMESPACE).build();
+    KEY4 = Key.newBuilder(KEY2).setName("newName1").setNamespace(NAMESPACE).build();
+    KEY5 = Key.newBuilder(KEY2).setName("newName2").setNamespace(NAMESPACE).build();
+
+    LIST_VALUE2 = ListValue.of(Collections.singletonList(KeyValue.of(KEY1)));
+
+    PARTIAL_ENTITY1 =
+        FullEntity.newBuilder(INCOMPLETE_KEY2)
+            .set("str", STR_VALUE)
+            .set("bool", BOOL_VALUE)
+            .set("list", LIST_VALUE1)
+            .build();
+    PARTIAL_ENTITY2 =
+        FullEntity.newBuilder(PARTIAL_ENTITY1)
+            .remove("str")
+            .set("bool", true)
+            .set("list", LIST_VALUE1.get())
+            .build();
+    PARTIAL_ENTITY3 =
+        FullEntity.newBuilder(PARTIAL_ENTITY1)
+            .setKey(
+                IncompleteKey.newBuilder(PROJECT_ID, KIND3)
+                    .setDatabaseId(options.getDatabaseId())
+                    .build())
+            .build();
+    ENTITY1 =
+        Entity.newBuilder(KEY1)
+            .set("str", STR_VALUE)
+            .set("date", TIMESTAMP_VALUE)
+            .set("latLng", LAT_LNG_VALUE)
+            .set("bool", BOOL_VALUE)
+            .set("partial1", EntityValue.of(PARTIAL_ENTITY1))
+            .set("list", LIST_VALUE2)
+            .set("emptyList", EMPTY_LIST_VALUE)
+            .build();
+    ENTITY2 =
+        Entity.newBuilder(ENTITY1)
+            .setKey(KEY2)
+            .remove("str")
+            .set("name", "Dan")
+            .setNull("null")
+            .set("age", 20)
+            .build();
+    ENTITY3 =
+        Entity.newBuilder(ENTITY1)
+            .setKey(KEY3)
+            .remove("str")
+            .set("null", NULL_VALUE)
+            .set("partial1", PARTIAL_ENTITY2)
+            .set("partial2", ENTITY2)
+            .build();
+
+    datastore.put(ENTITY1, ENTITY2);
   }
 
   @After
   public void tearDown() {
     EntityQuery allEntitiesQuery = Query.newEntityQueryBuilder().build();
-    QueryResults<Entity> allEntities = DATASTORE.run(allEntitiesQuery);
+    QueryResults<Entity> allEntities = datastore.run(allEntitiesQuery);
     Key[] keysToDelete =
         ImmutableList.copyOf(allEntities).stream().map(Entity::getKey).toArray(Key[]::new);
-    DATASTORE.delete(keysToDelete);
+    datastore.delete(keysToDelete);
   }
 
   private <T> Iterator<T> getStronglyConsistentResults(Query scQuery, Query query)
       throws InterruptedException {
     // scQuery is equivalent to query, but with an ancestor filter in it
     // this makes scQuery strongly consistent
-    QueryResults<T> scResults = DATASTORE.run(scQuery);
+    QueryResults<T> scResults = datastore.run(scQuery);
     List<T> scResultsCopy = makeResultsCopy(scResults);
     Set<T> scResultsSet = new HashSet<>(scResultsCopy);
     int maxAttempts = 20;
 
     while (maxAttempts > 0) {
       --maxAttempts;
-      QueryResults<T> results = DATASTORE.run(query);
+      QueryResults<T> results = datastore.run(query);
       List<T> resultsCopy = makeResultsCopy(results);
       Set<T> resultsSet = new HashSet<>(resultsCopy);
       if (scResultsSet.size() == resultsSet.size() && scResultsSet.containsAll(resultsSet)) {
@@ -225,7 +261,7 @@ public class ITDatastoreTest {
 
   @Test
   public void testNewTransactionCommit() {
-    Transaction transaction = DATASTORE.newTransaction();
+    Transaction transaction = datastore.newTransaction();
     transaction.add(ENTITY3);
     Entity entity2 = Entity.newBuilder(ENTITY2).clear().setNull("bla").build();
     transaction.update(entity2);
@@ -233,7 +269,7 @@ public class ITDatastoreTest {
     transaction.commit();
     assertFalse(transaction.isActive());
 
-    List<Entity> list = DATASTORE.fetch(KEY1, KEY2, KEY3);
+    List<Entity> list = datastore.fetch(KEY1, KEY2, KEY3);
     assertNull(list.get(0));
     assertEquals(entity2, list.get(1));
     assertEquals(ENTITY3, list.get(2));
@@ -256,16 +292,16 @@ public class ITDatastoreTest {
 
   @Test
   public void testTransactionWithRead() {
-    Transaction transaction = DATASTORE.newTransaction();
+    Transaction transaction = datastore.newTransaction();
     assertNull(transaction.get(KEY3));
     transaction.add(ENTITY3);
     transaction.commit();
-    assertEquals(ENTITY3, DATASTORE.get(KEY3));
+    assertEquals(ENTITY3, datastore.get(KEY3));
 
-    transaction = DATASTORE.newTransaction();
+    transaction = datastore.newTransaction();
     assertEquals(ENTITY3, transaction.get(KEY3));
     // update entity3 during the transaction
-    DATASTORE.put(Entity.newBuilder(ENTITY2).clear().set("from", "datastore").build());
+    datastore.put(Entity.newBuilder(ENTITY2).clear().set("from", "datastore").build());
     transaction.update(Entity.newBuilder(ENTITY2).clear().set("from", "transaction").build());
     try {
       transaction.commit();
@@ -283,23 +319,23 @@ public class ITDatastoreTest {
             .setFilter(PropertyFilter.hasAncestor(KEY2))
             .setNamespace(NAMESPACE)
             .build();
-    Transaction transaction = DATASTORE.newTransaction();
+    Transaction transaction = datastore.newTransaction();
     QueryResults<Entity> results = transaction.run(query);
     assertTrue(results.hasNext());
     assertEquals(ENTITY2, results.next());
     assertFalse(results.hasNext());
     transaction.add(ENTITY3);
     transaction.commit();
-    assertEquals(ENTITY3, DATASTORE.get(KEY3));
+    assertEquals(ENTITY3, datastore.get(KEY3));
 
-    transaction = DATASTORE.newTransaction();
+    transaction = datastore.newTransaction();
     results = transaction.run(query);
     assertTrue(results.hasNext());
     assertEquals(ENTITY2, results.next());
     assertFalse(results.hasNext());
     transaction.delete(ENTITY3.getKey());
     // update entity2 during the transaction
-    DATASTORE.put(Entity.newBuilder(ENTITY2).clear().build());
+    datastore.put(Entity.newBuilder(ENTITY2).clear().build());
     try {
       transaction.commit();
       fail("Expecting a failure");
@@ -310,7 +346,7 @@ public class ITDatastoreTest {
 
   @Test
   public void testNewTransactionRollback() {
-    Transaction transaction = DATASTORE.newTransaction();
+    Transaction transaction = datastore.newTransaction();
     transaction.add(ENTITY3);
     Entity entity2 =
         Entity.newBuilder(ENTITY2)
@@ -330,7 +366,7 @@ public class ITDatastoreTest {
       assertEquals("FAILED_PRECONDITION", expected.getReason());
     }
 
-    List<Entity> list = DATASTORE.fetch(KEY1, KEY2, KEY3);
+    List<Entity> list = datastore.fetch(KEY1, KEY2, KEY3);
     assertEquals(ENTITY1, list.get(0));
     assertEquals(ENTITY2, list.get(1));
     assertNull(list.get(2));
@@ -339,7 +375,7 @@ public class ITDatastoreTest {
 
   @Test
   public void testNewBatch() {
-    Batch batch = DATASTORE.newBatch();
+    Batch batch = datastore.newBatch();
     Entity entity1 = Entity.newBuilder(ENTITY1).clear().build();
     Entity entity2 = Entity.newBuilder(ENTITY2).clear().setNull("bla").build();
     Entity entity4 = Entity.newBuilder(KEY4).set("value", StringValue.of("value")).build();
@@ -362,7 +398,7 @@ public class ITDatastoreTest {
 
     Batch.Response response = batch.submit();
     entities =
-        DATASTORE.fetch(KEY1, KEY2, KEY3, entity4.getKey(), entity5.getKey(), entity6.getKey());
+        datastore.fetch(KEY1, KEY2, KEY3, entity4.getKey(), entity5.getKey(), entity6.getKey());
     assertEquals(entity1, entities.get(0));
     assertEquals(entity2, entities.get(1));
     assertEquals(ENTITY3, entities.get(2));
@@ -372,7 +408,7 @@ public class ITDatastoreTest {
     assertEquals(6, entities.size());
     List<Key> generatedKeys = response.getGeneratedKeys();
     assertEquals(1, generatedKeys.size());
-    assertEquals(PARTIAL_ENTITY3.getNames(), DATASTORE.get(generatedKeys.get(0)).getNames());
+    assertEquals(PARTIAL_ENTITY3.getNames(), datastore.get(generatedKeys.get(0)).getNames());
     assertEquals(PARTIAL_ENTITY3.getKey(), IncompleteKey.newBuilder(generatedKeys.get(0)).build());
 
     try {
@@ -382,12 +418,12 @@ public class ITDatastoreTest {
       assertEquals("FAILED_PRECONDITION", expected.getReason());
     }
 
-    batch = DATASTORE.newBatch();
+    batch = datastore.newBatch();
     batch.delete(entity4.getKey(), entity5.getKey(), entity6.getKey());
     batch.update(ENTITY1, ENTITY2, ENTITY3);
     batch.submit();
     entities =
-        DATASTORE.fetch(KEY1, KEY2, KEY3, entity4.getKey(), entity5.getKey(), entity6.getKey());
+        datastore.fetch(KEY1, KEY2, KEY3, entity4.getKey(), entity5.getKey(), entity6.getKey());
     assertEquals(ENTITY1, entities.get(0));
     assertEquals(ENTITY2, entities.get(1));
     assertEquals(ENTITY3, entities.get(2));
@@ -416,7 +452,7 @@ public class ITDatastoreTest {
     assertEquals(ENTITY1, results1.next());
     assertFalse(results1.hasNext());
 
-    DATASTORE.put(ENTITY3);
+    datastore.put(ENTITY3);
     Query<? extends Entity> query2 =
         Query.newGqlQueryBuilder(ResultType.ENTITY, "select * from " + KIND2 + " order by __key__")
             .setNamespace(NAMESPACE)
@@ -483,7 +519,7 @@ public class ITDatastoreTest {
     assertEquals(KEY1, projectionEntity.getKey());
     assertTrue(projectionEntity.getNames().isEmpty());
     assertFalse(keyProjectionResult.hasNext());
-    DATASTORE.delete(ENTITY3.getKey());
+    datastore.delete(ENTITY3.getKey());
   }
 
   @Test
@@ -506,7 +542,7 @@ public class ITDatastoreTest {
     Query<?> query2 =
         Query.newGqlQueryBuilder("select * from " + KIND1).setNamespace(NAMESPACE).build();
 
-    QueryResults<?> results2 = DATASTORE.run(query2);
+    QueryResults<?> results2 = datastore.run(query2);
 
     assertSame(Entity.class, results2.getResultClass());
 
@@ -661,11 +697,11 @@ public class ITDatastoreTest {
             .build();
 
     // original entity count is 2
-    assertThat(getOnlyElement(DATASTORE.runAggregation(aggregationQuery)).get("count"))
+    assertThat(getOnlyElement(datastore.runAggregation(aggregationQuery)).get("count"))
         .isEqualTo(2L);
 
     // FIRST TRANSACTION
-    DATASTORE.runInTransaction(
+    datastore.runInTransaction(
         (TransactionCallable<Void>)
             inFirstTransaction -> {
               // creating a new entity
@@ -678,16 +714,16 @@ public class ITDatastoreTest {
                       getOnlyElement(inFirstTransaction.runAggregation(aggregationQuery))
                           .get("count"))
                   .isEqualTo(2L);
-              assertThat(getOnlyElement(DATASTORE.runAggregation(aggregationQuery)).get("count"))
+              assertThat(getOnlyElement(datastore.runAggregation(aggregationQuery)).get("count"))
                   .isEqualTo(2L);
               return null;
             });
     // after first transaction is committed, count is updated to 3 now.
-    assertThat(getOnlyElement(DATASTORE.runAggregation(aggregationQuery)).get("count"))
+    assertThat(getOnlyElement(datastore.runAggregation(aggregationQuery)).get("count"))
         .isEqualTo(3L);
 
     // SECOND TRANSACTION
-    DATASTORE.runInTransaction(
+    datastore.runInTransaction(
         (TransactionCallable<Void>)
             inSecondTransaction -> {
               // deleting ENTITY2
@@ -698,14 +734,14 @@ public class ITDatastoreTest {
                       getOnlyElement(inSecondTransaction.runAggregation(aggregationQuery))
                           .get("count"))
                   .isEqualTo(3L);
-              assertThat(getOnlyElement(DATASTORE.runAggregation(aggregationQuery)).get("count"))
+              assertThat(getOnlyElement(datastore.runAggregation(aggregationQuery)).get("count"))
                   .isEqualTo(3L);
               return null;
             });
     // after second transaction is committed, count is updated to 2 now.
-    assertThat(getOnlyElement(DATASTORE.runAggregation(aggregationQuery)).get("count"))
+    assertThat(getOnlyElement(datastore.runAggregation(aggregationQuery)).get("count"))
         .isEqualTo(2L);
-    DATASTORE.delete(newEntityKey);
+    datastore.delete(newEntityKey);
   }
 
   @Test
@@ -726,7 +762,7 @@ public class ITDatastoreTest {
 
     TransactionOptions transactionOptions =
         TransactionOptions.newBuilder().setReadOnly(ReadOnly.newBuilder().build()).build();
-    Transaction readOnlyTransaction = DATASTORE.newTransaction(transactionOptions);
+    Transaction readOnlyTransaction = datastore.newTransaction(transactionOptions);
 
     // Executing query in transaction
     assertThat(getOnlyElement(readOnlyTransaction.runAggregation(aggregationQuery)).get("count"))
@@ -741,7 +777,7 @@ public class ITDatastoreTest {
                       .setKey(Key.newBuilder(KEY1, "newKind", "name-01").build())
                       .set("v_int", 10)
                       .build();
-              DATASTORE.put(aNewEntity);
+              datastore.put(aNewEntity);
               return null;
             });
 
@@ -752,7 +788,7 @@ public class ITDatastoreTest {
     readOnlyTransaction.commit();
     executor.shutdownNow();
 
-    assertThat(getOnlyElement(DATASTORE.runAggregation(aggregationQuery)).get("count"))
+    assertThat(getOnlyElement(datastore.runAggregation(aggregationQuery)).get("count"))
         .isEqualTo(3L);
   }
 
@@ -885,7 +921,7 @@ public class ITDatastoreTest {
             .setKey(Key.newBuilder(INCOMPLETE_KEY1, "e2").build())
             .set("v_int", 20)
             .build();
-    DATASTORE.put(e1, e2);
+    datastore.put(e1, e2);
 
     Query<Entity> queryIn =
         Query.newEntityQueryBuilder()
@@ -946,15 +982,15 @@ public class ITDatastoreTest {
     assertEquals(e2, resultNeq.next());
     assertFalse(resultNeq.hasNext());
 
-    DATASTORE.delete(e1.getKey());
-    DATASTORE.delete(e2.getKey());
+    datastore.delete(e1.getKey());
+    datastore.delete(e2.getKey());
   }
 
   @Test
   public void testAllocateId() {
-    KeyFactory keyFactory = DATASTORE.newKeyFactory().setKind(KIND1);
+    KeyFactory keyFactory = datastore.newKeyFactory().setKind(KIND1);
     IncompleteKey pk1 = keyFactory.newKey();
-    Key key1 = DATASTORE.allocateId(pk1);
+    Key key1 = datastore.allocateId(pk1);
     assertEquals(key1.getProjectId(), pk1.getProjectId());
     assertEquals(key1.getNamespace(), pk1.getNamespace());
     assertEquals(key1.getAncestors(), pk1.getAncestors());
@@ -963,27 +999,31 @@ public class ITDatastoreTest {
     assertFalse(key1.hasName());
     assertEquals(Key.newBuilder(pk1, key1.getId()).build(), key1);
 
-    Key key2 = DATASTORE.allocateId(pk1);
+    Key key2 = datastore.allocateId(pk1);
     assertNotEquals(key1, key2);
     assertEquals(Key.newBuilder(pk1, key2.getId()).build(), key2);
   }
 
   @Test
   public void testReserveIds() {
-    KeyFactory keyFactory = DATASTORE.newKeyFactory().setKind("MyKind");
+    KeyFactory keyFactory = datastore.newKeyFactory().setKind("MyKind");
     Key key1 = keyFactory.newKey(10);
     Key key2 = keyFactory.newKey("name");
-    List<Key> keyList = DATASTORE.reserveIds(key1, key2);
+    List<Key> keyList = datastore.reserveIds(key1, key2);
     assertEquals(2, keyList.size());
   }
 
   @Test
   public void testAllocateIdArray() {
-    KeyFactory keyFactory = DATASTORE.newKeyFactory().setKind(KIND1);
+    KeyFactory keyFactory = datastore.newKeyFactory().setKind(KIND1);
     IncompleteKey incompleteKey1 = keyFactory.newKey();
     IncompleteKey incompleteKey2 =
-        keyFactory.setKind(KIND2).addAncestors(PathElement.of(KIND1, 10)).newKey();
-    List<Key> result = DATASTORE.allocateId(incompleteKey1, incompleteKey2, incompleteKey1);
+        keyFactory
+            .setKind(KIND2)
+            .setDatabaseId(options.getDatabaseId())
+            .addAncestors(PathElement.of(KIND1, 10))
+            .newKey();
+    List<Key> result = datastore.allocateId(incompleteKey1, incompleteKey2, incompleteKey1);
     assertEquals(3, result.size());
     assertEquals(Key.newBuilder(incompleteKey1, result.get(0).getId()).build(), result.get(0));
     assertEquals(Key.newBuilder(incompleteKey1, result.get(2).getId()).build(), result.get(2));
@@ -992,10 +1032,10 @@ public class ITDatastoreTest {
 
   @Test
   public void testGet() {
-    Entity entity = DATASTORE.get(KEY3);
+    Entity entity = datastore.get(KEY3);
     assertNull(entity);
 
-    entity = DATASTORE.get(KEY1);
+    entity = datastore.get(KEY1);
     assertEquals(ENTITY1, entity);
     StringValue value1 = entity.getValue("str");
     assertEquals(STR_VALUE, value1);
@@ -1017,34 +1057,37 @@ public class ITDatastoreTest {
 
   @Test
   public void testGetWithReadTime() throws InterruptedException {
-    Key key = Key.newBuilder(PROJECT_ID, "new_kind", "name").setNamespace(NAMESPACE).build();
+    Key key =
+        Key.newBuilder(PROJECT_ID, "new_kind", "name", options.getDatabaseId())
+            .setNamespace(NAMESPACE)
+            .build();
 
     try {
-      DATASTORE.put(Entity.newBuilder(key).set("str", "old_str_value").build());
+      datastore.put(Entity.newBuilder(key).set("str", "old_str_value").build());
 
       Thread.sleep(1000);
       Timestamp now = Timestamp.now();
       Thread.sleep(1000);
 
-      DATASTORE.put(Entity.newBuilder(key).set("str", "new_str_value").build());
+      datastore.put(Entity.newBuilder(key).set("str", "new_str_value").build());
 
-      Entity entity = DATASTORE.get(key);
+      Entity entity = datastore.get(key);
       StringValue value1 = entity.getValue("str");
       assertEquals(StringValue.of("new_str_value"), value1);
 
-      entity = DATASTORE.get(key, ReadOption.readTime(now));
+      entity = datastore.get(key, ReadOption.readTime(now));
       value1 = entity.getValue("str");
       assertEquals(StringValue.of("old_str_value"), value1);
     } finally {
-      DATASTORE.delete(key);
+      datastore.delete(key);
     }
   }
 
   @Test
   public void testGetArrayNoDeferredResults() {
-    DATASTORE.put(ENTITY3);
+    datastore.put(ENTITY3);
     Iterator<Entity> result =
-        DATASTORE.fetch(KEY1, Key.newBuilder(KEY1).setName("bla").build(), KEY2, KEY3).iterator();
+        datastore.fetch(KEY1, Key.newBuilder(KEY1).setName("bla").build(), KEY2, KEY3).iterator();
     assertEquals(ENTITY1, result.next());
     assertNull(result.next());
     assertEquals(ENTITY2, result.next());
@@ -1069,93 +1112,93 @@ public class ITDatastoreTest {
       // expected - no such property
     }
     assertFalse(result.hasNext());
-    DATASTORE.delete(ENTITY3.getKey());
+    datastore.delete(ENTITY3.getKey());
   }
 
   @Test
   public void testAddEntity() {
-    List<Entity> keys = DATASTORE.fetch(ENTITY1.getKey(), ENTITY3.getKey());
+    List<Entity> keys = datastore.fetch(ENTITY1.getKey(), ENTITY3.getKey());
     assertEquals(ENTITY1, keys.get(0));
     assertNull(keys.get(1));
     assertEquals(2, keys.size());
 
     try {
-      DATASTORE.add(ENTITY1);
+      datastore.add(ENTITY1);
       fail("Expecting a failure");
     } catch (DatastoreException expected) {
       // expected;
     }
 
-    List<Entity> entities = DATASTORE.add(ENTITY3, PARTIAL_ENTITY1, PARTIAL_ENTITY2);
-    assertEquals(ENTITY3, DATASTORE.get(ENTITY3.getKey()));
+    List<Entity> entities = datastore.add(ENTITY3, PARTIAL_ENTITY1, PARTIAL_ENTITY2);
+    assertEquals(ENTITY3, datastore.get(ENTITY3.getKey()));
     assertEquals(ENTITY3, entities.get(0));
     assertEquals(PARTIAL_ENTITY1.getNames(), entities.get(1).getNames());
     assertEquals(PARTIAL_ENTITY1.getKey().getAncestors(), entities.get(1).getKey().getAncestors());
-    assertNotNull(DATASTORE.get(entities.get(1).getKey()));
+    assertNotNull(datastore.get(entities.get(1).getKey()));
     assertEquals(PARTIAL_ENTITY2.getNames(), entities.get(2).getNames());
     assertEquals(PARTIAL_ENTITY2.getKey().getAncestors(), entities.get(2).getKey().getAncestors());
-    assertNotNull(DATASTORE.get(entities.get(2).getKey()));
+    assertNotNull(datastore.get(entities.get(2).getKey()));
     for (Entity entity : entities) {
-      DATASTORE.delete(entity.getKey());
+      datastore.delete(entity.getKey());
     }
   }
 
   @Test
   public void testUpdate() {
-    List<Entity> keys = DATASTORE.fetch(ENTITY1.getKey(), ENTITY3.getKey());
+    List<Entity> keys = datastore.fetch(ENTITY1.getKey(), ENTITY3.getKey());
     assertEquals(ENTITY1, keys.get(0));
     assertNull(keys.get(1));
     assertEquals(2, keys.size());
 
     try {
-      DATASTORE.update(ENTITY3);
+      datastore.update(ENTITY3);
       fail("Expecting a failure");
     } catch (DatastoreException expected) {
       // expected;
     }
-    DATASTORE.add(ENTITY3);
-    assertEquals(ENTITY3, DATASTORE.get(ENTITY3.getKey()));
+    datastore.add(ENTITY3);
+    assertEquals(ENTITY3, datastore.get(ENTITY3.getKey()));
     Entity entity3 = Entity.newBuilder(ENTITY3).clear().set("bla", new NullValue()).build();
     assertNotEquals(ENTITY3, entity3);
-    DATASTORE.update(entity3);
-    assertEquals(entity3, DATASTORE.get(ENTITY3.getKey()));
-    DATASTORE.delete(ENTITY3.getKey());
+    datastore.update(entity3);
+    assertEquals(entity3, datastore.get(ENTITY3.getKey()));
+    datastore.delete(ENTITY3.getKey());
   }
 
   @Test
   public void testPut() {
     Entity updatedEntity = Entity.newBuilder(ENTITY1).set("new_property", 42L).build();
-    assertEquals(updatedEntity, DATASTORE.put(updatedEntity));
-    assertEquals(updatedEntity, DATASTORE.get(updatedEntity.getKey()));
+    assertEquals(updatedEntity, datastore.put(updatedEntity));
+    assertEquals(updatedEntity, datastore.get(updatedEntity.getKey()));
 
     Entity entity2 = Entity.newBuilder(ENTITY2).clear().set("bla", new NullValue()).build();
     assertNotEquals(ENTITY2, entity2);
-    List<Entity> entities = DATASTORE.put(ENTITY1, entity2, ENTITY3, PARTIAL_ENTITY1);
+    List<Entity> entities = datastore.put(ENTITY1, entity2, ENTITY3, PARTIAL_ENTITY1);
     assertEquals(ENTITY1, entities.get(0));
     assertEquals(entity2, entities.get(1));
     assertEquals(ENTITY3, entities.get(2));
     assertEquals(PARTIAL_ENTITY1.getNames(), entities.get(3).getNames());
     assertEquals(PARTIAL_ENTITY1.getKey().getAncestors(), entities.get(3).getKey().getAncestors());
-    assertEquals(ENTITY1, DATASTORE.get(ENTITY1.getKey()));
-    assertEquals(entity2, DATASTORE.get(entity2.getKey()));
-    assertEquals(ENTITY3, DATASTORE.get(ENTITY3.getKey()));
-    Entity entity = DATASTORE.get(entities.get(3).getKey());
+    assertEquals(ENTITY1, datastore.get(ENTITY1.getKey()));
+    assertEquals(entity2, datastore.get(entity2.getKey()));
+    assertEquals(ENTITY3, datastore.get(ENTITY3.getKey()));
+    Entity entity = datastore.get(entities.get(3).getKey());
     assertEquals(entities.get(3), entity);
     for (Entity entityToDelete : entities) {
-      DATASTORE.delete(entityToDelete.getKey());
+      datastore.delete(entityToDelete.getKey());
     }
   }
 
   @Test
   public void testDelete() {
     Iterator<Entity> keys =
-        DATASTORE.fetch(ENTITY1.getKey(), ENTITY2.getKey(), ENTITY3.getKey()).iterator();
+        datastore.fetch(ENTITY1.getKey(), ENTITY2.getKey(), ENTITY3.getKey()).iterator();
     assertEquals(ENTITY1, keys.next());
     assertEquals(ENTITY2, keys.next());
     assertNull(keys.next());
     assertFalse(keys.hasNext());
-    DATASTORE.delete(ENTITY1.getKey(), ENTITY2.getKey(), ENTITY3.getKey());
-    keys = DATASTORE.fetch(ENTITY1.getKey(), ENTITY2.getKey(), ENTITY3.getKey()).iterator();
+    datastore.delete(ENTITY1.getKey(), ENTITY2.getKey(), ENTITY3.getKey());
+    keys = datastore.fetch(ENTITY1.getKey(), ENTITY2.getKey(), ENTITY3.getKey()).iterator();
     assertNull(keys.next());
     assertNull(keys.next());
     assertNull(keys.next());
@@ -1180,7 +1223,7 @@ public class ITDatastoreTest {
           }
         };
 
-    int result = DATASTORE.runInTransaction(callable1);
+    int result = datastore.runInTransaction(callable1);
     assertEquals(result, 2);
 
     Datastore.TransactionCallable<Integer> callable2 =
@@ -1199,7 +1242,7 @@ public class ITDatastoreTest {
         };
 
     try {
-      DATASTORE.runInTransaction(callable2);
+      datastore.runInTransaction(callable2);
       fail("Expecting a failure");
     } catch (DatastoreException expected) {
       assertEquals(4, ((DatastoreException) expected.getCause()).getCode());
@@ -1227,7 +1270,7 @@ public class ITDatastoreTest {
           }
         };
 
-    int result = DATASTORE.runInTransaction(callable1);
+    int result = datastore.runInTransaction(callable1);
     assertEquals(result, 2);
 
     final Entity entity2 = Entity.newBuilder(ENTITY2).clear().setNull("bla").build();
@@ -1253,7 +1296,7 @@ public class ITDatastoreTest {
             .build();
 
     try {
-      DATASTORE.runInTransaction(callable2, readOnlyOptions);
+      datastore.runInTransaction(callable2, readOnlyOptions);
       fail("Expecting a failure");
     } catch (DatastoreException expected) {
       assertEquals(3, ((DatastoreException) expected.getCause()).getCode());
@@ -1263,20 +1306,20 @@ public class ITDatastoreTest {
   @Test
   public void testSkippedResults() {
     Query<Key> query = Query.newKeyQueryBuilder().setOffset(Integer.MAX_VALUE).build();
-    int numberOfEntities = DATASTORE.run(query).getSkippedResults();
+    int numberOfEntities = datastore.run(query).getSkippedResults();
     assertEquals(2, numberOfEntities);
   }
 
   @Test
   public void testSetLimit() {
-    DATASTORE.put(ENTITY1);
+    datastore.put(ENTITY1);
     Query<Key> keyQuery = Query.newKeyQueryBuilder().setLimit(1).build();
-    QueryResults<?> queryResults = DATASTORE.run(keyQuery);
+    QueryResults<?> queryResults = datastore.run(keyQuery);
     assertTrue(queryResults.hasNext());
     assertEquals(KEY1, queryResults.next());
 
     Query<?> query = Query.newEntityQueryBuilder().setLimit(0).build();
-    QueryResults<?> results = DATASTORE.run(query);
+    QueryResults<?> results = datastore.run(query);
     assertFalse(results.hasNext());
   }
 
@@ -1287,7 +1330,7 @@ public class ITDatastoreTest {
             .setNamespace(NAMESPACE)
             .setNullBinding("name")
             .build();
-    Iterator<Entity> results = DATASTORE.run(query);
+    Iterator<Entity> results = datastore.run(query);
     assertTrue(results.hasNext());
     assertEquals(ENTITY1, results.next());
     assertFalse(results.hasNext());
@@ -1296,49 +1339,64 @@ public class ITDatastoreTest {
   @Test
   public void testQueryWithStartCursor() {
     Entity entity1 =
-        Entity.newBuilder(Key.newBuilder(PROJECT_ID, KIND1, "name-01").build()).build();
+        Entity.newBuilder(
+                Key.newBuilder(PROJECT_ID, KIND1, "name-01", options.getDatabaseId()).build())
+            .build();
     Entity entity2 =
-        Entity.newBuilder(Key.newBuilder(PROJECT_ID, KIND1, "name-02").build()).build();
+        Entity.newBuilder(
+                Key.newBuilder(PROJECT_ID, KIND1, "name-02", options.getDatabaseId()).build())
+            .build();
     Entity entity3 =
-        Entity.newBuilder(Key.newBuilder(PROJECT_ID, KIND1, "name-03").build()).build();
-    DATASTORE.put(entity1, entity2, entity3);
-    QueryResults<Entity> run1 = DATASTORE.run(Query.newEntityQueryBuilder().setKind(KIND1).build());
+        Entity.newBuilder(
+                Key.newBuilder(PROJECT_ID, KIND1, "name-03", options.getDatabaseId()).build())
+            .build();
+    datastore.put(entity1, entity2, entity3);
+    QueryResults<Entity> run1 = datastore.run(Query.newEntityQueryBuilder().setKind(KIND1).build());
     run1.next();
     Cursor cursor1 = run1.getCursorAfter();
     assertNotNull(cursor1);
     QueryResults<Entity> run2 =
-        DATASTORE.run(Query.newEntityQueryBuilder().setKind(KIND1).setStartCursor(cursor1).build());
+        datastore.run(Query.newEntityQueryBuilder().setKind(KIND1).setStartCursor(cursor1).build());
     Cursor cursor2 = run2.getCursorAfter();
     assertNotNull(cursor2);
     assertEquals(cursor2, cursor1);
-    DATASTORE.delete(entity1.getKey(), entity2.getKey(), entity3.getKey());
+    datastore.delete(entity1.getKey(), entity2.getKey(), entity3.getKey());
   }
 
   @Test
   public void testQueryWithReadTime() throws InterruptedException {
     Entity entity1 =
         Entity.newBuilder(
-                Key.newBuilder(PROJECT_ID, "new_kind", "name-01").setNamespace(NAMESPACE).build())
+                Key.newBuilder(PROJECT_ID, "new_kind", "name-01")
+                    .setDatabaseId(options.getDatabaseId())
+                    .setNamespace(NAMESPACE)
+                    .build())
             .build();
     Entity entity2 =
         Entity.newBuilder(
-                Key.newBuilder(PROJECT_ID, "new_kind", "name-02").setNamespace(NAMESPACE).build())
+                Key.newBuilder(PROJECT_ID, "new_kind", "name-02")
+                    .setDatabaseId(options.getDatabaseId())
+                    .setNamespace(NAMESPACE)
+                    .build())
             .build();
     Entity entity3 =
         Entity.newBuilder(
-                Key.newBuilder(PROJECT_ID, "new_kind", "name-03").setNamespace(NAMESPACE).build())
+                Key.newBuilder(PROJECT_ID, "new_kind", "name-03")
+                    .setDatabaseId(options.getDatabaseId())
+                    .setNamespace(NAMESPACE)
+                    .build())
             .build();
 
-    DATASTORE.put(entity1, entity2);
+    datastore.put(entity1, entity2);
     Thread.sleep(1000);
     Timestamp now = Timestamp.now();
     Thread.sleep(1000);
-    DATASTORE.put(entity3);
+    datastore.put(entity3);
 
     try {
       Query<Entity> query = Query.newEntityQueryBuilder().setKind("new_kind").build();
 
-      QueryResults<Entity> withoutReadTime = DATASTORE.run(query);
+      QueryResults<Entity> withoutReadTime = datastore.run(query);
       assertTrue(withoutReadTime.hasNext());
       assertEquals(entity1, withoutReadTime.next());
       assertTrue(withoutReadTime.hasNext());
@@ -1347,14 +1405,14 @@ public class ITDatastoreTest {
       assertEquals(entity3, withoutReadTime.next());
       assertFalse(withoutReadTime.hasNext());
 
-      QueryResults<Entity> withReadTime = DATASTORE.run(query, ReadOption.readTime(now));
+      QueryResults<Entity> withReadTime = datastore.run(query, ReadOption.readTime(now));
       assertTrue(withReadTime.hasNext());
       assertEquals(entity1, withReadTime.next());
       assertTrue(withReadTime.hasNext());
       assertEquals(entity2, withReadTime.next());
       assertFalse(withReadTime.hasNext());
     } finally {
-      DATASTORE.delete(entity1.getKey(), entity2.getKey(), entity3.getKey());
+      datastore.delete(entity1.getKey(), entity2.getKey(), entity3.getKey());
     }
   }
 
@@ -1364,7 +1422,7 @@ public class ITDatastoreTest {
     AggregationQuery aggregationQuery = builder.build();
     String alias = "total_count";
 
-    Long countBeforeAdd = getOnlyElement(DATASTORE.runAggregation(aggregationQuery)).get(alias);
+    Long countBeforeAdd = getOnlyElement(datastore.runAggregation(aggregationQuery)).get(alias);
     long expectedCount = countBeforeAdd + 1;
 
     Entity newEntity =
@@ -1374,12 +1432,12 @@ public class ITDatastoreTest {
             .set("partial1", PARTIAL_ENTITY2)
             .set("partial2", ENTITY2)
             .build();
-    DATASTORE.put(newEntity);
+    datastore.put(newEntity);
 
-    Long countAfterAdd = getOnlyElement(DATASTORE.runAggregation(aggregationQuery)).get(alias);
+    Long countAfterAdd = getOnlyElement(datastore.runAggregation(aggregationQuery)).get(alias);
     assertThat(countAfterAdd).isEqualTo(expectedCount);
 
-    DATASTORE.delete(newEntity.getKey());
+    datastore.delete(newEntity.getKey());
   }
 
   private void testCountAggregationWithLimit(
@@ -1392,7 +1450,7 @@ public class ITDatastoreTest {
     withoutLimitConfigurer.accept(withoutLimitBuilder);
 
     Long currentCount =
-        getOnlyElement(DATASTORE.runAggregation(withoutLimitBuilder.build())).get(alias);
+        getOnlyElement(datastore.runAggregation(withoutLimitBuilder.build())).get(alias);
     long limit = currentCount - 1;
 
     AggregationQuery.Builder withLimitBuilder =
@@ -1400,7 +1458,7 @@ public class ITDatastoreTest {
     withLimitConfigurer.accept(withLimitBuilder, limit);
 
     Long countWithLimit =
-        getOnlyElement(DATASTORE.runAggregation(withLimitBuilder.build())).get(alias);
+        getOnlyElement(datastore.runAggregation(withLimitBuilder.build())).get(alias);
     assertThat(countWithLimit).isEqualTo(limit);
   }
 
@@ -1408,25 +1466,34 @@ public class ITDatastoreTest {
       throws InterruptedException {
     Entity entity1 =
         Entity.newBuilder(
-                Key.newBuilder(PROJECT_ID, "new_kind", "name-01").setNamespace(NAMESPACE).build())
+                Key.newBuilder(PROJECT_ID, "new_kind", "name-01")
+                    .setDatabaseId(options.getDatabaseId())
+                    .setNamespace(NAMESPACE)
+                    .build())
             .set("name", "Tyrion Lannister")
             .build();
     Entity entity2 =
         Entity.newBuilder(
-                Key.newBuilder(PROJECT_ID, "new_kind", "name-02").setNamespace(NAMESPACE).build())
+                Key.newBuilder(PROJECT_ID, "new_kind", "name-02")
+                    .setDatabaseId(options.getDatabaseId())
+                    .setNamespace(NAMESPACE)
+                    .build())
             .set("name", "Jaime Lannister")
             .build();
     Entity entity3 =
         Entity.newBuilder(
-                Key.newBuilder(PROJECT_ID, "new_kind", "name-03").setNamespace(NAMESPACE).build())
+                Key.newBuilder(PROJECT_ID, "new_kind", "name-03")
+                    .setDatabaseId(options.getDatabaseId())
+                    .setNamespace(NAMESPACE)
+                    .build())
             .set("name", "Cersei Lannister")
             .build();
 
-    DATASTORE.put(entity1, entity2);
+    datastore.put(entity1, entity2);
     Thread.sleep(1000);
     Timestamp now = Timestamp.now();
     Thread.sleep(1000);
-    DATASTORE.put(entity3);
+    datastore.put(entity3);
 
     try {
       AggregationQuery.Builder builder = Query.newAggregationQueryBuilder().setNamespace(NAMESPACE);
@@ -1434,15 +1501,15 @@ public class ITDatastoreTest {
       AggregationQuery countAggregationQuery = builder.build();
 
       Long latestCount =
-          getOnlyElement(DATASTORE.runAggregation(countAggregationQuery)).get("total_count");
+          getOnlyElement(datastore.runAggregation(countAggregationQuery)).get("total_count");
       assertThat(latestCount).isEqualTo(3L);
 
       Long oldCount =
-          getOnlyElement(DATASTORE.runAggregation(countAggregationQuery, ReadOption.readTime(now)))
+          getOnlyElement(datastore.runAggregation(countAggregationQuery, ReadOption.readTime(now)))
               .get("total_count");
       assertThat(oldCount).isEqualTo(2L);
     } finally {
-      DATASTORE.delete(entity1.getKey(), entity2.getKey(), entity3.getKey());
+      datastore.delete(entity1.getKey(), entity2.getKey(), entity3.getKey());
     }
   }
 }
diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITDatastoreTestWithNamedDb.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITDatastoreTestWithNamedDb.java
new file mode 100644
index 00000000..ae74dd4c
--- /dev/null
+++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITDatastoreTestWithNamedDb.java
@@ -0,0 +1,1512 @@
+/*
+ * 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.cloud.datastore.it;
+
+import static com.google.cloud.datastore.aggregation.Aggregation.count;
+import static com.google.common.collect.Iterables.getOnlyElement;
+import static com.google.common.truth.Truth.assertThat;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNotEquals;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertSame;
+import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.fail;
+
+import com.google.cloud.Timestamp;
+import com.google.cloud.datastore.AggregationQuery;
+import com.google.cloud.datastore.Batch;
+import com.google.cloud.datastore.BooleanValue;
+import com.google.cloud.datastore.Cursor;
+import com.google.cloud.datastore.Datastore;
+import com.google.cloud.datastore.Datastore.TransactionCallable;
+import com.google.cloud.datastore.DatastoreException;
+import com.google.cloud.datastore.DatastoreOptions;
+import com.google.cloud.datastore.DatastoreReaderWriter;
+import com.google.cloud.datastore.Entity;
+import com.google.cloud.datastore.EntityQuery;
+import com.google.cloud.datastore.EntityValue;
+import com.google.cloud.datastore.FullEntity;
+import com.google.cloud.datastore.GqlQuery;
+import com.google.cloud.datastore.IncompleteKey;
+import com.google.cloud.datastore.Key;
+import com.google.cloud.datastore.KeyFactory;
+import com.google.cloud.datastore.KeyValue;
+import com.google.cloud.datastore.LatLng;
+import com.google.cloud.datastore.LatLngValue;
+import com.google.cloud.datastore.ListValue;
+import com.google.cloud.datastore.NullValue;
+import com.google.cloud.datastore.PathElement;
+import com.google.cloud.datastore.ProjectionEntity;
+import com.google.cloud.datastore.Query;
+import com.google.cloud.datastore.Query.ResultType;
+import com.google.cloud.datastore.QueryResults;
+import com.google.cloud.datastore.ReadOption;
+import com.google.cloud.datastore.StringValue;
+import com.google.cloud.datastore.StructuredQuery;
+import com.google.cloud.datastore.StructuredQuery.OrderBy;
+import com.google.cloud.datastore.StructuredQuery.PropertyFilter;
+import com.google.cloud.datastore.TimestampValue;
+import com.google.cloud.datastore.Transaction;
+import com.google.cloud.datastore.ValueType;
+import com.google.cloud.datastore.testing.RemoteDatastoreHelper;
+import com.google.common.base.Preconditions;
+import com.google.common.collect.ImmutableList;
+import com.google.datastore.v1.TransactionOptions;
+import com.google.datastore.v1.TransactionOptions.ReadOnly;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.Iterator;
+import java.util.List;
+import java.util.Set;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.function.BiConsumer;
+import java.util.function.Consumer;
+import org.junit.After;
+import org.junit.AfterClass;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.Timeout;
+
+// todo: combine this with ITDatastoreTest once we resolve issues with parameterized graalvm tests
+public class ITDatastoreTestWithNamedDb {
+
+  private static final RemoteDatastoreHelper HELPER = RemoteDatastoreHelper.create();
+  private static final DatastoreOptions OPTIONS_1 = HELPER.getOptions();
+  private static final Datastore DATASTORE_1 = OPTIONS_1.getService();
+
+  private static final String CUSTOM_DB_ID = "test-db";
+  private static final RemoteDatastoreHelper HELPER2 = RemoteDatastoreHelper.create(CUSTOM_DB_ID);
+  private static final DatastoreOptions OPTIONS_2 = HELPER2.getOptions();
+  private static final Datastore DATASTORE_2 = OPTIONS_2.getService();
+
+  private final DatastoreOptions options = OPTIONS_2;
+  private final Datastore datastore = DATASTORE_2;
+
+  private static final String KIND1 = "kind1";
+  private static final String KIND2 = "kind2";
+  private static final String KIND3 = "kind3";
+  private static final NullValue NULL_VALUE = NullValue.of();
+  private static final StringValue STR_VALUE = StringValue.of("str");
+  private static final BooleanValue BOOL_VALUE =
+      BooleanValue.newBuilder(false).setExcludeFromIndexes(true).build();
+  private static final ListValue EMPTY_LIST_VALUE = ListValue.of(Collections.emptyList());
+  private static final ListValue LIST_VALUE1 =
+      ListValue.newBuilder().addValue(NULL_VALUE).addValue(STR_VALUE, BOOL_VALUE).build();
+  private static final TimestampValue TIMESTAMP_VALUE = new TimestampValue(Timestamp.now());
+  private static final LatLngValue LAT_LNG_VALUE =
+      new LatLngValue(LatLng.of(37.422035, -122.084124));
+
+  private static String PROJECT_ID;
+  private static String NAMESPACE;
+  private static Key ROOT_KEY;
+  private static IncompleteKey INCOMPLETE_KEY1;
+  private static ListValue LIST_VALUE2;
+  private static Key KEY1;
+  private static Key KEY2;
+  private static Key KEY3;
+  private static Key KEY4;
+  private static Key KEY5;
+  private static FullEntity<IncompleteKey> PARTIAL_ENTITY1;
+  private static FullEntity<IncompleteKey> PARTIAL_ENTITY2;
+  private static FullEntity<IncompleteKey> PARTIAL_ENTITY3;
+  private static Entity ENTITY1;
+  private static Entity ENTITY2;
+  private static Entity ENTITY3;
+
+  @Rule public Timeout globalTimeout = Timeout.seconds(100);
+
+  @Rule public MultipleAttemptsRule multipleAttemptsRule = new MultipleAttemptsRule(3);
+
+  @AfterClass
+  public static void afterClass() {
+    HELPER.deleteNamespace();
+  }
+
+  @Before
+  public void setUp() {
+    PROJECT_ID = this.options.getProjectId();
+    NAMESPACE = this.options.getNamespace();
+
+    ROOT_KEY =
+        Key.newBuilder(PROJECT_ID, "rootkey", "default", options.getDatabaseId())
+            .setNamespace(NAMESPACE)
+            .build();
+    INCOMPLETE_KEY1 = IncompleteKey.newBuilder(ROOT_KEY, KIND1).setNamespace(NAMESPACE).build();
+
+    IncompleteKey INCOMPLETE_KEY2 =
+        IncompleteKey.newBuilder(PROJECT_ID, KIND2)
+            .setDatabaseId(options.getDatabaseId())
+            .setNamespace(NAMESPACE)
+            .build();
+
+    KEY1 = Key.newBuilder(INCOMPLETE_KEY1, "name").build();
+    KEY2 = Key.newBuilder(KEY1, KIND2, 1).build();
+    KEY3 = Key.newBuilder(KEY2).setName("bla").setNamespace(NAMESPACE).build();
+    KEY4 = Key.newBuilder(KEY2).setName("newName1").setNamespace(NAMESPACE).build();
+    KEY5 = Key.newBuilder(KEY2).setName("newName2").setNamespace(NAMESPACE).build();
+
+    LIST_VALUE2 = ListValue.of(Collections.singletonList(KeyValue.of(KEY1)));
+
+    PARTIAL_ENTITY1 =
+        FullEntity.newBuilder(INCOMPLETE_KEY2)
+            .set("str", STR_VALUE)
+            .set("bool", BOOL_VALUE)
+            .set("list", LIST_VALUE1)
+            .build();
+    PARTIAL_ENTITY2 =
+        FullEntity.newBuilder(PARTIAL_ENTITY1)
+            .remove("str")
+            .set("bool", true)
+            .set("list", LIST_VALUE1.get())
+            .build();
+    PARTIAL_ENTITY3 =
+        FullEntity.newBuilder(PARTIAL_ENTITY1)
+            .setKey(
+                IncompleteKey.newBuilder(PROJECT_ID, KIND3)
+                    .setDatabaseId(options.getDatabaseId())
+                    .build())
+            .build();
+    ENTITY1 =
+        Entity.newBuilder(KEY1)
+            .set("str", STR_VALUE)
+            .set("date", TIMESTAMP_VALUE)
+            .set("latLng", LAT_LNG_VALUE)
+            .set("bool", BOOL_VALUE)
+            .set("partial1", EntityValue.of(PARTIAL_ENTITY1))
+            .set("list", LIST_VALUE2)
+            .set("emptyList", EMPTY_LIST_VALUE)
+            .build();
+    ENTITY2 =
+        Entity.newBuilder(ENTITY1)
+            .setKey(KEY2)
+            .remove("str")
+            .set("name", "Dan")
+            .setNull("null")
+            .set("age", 20)
+            .build();
+    ENTITY3 =
+        Entity.newBuilder(ENTITY1)
+            .setKey(KEY3)
+            .remove("str")
+            .set("null", NULL_VALUE)
+            .set("partial1", PARTIAL_ENTITY2)
+            .set("partial2", ENTITY2)
+            .build();
+
+    datastore.put(ENTITY1, ENTITY2);
+  }
+
+  @After
+  public void tearDown() {
+    EntityQuery allEntitiesQuery = Query.newEntityQueryBuilder().build();
+    QueryResults<Entity> allEntities = datastore.run(allEntitiesQuery);
+    Key[] keysToDelete =
+        ImmutableList.copyOf(allEntities).stream().map(Entity::getKey).toArray(Key[]::new);
+    datastore.delete(keysToDelete);
+  }
+
+  private <T> Iterator<T> getStronglyConsistentResults(Query scQuery, Query query)
+      throws InterruptedException {
+    // scQuery is equivalent to query, but with an ancestor filter in it
+    // this makes scQuery strongly consistent
+    QueryResults<T> scResults = datastore.run(scQuery);
+    List<T> scResultsCopy = makeResultsCopy(scResults);
+    Set<T> scResultsSet = new HashSet<>(scResultsCopy);
+    int maxAttempts = 20;
+
+    while (maxAttempts > 0) {
+      --maxAttempts;
+      QueryResults<T> results = datastore.run(query);
+      List<T> resultsCopy = makeResultsCopy(results);
+      Set<T> resultsSet = new HashSet<>(resultsCopy);
+      if (scResultsSet.size() == resultsSet.size() && scResultsSet.containsAll(resultsSet)) {
+        return resultsCopy.iterator();
+      }
+      Thread.sleep(500);
+    }
+
+    throw new RuntimeException(
+        "reached max number of attempts to get strongly consistent results.");
+  }
+
+  private <T> List<T> makeResultsCopy(QueryResults<T> scResults) {
+    Preconditions.checkNotNull(scResults);
+    List<T> results = new ArrayList<>();
+    while (scResults.hasNext()) {
+      results.add(scResults.next());
+    }
+    return results;
+  }
+
+  @Test
+  public void testNewTransactionCommit() {
+    Transaction transaction = datastore.newTransaction();
+    transaction.add(ENTITY3);
+    Entity entity2 = Entity.newBuilder(ENTITY2).clear().setNull("bla").build();
+    transaction.update(entity2);
+    transaction.delete(KEY1);
+    transaction.commit();
+    assertFalse(transaction.isActive());
+
+    List<Entity> list = datastore.fetch(KEY1, KEY2, KEY3);
+    assertNull(list.get(0));
+    assertEquals(entity2, list.get(1));
+    assertEquals(ENTITY3, list.get(2));
+    assertEquals(3, list.size());
+
+    try {
+      transaction.commit();
+      fail("Expecting a failure");
+    } catch (DatastoreException expected) {
+      assertEquals("FAILED_PRECONDITION", expected.getReason());
+    }
+
+    try {
+      transaction.rollback();
+      fail("Expecting a failure");
+    } catch (DatastoreException expected) {
+      assertEquals("FAILED_PRECONDITION", expected.getReason());
+    }
+  }
+
+  @Test
+  public void testTransactionWithRead() {
+    Transaction transaction = datastore.newTransaction();
+    assertNull(transaction.get(KEY3));
+    transaction.add(ENTITY3);
+    transaction.commit();
+    assertEquals(ENTITY3, datastore.get(KEY3));
+
+    transaction = datastore.newTransaction();
+    assertEquals(ENTITY3, transaction.get(KEY3));
+    // update entity3 during the transaction
+    datastore.put(Entity.newBuilder(ENTITY2).clear().set("from", "datastore").build());
+    transaction.update(Entity.newBuilder(ENTITY2).clear().set("from", "transaction").build());
+    try {
+      transaction.commit();
+      fail("Expecting a failure");
+    } catch (DatastoreException expected) {
+      assertEquals("ABORTED", expected.getReason());
+    }
+  }
+
+  @Test
+  public void testTransactionWithQuery() {
+    Query<Entity> query =
+        Query.newEntityQueryBuilder()
+            .setKind(KIND2)
+            .setFilter(PropertyFilter.hasAncestor(KEY2))
+            .setNamespace(NAMESPACE)
+            .build();
+    Transaction transaction = datastore.newTransaction();
+    QueryResults<Entity> results = transaction.run(query);
+    assertTrue(results.hasNext());
+    assertEquals(ENTITY2, results.next());
+    assertFalse(results.hasNext());
+    transaction.add(ENTITY3);
+    transaction.commit();
+    assertEquals(ENTITY3, datastore.get(KEY3));
+
+    transaction = datastore.newTransaction();
+    results = transaction.run(query);
+    assertTrue(results.hasNext());
+    assertEquals(ENTITY2, results.next());
+    assertFalse(results.hasNext());
+    transaction.delete(ENTITY3.getKey());
+    // update entity2 during the transaction
+    datastore.put(Entity.newBuilder(ENTITY2).clear().build());
+    try {
+      transaction.commit();
+      fail("Expecting a failure");
+    } catch (DatastoreException expected) {
+      assertEquals("ABORTED", expected.getReason());
+    }
+  }
+
+  @Test
+  public void testNewTransactionRollback() {
+    Transaction transaction = datastore.newTransaction();
+    transaction.add(ENTITY3);
+    Entity entity2 =
+        Entity.newBuilder(ENTITY2)
+            .clear()
+            .setNull("bla")
+            .set("list3", StringValue.of("bla"), StringValue.newBuilder("bla").build())
+            .build();
+    transaction.update(entity2);
+    transaction.delete(KEY1);
+    transaction.rollback();
+    transaction.rollback(); // should be safe to repeat rollback calls
+
+    try {
+      transaction.commit();
+      fail("Expecting a failure");
+    } catch (DatastoreException expected) {
+      assertEquals("FAILED_PRECONDITION", expected.getReason());
+    }
+
+    List<Entity> list = datastore.fetch(KEY1, KEY2, KEY3);
+    assertEquals(ENTITY1, list.get(0));
+    assertEquals(ENTITY2, list.get(1));
+    assertNull(list.get(2));
+    assertEquals(3, list.size());
+  }
+
+  @Test
+  public void testNewBatch() {
+    Batch batch = datastore.newBatch();
+    Entity entity1 = Entity.newBuilder(ENTITY1).clear().build();
+    Entity entity2 = Entity.newBuilder(ENTITY2).clear().setNull("bla").build();
+    Entity entity4 = Entity.newBuilder(KEY4).set("value", StringValue.of("value")).build();
+    Entity entity5 = Entity.newBuilder(KEY5).set("value", "value").build();
+
+    List<Entity> entities = batch.add(entity4, PARTIAL_ENTITY2, entity5);
+    Entity entity6 = entities.get(1);
+    assertSame(entity4, entities.get(0));
+    assertEquals(PARTIAL_ENTITY2.getNames(), entity6.getNames());
+    assertEquals(PARTIAL_ENTITY2.getKey().getProjectId(), entity6.getKey().getProjectId());
+    assertEquals(PARTIAL_ENTITY2.getKey().getNamespace(), entity6.getKey().getNamespace());
+    assertEquals(PARTIAL_ENTITY2.getKey().getAncestors(), entity6.getKey().getAncestors());
+    assertEquals(PARTIAL_ENTITY2.getKey().getKind(), entity6.getKey().getKind());
+    assertEquals(PARTIAL_ENTITY2.getKey(), IncompleteKey.newBuilder(entity6.getKey()).build());
+    assertEquals(PARTIAL_ENTITY2.getKey().getAncestors(), entity6.getKey().getAncestors());
+    assertNotEquals(PARTIAL_ENTITY2.getKey(), entity6.getKey());
+    assertSame(entity5, entities.get(2));
+    batch.addWithDeferredIdAllocation(PARTIAL_ENTITY3);
+    batch.put(ENTITY3, entity1, entity2);
+
+    Batch.Response response = batch.submit();
+    entities =
+        datastore.fetch(KEY1, KEY2, KEY3, entity4.getKey(), entity5.getKey(), entity6.getKey());
+    assertEquals(entity1, entities.get(0));
+    assertEquals(entity2, entities.get(1));
+    assertEquals(ENTITY3, entities.get(2));
+    assertEquals(entity4, entities.get(3));
+    assertEquals(entity5, entities.get(4));
+    assertEquals(entity6, entities.get(5));
+    assertEquals(6, entities.size());
+    List<Key> generatedKeys = response.getGeneratedKeys();
+    assertEquals(1, generatedKeys.size());
+    assertEquals(PARTIAL_ENTITY3.getNames(), datastore.get(generatedKeys.get(0)).getNames());
+    assertEquals(PARTIAL_ENTITY3.getKey(), IncompleteKey.newBuilder(generatedKeys.get(0)).build());
+
+    try {
+      batch.submit();
+      fail("Expecting a failure");
+    } catch (DatastoreException expected) {
+      assertEquals("FAILED_PRECONDITION", expected.getReason());
+    }
+
+    batch = datastore.newBatch();
+    batch.delete(entity4.getKey(), entity5.getKey(), entity6.getKey());
+    batch.update(ENTITY1, ENTITY2, ENTITY3);
+    batch.submit();
+    entities =
+        datastore.fetch(KEY1, KEY2, KEY3, entity4.getKey(), entity5.getKey(), entity6.getKey());
+    assertEquals(ENTITY1, entities.get(0));
+    assertEquals(ENTITY2, entities.get(1));
+    assertEquals(ENTITY3, entities.get(2));
+    assertNull(entities.get(3));
+    assertNull(entities.get(4));
+    assertNull(entities.get(5));
+    assertEquals(6, entities.size());
+  }
+
+  @Test
+  public void testRunGqlQueryNoCasting() throws InterruptedException {
+    Query<Entity> query1 =
+        Query.newGqlQueryBuilder(ResultType.ENTITY, "select * from " + KIND1)
+            .setNamespace(NAMESPACE)
+            .build();
+    Query<Entity> scQuery1 =
+        Query.newEntityQueryBuilder()
+            .setNamespace(NAMESPACE)
+            .setKind(KIND1)
+            .setFilter(PropertyFilter.hasAncestor(ROOT_KEY))
+            .build();
+
+    Iterator<Entity> results1 = getStronglyConsistentResults(scQuery1, query1);
+
+    assertTrue(results1.hasNext());
+    assertEquals(ENTITY1, results1.next());
+    assertFalse(results1.hasNext());
+
+    datastore.put(ENTITY3);
+    Query<? extends Entity> query2 =
+        Query.newGqlQueryBuilder(ResultType.ENTITY, "select * from " + KIND2 + " order by __key__")
+            .setNamespace(NAMESPACE)
+            .build();
+    Query<? extends Entity> scQuery2 =
+        Query.newEntityQueryBuilder()
+            .setNamespace(NAMESPACE)
+            .setKind(KIND2)
+            .setFilter(PropertyFilter.hasAncestor(ROOT_KEY))
+            .setOrderBy(OrderBy.asc("__key__"))
+            .build();
+
+    Iterator<Entity> results2 = getStronglyConsistentResults(scQuery2, query2);
+    assertTrue(results2.hasNext());
+    assertEquals(ENTITY2, results2.next());
+    assertTrue(results2.hasNext());
+    assertEquals(ENTITY3, results2.next());
+    assertFalse(results2.hasNext());
+
+    query1 =
+        Query.newGqlQueryBuilder(ResultType.ENTITY, "select * from bla")
+            .setNamespace(NAMESPACE)
+            .build();
+    scQuery1 =
+        Query.newEntityQueryBuilder()
+            .setNamespace(NAMESPACE)
+            .setFilter(PropertyFilter.hasAncestor(ROOT_KEY))
+            .setKind("bla")
+            .build();
+    results1 = getStronglyConsistentResults(scQuery1, query1);
+    assertFalse(results1.hasNext());
+
+    Query<Key> keyOnlyQuery =
+        Query.newGqlQueryBuilder(ResultType.KEY, "select __key__ from " + KIND1)
+            .setNamespace(NAMESPACE)
+            .build();
+    Query<Key> scKeyOnlyQuery =
+        Query.newKeyQueryBuilder()
+            .setNamespace(NAMESPACE)
+            .setFilter(PropertyFilter.hasAncestor(ROOT_KEY))
+            .setKind(KIND1)
+            .build();
+    Iterator<Key> keyOnlyResults = getStronglyConsistentResults(scKeyOnlyQuery, keyOnlyQuery);
+    assertTrue(keyOnlyResults.hasNext());
+    assertEquals(KEY1, keyOnlyResults.next());
+    assertFalse(keyOnlyResults.hasNext());
+
+    GqlQuery<ProjectionEntity> keyProjectionQuery =
+        Query.newGqlQueryBuilder(ResultType.PROJECTION_ENTITY, "select __key__ from " + KIND1)
+            .setNamespace(NAMESPACE)
+            .build();
+    Query<ProjectionEntity> scKeyProjectionQuery =
+        Query.newProjectionEntityQueryBuilder()
+            .addProjection("__key__")
+            .setNamespace(NAMESPACE)
+            .setKind(KIND1)
+            .setFilter(PropertyFilter.hasAncestor(ROOT_KEY))
+            .build();
+
+    Iterator<ProjectionEntity> keyProjectionResult =
+        getStronglyConsistentResults(scKeyProjectionQuery, keyProjectionQuery);
+    assertTrue(keyProjectionResult.hasNext());
+    ProjectionEntity projectionEntity = keyProjectionResult.next();
+    assertEquals(KEY1, projectionEntity.getKey());
+    assertTrue(projectionEntity.getNames().isEmpty());
+    assertFalse(keyProjectionResult.hasNext());
+    datastore.delete(ENTITY3.getKey());
+  }
+
+  @Test
+  public void testRunGqlQueryWithCasting() throws InterruptedException {
+    @SuppressWarnings("unchecked")
+    Query<Entity> query1 =
+        (Query<Entity>)
+            Query.newGqlQueryBuilder("select * from " + KIND1).setNamespace(NAMESPACE).build();
+    Query<Entity> scQuery1 =
+        Query.newEntityQueryBuilder()
+            .setNamespace(NAMESPACE)
+            .setKind(KIND1)
+            .setFilter(PropertyFilter.hasAncestor(ROOT_KEY))
+            .build();
+    Iterator<Entity> results1 = getStronglyConsistentResults(scQuery1, query1);
+    assertTrue(results1.hasNext());
+    assertEquals(ENTITY1, results1.next());
+    assertFalse(results1.hasNext());
+
+    Query<?> query2 =
+        Query.newGqlQueryBuilder("select * from " + KIND1).setNamespace(NAMESPACE).build();
+
+    QueryResults<?> results2 = datastore.run(query2);
+
+    assertSame(Entity.class, results2.getResultClass());
+
+    Query<?> scQuery2 =
+        Query.newEntityQueryBuilder()
+            .setNamespace(NAMESPACE)
+            .setKind(KIND1)
+            .setFilter(PropertyFilter.hasAncestor(ROOT_KEY))
+            .build();
+
+    Iterator<Entity> results3 = getStronglyConsistentResults(scQuery2, query2);
+    assertTrue(results3.hasNext());
+    assertEquals(ENTITY1, results3.next());
+    assertFalse(results3.hasNext());
+  }
+
+  @Test
+  public void testRunAggregationQuery() {
+    // verifying aggregation with an entity query
+    testCountAggregationWith(
+        builder ->
+            builder
+                .addAggregation(count().as("total_count"))
+                .over(
+                    Query.newEntityQueryBuilder().setNamespace(NAMESPACE).setKind(KIND1).build()));
+
+    // verifying aggregation with a projection query
+    testCountAggregationWith(
+        builder ->
+            builder
+                .addAggregation(count().as("total_count"))
+                .over(
+                    Query.newProjectionEntityQueryBuilder()
+                        .setProjection("str")
+                        .setNamespace(NAMESPACE)
+                        .setKind(KIND1)
+                        .build()));
+
+    // verifying aggregation with a key query
+    testCountAggregationWith(
+        builder ->
+            builder
+                .addAggregation(count().as("total_count"))
+                .over(Query.newKeyQueryBuilder().setNamespace(NAMESPACE).setKind(KIND1).build()));
+
+    // verifying aggregation with a GQL query
+    testCountAggregationWith(
+        builder ->
+            builder.over(
+                Query.newGqlQueryBuilder(
+                        "AGGREGATE COUNT(*) AS total_count OVER (SELECT * FROM kind1)")
+                    .setNamespace(NAMESPACE)
+                    .build()));
+  }
+
+  @Test
+  public void testRunAggregationQueryWithLimit() {
+    // verifying aggregation with an entity query
+    testCountAggregationWithLimit(
+        builder ->
+            builder
+                .addAggregation(count().as("total_count"))
+                .over(Query.newEntityQueryBuilder().setNamespace(NAMESPACE).setKind(KIND1).build()),
+        ((builder, limit) ->
+            builder
+                .addAggregation(count().as("total_count"))
+                .over(
+                    Query.newEntityQueryBuilder()
+                        .setNamespace(NAMESPACE)
+                        .setKind(KIND1)
+                        .setLimit(limit.intValue())
+                        .build())));
+
+    // verifying aggregation with a projection query
+    testCountAggregationWithLimit(
+        builder ->
+            builder
+                .addAggregation(count().as("total_count"))
+                .over(
+                    Query.newProjectionEntityQueryBuilder()
+                        .setProjection("str")
+                        .setNamespace(NAMESPACE)
+                        .setKind(KIND1)
+                        .build()),
+        ((builder, limit) ->
+            builder
+                .addAggregation(count().as("total_count"))
+                .over(
+                    Query.newProjectionEntityQueryBuilder()
+                        .setProjection("str")
+                        .setNamespace(NAMESPACE)
+                        .setKind(KIND1)
+                        .setLimit(limit.intValue())
+                        .build())));
+
+    // verifying aggregation with a key query
+    testCountAggregationWithLimit(
+        builder ->
+            builder
+                .addAggregation(count().as("total_count"))
+                .over(Query.newKeyQueryBuilder().setNamespace(NAMESPACE).setKind(KIND1).build()),
+        (builder, limit) ->
+            builder
+                .addAggregation(count().as("total_count"))
+                .over(
+                    Query.newKeyQueryBuilder()
+                        .setNamespace(NAMESPACE)
+                        .setKind(KIND1)
+                        .setLimit(limit.intValue())
+                        .build()));
+
+    // verifying aggregation with a GQL query
+    testCountAggregationWithLimit(
+        builder ->
+            builder.over(
+                Query.newGqlQueryBuilder(
+                        "AGGREGATE COUNT(*) AS total_count OVER (SELECT * FROM kind1)")
+                    .setNamespace(NAMESPACE)
+                    .build()),
+        (builder, limit) ->
+            builder.over(
+                Query.newGqlQueryBuilder(
+                        "AGGREGATE COUNT(*) AS total_count OVER (SELECT * FROM kind1 LIMIT @limit)")
+                    .setNamespace(NAMESPACE)
+                    .setBinding("limit", limit)
+                    .build()));
+  }
+
+  /**
+   * if an entity is modified or deleted within a transaction, a query or lookup returns the
+   * original version of the entity as of the beginning of the transaction, or nothing if the entity
+   * did not exist then.
+   *
+   * @see <a
+   *     href="https://cloud.google.com/datastore/docs/concepts/transactions#isolation_and_consistency">
+   *     Source</a>
+   */
+  @Test
+  public void testRunAggregationQueryInTransactionShouldReturnAConsistentSnapshot() {
+    Key newEntityKey = Key.newBuilder(KEY1, "newKind", "name-01").build();
+    EntityQuery entityQuery =
+        Query.newEntityQueryBuilder()
+            .setNamespace(NAMESPACE)
+            .setFilter(PropertyFilter.hasAncestor(KEY1))
+            .build();
+
+    AggregationQuery aggregationQuery =
+        Query.newAggregationQueryBuilder()
+            .setNamespace(NAMESPACE)
+            .over(entityQuery)
+            .addAggregation(count().as("count"))
+            .build();
+
+    // original entity count is 2
+    assertThat(getOnlyElement(datastore.runAggregation(aggregationQuery)).get("count"))
+        .isEqualTo(2L);
+
+    // FIRST TRANSACTION
+    datastore.runInTransaction(
+        (TransactionCallable<Void>)
+            inFirstTransaction -> {
+              // creating a new entity
+              Entity aNewEntity =
+                  Entity.newBuilder(ENTITY2).setKey(newEntityKey).set("v_int", 10).build();
+              inFirstTransaction.put(aNewEntity);
+
+              // count remains 2
+              assertThat(
+                      getOnlyElement(inFirstTransaction.runAggregation(aggregationQuery))
+                          .get("count"))
+                  .isEqualTo(2L);
+              assertThat(getOnlyElement(datastore.runAggregation(aggregationQuery)).get("count"))
+                  .isEqualTo(2L);
+              return null;
+            });
+    // after first transaction is committed, count is updated to 3 now.
+    assertThat(getOnlyElement(datastore.runAggregation(aggregationQuery)).get("count"))
+        .isEqualTo(3L);
+
+    // SECOND TRANSACTION
+    datastore.runInTransaction(
+        (TransactionCallable<Void>)
+            inSecondTransaction -> {
+              // deleting ENTITY2
+              inSecondTransaction.delete(ENTITY2.getKey());
+
+              // count remains 3
+              assertThat(
+                      getOnlyElement(inSecondTransaction.runAggregation(aggregationQuery))
+                          .get("count"))
+                  .isEqualTo(3L);
+              assertThat(getOnlyElement(datastore.runAggregation(aggregationQuery)).get("count"))
+                  .isEqualTo(3L);
+              return null;
+            });
+    // after second transaction is committed, count is updated to 2 now.
+    assertThat(getOnlyElement(datastore.runAggregation(aggregationQuery)).get("count"))
+        .isEqualTo(2L);
+    datastore.delete(newEntityKey);
+  }
+
+  @Test
+  public void testRunAggregationQueryInAReadOnlyTransactionShouldNotLockTheCountedDocuments()
+      throws Exception {
+    ExecutorService executor = Executors.newSingleThreadExecutor();
+    EntityQuery entityQuery =
+        Query.newEntityQueryBuilder()
+            .setNamespace(NAMESPACE)
+            .setFilter(PropertyFilter.hasAncestor(KEY1))
+            .build();
+    AggregationQuery aggregationQuery =
+        Query.newAggregationQueryBuilder()
+            .setNamespace(NAMESPACE)
+            .over(entityQuery)
+            .addAggregation(count().as("count"))
+            .build();
+
+    TransactionOptions transactionOptions =
+        TransactionOptions.newBuilder().setReadOnly(ReadOnly.newBuilder().build()).build();
+    Transaction readOnlyTransaction = datastore.newTransaction(transactionOptions);
+
+    // Executing query in transaction
+    assertThat(getOnlyElement(readOnlyTransaction.runAggregation(aggregationQuery)).get("count"))
+        .isEqualTo(2L);
+
+    // Concurrent write task.
+    Future<Void> addNewEntityTaskOutsideTransaction =
+        executor.submit(
+            () -> {
+              Entity aNewEntity =
+                  Entity.newBuilder(ENTITY2)
+                      .setKey(Key.newBuilder(KEY1, "newKind", "name-01").build())
+                      .set("v_int", 10)
+                      .build();
+              datastore.put(aNewEntity);
+              return null;
+            });
+
+    // should not throw exception and complete successfully as the ongoing transaction is read-only.
+    addNewEntityTaskOutsideTransaction.get();
+
+    // cleanup
+    readOnlyTransaction.commit();
+    executor.shutdownNow();
+
+    assertThat(getOnlyElement(datastore.runAggregation(aggregationQuery)).get("count"))
+        .isEqualTo(3L);
+  }
+
+  @Test
+  public void testRunAggregationQueryWithReadTime() throws InterruptedException {
+    String alias = "total_count";
+
+    // verifying aggregation readTime with an entity query
+    testCountAggregationReadTimeWith(
+        builder ->
+            builder
+                .over(Query.newEntityQueryBuilder().setKind("new_kind").build())
+                .addAggregation(count().as(alias)));
+
+    // verifying aggregation readTime with a projection query
+    testCountAggregationReadTimeWith(
+        builder ->
+            builder
+                .over(
+                    Query.newProjectionEntityQueryBuilder()
+                        .setProjection("name")
+                        .setKind("new_kind")
+                        .build())
+                .addAggregation(count().as(alias)));
+
+    // verifying aggregation readTime with a key query
+    testCountAggregationReadTimeWith(
+        builder ->
+            builder
+                .over(Query.newKeyQueryBuilder().setKind("new_kind").build())
+                .addAggregation(count().as(alias)));
+
+    // verifying aggregation readTime with a GQL query
+    testCountAggregationReadTimeWith(
+        builder ->
+            builder
+                .over(
+                    Query.newGqlQueryBuilder(
+                            "AGGREGATE COUNT(*) AS total_count OVER (SELECT * FROM new_kind)")
+                        .build())
+                .addAggregation(count().as(alias)));
+  }
+
+  @Test
+  public void testRunStructuredQuery() throws InterruptedException {
+    Query<Entity> query =
+        Query.newEntityQueryBuilder().setKind(KIND1).setOrderBy(OrderBy.asc("__key__")).build();
+
+    Query<Entity> scQuery =
+        Query.newEntityQueryBuilder()
+            .setKind(KIND1)
+            .setFilter(PropertyFilter.hasAncestor(ROOT_KEY))
+            .setOrderBy(OrderBy.asc("__key__"))
+            .build();
+
+    Iterator<Entity> results1 = getStronglyConsistentResults(scQuery, query);
+
+    assertTrue(results1.hasNext());
+    assertEquals(ENTITY1, results1.next());
+    assertFalse(results1.hasNext());
+
+    Query<Key> keyOnlyQuery = Query.newKeyQueryBuilder().setKind(KIND1).build();
+    Query<Key> scKeyOnlyQuery =
+        Query.newKeyQueryBuilder()
+            .setKind(KIND1)
+            .setFilter(PropertyFilter.hasAncestor(ROOT_KEY))
+            .build();
+
+    Iterator<Key> results2 = getStronglyConsistentResults(scKeyOnlyQuery, keyOnlyQuery);
+    assertTrue(results2.hasNext());
+    assertEquals(ENTITY1.getKey(), results2.next());
+    assertFalse(results2.hasNext());
+
+    StructuredQuery<ProjectionEntity> keyOnlyProjectionQuery =
+        Query.newProjectionEntityQueryBuilder().setKind(KIND1).setProjection("__key__").build();
+    StructuredQuery<ProjectionEntity> scKeyOnlyProjectionQuery =
+        Query.newProjectionEntityQueryBuilder()
+            .setKind(KIND1)
+            .setFilter(PropertyFilter.hasAncestor(ROOT_KEY))
+            .setProjection("__key__")
+            .build();
+    Iterator<ProjectionEntity> results3 =
+        getStronglyConsistentResults(scKeyOnlyProjectionQuery, keyOnlyProjectionQuery);
+    assertTrue(results3.hasNext());
+    ProjectionEntity projectionEntity = results3.next();
+    assertEquals(ENTITY1.getKey(), projectionEntity.getKey());
+    assertTrue(projectionEntity.getNames().isEmpty());
+    assertFalse(results3.hasNext());
+
+    StructuredQuery<ProjectionEntity> projectionQuery =
+        Query.newProjectionEntityQueryBuilder()
+            .setKind(KIND2)
+            .setProjection("age")
+            .setFilter(PropertyFilter.gt("age", 18))
+            .setDistinctOn("age")
+            .setOrderBy(OrderBy.asc("age"))
+            .setLimit(10)
+            .build();
+
+    StructuredQuery<ProjectionEntity> scProjectionQuery =
+        Query.newProjectionEntityQueryBuilder()
+            .setKind(KIND2)
+            .setFilter(PropertyFilter.hasAncestor(ROOT_KEY))
+            .setProjection("age")
+            .setFilter(PropertyFilter.gt("age", 18))
+            .setDistinctOn("age")
+            .setOrderBy(OrderBy.asc("age"))
+            .setLimit(10)
+            .build();
+
+    Iterator<ProjectionEntity> results4 =
+        getStronglyConsistentResults(scProjectionQuery, projectionQuery);
+    assertTrue(results4.hasNext());
+    ProjectionEntity entity = results4.next();
+    assertEquals(ENTITY2.getKey(), entity.getKey());
+    assertEquals(20, entity.getLong("age"));
+    assertEquals(1, entity.getNames().size());
+    assertFalse(results4.hasNext());
+  }
+
+  @Test
+  public void testInNotInNeqFilters() throws InterruptedException {
+    Entity e1 =
+        Entity.newBuilder(ENTITY1)
+            .setKey(Key.newBuilder(INCOMPLETE_KEY1, "e1").build())
+            .set("v_int", 10)
+            .build();
+    Entity e2 =
+        Entity.newBuilder(ENTITY1)
+            .setKey(Key.newBuilder(INCOMPLETE_KEY1, "e2").build())
+            .set("v_int", 20)
+            .build();
+    datastore.put(e1, e2);
+
+    Query<Entity> queryIn =
+        Query.newEntityQueryBuilder()
+            .setKind(KIND1)
+            .setFilter(PropertyFilter.in("v_int", ListValue.of(10, 20)))
+            .build();
+
+    Query<Entity> scQueryIn =
+        Query.newEntityQueryBuilder()
+            .setKind(KIND1)
+            .setFilter(PropertyFilter.hasAncestor(ROOT_KEY))
+            .setFilter(PropertyFilter.in("v_int", ListValue.of(10, 20)))
+            .build();
+
+    Iterator<Entity> resultIn = getStronglyConsistentResults(scQueryIn, queryIn);
+
+    assertTrue(resultIn.hasNext());
+    assertEquals(e1, resultIn.next());
+    assertTrue(resultIn.hasNext());
+    assertEquals(e2, resultIn.next());
+    assertFalse(resultIn.hasNext());
+
+    Query<Entity> queryNotIn =
+        Query.newEntityQueryBuilder()
+            .setKind(KIND1)
+            .setFilter(PropertyFilter.not_in("v_int", ListValue.of(20, 30)))
+            .build();
+
+    Query<Entity> scQueryNotIn =
+        Query.newEntityQueryBuilder()
+            .setKind(KIND1)
+            .setFilter(PropertyFilter.hasAncestor(ROOT_KEY))
+            .setFilter(PropertyFilter.not_in("v_int", ListValue.of(20, 30)))
+            .build();
+
+    Iterator<Entity> resultNotIn = getStronglyConsistentResults(scQueryNotIn, queryNotIn);
+
+    assertTrue(resultNotIn.hasNext());
+    assertEquals(e1, resultNotIn.next());
+    assertFalse(resultNotIn.hasNext());
+
+    Query<Entity> queryNeq =
+        Query.newEntityQueryBuilder()
+            .setKind(KIND1)
+            .setFilter(PropertyFilter.neq("v_int", 10))
+            .build();
+
+    Query<Entity> scQueryNeq =
+        Query.newEntityQueryBuilder()
+            .setKind(KIND1)
+            .setFilter(PropertyFilter.hasAncestor(ROOT_KEY))
+            .setFilter(PropertyFilter.neq("v_int", 10))
+            .build();
+
+    Iterator<Entity> resultNeq = getStronglyConsistentResults(scQueryNeq, queryNeq);
+
+    assertTrue(resultNeq.hasNext());
+    assertEquals(e2, resultNeq.next());
+    assertFalse(resultNeq.hasNext());
+
+    datastore.delete(e1.getKey());
+    datastore.delete(e2.getKey());
+  }
+
+  @Test
+  public void testAllocateId() {
+    KeyFactory keyFactory = datastore.newKeyFactory().setKind(KIND1);
+    IncompleteKey pk1 = keyFactory.newKey();
+    Key key1 = datastore.allocateId(pk1);
+    assertEquals(key1.getProjectId(), pk1.getProjectId());
+    assertEquals(key1.getNamespace(), pk1.getNamespace());
+    assertEquals(key1.getAncestors(), pk1.getAncestors());
+    assertEquals(key1.getKind(), pk1.getKind());
+    assertTrue(key1.hasId());
+    assertFalse(key1.hasName());
+    assertEquals(Key.newBuilder(pk1, key1.getId()).build(), key1);
+
+    Key key2 = datastore.allocateId(pk1);
+    assertNotEquals(key1, key2);
+    assertEquals(Key.newBuilder(pk1, key2.getId()).build(), key2);
+  }
+
+  @Test
+  public void testReserveIds() {
+    KeyFactory keyFactory = datastore.newKeyFactory().setKind("MyKind");
+    Key key1 = keyFactory.newKey(10);
+    Key key2 = keyFactory.newKey("name");
+    List<Key> keyList = datastore.reserveIds(key1, key2);
+    assertEquals(2, keyList.size());
+  }
+
+  @Test
+  public void testAllocateIdArray() {
+    KeyFactory keyFactory = datastore.newKeyFactory().setKind(KIND1);
+    IncompleteKey incompleteKey1 = keyFactory.newKey();
+    IncompleteKey incompleteKey2 =
+        keyFactory
+            .setKind(KIND2)
+            .setDatabaseId(options.getDatabaseId())
+            .addAncestors(PathElement.of(KIND1, 10))
+            .newKey();
+    List<Key> result = datastore.allocateId(incompleteKey1, incompleteKey2, incompleteKey1);
+    assertEquals(3, result.size());
+    assertEquals(Key.newBuilder(incompleteKey1, result.get(0).getId()).build(), result.get(0));
+    assertEquals(Key.newBuilder(incompleteKey1, result.get(2).getId()).build(), result.get(2));
+    assertEquals(Key.newBuilder(incompleteKey2, result.get(1).getId()).build(), result.get(1));
+  }
+
+  @Test
+  public void testGet() {
+    Entity entity = datastore.get(KEY3);
+    assertNull(entity);
+
+    entity = datastore.get(KEY1);
+    assertEquals(ENTITY1, entity);
+    StringValue value1 = entity.getValue("str");
+    assertEquals(STR_VALUE, value1);
+    BooleanValue value2 = entity.getValue("bool");
+    assertEquals(BOOL_VALUE, value2);
+    ListValue value3 = entity.getValue("list");
+    assertEquals(LIST_VALUE2, value3);
+    TimestampValue value4 = entity.getValue("date");
+    assertEquals(TIMESTAMP_VALUE, value4);
+    LatLngValue value5 = entity.getValue("latLng");
+    assertEquals(LAT_LNG_VALUE, value5);
+    FullEntity<IncompleteKey> value6 = entity.getEntity("partial1");
+    assertEquals(PARTIAL_ENTITY1, value6);
+    ListValue value7 = entity.getValue("emptyList");
+    assertEquals(EMPTY_LIST_VALUE, value7);
+    assertEquals(7, entity.getNames().size());
+    assertFalse(entity.contains("bla"));
+  }
+
+  @Test
+  public void testGetWithReadTime() throws InterruptedException {
+    Key key =
+        Key.newBuilder(PROJECT_ID, "new_kind", "name", options.getDatabaseId())
+            .setNamespace(NAMESPACE)
+            .build();
+
+    try {
+      datastore.put(Entity.newBuilder(key).set("str", "old_str_value").build());
+
+      Thread.sleep(1000);
+      Timestamp now = Timestamp.now();
+      Thread.sleep(1000);
+
+      datastore.put(Entity.newBuilder(key).set("str", "new_str_value").build());
+
+      Entity entity = datastore.get(key);
+      StringValue value1 = entity.getValue("str");
+      assertEquals(StringValue.of("new_str_value"), value1);
+
+      entity = datastore.get(key, ReadOption.readTime(now));
+      value1 = entity.getValue("str");
+      assertEquals(StringValue.of("old_str_value"), value1);
+    } finally {
+      datastore.delete(key);
+    }
+  }
+
+  @Test
+  public void testGetArrayNoDeferredResults() {
+    datastore.put(ENTITY3);
+    Iterator<Entity> result =
+        datastore.fetch(KEY1, Key.newBuilder(KEY1).setName("bla").build(), KEY2, KEY3).iterator();
+    assertEquals(ENTITY1, result.next());
+    assertNull(result.next());
+    assertEquals(ENTITY2, result.next());
+    Entity entity3 = result.next();
+    assertEquals(ENTITY3, entity3);
+    assertTrue(entity3.isNull("null"));
+    assertFalse(entity3.getBoolean("bool"));
+    assertEquals(LIST_VALUE2.get(), entity3.getList("list"));
+    FullEntity<IncompleteKey> partial1 = entity3.getEntity("partial1");
+    FullEntity<IncompleteKey> partial2 = entity3.getEntity("partial2");
+    assertEquals(PARTIAL_ENTITY2, partial1);
+    assertEquals(ENTITY2, partial2);
+    assertEquals(ValueType.BOOLEAN, entity3.getValue("bool").getType());
+    assertEquals(LAT_LNG_VALUE, entity3.getValue("latLng"));
+    assertEquals(EMPTY_LIST_VALUE, entity3.getValue("emptyList"));
+    assertEquals(8, entity3.getNames().size());
+    assertFalse(entity3.contains("bla"));
+    try {
+      entity3.getString("str");
+      fail("Expecting a failure");
+    } catch (DatastoreException expected) {
+      // expected - no such property
+    }
+    assertFalse(result.hasNext());
+    datastore.delete(ENTITY3.getKey());
+  }
+
+  @Test
+  public void testAddEntity() {
+    List<Entity> keys = datastore.fetch(ENTITY1.getKey(), ENTITY3.getKey());
+    assertEquals(ENTITY1, keys.get(0));
+    assertNull(keys.get(1));
+    assertEquals(2, keys.size());
+
+    try {
+      datastore.add(ENTITY1);
+      fail("Expecting a failure");
+    } catch (DatastoreException expected) {
+      // expected;
+    }
+
+    List<Entity> entities = datastore.add(ENTITY3, PARTIAL_ENTITY1, PARTIAL_ENTITY2);
+    assertEquals(ENTITY3, datastore.get(ENTITY3.getKey()));
+    assertEquals(ENTITY3, entities.get(0));
+    assertEquals(PARTIAL_ENTITY1.getNames(), entities.get(1).getNames());
+    assertEquals(PARTIAL_ENTITY1.getKey().getAncestors(), entities.get(1).getKey().getAncestors());
+    assertNotNull(datastore.get(entities.get(1).getKey()));
+    assertEquals(PARTIAL_ENTITY2.getNames(), entities.get(2).getNames());
+    assertEquals(PARTIAL_ENTITY2.getKey().getAncestors(), entities.get(2).getKey().getAncestors());
+    assertNotNull(datastore.get(entities.get(2).getKey()));
+    for (Entity entity : entities) {
+      datastore.delete(entity.getKey());
+    }
+  }
+
+  @Test
+  public void testUpdate() {
+    List<Entity> keys = datastore.fetch(ENTITY1.getKey(), ENTITY3.getKey());
+    assertEquals(ENTITY1, keys.get(0));
+    assertNull(keys.get(1));
+    assertEquals(2, keys.size());
+
+    try {
+      datastore.update(ENTITY3);
+      fail("Expecting a failure");
+    } catch (DatastoreException expected) {
+      // expected;
+    }
+    datastore.add(ENTITY3);
+    assertEquals(ENTITY3, datastore.get(ENTITY3.getKey()));
+    Entity entity3 = Entity.newBuilder(ENTITY3).clear().set("bla", new NullValue()).build();
+    assertNotEquals(ENTITY3, entity3);
+    datastore.update(entity3);
+    assertEquals(entity3, datastore.get(ENTITY3.getKey()));
+    datastore.delete(ENTITY3.getKey());
+  }
+
+  @Test
+  public void testPut() {
+    Entity updatedEntity = Entity.newBuilder(ENTITY1).set("new_property", 42L).build();
+    assertEquals(updatedEntity, datastore.put(updatedEntity));
+    assertEquals(updatedEntity, datastore.get(updatedEntity.getKey()));
+
+    Entity entity2 = Entity.newBuilder(ENTITY2).clear().set("bla", new NullValue()).build();
+    assertNotEquals(ENTITY2, entity2);
+    List<Entity> entities = datastore.put(ENTITY1, entity2, ENTITY3, PARTIAL_ENTITY1);
+    assertEquals(ENTITY1, entities.get(0));
+    assertEquals(entity2, entities.get(1));
+    assertEquals(ENTITY3, entities.get(2));
+    assertEquals(PARTIAL_ENTITY1.getNames(), entities.get(3).getNames());
+    assertEquals(PARTIAL_ENTITY1.getKey().getAncestors(), entities.get(3).getKey().getAncestors());
+    assertEquals(ENTITY1, datastore.get(ENTITY1.getKey()));
+    assertEquals(entity2, datastore.get(entity2.getKey()));
+    assertEquals(ENTITY3, datastore.get(ENTITY3.getKey()));
+    Entity entity = datastore.get(entities.get(3).getKey());
+    assertEquals(entities.get(3), entity);
+    for (Entity entityToDelete : entities) {
+      datastore.delete(entityToDelete.getKey());
+    }
+  }
+
+  @Test
+  public void testDelete() {
+    Iterator<Entity> keys =
+        datastore.fetch(ENTITY1.getKey(), ENTITY2.getKey(), ENTITY3.getKey()).iterator();
+    assertEquals(ENTITY1, keys.next());
+    assertEquals(ENTITY2, keys.next());
+    assertNull(keys.next());
+    assertFalse(keys.hasNext());
+    datastore.delete(ENTITY1.getKey(), ENTITY2.getKey(), ENTITY3.getKey());
+    keys = datastore.fetch(ENTITY1.getKey(), ENTITY2.getKey(), ENTITY3.getKey()).iterator();
+    assertNull(keys.next());
+    assertNull(keys.next());
+    assertNull(keys.next());
+    assertFalse(keys.hasNext());
+  }
+
+  @Test
+  public void testRunInTransaction() {
+    TransactionCallable<Integer> callable1 =
+        new TransactionCallable<Integer>() {
+          private Integer attempts = 1;
+
+          @Override
+          public Integer run(DatastoreReaderWriter transaction) {
+            transaction.get(KEY1);
+            if (attempts < 2) {
+              ++attempts;
+              throw new DatastoreException(10, "", "ABORTED", false, null);
+            }
+
+            return attempts;
+          }
+        };
+
+    int result = datastore.runInTransaction(callable1);
+    assertEquals(result, 2);
+
+    TransactionCallable<Integer> callable2 =
+        new TransactionCallable<Integer>() {
+          private Integer attempts = 1;
+
+          @Override
+          public Integer run(DatastoreReaderWriter transaction) {
+            transaction.get(KEY1);
+            if (attempts < 2) {
+              ++attempts;
+              throw new DatastoreException(4, "", "DEADLINE_EXCEEDED", false, null);
+            }
+            return attempts;
+          }
+        };
+
+    try {
+      datastore.runInTransaction(callable2);
+      fail("Expecting a failure");
+    } catch (DatastoreException expected) {
+      assertEquals(4, ((DatastoreException) expected.getCause()).getCode());
+    }
+  }
+
+  @Test
+  public void testRunInTransactionReadWrite() {
+
+    final Entity entity1 = Entity.newBuilder(ENTITY1).clear().setNull("bla").build();
+
+    TransactionCallable<Integer> callable1 =
+        new TransactionCallable<Integer>() {
+          private Integer attempts = 1;
+
+          @Override
+          public Integer run(DatastoreReaderWriter transaction) {
+            transaction.update(entity1);
+            if (attempts < 2) {
+              ++attempts;
+              throw new DatastoreException(10, "", "ABORTED", false, null);
+            }
+
+            return attempts;
+          }
+        };
+
+    int result = datastore.runInTransaction(callable1);
+    assertEquals(result, 2);
+
+    final Entity entity2 = Entity.newBuilder(ENTITY2).clear().setNull("bla").build();
+    TransactionCallable<Integer> callable2 =
+        new TransactionCallable<Integer>() {
+          private Integer attempts = 1;
+
+          @Override
+          public Integer run(DatastoreReaderWriter transaction) {
+            transaction.update(entity2);
+            if (attempts < 2) {
+              ++attempts;
+              throw new DatastoreException(10, "", "ABORTED", false, null);
+            }
+
+            return attempts;
+          }
+        };
+
+    TransactionOptions readOnlyOptions =
+        TransactionOptions.newBuilder().setReadOnly(ReadOnly.getDefaultInstance()).build();
+
+    try {
+      datastore.runInTransaction(callable2, readOnlyOptions);
+      fail("Expecting a failure");
+    } catch (DatastoreException expected) {
+      assertEquals(3, ((DatastoreException) expected.getCause()).getCode());
+    }
+  }
+
+  @Test
+  public void testSkippedResults() {
+    Query<Key> query = Query.newKeyQueryBuilder().setOffset(Integer.MAX_VALUE).build();
+    int numberOfEntities = datastore.run(query).getSkippedResults();
+    assertEquals(2, numberOfEntities);
+  }
+
+  @Test
+  public void testSetLimit() {
+    datastore.put(ENTITY1);
+    Query<Key> keyQuery = Query.newKeyQueryBuilder().setLimit(1).build();
+    QueryResults<?> queryResults = datastore.run(keyQuery);
+    assertTrue(queryResults.hasNext());
+    assertEquals(KEY1, queryResults.next());
+
+    Query<?> query = Query.newEntityQueryBuilder().setLimit(0).build();
+    QueryResults<?> results = datastore.run(query);
+    assertFalse(results.hasNext());
+  }
+
+  @Test
+  public void testGqlQueryWithNullBinding() {
+    Query<Entity> query =
+        Query.newGqlQueryBuilder(ResultType.ENTITY, "select * from " + KIND1)
+            .setNamespace(NAMESPACE)
+            .setNullBinding("name")
+            .build();
+    Iterator<Entity> results = datastore.run(query);
+    assertTrue(results.hasNext());
+    assertEquals(ENTITY1, results.next());
+    assertFalse(results.hasNext());
+  }
+
+  @Test
+  public void testQueryWithStartCursor() {
+    Entity entity1 =
+        Entity.newBuilder(
+                Key.newBuilder(PROJECT_ID, KIND1, "name-01", options.getDatabaseId()).build())
+            .build();
+    Entity entity2 =
+        Entity.newBuilder(
+                Key.newBuilder(PROJECT_ID, KIND1, "name-02", options.getDatabaseId()).build())
+            .build();
+    Entity entity3 =
+        Entity.newBuilder(
+                Key.newBuilder(PROJECT_ID, KIND1, "name-03", options.getDatabaseId()).build())
+            .build();
+    datastore.put(entity1, entity2, entity3);
+    QueryResults<Entity> run1 = datastore.run(Query.newEntityQueryBuilder().setKind(KIND1).build());
+    run1.next();
+    Cursor cursor1 = run1.getCursorAfter();
+    assertNotNull(cursor1);
+    QueryResults<Entity> run2 =
+        datastore.run(Query.newEntityQueryBuilder().setKind(KIND1).setStartCursor(cursor1).build());
+    Cursor cursor2 = run2.getCursorAfter();
+    assertNotNull(cursor2);
+    assertEquals(cursor2, cursor1);
+    datastore.delete(entity1.getKey(), entity2.getKey(), entity3.getKey());
+  }
+
+  @Test
+  public void testQueryWithReadTime() throws InterruptedException {
+    Entity entity1 =
+        Entity.newBuilder(
+                Key.newBuilder(PROJECT_ID, "new_kind", "name-01")
+                    .setDatabaseId(options.getDatabaseId())
+                    .setNamespace(NAMESPACE)
+                    .build())
+            .build();
+    Entity entity2 =
+        Entity.newBuilder(
+                Key.newBuilder(PROJECT_ID, "new_kind", "name-02")
+                    .setDatabaseId(options.getDatabaseId())
+                    .setNamespace(NAMESPACE)
+                    .build())
+            .build();
+    Entity entity3 =
+        Entity.newBuilder(
+                Key.newBuilder(PROJECT_ID, "new_kind", "name-03")
+                    .setDatabaseId(options.getDatabaseId())
+                    .setNamespace(NAMESPACE)
+                    .build())
+            .build();
+
+    datastore.put(entity1, entity2);
+    Thread.sleep(1000);
+    Timestamp now = Timestamp.now();
+    Thread.sleep(1000);
+    datastore.put(entity3);
+
+    try {
+      Query<Entity> query = Query.newEntityQueryBuilder().setKind("new_kind").build();
+
+      QueryResults<Entity> withoutReadTime = datastore.run(query);
+      assertTrue(withoutReadTime.hasNext());
+      assertEquals(entity1, withoutReadTime.next());
+      assertTrue(withoutReadTime.hasNext());
+      assertEquals(entity2, withoutReadTime.next());
+      assertTrue(withoutReadTime.hasNext());
+      assertEquals(entity3, withoutReadTime.next());
+      assertFalse(withoutReadTime.hasNext());
+
+      QueryResults<Entity> withReadTime = datastore.run(query, ReadOption.readTime(now));
+      assertTrue(withReadTime.hasNext());
+      assertEquals(entity1, withReadTime.next());
+      assertTrue(withReadTime.hasNext());
+      assertEquals(entity2, withReadTime.next());
+      assertFalse(withReadTime.hasNext());
+    } finally {
+      datastore.delete(entity1.getKey(), entity2.getKey(), entity3.getKey());
+    }
+  }
+
+  private void testCountAggregationWith(Consumer<AggregationQuery.Builder> configurer) {
+    AggregationQuery.Builder builder = Query.newAggregationQueryBuilder().setNamespace(NAMESPACE);
+    configurer.accept(builder);
+    AggregationQuery aggregationQuery = builder.build();
+    String alias = "total_count";
+
+    Long countBeforeAdd = getOnlyElement(datastore.runAggregation(aggregationQuery)).get(alias);
+    long expectedCount = countBeforeAdd + 1;
+
+    Entity newEntity =
+        Entity.newBuilder(ENTITY1)
+            .setKey(Key.newBuilder(KEY3, KIND1, 1).build())
+            .set("null", NULL_VALUE)
+            .set("partial1", PARTIAL_ENTITY2)
+            .set("partial2", ENTITY2)
+            .build();
+    datastore.put(newEntity);
+
+    Long countAfterAdd = getOnlyElement(datastore.runAggregation(aggregationQuery)).get(alias);
+    assertThat(countAfterAdd).isEqualTo(expectedCount);
+
+    datastore.delete(newEntity.getKey());
+  }
+
+  private void testCountAggregationWithLimit(
+      Consumer<AggregationQuery.Builder> withoutLimitConfigurer,
+      BiConsumer<AggregationQuery.Builder, Long> withLimitConfigurer) {
+    String alias = "total_count";
+
+    AggregationQuery.Builder withoutLimitBuilder =
+        Query.newAggregationQueryBuilder().setNamespace(NAMESPACE);
+    withoutLimitConfigurer.accept(withoutLimitBuilder);
+
+    Long currentCount =
+        getOnlyElement(datastore.runAggregation(withoutLimitBuilder.build())).get(alias);
+    long limit = currentCount - 1;
+
+    AggregationQuery.Builder withLimitBuilder =
+        Query.newAggregationQueryBuilder().setNamespace(NAMESPACE);
+    withLimitConfigurer.accept(withLimitBuilder, limit);
+
+    Long countWithLimit =
+        getOnlyElement(datastore.runAggregation(withLimitBuilder.build())).get(alias);
+    assertThat(countWithLimit).isEqualTo(limit);
+  }
+
+  private void testCountAggregationReadTimeWith(Consumer<AggregationQuery.Builder> configurer)
+      throws InterruptedException {
+    Entity entity1 =
+        Entity.newBuilder(
+                Key.newBuilder(PROJECT_ID, "new_kind", "name-01")
+                    .setDatabaseId(options.getDatabaseId())
+                    .setNamespace(NAMESPACE)
+                    .build())
+            .set("name", "Tyrion Lannister")
+            .build();
+    Entity entity2 =
+        Entity.newBuilder(
+                Key.newBuilder(PROJECT_ID, "new_kind", "name-02")
+                    .setDatabaseId(options.getDatabaseId())
+                    .setNamespace(NAMESPACE)
+                    .build())
+            .set("name", "Jaime Lannister")
+            .build();
+    Entity entity3 =
+        Entity.newBuilder(
+                Key.newBuilder(PROJECT_ID, "new_kind", "name-03")
+                    .setDatabaseId(options.getDatabaseId())
+                    .setNamespace(NAMESPACE)
+                    .build())
+            .set("name", "Cersei Lannister")
+            .build();
+
+    datastore.put(entity1, entity2);
+    Thread.sleep(1000);
+    Timestamp now = Timestamp.now();
+    Thread.sleep(1000);
+    datastore.put(entity3);
+
+    try {
+      AggregationQuery.Builder builder = Query.newAggregationQueryBuilder().setNamespace(NAMESPACE);
+      configurer.accept(builder);
+      AggregationQuery countAggregationQuery = builder.build();
+
+      Long latestCount =
+          getOnlyElement(datastore.runAggregation(countAggregationQuery)).get("total_count");
+      assertThat(latestCount).isEqualTo(3L);
+
+      Long oldCount =
+          getOnlyElement(datastore.runAggregation(countAggregationQuery, ReadOption.readTime(now)))
+              .get("total_count");
+      assertThat(oldCount).isEqualTo(2L);
+    } finally {
+      datastore.delete(entity1.getKey(), entity2.getKey(), entity3.getKey());
+    }
+  }
+}
diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/testing/ITLocalDatastoreHelperTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/testing/ITLocalDatastoreHelperTest.java
index 6cc236a8..615ba709 100644
--- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/testing/ITLocalDatastoreHelperTest.java
+++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/testing/ITLocalDatastoreHelperTest.java
@@ -123,20 +123,28 @@ public class ITLocalDatastoreHelperTest {
   @Test
   public void testCreatePort() {
     LocalDatastoreHelper helper = LocalDatastoreHelper.create(0.75, 8888);
-    DatastoreOptions options = helper.getOptions(NAMESPACE);
+    DatastoreOptions options = helper.setNamespace(NAMESPACE).build();
     assertTrue(options.getHost().endsWith("8888"));
     assertTrue(Math.abs(0.75 - helper.getConsistency()) < TOLERANCE);
     helper = LocalDatastoreHelper.create();
-    options = helper.getOptions(NAMESPACE);
+    options = helper.setNamespace(NAMESPACE).build();
     assertTrue(Math.abs(0.9 - helper.getConsistency()) < TOLERANCE);
     assertFalse(options.getHost().endsWith("8888"));
     assertFalse(options.getHost().endsWith("8080"));
     helper = LocalDatastoreHelper.create(9999);
-    options = helper.getOptions(NAMESPACE);
+    options = helper.setNamespace(NAMESPACE).build();
     assertTrue(Math.abs(0.9 - helper.getConsistency()) < TOLERANCE);
     assertTrue(options.getHost().endsWith("9999"));
   }
 
+  @Test
+  public void testSetDatabaseId() {
+    LocalDatastoreHelper helper = LocalDatastoreHelper.create(0.75, 8888);
+    DatastoreOptions options =
+        helper.setNamespace(NAMESPACE).setDatabaseId("new-database-id").build();
+    assertEquals("new-database-id", options.getDatabaseId());
+  }
+
   @Test
   public void testOptions() {
     LocalDatastoreHelper helper = LocalDatastoreHelper.create();
@@ -144,7 +152,7 @@ public class ITLocalDatastoreHelperTest {
     assertTrue(options.getProjectId().startsWith(PROJECT_ID_PREFIX));
     assertTrue(options.getHost().startsWith("localhost:"));
     assertSame(NoCredentials.getInstance(), options.getCredentials());
-    options = helper.getOptions(NAMESPACE);
+    options = helper.setNamespace(NAMESPACE).build();
     assertTrue(options.getProjectId().startsWith(PROJECT_ID_PREFIX));
     assertTrue(options.getHost().startsWith("localhost:"));
     assertSame(NoCredentials.getInstance(), options.getCredentials());