Repository: googleapis/java-datastore
Pull request #1674 (2024-11-21 18:27:25)

Refactor: refactor grpc query splitter logic
With the new Gapic upgrade change, we are deprecating handwritten HTTP Datastore module [datastore-v1-proto-client](https://github.com/googleapis/java-datastore/tree/main/datastore-v1-proto-client) since it will be replaced with auto generated gRPC code. `QuerySplitter` is a class in [datastore-v1-proto-client](https://github.com/googleapis/java-datastore/tree/main/datastore-v1-proto-client) which is still used by [Beam team](https://github.com/apache/beam/blob/2e4a152cecd269e5c6de46edaf51bd51853fc0cd/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/datastore/DatastoreV1.java#L2057). We want to refactor QuerySplitter and its related helper classes including DatatoreHelper.java into a new module `google-cloud-datastore-utils` in google-cloud-datastore so they will continue to be available for our customers while we deprecate [datastore-v1-proto-client](https://github.com/googleapis/java-datastore/tree/main/datastore-v1-proto-client) module.

For more details about the design decision, please refer [here](https://docs.google.com/document/d/16Dz1AWbAvzADA7fGHQKlpYC0JhnnxU8fORNU1MlHG6s/edit?resourcekey=0-G-goAhwkWDfG0zsm8qARRw&tab=t.0#heading=h.tnhuplgtldlk).



Patch:
diff --git a/.readme-partials.yaml b/.readme-partials.yaml
index 965a0bdd..7960c687 100644
--- a/.readme-partials.yaml
+++ b/.readme-partials.yaml
@@ -226,7 +226,7 @@ custom_content: |
   <dependency>
     <groupId>com.google.cloud</groupId>
     <artifactId>google-cloud-datastore</artifactId>
-    <version>2.20.0-grpc-experimental-1-SNAPSHOT</version>
+    <version>2.22.0-grpc-experimental-1-SNAPSHOT</version>
   </dependency>
   ```
   
diff --git a/google-cloud-datastore-utils/pom.xml b/google-cloud-datastore-utils/pom.xml
new file mode 100644
index 00000000..b45e16ce
--- /dev/null
+++ b/google-cloud-datastore-utils/pom.xml
@@ -0,0 +1,111 @@
+<?xml version='1.0' encoding='UTF-8'?>
+<project xmlns="http://maven.apache.org/POM/4.0.0" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
+    <modelVersion>4.0.0</modelVersion>
+    <groupId>com.google.cloud</groupId>
+    <artifactId>google-cloud-datastore-utils</artifactId>
+    <version>2.22.0-grpc-experimental-1-SNAPSHOT</version><!-- {x-version-update:google-cloud-datastore:current} -->
+    <packaging>jar</packaging>
+    <name>Google Cloud Datastore Utilities</name>
+    <url>https://github.com/googleapis/java-datastore</url>
+    <description>
+        Java datastore client utility library.
+    </description>
+    <parent>
+        <groupId>com.google.cloud</groupId>
+        <artifactId>google-cloud-datastore-parent</artifactId>
+        <version>2.22.0-grpc-experimental-1-SNAPSHOT</version><!-- {x-version-update:google-cloud-datastore:current} -->
+    </parent>
+    <properties>
+        <site.installationModule>google-cloud-datastore-utils</site.installationModule>
+    </properties>
+    <dependencies>
+        <dependency>
+            <groupId>com.google.api-client</groupId>
+            <artifactId>google-api-client</artifactId>
+        </dependency>
+        <dependency>
+            <groupId>com.google.http-client</groupId>
+            <artifactId>google-http-client-protobuf</artifactId>
+        </dependency>
+        <dependency>
+            <groupId>com.google.http-client</groupId>
+            <artifactId>google-http-client-gson</artifactId>
+        </dependency>
+        <dependency>
+            <groupId>com.google.api.grpc</groupId>
+            <artifactId>proto-google-cloud-datastore-v1</artifactId>
+        </dependency>
+        <dependency>
+            <groupId>com.google.api</groupId>
+            <artifactId>api-common</artifactId>
+        </dependency>
+        <dependency>
+            <groupId>com.google.protobuf</groupId>
+            <artifactId>protobuf-java</artifactId>
+        </dependency>
+        <dependency>
+            <groupId>com.google.guava</groupId>
+            <artifactId>guava</artifactId>
+        </dependency>
+        <dependency>
+            <groupId>com.google.api.grpc</groupId>
+            <artifactId>proto-google-common-protos</artifactId>
+        </dependency>
+        <dependency>
+            <groupId>com.google.http-client</groupId>
+            <artifactId>google-http-client</artifactId>
+        </dependency>
+        <dependency>
+            <groupId>com.google.http-client</groupId>
+            <artifactId>google-http-client-jackson2</artifactId>
+        </dependency>
+        <dependency>
+            <groupId>com.google.oauth-client</groupId>
+            <artifactId>google-oauth-client</artifactId>
+        </dependency>
+        <dependency>
+            <groupId>com.google.code.findbugs</groupId>
+            <artifactId>jsr305</artifactId>
+        </dependency>
+        <!-- Test dependencies -->
+        <dependency>
+            <groupId>junit</groupId>
+            <artifactId>junit</artifactId>
+            <scope>test</scope>
+        </dependency>
+        <dependency>
+            <groupId>com.google.truth</groupId>
+            <artifactId>truth</artifactId>
+            <version>1.4.2</version>
+            <scope>test</scope>
+            <exclusions>
+                <exclusion>
+                    <groupId>org.checkerframework</groupId>
+                    <artifactId>checker-qual</artifactId>
+                </exclusion>
+            </exclusions>
+        </dependency>
+    </dependencies>
+
+    <build>
+        <plugins>
+            <plugin>
+                <groupId>org.apache.maven.plugins</groupId>
+                <artifactId>maven-surefire-plugin</artifactId>
+                <configuration>
+                    <argLine>-Xmx2048m</argLine>
+                </configuration>
+            </plugin>
+        </plugins>
+    </build>
+
+    <profiles>
+        <profile>
+            <id>native</id>
+            <properties>
+                <!-- skip native tests for this module, it will be tested in google-cloud-datastore-->
+                <skipTests>true</skipTests>
+            </properties>
+        </profile>
+    </profiles>
+</project>
diff --git a/google-cloud-datastore-utils/src/main/java/com/google/datastore/utils/Datastore.java b/google-cloud-datastore-utils/src/main/java/com/google/datastore/utils/Datastore.java
new file mode 100644
index 00000000..d66e9ce6
--- /dev/null
+++ b/google-cloud-datastore-utils/src/main/java/com/google/datastore/utils/Datastore.java
@@ -0,0 +1,136 @@
+/*
+ * Copyright 2024 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.utils;
+
+import com.google.datastore.v1.*;
+import com.google.rpc.Code;
+import java.io.IOException;
+import java.io.InputStream;
+
+/**
+ * Provides access to Cloud Datastore.
+ *
+ * <p>This class is thread-safe.
+ */
+public class Datastore {
+
+  final RemoteRpc remoteRpc;
+
+  Datastore(RemoteRpc remoteRpc) {
+    this.remoteRpc = remoteRpc;
+  }
+
+  /** Reset the RPC count. */
+  public void resetRpcCount() {
+    remoteRpc.resetRpcCount();
+  }
+
+  /**
+   * Returns the number of RPC calls made since the client was created or {@link #resetRpcCount} was
+   * called.
+   */
+  public int getRpcCount() {
+    return remoteRpc.getRpcCount();
+  }
+
+  private com.google.datastore.utils.DatastoreException invalidResponseException(
+      String method, IOException exception) {
+    return RemoteRpc.makeException(
+        remoteRpc.getUrl(), method, Code.UNAVAILABLE, "Invalid response", exception);
+  }
+
+  public AllocateIdsResponse allocateIds(AllocateIdsRequest request)
+      throws com.google.datastore.utils.DatastoreException {
+    try (InputStream is =
+        remoteRpc.call("allocateIds", request, request.getProjectId(), request.getDatabaseId())) {
+      return AllocateIdsResponse.parseFrom(is);
+    } catch (IOException exception) {
+      throw invalidResponseException("allocateIds", exception);
+    }
+  }
+
+  public BeginTransactionResponse beginTransaction(BeginTransactionRequest request)
+      throws com.google.datastore.utils.DatastoreException {
+    try (InputStream is =
+        remoteRpc.call(
+            "beginTransaction", request, request.getProjectId(), request.getDatabaseId())) {
+      return BeginTransactionResponse.parseFrom(is);
+    } catch (IOException exception) {
+      throw invalidResponseException("beginTransaction", exception);
+    }
+  }
+
+  public CommitResponse commit(CommitRequest request)
+      throws com.google.datastore.utils.DatastoreException {
+    try (InputStream is =
+        remoteRpc.call("commit", request, request.getProjectId(), request.getDatabaseId())) {
+      return CommitResponse.parseFrom(is);
+    } catch (IOException exception) {
+      throw invalidResponseException("commit", exception);
+    }
+  }
+
+  public LookupResponse lookup(LookupRequest request)
+      throws com.google.datastore.utils.DatastoreException {
+    try (InputStream is =
+        remoteRpc.call("lookup", request, request.getProjectId(), request.getDatabaseId())) {
+      return LookupResponse.parseFrom(is);
+    } catch (IOException exception) {
+      throw invalidResponseException("lookup", exception);
+    }
+  }
+
+  public ReserveIdsResponse reserveIds(ReserveIdsRequest request)
+      throws com.google.datastore.utils.DatastoreException {
+    try (InputStream is =
+        remoteRpc.call("reserveIds", request, request.getProjectId(), request.getDatabaseId())) {
+      return ReserveIdsResponse.parseFrom(is);
+    } catch (IOException exception) {
+      throw invalidResponseException("reserveIds", exception);
+    }
+  }
+
+  public RollbackResponse rollback(RollbackRequest request)
+      throws com.google.datastore.utils.DatastoreException {
+    try (InputStream is =
+        remoteRpc.call("rollback", request, request.getProjectId(), request.getDatabaseId())) {
+      return RollbackResponse.parseFrom(is);
+    } catch (IOException exception) {
+      throw invalidResponseException("rollback", exception);
+    }
+  }
+
+  public RunQueryResponse runQuery(RunQueryRequest request)
+      throws com.google.datastore.utils.DatastoreException {
+    try (InputStream is =
+        remoteRpc.call("runQuery", request, request.getProjectId(), request.getDatabaseId())) {
+      return RunQueryResponse.parseFrom(is);
+    } catch (IOException exception) {
+      throw invalidResponseException("runQuery", exception);
+    }
+  }
+
+  public RunAggregationQueryResponse runAggregationQuery(RunAggregationQueryRequest request)
+      throws DatastoreException {
+    try (InputStream is =
+        remoteRpc.call(
+            "runAggregationQuery", request, request.getProjectId(), request.getDatabaseId())) {
+      return RunAggregationQueryResponse.parseFrom(is);
+    } catch (IOException exception) {
+      throw invalidResponseException("runAggregationQuery", exception);
+    }
+  }
+}
diff --git a/google-cloud-datastore-utils/src/main/java/com/google/datastore/utils/DatastoreException.java b/google-cloud-datastore-utils/src/main/java/com/google/datastore/utils/DatastoreException.java
new file mode 100644
index 00000000..48a5dac3
--- /dev/null
+++ b/google-cloud-datastore-utils/src/main/java/com/google/datastore/utils/DatastoreException.java
@@ -0,0 +1,45 @@
+/*
+ * Copyright 2024 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.utils;
+
+import com.google.rpc.Code;
+
+/** Indicates an error in a {@link Datastore} call. */
+public class DatastoreException extends Exception {
+  private final String methodName;
+  private final Code code;
+
+  public DatastoreException(String methodName, Code code, String message, Throwable cause) {
+    super(message, cause);
+    this.methodName = methodName;
+    this.code = code;
+  }
+
+  /** @return the canonical error code */
+  public Code getCode() {
+    return code;
+  }
+
+  /** @return the datastore method name */
+  public String getMethodName() {
+    return methodName;
+  }
+
+  @Override
+  public String toString() {
+    return String.format("%s, code=%s", super.toString(), code);
+  }
+}
diff --git a/google-cloud-datastore-utils/src/main/java/com/google/datastore/utils/DatastoreFactory.java b/google-cloud-datastore-utils/src/main/java/com/google/datastore/utils/DatastoreFactory.java
new file mode 100644
index 00000000..2befe276
--- /dev/null
+++ b/google-cloud-datastore-utils/src/main/java/com/google/datastore/utils/DatastoreFactory.java
@@ -0,0 +1,127 @@
+/*
+ * Copyright 2024 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.utils;
+
+import static com.google.common.base.Preconditions.checkNotNull;
+
+import com.google.api.client.auth.oauth2.Credential;
+import com.google.api.client.http.HttpRequestFactory;
+import com.google.api.client.http.HttpTransport;
+import com.google.api.client.http.javanet.NetHttpTransport;
+import java.net.URI;
+import java.net.URISyntaxException;
+import java.util.Arrays;
+import java.util.logging.*;
+
+/** Client factory for {@link com.google.datastore.utils.Datastore}. */
+public class DatastoreFactory {
+
+  // Lazy load this because we might be running inside App Engine and this
+  // class isn't on the whitelist.
+  private static ConsoleHandler methodHandler;
+
+  /** API version. */
+  public static final String VERSION = "v1";
+
+  public static final String DEFAULT_HOST = "https://datastore.googleapis.com";
+
+  /** Singleton factory instance. */
+  private static final DatastoreFactory INSTANCE = new DatastoreFactory();
+
+  public static DatastoreFactory get() {
+    return INSTANCE;
+  }
+
+  /**
+   * Provides access to a datastore using the provided options. Logs into the application using the
+   * credentials available via these options.
+   *
+   * @throws IllegalArgumentException if the server or credentials weren't provided.
+   */
+  public com.google.datastore.utils.Datastore create(
+      com.google.datastore.utils.DatastoreOptions options) {
+    return new com.google.datastore.utils.Datastore(newRemoteRpc(options));
+  }
+
+  /** Constructs a Google APIs HTTP client with the associated credentials. */
+  public HttpRequestFactory makeClient(com.google.datastore.utils.DatastoreOptions options) {
+    Credential credential = options.getCredential();
+    HttpTransport transport = options.getTransport();
+    if (transport == null) {
+      transport = credential == null ? new NetHttpTransport() : credential.getTransport();
+      transport = transport == null ? new NetHttpTransport() : transport;
+    }
+    return transport.createRequestFactory(credential);
+  }
+
+  /** Starts logging datastore method calls to the console. (Useful within tests.) */
+  public static void logMethodCalls() {
+    Logger logger = Logger.getLogger(Datastore.class.getName());
+    logger.setLevel(Level.FINE);
+    if (!Arrays.asList(logger.getHandlers()).contains(getStreamHandler())) {
+      logger.addHandler(getStreamHandler());
+    }
+  }
+
+  /** Build a valid datastore URL. */
+  String buildProjectEndpoint(com.google.datastore.utils.DatastoreOptions options) {
+    if (options.getProjectEndpoint() != null) {
+      return options.getProjectEndpoint();
+    }
+    // DatastoreOptions ensures either project endpoint or project ID is set.
+    String projectId = checkNotNull(options.getProjectId());
+    if (options.getHost() != null) {
+      return validateUrl(
+          String.format("https://%s/%s/projects/%s", options.getHost(), VERSION, projectId));
+    } else if (options.getLocalHost() != null) {
+      return validateUrl(
+          String.format("http://%s/%s/projects/%s", options.getLocalHost(), VERSION, projectId));
+    }
+    return validateUrl(String.format("%s/%s/projects/%s", DEFAULT_HOST, VERSION, projectId));
+  }
+
+  protected com.google.datastore.utils.RemoteRpc newRemoteRpc(DatastoreOptions options) {
+    checkNotNull(options);
+    HttpRequestFactory client = makeClient(options);
+    return new com.google.datastore.utils.RemoteRpc(
+        client, options.getInitializer(), buildProjectEndpoint(options));
+  }
+
+  private static String validateUrl(String url) {
+    try {
+      return new URI(url).toString();
+    } catch (URISyntaxException e) {
+      throw new IllegalArgumentException(e);
+    }
+  }
+
+  // TODO: Support something other than console handler for when we're
+  // running in App Engine
+  private static synchronized StreamHandler getStreamHandler() {
+    if (methodHandler == null) {
+      methodHandler = new ConsoleHandler();
+      methodHandler.setFormatter(
+          new Formatter() {
+            @Override
+            public String format(LogRecord record) {
+              return record.getMessage() + "\n";
+            }
+          });
+      methodHandler.setLevel(Level.FINE);
+    }
+    return methodHandler;
+  }
+}
diff --git a/google-cloud-datastore-utils/src/main/java/com/google/datastore/utils/DatastoreHelper.java b/google-cloud-datastore-utils/src/main/java/com/google/datastore/utils/DatastoreHelper.java
new file mode 100644
index 00000000..6480de13
--- /dev/null
+++ b/google-cloud-datastore-utils/src/main/java/com/google/datastore/utils/DatastoreHelper.java
@@ -0,0 +1,729 @@
+/*
+ * Copyright 2024 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.utils;
+
+import com.google.api.client.auth.oauth2.Credential;
+import com.google.api.client.googleapis.auth.oauth2.GoogleCredential;
+import com.google.api.client.googleapis.javanet.GoogleNetHttpTransport;
+import com.google.api.client.http.GenericUrl;
+import com.google.api.client.http.HttpRequest;
+import com.google.api.client.http.HttpTransport;
+import com.google.api.client.json.JsonFactory;
+import com.google.api.client.json.gson.GsonFactory;
+import com.google.datastore.v1.ArrayValue;
+import com.google.datastore.v1.CompositeFilter;
+import com.google.datastore.v1.Entity;
+import com.google.datastore.v1.Filter;
+import com.google.datastore.v1.Key;
+import com.google.datastore.v1.Key.PathElement;
+import com.google.datastore.v1.Key.PathElement.IdTypeCase;
+import com.google.datastore.v1.Mutation;
+import com.google.datastore.v1.PartitionId;
+import com.google.datastore.v1.PropertyFilter;
+import com.google.datastore.v1.PropertyOrder;
+import com.google.datastore.v1.PropertyReference;
+import com.google.datastore.v1.Value;
+import com.google.datastore.v1.Value.ValueTypeCase;
+import com.google.protobuf.ByteString;
+import com.google.protobuf.Timestamp;
+import com.google.protobuf.TimestampOrBuilder;
+import com.google.type.LatLng;
+import java.io.File;
+import java.io.IOException;
+import java.security.GeneralSecurityException;
+import java.security.PrivateKey;
+import java.util.Arrays;
+import java.util.Collection;
+import java.util.Comparator;
+import java.util.Date;
+import java.util.Iterator;
+import java.util.List;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.logging.Level;
+import java.util.logging.Logger;
+import javax.annotation.Nullable;
+
+/** Helper methods for {@link Datastore}. */
+// TODO: Accept OrBuilders when possible.
+public final class DatastoreHelper {
+  private static final Logger logger =
+      Logger.getLogger(com.google.datastore.utils.DatastoreHelper.class.getName());
+
+  private static final int MICROSECONDS_PER_SECOND = 1000 * 1000;
+  private static final int NANOSECONDS_PER_MICROSECOND = 1000;
+
+  /** The property used in the Datastore to give us a random distribution. * */
+  public static final String SCATTER_PROPERTY_NAME = "__scatter__";
+
+  /** The property used in the Datastore to get the key of the entity. * */
+  public static final String KEY_PROPERTY_NAME = "__key__";
+
+  /** Name of the environment variable used to set the project ID. */
+  public static final String PROJECT_ID_ENV_VAR = "DATASTORE_PROJECT_ID";
+
+  /** Name of the environment variable used to set the local host. */
+  public static final String LOCAL_HOST_ENV_VAR = "DATASTORE_EMULATOR_HOST";
+
+  /** Name of the environment variable used to set the service account. */
+  public static final String SERVICE_ACCOUNT_ENV_VAR = "DATASTORE_SERVICE_ACCOUNT";
+
+  /** Name of the environment variable used to set the private key file. */
+  public static final String PRIVATE_KEY_FILE_ENV_VAR = "DATASTORE_PRIVATE_KEY_FILE";
+
+  private static final String URL_OVERRIDE_ENV_VAR = "__DATASTORE_URL_OVERRIDE";
+
+  private static final AtomicReference<String> projectIdFromComputeEngine = new AtomicReference<>();
+
+  /** Comparator for Keys */
+  private static final class KeyComparator implements Comparator<Key> {
+
+    static final com.google.datastore.utils.DatastoreHelper.KeyComparator INSTANCE =
+        new com.google.datastore.utils.DatastoreHelper.KeyComparator();
+
+    private int comparePathElement(PathElement thisElement, PathElement otherElement) {
+      int result = thisElement.getKind().compareTo(otherElement.getKind());
+      if (result != 0) {
+        return result;
+      }
+      if (thisElement.getIdTypeCase() == IdTypeCase.ID) {
+        if (otherElement.getIdTypeCase() != IdTypeCase.ID) {
+          return -1;
+        }
+        return Long.valueOf(thisElement.getId()).compareTo(otherElement.getId());
+      }
+      if (otherElement.getIdTypeCase() == IdTypeCase.ID) {
+        return 1;
+      }
+
+      return thisElement.getName().compareTo(otherElement.getName());
+    }
+
+    @Override
+    public int compare(Key thisKey, Key otherKey) {
+      if (!thisKey.getPartitionId().equals(otherKey.getPartitionId())) {
+        throw new IllegalArgumentException("Cannot compare keys with different partition ids.");
+      }
+
+      Iterator<PathElement> thisPath = thisKey.getPathList().iterator();
+      Iterator<PathElement> otherPath = otherKey.getPathList().iterator();
+      while (thisPath.hasNext()) {
+        if (!otherPath.hasNext()) {
+          return 1;
+        }
+        int result = comparePathElement(thisPath.next(), otherPath.next());
+        if (result != 0) {
+          return result;
+        }
+      }
+
+      return otherPath.hasNext() ? -1 : 0;
+    }
+  }
+
+  private DatastoreHelper() {}
+
+  private static HttpTransport newTransport() throws GeneralSecurityException, IOException {
+    return GoogleNetHttpTransport.newTrustedTransport();
+  }
+
+  static JsonFactory newJsonFactory() {
+    return new GsonFactory();
+  }
+
+  /**
+   * Constructs credentials for the given account and key.
+   *
+   * @param serviceAccountId service account ID (typically an e-mail address).
+   * @param privateKeyFile the file name from which to get the private key.
+   * @return valid credentials or {@code null}
+   */
+  public static Credential getServiceAccountCredential(
+      String serviceAccountId, String privateKeyFile) throws GeneralSecurityException, IOException {
+    return getServiceAccountCredential(serviceAccountId, privateKeyFile, DatastoreOptions.SCOPES);
+  }
+
+  /**
+   * Constructs credentials for the given account and key file.
+   *
+   * @param serviceAccountId service account ID (typically an e-mail address).
+   * @param privateKeyFile the file name from which to get the private key.
+   * @param serviceAccountScopes Collection of OAuth scopes to use with the the service account flow
+   *     or {@code null} if not.
+   * @return valid credentials or {@code null}
+   */
+  public static Credential getServiceAccountCredential(
+      String serviceAccountId, String privateKeyFile, Collection<String> serviceAccountScopes)
+      throws GeneralSecurityException, IOException {
+    return getCredentialBuilderWithoutPrivateKey(serviceAccountId, serviceAccountScopes)
+        .setServiceAccountPrivateKeyFromP12File(new File(privateKeyFile))
+        .build();
+  }
+
+  /**
+   * Constructs credentials for the given account and key.
+   *
+   * @param serviceAccountId service account ID (typically an e-mail address).
+   * @param privateKey the private key for the given account.
+   * @param serviceAccountScopes Collection of OAuth scopes to use with the the service account flow
+   *     or {@code null} if not.
+   * @return valid credentials or {@code null}
+   */
+  public static Credential getServiceAccountCredential(
+      String serviceAccountId, PrivateKey privateKey, Collection<String> serviceAccountScopes)
+      throws GeneralSecurityException, IOException {
+    return getCredentialBuilderWithoutPrivateKey(serviceAccountId, serviceAccountScopes)
+        .setServiceAccountPrivateKey(privateKey)
+        .build();
+  }
+
+  private static GoogleCredential.Builder getCredentialBuilderWithoutPrivateKey(
+      String serviceAccountId, Collection<String> serviceAccountScopes)
+      throws GeneralSecurityException, IOException {
+    HttpTransport transport = newTransport();
+    JsonFactory jsonFactory = newJsonFactory();
+    return new GoogleCredential.Builder()
+        .setTransport(transport)
+        .setJsonFactory(jsonFactory)
+        .setServiceAccountId(serviceAccountId)
+        .setServiceAccountScopes(serviceAccountScopes);
+  }
+
+  /**
+   * Constructs a {@link Datastore} from environment variables and/or the Compute Engine metadata
+   * server.
+   *
+   * <p>The project ID is determined from, in order of preference:
+   *
+   * <ul>
+   *   <li>DATASTORE_PROJECT_ID environment variable
+   *   <li>Compute Engine
+   * </ul>
+   *
+   * <p>Credentials are taken from, in order of preference:
+   *
+   * <ol>
+   *   <li>No credentials (if the DATASTORE_EMULATOR_HOST environment variable is set)
+   *   <li>Service Account specified by the DATASTORE_SERVICE_ACCOUNT and DATASTORE_PRIVATE_KEY_FILE
+   *       environment variables
+   *   <li>Google Application Default as described <a
+   *       href="https://cloud.google.com/docs/authentication/production">here</a>.
+   * </ol>
+   */
+  public static DatastoreOptions.Builder getOptionsFromEnv()
+      throws GeneralSecurityException, IOException {
+    DatastoreOptions.Builder options = new DatastoreOptions.Builder();
+    setProjectEndpointFromEnv(options);
+    options.credential(getCredentialFromEnv());
+    return options;
+  }
+
+  private static Credential getCredentialFromEnv() throws GeneralSecurityException, IOException {
+    if (System.getenv(LOCAL_HOST_ENV_VAR) != null) {
+      logger.log(
+          Level.INFO,
+          "{0} environment variable was set. Not using credentials.",
+          new Object[] {LOCAL_HOST_ENV_VAR});
+      return null;
+    }
+    String serviceAccount = System.getenv(SERVICE_ACCOUNT_ENV_VAR);
+    String privateKeyFile = System.getenv(PRIVATE_KEY_FILE_ENV_VAR);
+    if (serviceAccount != null && privateKeyFile != null) {
+      logger.log(
+          Level.INFO,
+          "{0} and {1} environment variables were set. " + "Using service account credential.",
+          new Object[] {SERVICE_ACCOUNT_ENV_VAR, PRIVATE_KEY_FILE_ENV_VAR});
+      return getServiceAccountCredential(serviceAccount, privateKeyFile);
+    }
+    return GoogleCredential.getApplicationDefault().createScoped(DatastoreOptions.SCOPES);
+  }
+
+  /**
+   * Determines the project id from the environment. Uses the following sources in order of
+   * preference:
+   *
+   * <ol>
+   *   <li>Value of the DATASTORE_PROJECT_ID environment variable
+   *   <li>Compute Engine
+   * </ol>
+   *
+   * @throws IllegalStateException if the project ID cannot be determined
+   */
+  private static String getProjectIdFromEnv() {
+    if (System.getenv(PROJECT_ID_ENV_VAR) != null) {
+      return System.getenv(PROJECT_ID_ENV_VAR);
+    }
+    String projectIdFromComputeEngine = getProjectIdFromComputeEngine();
+    if (projectIdFromComputeEngine != null) {
+      return projectIdFromComputeEngine;
+    }
+    throw new IllegalStateException(
+        String.format(
+            "Could not determine project ID."
+                + " If you are not running on Compute Engine, set the"
+                + " %s environment variable.",
+            PROJECT_ID_ENV_VAR));
+  }
+
+  /**
+   * Gets the project ID from the Compute Engine metadata server. Returns {@code null} if the
+   * project ID cannot be determined (because, for instance, the code is not running on Compute
+   * Engine).
+   */
+  @Nullable
+  public static String getProjectIdFromComputeEngine() {
+    String cachedProjectId = projectIdFromComputeEngine.get();
+    return cachedProjectId != null ? cachedProjectId : queryProjectIdFromComputeEngine();
+  }
+
+  @Nullable
+  private static String queryProjectIdFromComputeEngine() {
+    HttpTransport transport;
+
+    try {
+      transport = newTransport();
+    } catch (GeneralSecurityException | IOException e) {
+      logger.log(Level.WARNING, "Failed to create HttpTransport.", e);
+      return null;
+    }
+
+    try {
+      GenericUrl projectIdUrl =
+          new GenericUrl("http://metadata/computeMetadata/v1/project/project-id");
+      HttpRequest request = transport.createRequestFactory().buildGetRequest(projectIdUrl);
+      request.getHeaders().set("Metadata-Flavor", "Google");
+      String result = request.execute().parseAsString();
+      projectIdFromComputeEngine.set(result);
+      return result;
+    } catch (IOException e) {
+      logger.log(Level.INFO, "Could not determine project ID from Compute Engine.", e);
+      return null;
+    }
+  }
+
+  private static void setProjectEndpointFromEnv(DatastoreOptions.Builder options) {
+    // DATASTORE_HOST is deprecated.
+    if (System.getenv("DATASTORE_HOST") != null) {
+      logger.warning(
+          String.format(
+              "Ignoring value of environment variable DATASTORE_HOST. "
+                  + "To point datastore to a host running locally, use "
+                  + "the environment variable %s.",
+              LOCAL_HOST_ENV_VAR));
+    }
+    String projectId = getProjectIdFromEnv();
+    if (System.getenv(URL_OVERRIDE_ENV_VAR) != null) {
+      options.projectEndpoint(
+          String.format("%s/projects/%s", System.getenv(URL_OVERRIDE_ENV_VAR), projectId));
+      return;
+    }
+    if (System.getenv(LOCAL_HOST_ENV_VAR) != null) {
+      options.projectId(projectId);
+      options.localHost(System.getenv(LOCAL_HOST_ENV_VAR));
+      return;
+    }
+    options.projectId(projectId);
+    return;
+  }
+
+  /** @see #getOptionsFromEnv() */
+  public static Datastore getDatastoreFromEnv() throws GeneralSecurityException, IOException {
+    return DatastoreFactory.get().create(getOptionsFromEnv().build());
+  }
+
+  /**
+   * Gets a {@link com.google.datastore.utils.QuerySplitter}.
+   *
+   * <p>The returned {@link com.google.datastore.utils.QuerySplitter#getSplits} cannot accept a
+   * query that contains inequality filters, a sort filter, or a missing kind.
+   */
+  public static QuerySplitter getQuerySplitter() {
+    return com.google.datastore.utils.QuerySplitterImpl.INSTANCE;
+  }
+
+  public static Comparator<Key> getKeyComparator() {
+    return com.google.datastore.utils.DatastoreHelper.KeyComparator.INSTANCE;
+  }
+
+  /** Make a sort order for use in a query. */
+  public static PropertyOrder.Builder makeOrder(
+      String property, PropertyOrder.Direction direction) {
+    return PropertyOrder.newBuilder()
+        .setProperty(makePropertyReference(property))
+        .setDirection(direction);
+  }
+
+  /** Makes an ancestor filter. */
+  public static Filter.Builder makeAncestorFilter(Key ancestor) {
+    return makeFilter(
+        com.google.datastore.utils.DatastoreHelper.KEY_PROPERTY_NAME,
+        PropertyFilter.Operator.HAS_ANCESTOR,
+        makeValue(ancestor));
+  }
+
+  /** Make a filter on a property for use in a query. */
+  public static Filter.Builder makeFilter(
+      String property, PropertyFilter.Operator operator, Value value) {
+    return Filter.newBuilder()
+        .setPropertyFilter(
+            PropertyFilter.newBuilder()
+                .setProperty(makePropertyReference(property))
+                .setOp(operator)
+                .setValue(value));
+  }
+
+  /** Make a filter on a property for use in a query. */
+  public static Filter.Builder makeFilter(
+      String property, PropertyFilter.Operator operator, Value.Builder value) {
+    return makeFilter(property, operator, value.build());
+  }
+
+  /** Make a composite filter from the given sub-filters using AND to combine filters. */
+  public static Filter.Builder makeAndFilter(Filter... subfilters) {
+    return makeAndFilter(Arrays.asList(subfilters));
+  }
+
+  /** Make a composite filter from the given sub-filters using AND to combine filters. */
+  public static Filter.Builder makeAndFilter(Iterable<Filter> subfilters) {
+    return Filter.newBuilder()
+        .setCompositeFilter(
+            CompositeFilter.newBuilder()
+                .addAllFilters(subfilters)
+                .setOp(CompositeFilter.Operator.AND));
+  }
+
+  /** Make a property reference for use in a query. */
+  public static PropertyReference.Builder makePropertyReference(String propertyName) {
+    return PropertyReference.newBuilder().setName(propertyName);
+  }
+
+  /** Make an array value containing the specified values. */
+  public static Value.Builder makeValue(Iterable<Value> values) {
+    return Value.newBuilder().setArrayValue(ArrayValue.newBuilder().addAllValues(values));
+  }
+
+  /** Make a list value containing the specified values. */
+  public static Value.Builder makeValue(Value value1, Value value2, Value... rest) {
+    ArrayValue.Builder arrayValue = ArrayValue.newBuilder();
+    arrayValue.addValues(value1);
+    arrayValue.addValues(value2);
+    arrayValue.addAllValues(Arrays.asList(rest));
+    return Value.newBuilder().setArrayValue(arrayValue);
+  }
+
+  /** Make an array value containing the specified values. */
+  public static Value.Builder makeValue(
+      Value.Builder value1, Value.Builder value2, Value.Builder... rest) {
+    ArrayValue.Builder arrayValue = ArrayValue.newBuilder();
+    arrayValue.addValues(value1);
+    arrayValue.addValues(value2);
+    for (Value.Builder builder : rest) {
+      arrayValue.addValues(builder);
+    }
+    return Value.newBuilder().setArrayValue(arrayValue);
+  }
+
+  /** Make a key value. */
+  public static Value.Builder makeValue(Key key) {
+    return Value.newBuilder().setKeyValue(key);
+  }
+
+  /** Make a key value. */
+  public static Value.Builder makeValue(Key.Builder key) {
+    return makeValue(key.build());
+  }
+
+  /** Make an integer value. */
+  public static Value.Builder makeValue(long key) {
+    return Value.newBuilder().setIntegerValue(key);
+  }
+
+  /** Make a floating point value. */
+  public static Value.Builder makeValue(double value) {
+    return Value.newBuilder().setDoubleValue(value);
+  }
+
+  /** Make a boolean value. */
+  public static Value.Builder makeValue(boolean value) {
+    return Value.newBuilder().setBooleanValue(value);
+  }
+
+  /** Make a string value. */
+  public static Value.Builder makeValue(String value) {
+    return Value.newBuilder().setStringValue(value);
+  }
+
+  /** Make an entity value. */
+  public static Value.Builder makeValue(Entity entity) {
+    return Value.newBuilder().setEntityValue(entity);
+  }
+
+  /** Make a entity value. */
+  public static Value.Builder makeValue(Entity.Builder entity) {
+    return makeValue(entity.build());
+  }
+
+  /** Make a ByteString value. */
+  public static Value.Builder makeValue(ByteString blob) {
+    return Value.newBuilder().setBlobValue(blob);
+  }
+
+  /** Make a timestamp value given a date. */
+  public static Value.Builder makeValue(Date date) {
+    return Value.newBuilder().setTimestampValue(toTimestamp(date.getTime() * 1000L));
+  }
+
+  /** Makes a GeoPoint value. */
+  public static Value.Builder makeValue(LatLng value) {
+    return Value.newBuilder().setGeoPointValue(value);
+  }
+
+  /** Makes a GeoPoint value. */
+  public static Value.Builder makeValue(LatLng.Builder value) {
+    return makeValue(value.build());
+  }
+
+  private static Timestamp.Builder toTimestamp(long microseconds) {
+    long seconds = microseconds / MICROSECONDS_PER_SECOND;
+    long microsecondsRemainder = microseconds % MICROSECONDS_PER_SECOND;
+    if (microsecondsRemainder < 0) {
+      // Nanos must be positive even if microseconds is negative.
+      // Java modulus doesn't take care of this for us.
+      microsecondsRemainder += MICROSECONDS_PER_SECOND;
+      seconds -= 1;
+    }
+    return Timestamp.newBuilder()
+        .setSeconds(seconds)
+        .setNanos((int) microsecondsRemainder * NANOSECONDS_PER_MICROSECOND);
+  }
+
+  /**
+   * Make a key from the specified path of kind/id-or-name pairs and/or Keys.
+   *
+   * <p>The id-or-name values must be either String, Long, Integer or Short.
+   *
+   * <p>The last id-or-name value may be omitted, in which case an entity without an id is created
+   * (for use with automatic id allocation).
+   *
+   * <p>The PartitionIds of all Keys in the path must be equal. The returned Key.Builder will use
+   * this PartitionId.
+   */
+  public static Key.Builder makeKey(Object... elements) {
+    Key.Builder key = Key.newBuilder();
+    PartitionId partitionId = null;
+    for (int pathIndex = 0; pathIndex < elements.length; pathIndex += 2) {
+      PathElement.Builder pathElement = PathElement.newBuilder();
+      Object element = elements[pathIndex];
+      if (element instanceof Key) {
+        Key subKey = (Key) element;
+        if (partitionId == null) {
+          partitionId = subKey.getPartitionId();
+        } else if (!partitionId.equals(subKey.getPartitionId())) {
+          throw new IllegalArgumentException(
+              "Partition IDs did not match, found: "
+                  + partitionId
+                  + " and "
+                  + subKey.getPartitionId());
+        }
+        key.addAllPath(((Key) element).getPathList());
+        // We increment by 2, but since we got a Key argument we're only consuming 1 element in this
+        // iteration of the loop. Decrement the index so that when we jump by 2 we end up in the
+        // right spot.
+        pathIndex--;
+      } else {
+        String kind;
+        try {
+          kind = (String) element;
+        } catch (ClassCastException e) {
+          throw new IllegalArgumentException(
+              "Expected string or Key, got: " + element.getClass(), e);
+        }
+        pathElement.setKind(kind);
+        if (pathIndex + 1 < elements.length) {
+          Object value = elements[pathIndex + 1];
+          if (value instanceof String) {
+            pathElement.setName((String) value);
+          } else if (value instanceof Long) {
+            pathElement.setId((Long) value);
+          } else if (value instanceof Integer) {
+            pathElement.setId((Integer) value);
+          } else if (value instanceof Short) {
+            pathElement.setId((Short) value);
+          } else {
+            throw new IllegalArgumentException(
+                "Expected string or integer, got: " + value.getClass());
+          }
+        }
+        key.addPath(pathElement);
+      }
+    }
+    if (partitionId != null && !partitionId.equals(PartitionId.getDefaultInstance())) {
+      key.setPartitionId(partitionId);
+    }
+    return key;
+  }
+
+  /**
+   * @return the double contained in value
+   * @throws IllegalArgumentException if the value does not contain a double.
+   */
+  public static double getDouble(Value value) {
+    if (value.getValueTypeCase() != ValueTypeCase.DOUBLE_VALUE) {
+      throw new IllegalArgumentException("Value does not contain a double.");
+    }
+    return value.getDoubleValue();
+  }
+
+  /**
+   * @return the key contained in value
+   * @throws IllegalArgumentException if the value does not contain a key.
+   */
+  public static Key getKey(Value value) {
+    if (value.getValueTypeCase() != ValueTypeCase.KEY_VALUE) {
+      throw new IllegalArgumentException("Value does not contain a key.");
+    }
+    return value.getKeyValue();
+  }
+
+  /**
+   * @return the blob contained in value
+   * @throws IllegalArgumentException if the value does not contain a blob.
+   */
+  public static ByteString getByteString(Value value) {
+    if (value.getMeaning() == 18 && value.getValueTypeCase() == ValueTypeCase.STRING_VALUE) {
+      return value.getStringValueBytes();
+    } else if (value.getValueTypeCase() == ValueTypeCase.BLOB_VALUE) {
+      return value.getBlobValue();
+    }
+    throw new IllegalArgumentException("Value does not contain a blob.");
+  }
+
+  /**
+   * @return the entity contained in value
+   * @throws IllegalArgumentException if the value does not contain an entity.
+   */
+  public static Entity getEntity(Value value) {
+    if (value.getValueTypeCase() != ValueTypeCase.ENTITY_VALUE) {
+      throw new IllegalArgumentException("Value does not contain an Entity.");
+    }
+    return value.getEntityValue();
+  }
+
+  /**
+   * @return the string contained in value
+   * @throws IllegalArgumentException if the value does not contain a string.
+   */
+  public static String getString(Value value) {
+    if (value.getValueTypeCase() != ValueTypeCase.STRING_VALUE) {
+      throw new IllegalArgumentException("Value does not contain a string.");
+    }
+    return value.getStringValue();
+  }
+
+  /**
+   * @return the boolean contained in value
+   * @throws IllegalArgumentException if the value does not contain a boolean.
+   */
+  public static boolean getBoolean(Value value) {
+    if (value.getValueTypeCase() != ValueTypeCase.BOOLEAN_VALUE) {
+      throw new IllegalArgumentException("Value does not contain a boolean.");
+    }
+    return value.getBooleanValue();
+  }
+
+  /**
+   * @return the long contained in value
+   * @throws IllegalArgumentException if the value does not contain a long.
+   */
+  public static long getLong(Value value) {
+    if (value.getValueTypeCase() != ValueTypeCase.INTEGER_VALUE) {
+      throw new IllegalArgumentException("Value does not contain an integer.");
+    }
+    return value.getIntegerValue();
+  }
+
+  /**
+   * @return the timestamp in microseconds contained in value
+   * @throws IllegalArgumentException if the value does not contain a timestamp.
+   */
+  public static long getTimestamp(Value value) {
+    if (value.getMeaning() == 18 && value.getValueTypeCase() == ValueTypeCase.INTEGER_VALUE) {
+      return value.getIntegerValue();
+    } else if (value.getValueTypeCase() == ValueTypeCase.TIMESTAMP_VALUE) {
+      return toMicroseconds(value.getTimestampValue());
+    }
+    throw new IllegalArgumentException("Value does not contain a timestamp.");
+  }
+
+  private static long toMicroseconds(TimestampOrBuilder timestamp) {
+    // Nanosecond precision is lost.
+    return timestamp.getSeconds() * MICROSECONDS_PER_SECOND
+        + timestamp.getNanos() / NANOSECONDS_PER_MICROSECOND;
+  }
+
+  /**
+   * @return the array contained in value as a list.
+   * @throws IllegalArgumentException if the value does not contain an array.
+   */
+  public static List<Value> getList(Value value) {
+    if (value.getValueTypeCase() != ValueTypeCase.ARRAY_VALUE) {
+      throw new IllegalArgumentException("Value does not contain an array.");
+    }
+    return value.getArrayValue().getValuesList();
+  }
+
+  /**
+   * Convert a timestamp value into a {@link Date} clipping off the microseconds.
+   *
+   * @param value a timestamp value to convert
+   * @return the resulting {@link Date}
+   * @throws IllegalArgumentException if the value does not contain a timestamp.
+   */
+  public static Date toDate(Value value) {
+    return new Date(getTimestamp(value) / 1000);
+  }
+
+  /**
+   * @param entity the entity to insert
+   * @return a mutation that will insert an entity
+   */
+  public static Mutation.Builder makeInsert(Entity entity) {
+    return Mutation.newBuilder().setInsert(entity);
+  }
+
+  /**
+   * @param entity the entity to update
+   * @return a mutation that will update an entity
+   */
+  public static Mutation.Builder makeUpdate(Entity entity) {
+    return Mutation.newBuilder().setUpdate(entity);
+  }
+
+  /**
+   * @param entity the entity to upsert
+   * @return a mutation that will upsert an entity
+   */
+  public static Mutation.Builder makeUpsert(Entity entity) {
+    return Mutation.newBuilder().setUpsert(entity);
+  }
+
+  /**
+   * @param key the key of the entity to delete
+   * @return a mutation that will delete an entity
+   */
+  public static Mutation.Builder makeDelete(Key key) {
+    return Mutation.newBuilder().setDelete(key);
+  }
+}
diff --git a/google-cloud-datastore-utils/src/main/java/com/google/datastore/utils/DatastoreOptions.java b/google-cloud-datastore-utils/src/main/java/com/google/datastore/utils/DatastoreOptions.java
new file mode 100644
index 00000000..f6e91a41
--- /dev/null
+++ b/google-cloud-datastore-utils/src/main/java/com/google/datastore/utils/DatastoreOptions.java
@@ -0,0 +1,204 @@
+/*
+ * Copyright 2024 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.utils;
+
+import static com.google.common.base.Preconditions.checkArgument;
+
+import com.google.api.client.auth.oauth2.Credential;
+import com.google.api.client.http.HttpRequestInitializer;
+import com.google.api.client.http.HttpTransport;
+import java.util.Arrays;
+import java.util.List;
+
+/**
+ * An immutable object containing settings for the datastore.
+ *
+ * <p>Example for connecting to a datastore:
+ *
+ * <pre>
+ * DatastoreOptions options = new DatastoreOptions.Builder()
+ *     .projectId("my-project-id")
+ *     .credential(DatastoreHelper.getComputeEngineCredential())
+ *     .build();
+ * DatastoreFactory.get().create(options);
+ * </pre>
+ *
+ * <p>The options should be passed to {@link DatastoreFactory#create}.
+ */
+public class DatastoreOptions {
+  private final String projectId;
+
+  private final String projectEndpoint;
+  private final String host;
+  private final String localHost;
+
+  private final HttpRequestInitializer initializer;
+
+  private final Credential credential;
+  private final HttpTransport transport;
+  public static final List<String> SCOPES =
+      Arrays.asList("https://www.googleapis.com/auth/datastore");
+
+  DatastoreOptions(Builder b) {
+    checkArgument(
+        b.projectId != null || b.projectEndpoint != null,
+        "Either project ID or project endpoint must be provided.");
+    this.projectId = b.projectId;
+    this.projectEndpoint = b.projectEndpoint;
+    this.host = b.host;
+    this.localHost = b.localHost;
+    this.initializer = b.initializer;
+    this.credential = b.credential;
+    this.transport = b.transport;
+  }
+
+  /** Builder for {@link DatastoreOptions}. */
+  public static class Builder {
+    private static final String PROJECT_ENDPOINT_AND_PROJECT_ID_ERROR =
+        "Cannot set both project endpoint and project ID.";
+    private static final String PROJECT_ENDPOINT_AND_HOST_ERROR =
+        "Can set at most one of project endpoint, host, and local host.";
+
+    private String projectId;
+
+    private String projectEndpoint;
+    private String host;
+    private String localHost;
+    private HttpRequestInitializer initializer;
+    private Credential credential;
+    private HttpTransport transport;
+
+    public Builder() {}
+
+    public Builder(DatastoreOptions options) {
+      this.projectId = options.projectId;
+      this.projectEndpoint = options.projectEndpoint;
+      this.host = options.host;
+      this.localHost = options.localHost;
+      this.initializer = options.initializer;
+      this.credential = options.credential;
+      this.transport = options.transport;
+    }
+
+    public DatastoreOptions build() {
+      return new DatastoreOptions(this);
+    }
+
+    /** Sets the project ID used to access Cloud Datastore. */
+    public Builder projectId(String projectId) {
+      checkArgument(projectEndpoint == null, PROJECT_ENDPOINT_AND_PROJECT_ID_ERROR);
+      this.projectId = projectId;
+      return this;
+    }
+
+    /**
+     * Sets the host used to access Cloud Datastore. To connect to the Cloud Datastore Emulator, use
+     * {@link #localHost} instead.
+     */
+    public Builder host(String host) {
+      checkArgument(projectEndpoint == null && localHost == null, PROJECT_ENDPOINT_AND_HOST_ERROR);
+      if (includesScheme(host)) {
+        throw new IllegalArgumentException(
+            String.format("Host \"%s\" must not include scheme.", host));
+      }
+      this.host = host;
+      return this;
+    }
+
+    /**
+     * Configures the client to access Cloud Datastore on a local host (typically a Cloud Datastore
+     * Emulator instance). Call this method also configures the client not to attach credentials to
+     * requests.
+     */
+    public Builder localHost(String localHost) {
+      checkArgument(projectEndpoint == null && host == null, PROJECT_ENDPOINT_AND_HOST_ERROR);
+      if (includesScheme(localHost)) {
+        throw new IllegalArgumentException(
+            String.format("Local host \"%s\" must not include scheme.", localHost));
+      }
+      this.localHost = localHost;
+      return this;
+    }
+
+    /**
+     * Sets the project endpoint used to access Cloud Datastore. Prefer using {@link #projectId}
+     * and/or {@link #host}/{@link #localHost} when possible.
+     *
+     * @deprecated Use {@link #projectId} and/or {@link #host}/{@link #localHost} instead.
+     */
+    @Deprecated
+    public Builder projectEndpoint(String projectEndpoint) {
+      checkArgument(projectId == null, PROJECT_ENDPOINT_AND_PROJECT_ID_ERROR);
+      checkArgument(localHost == null && host == null, PROJECT_ENDPOINT_AND_HOST_ERROR);
+      if (!includesScheme(projectEndpoint)) {
+        throw new IllegalArgumentException(
+            String.format("Project endpoint \"%s\" must include scheme.", projectEndpoint));
+      }
+      this.projectEndpoint = projectEndpoint;
+      return this;
+    }
+
+    /** Sets the (optional) initializer to run on HTTP requests to Cloud Datastore. */
+    public Builder initializer(HttpRequestInitializer initializer) {
+      this.initializer = initializer;
+      return this;
+    }
+
+    /** Sets the Google APIs {@link Credential} used to access Cloud Datastore. */
+    public Builder credential(Credential credential) {
+      this.credential = credential;
+      return this;
+    }
+
+    /** Sets the transport used to access Cloud Datastore. */
+    public Builder transport(HttpTransport transport) {
+      this.transport = transport;
+      return this;
+    }
+
+    private static boolean includesScheme(String url) {
+      return url.startsWith("http://") || url.startsWith("https://");
+    }
+  }
+
+  public String getProjectId() {
+    return projectId;
+  }
+
+  public String getProjectEndpoint() {
+    return projectEndpoint;
+  }
+
+  public String getHost() {
+    return host;
+  }
+
+  public String getLocalHost() {
+    return localHost;
+  }
+
+  public HttpRequestInitializer getInitializer() {
+    return initializer;
+  }
+
+  public Credential getCredential() {
+    return credential;
+  }
+
+  public HttpTransport getTransport() {
+    return transport;
+  }
+}
diff --git a/google-cloud-datastore-utils/src/main/java/com/google/datastore/utils/QuerySplitter.java b/google-cloud-datastore-utils/src/main/java/com/google/datastore/utils/QuerySplitter.java
new file mode 100644
index 00000000..31d1fd7d
--- /dev/null
+++ b/google-cloud-datastore-utils/src/main/java/com/google/datastore/utils/QuerySplitter.java
@@ -0,0 +1,56 @@
+/*
+ * Copyright 2024 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.utils;
+
+import com.google.api.core.BetaApi;
+import com.google.datastore.v1.PartitionId;
+import com.google.datastore.v1.Query;
+import com.google.protobuf.Timestamp;
+import java.util.List;
+
+/** Provides the ability to split a query into multiple shards. */
+public interface QuerySplitter {
+
+  /**
+   * Returns a list of sharded {@link Query}s for the given query.
+   *
+   * <p>This will create up to the desired number of splits, however it may return less splits if
+   * the desired number of splits is unavailable. This will happen if the number of split points
+   * provided by the underlying Datastore is less than the desired number, which will occur if the
+   * number of results for the query is too small.
+   *
+   * @param query the query to split.
+   * @param partition the partition the query is running in.
+   * @param numSplits the desired number of splits.
+   * @param datastore the datastore to run on.
+   * @throws DatastoreException if there was a datastore error while generating query splits.
+   * @throws IllegalArgumentException if the given query or numSplits was invalid.
+   */
+  List<Query> getSplits(Query query, PartitionId partition, int numSplits, Datastore datastore)
+      throws DatastoreException;
+
+  /**
+   * Same as {@link #getSplits(Query, PartitionId, int, Datastore)} but the splits are based on
+   * {@code readTime}, and the returned sharded {@link Query}s should also be executed with {@code
+   * readTime}. Reading from a timestamp is currently a private preview feature in Datastore.
+   */
+  @BetaApi
+  default List<Query> getSplits(
+      Query query, PartitionId partition, int numSplits, Datastore datastore, Timestamp readTime)
+      throws DatastoreException {
+    throw new UnsupportedOperationException("Not implemented.");
+  }
+}
diff --git a/google-cloud-datastore-utils/src/main/java/com/google/datastore/utils/QuerySplitterImpl.java b/google-cloud-datastore-utils/src/main/java/com/google/datastore/utils/QuerySplitterImpl.java
new file mode 100644
index 00000000..ac2a6557
--- /dev/null
+++ b/google-cloud-datastore-utils/src/main/java/com/google/datastore/utils/QuerySplitterImpl.java
@@ -0,0 +1,309 @@
+/*
+ * Copyright 2024 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.utils;
+
+import static com.google.datastore.utils.DatastoreHelper.makeAndFilter;
+
+import com.google.api.core.BetaApi;
+import com.google.datastore.v1.EntityResult;
+import com.google.datastore.v1.Filter;
+import com.google.datastore.v1.Key;
+import com.google.datastore.v1.PartitionId;
+import com.google.datastore.v1.Projection;
+import com.google.datastore.v1.PropertyFilter;
+import com.google.datastore.v1.PropertyFilter.Operator;
+import com.google.datastore.v1.PropertyOrder.Direction;
+import com.google.datastore.v1.PropertyReference;
+import com.google.datastore.v1.Query;
+import com.google.datastore.v1.QueryResultBatch;
+import com.google.datastore.v1.QueryResultBatch.MoreResultsType;
+import com.google.datastore.v1.ReadOptions;
+import com.google.datastore.v1.RunQueryRequest;
+import com.google.protobuf.Timestamp;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.EnumSet;
+import java.util.List;
+import javax.annotation.Nullable;
+
+/**
+ * Provides the ability to split a query into multiple shards using Cloud Datastore.
+ *
+ * <p>This implementation of the QuerySplitter uses the __scatter__ property to gather random split
+ * points for a query.
+ */
+final class QuerySplitterImpl implements QuerySplitter {
+
+  /** The number of keys to sample for each split. * */
+  private static final int KEYS_PER_SPLIT = 32;
+
+  private static final EnumSet<Operator> UNSUPPORTED_OPERATORS =
+      EnumSet.of(
+          Operator.LESS_THAN,
+          Operator.LESS_THAN_OR_EQUAL,
+          Operator.GREATER_THAN,
+          Operator.GREATER_THAN_OR_EQUAL);
+
+  static final QuerySplitter INSTANCE = new QuerySplitterImpl();
+
+  private QuerySplitterImpl() {
+    // No initialization required.
+  }
+
+  @Override
+  public List<Query> getSplits(
+      Query query, PartitionId partition, int numSplits, Datastore datastore)
+      throws DatastoreException, IllegalArgumentException {
+    return getSplitsInternal(query, partition, numSplits, datastore, null);
+  }
+
+  @BetaApi
+  @Override
+  public List<Query> getSplits(
+      Query query, PartitionId partition, int numSplits, Datastore datastore, Timestamp readTime)
+      throws DatastoreException, IllegalArgumentException {
+    return getSplitsInternal(query, partition, numSplits, datastore, readTime);
+  }
+
+  private List<Query> getSplitsInternal(
+      Query query,
+      PartitionId partition,
+      int numSplits,
+      Datastore datastore,
+      @Nullable Timestamp readTime)
+      throws DatastoreException, IllegalArgumentException {
+    List<Query> splits = new ArrayList<Query>(numSplits);
+    if (numSplits == 1) {
+      splits.add(query);
+      return splits;
+    }
+    validateQuery(query);
+    validateSplitSize(numSplits);
+
+    List<Key> scatterKeys = getScatterKeys(numSplits, query, partition, datastore, readTime);
+    Key lastKey = null;
+    for (Key nextKey : getSplitKey(scatterKeys, numSplits)) {
+      splits.add(createSplit(lastKey, nextKey, query));
+      lastKey = nextKey;
+    }
+    splits.add(createSplit(lastKey, null, query));
+    return splits;
+  }
+
+  /**
+   * Verify that the given number of splits is not out of bounds.
+   *
+   * @param numSplits the number of splits.
+   * @throws IllegalArgumentException if the split size is invalid.
+   */
+  private void validateSplitSize(int numSplits) throws IllegalArgumentException {
+    if (numSplits < 1) {
+      throw new IllegalArgumentException("The number of splits must be greater than 0.");
+    }
+  }
+
+  /**
+   * Validates that we only have allowable filters.
+   *
+   * <p>Note that equality and ancestor filters are allowed, however they may result in inefficient
+   * sharding.
+   */
+  private void validateFilter(Filter filter) throws IllegalArgumentException {
+    switch (filter.getFilterTypeCase()) {
+      case COMPOSITE_FILTER:
+        for (Filter subFilter : filter.getCompositeFilter().getFiltersList()) {
+          validateFilter(subFilter);
+        }
+        break;
+      case PROPERTY_FILTER:
+        if (UNSUPPORTED_OPERATORS.contains(filter.getPropertyFilter().getOp())) {
+          throw new IllegalArgumentException("Query cannot have any inequality filters.");
+        }
+        break;
+      default:
+        throw new IllegalArgumentException(
+            "Unsupported filter type: " + filter.getFilterTypeCase());
+    }
+  }
+
+  /**
+   * Verifies that the given query can be properly scattered.
+   *
+   * @param query the query to verify
+   * @throws IllegalArgumentException if the query is invalid.
+   */
+  private void validateQuery(Query query) throws IllegalArgumentException {
+    if (query.getKindCount() != 1) {
+      throw new IllegalArgumentException("Query must have exactly one kind.");
+    }
+    if (query.getOrderCount() != 0) {
+      throw new IllegalArgumentException("Query cannot have any sort orders.");
+    }
+    if (query.hasFilter()) {
+      validateFilter(query.getFilter());
+    }
+  }
+
+  /**
+   * Create a new {@link Query} given the query and range.
+   *
+   * @param lastKey the previous key. If null then assumed to be the beginning.
+   * @param nextKey the next key. If null then assumed to be the end.
+   * @param query the desired query.
+   */
+  private Query createSplit(Key lastKey, Key nextKey, Query query) {
+    if (lastKey == null && nextKey == null) {
+      return query;
+    }
+    List<Filter> keyFilters = new ArrayList<Filter>();
+    if (query.hasFilter()) {
+      keyFilters.add(query.getFilter());
+    }
+    if (lastKey != null) {
+      Filter lowerBound =
+          DatastoreHelper.makeFilter(
+                  DatastoreHelper.KEY_PROPERTY_NAME,
+                  PropertyFilter.Operator.GREATER_THAN_OR_EQUAL,
+                  DatastoreHelper.makeValue(lastKey))
+              .build();
+      keyFilters.add(lowerBound);
+    }
+    if (nextKey != null) {
+      Filter upperBound =
+          DatastoreHelper.makeFilter(
+                  DatastoreHelper.KEY_PROPERTY_NAME,
+                  PropertyFilter.Operator.LESS_THAN,
+                  DatastoreHelper.makeValue(nextKey))
+              .build();
+      keyFilters.add(upperBound);
+    }
+    return Query.newBuilder(query).setFilter(makeAndFilter(keyFilters)).build();
+  }
+
+  /**
+   * Gets a list of split keys given a desired number of splits.
+   *
+   * <p>This list will contain multiple split keys for each split. Only a single split key will be
+   * chosen as the split point, however providing multiple keys allows for more uniform sharding.
+   *
+   * @param numSplits the number of desired splits.
+   * @param query the user query.
+   * @param partition the partition to run the query in.
+   * @param datastore the datastore containing the data.
+   * @param readTime read time at which to get the split keys from the datastore.
+   * @throws com.google.datastore.utils.DatastoreException if there was an error when executing the
+   *     datastore query.
+   */
+  private List<Key> getScatterKeys(
+      int numSplits,
+      Query query,
+      PartitionId partition,
+      Datastore datastore,
+      @Nullable Timestamp readTime)
+      throws DatastoreException {
+    Query.Builder scatterPointQuery = createScatterQuery(query, numSplits);
+
+    List<Key> keySplits = new ArrayList<Key>();
+
+    QueryResultBatch batch;
+    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());
+      }
+      batch = datastore.runQuery(scatterRequest.build()).getBatch();
+      for (EntityResult result : batch.getEntityResultsList()) {
+        keySplits.add(result.getEntity().getKey());
+      }
+      scatterPointQuery.setStartCursor(batch.getEndCursor());
+      scatterPointQuery
+          .getLimitBuilder()
+          .setValue(scatterPointQuery.getLimit().getValue() - batch.getEntityResultsCount());
+    } while (batch.getMoreResults() == MoreResultsType.NOT_FINISHED);
+    Collections.sort(keySplits, DatastoreHelper.getKeyComparator());
+    return keySplits;
+  }
+
+  /**
+   * Creates a scatter query from the given user query
+   *
+   * @param query the user's query.
+   * @param numSplits the number of splits to create.
+   */
+  private Query.Builder createScatterQuery(Query query, int numSplits) {
+    // TODO(pcostello): We can potentially support better splits with equality filters in our query
+    // if there exists a composite index on property, __scatter__, __key__. Until an API for
+    // metadata exists, this isn't possible. Note that ancestor and inequality queries fall into
+    // the same category.
+    Query.Builder scatterPointQuery = Query.newBuilder();
+    scatterPointQuery.addAllKind(query.getKindList());
+    scatterPointQuery.addOrder(
+        DatastoreHelper.makeOrder(DatastoreHelper.SCATTER_PROPERTY_NAME, Direction.ASCENDING));
+    // There is a split containing entities before and after each scatter entity:
+    // ||---*------*------*------*------*------*------*---||  = scatter entity
+    // If we represent each split as a region before a scatter entity, there is an extra region
+    // following the last scatter point. Thus, we do not need the scatter entities for the last
+    // region.
+    scatterPointQuery.getLimitBuilder().setValue((numSplits - 1) * KEYS_PER_SPLIT);
+    scatterPointQuery.addProjection(
+        Projection.newBuilder().setProperty(PropertyReference.newBuilder().setName("__key__")));
+    return scatterPointQuery;
+  }
+
+  /**
+   * Given a list of keys and a number of splits find the keys to split on.
+   *
+   * @param keys the list of keys.
+   * @param numSplits the number of splits.
+   */
+  private Iterable<Key> getSplitKey(List<Key> keys, int numSplits) {
+    // If the number of keys is less than the number of splits, we are limited in the number of
+    // splits we can make.
+    if (keys.size() < numSplits - 1) {
+      return keys;
+    }
+
+    // Calculate the number of keys per split. This should be KEYS_PER_SPLIT, but may
+    // be less if there are not KEYS_PER_SPLIT * (numSplits - 1) scatter entities.
+    //
+    // Consider the following dataset, where - represents an entity and * represents an entity
+    // that is returned as a scatter entity:
+    // ||---*-----*----*-----*-----*------*----*----||
+    // If we want 4 splits in this data, the optimal split would look like:
+    // ||---*-----*----*-----*-----*------*----*----||
+    //            |          |            |
+    // The scatter keys in the last region are not useful to us, so we never request them:
+    // ||---*-----*----*-----*-----*------*---------||
+    //            |          |            |
+    // With 6 scatter keys we want to set scatter points at indexes: 1, 3, 5.
+    //
+    // We keep this as a double so that any "fractional" keys per split get distributed throughout
+    // the splits and don't make the last split significantly larger than the rest.
+    double numKeysPerSplit = Math.max(1.0, ((double) keys.size()) / (numSplits - 1));
+
+    List<Key> keysList = new ArrayList<Key>(numSplits - 1);
+    // Grab the last sample for each split, otherwise the first split will be too small.
+    for (int i = 1; i < numSplits; i++) {
+      int splitIndex = (int) Math.round(i * numKeysPerSplit) - 1;
+      keysList.add(keys.get(splitIndex));
+    }
+
+    return keysList;
+  }
+}
diff --git a/google-cloud-datastore-utils/src/main/java/com/google/datastore/utils/RemoteRpc.java b/google-cloud-datastore-utils/src/main/java/com/google/datastore/utils/RemoteRpc.java
new file mode 100644
index 00000000..492936e1
--- /dev/null
+++ b/google-cloud-datastore-utils/src/main/java/com/google/datastore/utils/RemoteRpc.java
@@ -0,0 +1,239 @@
+/*
+ * Copyright 2024 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.utils;
+
+import com.google.api.client.http.*;
+import com.google.api.client.http.protobuf.ProtoHttpContent;
+import com.google.api.client.util.IOUtils;
+import com.google.common.annotations.VisibleForTesting;
+import com.google.common.base.Strings;
+import com.google.protobuf.MessageLite;
+import com.google.rpc.Code;
+import com.google.rpc.Status;
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.io.InputStream;
+import java.net.SocketTimeoutException;
+import java.nio.charset.Charset;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.logging.Logger;
+
+/**
+ * An RPC transport that sends protocol buffers over HTTP.
+ *
+ * <p>This class is thread-safe.
+ */
+class RemoteRpc {
+  private static final Logger logger = Logger.getLogger(RemoteRpc.class.getName());
+
+  @VisibleForTesting static final String API_FORMAT_VERSION_HEADER = "X-Goog-Api-Format-Version";
+  private static final String API_FORMAT_VERSION = "2";
+
+  @VisibleForTesting static final String X_GOOG_REQUEST_PARAMS_HEADER = "x-goog-request-params";
+
+  private final HttpRequestFactory client;
+  private final HttpRequestInitializer initializer;
+  private final String url;
+  private final AtomicInteger rpcCount = new AtomicInteger(0);
+  // Not final - so it can be set/reset in Unittests
+  private static boolean enableE2EChecksum =
+      Boolean.parseBoolean(System.getenv("GOOGLE_CLOUD_DATASTORE_HTTP_ENABLE_E2E_CHECKSUM"));
+
+  RemoteRpc(HttpRequestFactory client, HttpRequestInitializer initializer, String url) {
+    this.client = client;
+    this.initializer = initializer;
+    this.url = url;
+    try {
+      resolveURL("dummyRpc");
+    } catch (Exception e) {
+      throw new IllegalArgumentException(
+          "Unable to construct RemoteRpc due to unsupported url: <" + url + ">", e);
+    }
+  }
+
+  /**
+   * Makes an RPC call using the client. Logs how long it took and any exceptions.
+   *
+   * <p>NOTE: The request could be an InputStream too, but the http client will need to find its
+   * length, which will require buffering it anyways.
+   *
+   * @throws com.google.datastore.utils.DatastoreException if the RPC fails.
+   */
+  public InputStream call(
+      String methodName, MessageLite request, String projectId, String databaseId)
+      throws com.google.datastore.utils.DatastoreException {
+    logger.fine("remote datastore call " + methodName);
+
+    long startTime = System.currentTimeMillis();
+    try {
+      HttpResponse httpResponse;
+      try {
+        rpcCount.incrementAndGet();
+        ProtoHttpContent payload = new ProtoHttpContent(request);
+        HttpRequest httpRequest = client.buildPostRequest(resolveURL(methodName), payload);
+        setHeaders(request, httpRequest, projectId, databaseId);
+        // Don't throw an HTTPResponseException on error. It converts the response to a String and
+        // throws away the original, whereas we need the raw bytes to parse it as a proto.
+        httpRequest.setThrowExceptionOnExecuteError(false);
+        // Datastore requests typically time out after 60s; set the read timeout to slightly longer
+        // than that by default (can be overridden via the HttpRequestInitializer).
+        httpRequest.setReadTimeout(65 * 1000);
+        if (initializer != null) {
+          initializer.initialize(httpRequest);
+        }
+        httpResponse = httpRequest.execute();
+        if (!httpResponse.isSuccessStatusCode()) {
+          try (InputStream content = httpResponse.getContent()) {
+            throw makeException(
+                url,
+                methodName,
+                content,
+                httpResponse.getContentType(),
+                httpResponse.getContentCharset(),
+                null,
+                httpResponse.getStatusCode());
+          }
+        }
+        InputStream inputStream = httpResponse.getContent();
+        return inputStream;
+      } catch (SocketTimeoutException e) {
+        throw makeException(url, methodName, Code.DEADLINE_EXCEEDED, "Deadline exceeded", e);
+      } catch (IOException e) {
+        throw makeException(url, methodName, Code.UNAVAILABLE, "I/O error", e);
+      }
+    } finally {
+      long elapsedTime = System.currentTimeMillis() - startTime;
+      logger.fine("remote datastore call " + methodName + " took " + elapsedTime + " ms");
+    }
+  }
+
+  @VisibleForTesting
+  void setHeaders(
+      MessageLite request, HttpRequest httpRequest, String projectId, String databaseId) {
+    httpRequest.getHeaders().put(API_FORMAT_VERSION_HEADER, API_FORMAT_VERSION);
+    StringBuilder builder = new StringBuilder("project_id=");
+    builder.append(projectId);
+    if (!Strings.isNullOrEmpty(databaseId)) {
+      builder.append("&database_id=");
+      builder.append(databaseId);
+    }
+    httpRequest.getHeaders().put(X_GOOG_REQUEST_PARAMS_HEADER, builder.toString());
+  }
+
+  @VisibleForTesting
+  HttpRequestFactory getClient() {
+    return client;
+  }
+
+  @VisibleForTesting
+  static void setSystemEnvE2EChecksum(boolean enableE2EChecksum) {
+    RemoteRpc.enableE2EChecksum = enableE2EChecksum;
+  }
+
+  void resetRpcCount() {
+    rpcCount.set(0);
+  }
+
+  int getRpcCount() {
+    return rpcCount.get();
+  }
+
+  public String getUrl() {
+    return url;
+  }
+
+  GenericUrl resolveURL(String path) {
+    return new GenericUrl(url + ":" + path);
+  }
+
+  HttpRequestFactory getHttpRequestFactory() {
+    return client;
+  }
+
+  public static com.google.datastore.utils.DatastoreException makeException(
+      String url, String methodName, Code code, String message, Throwable cause) {
+    logger.fine("remote datastore call " + methodName + " against " + url + " failed: " + message);
+    return new com.google.datastore.utils.DatastoreException(methodName, code, message, cause);
+  }
+
+  static DatastoreException makeException(
+      String url,
+      String methodName,
+      InputStream content,
+      String contentType,
+      Charset contentCharset,
+      Throwable cause,
+      int httpStatusCode) {
+    if (!contentType.equals("application/x-protobuf")) {
+      String responseContent;
+      try {
+        ByteArrayOutputStream out = new ByteArrayOutputStream();
+        IOUtils.copy(content, out, false);
+        responseContent = out.toString(contentCharset.name());
+      } catch (IOException e) {
+        responseContent = "";
+      }
+      return makeException(
+          url,
+          methodName,
+          Code.INTERNAL,
+          String.format(
+              "Non-protobuf error: %s. HTTP status code was %d.", responseContent, httpStatusCode),
+          cause);
+    }
+
+    Status rpcStatus;
+    try {
+      rpcStatus = Status.parseFrom(content);
+    } catch (IOException e) {
+      return makeException(
+          url,
+          methodName,
+          Code.INTERNAL,
+          String.format(
+              "Unable to parse Status protocol buffer: HTTP status code was %s.", httpStatusCode),
+          e);
+    }
+
+    Code code = Code.forNumber(rpcStatus.getCode());
+    if (code == null) {
+      return makeException(
+          url,
+          methodName,
+          Code.INTERNAL,
+          String.format(
+              "Invalid error code: %d. Message: %s.", rpcStatus.getCode(), rpcStatus.getMessage()),
+          cause);
+    } else if (code == Code.OK) {
+      // We can end up here because there was no response body (and we successfully parsed an
+      // empty Status message). This may happen for 401s in particular due to special handling
+      // in low-level HTTP libraries.
+      if (httpStatusCode == HttpStatusCodes.STATUS_CODE_UNAUTHORIZED) {
+        return makeException(url, methodName, Code.UNAUTHENTICATED, "Unauthenticated.", cause);
+      }
+      return makeException(
+          url,
+          methodName,
+          Code.INTERNAL,
+          String.format(
+              "Unexpected OK error code with HTTP status code of %d. Message: %s.",
+              httpStatusCode, rpcStatus.getMessage()),
+          cause);
+    }
+
+    return makeException(url, methodName, code, rpcStatus.getMessage(), cause);
+  }
+}
diff --git a/pom.xml b/pom.xml
index aec03ebc..549537da 100644
--- a/pom.xml
+++ b/pom.xml
@@ -181,6 +181,11 @@
         <artifactId>grpc-google-cloud-datastore-v1</artifactId>
         <version>2.22.0-grpc-experimental-1-SNAPSHOT</version><!-- {x-version-update:grpc-google-cloud-datastore-v1:current} -->
       </dependency>
+      <dependency>
+        <groupId>com.google.cloud</groupId>
+        <artifactId>google-cloud-datastore-utils</artifactId>
+        <version>2.22.0-grpc-experimental-1-SNAPSHOT</version><!-- {x-version-update:google-cloud-datastore-utils:current} -->
+      </dependency>
       <dependency>
         <groupId>com.google.cloud.datastore</groupId>
         <artifactId>datastore-v1-proto-client</artifactId>
@@ -277,6 +282,7 @@
     <module>grpc-google-cloud-datastore-v1</module>
     <module>datastore-v1-proto-client</module>
     <module>google-cloud-datastore-bom</module>
+    <module>google-cloud-datastore-utils</module>
   </modules>
 
   </project>


Test patch:
diff --git a/google-cloud-datastore-utils/src/main/java/com/google/datastore/utils/testing/MockCredential.java b/google-cloud-datastore-utils/src/main/java/com/google/datastore/utils/testing/MockCredential.java
new file mode 100644
index 00000000..d5d16bb6
--- /dev/null
+++ b/google-cloud-datastore-utils/src/main/java/com/google/datastore/utils/testing/MockCredential.java
@@ -0,0 +1,36 @@
+/*
+ * 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.utils.testing;
+
+import com.google.api.client.auth.oauth2.Credential;
+import com.google.api.client.http.HttpRequest;
+import java.io.IOException;
+
+/** Fake credential used for testing purpose. */
+public class MockCredential extends Credential {
+  public MockCredential() {
+    super(
+        new AccessMethod() {
+          @Override
+          public void intercept(HttpRequest request, String accessToken) throws IOException {}
+
+          @Override
+          public String getAccessTokenFromRequest(HttpRequest request) {
+            return "MockAccessToken";
+          }
+        });
+  }
+}
diff --git a/google-cloud-datastore-utils/src/main/java/com/google/datastore/utils/testing/MockDatastoreFactory.java b/google-cloud-datastore-utils/src/main/java/com/google/datastore/utils/testing/MockDatastoreFactory.java
new file mode 100644
index 00000000..d4dd5cae
--- /dev/null
+++ b/google-cloud-datastore-utils/src/main/java/com/google/datastore/utils/testing/MockDatastoreFactory.java
@@ -0,0 +1,132 @@
+/*
+ * 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.utils.testing;
+
+import static com.google.common.base.Preconditions.checkState;
+
+import com.google.api.client.auth.oauth2.Credential;
+import com.google.api.client.http.*;
+import com.google.api.client.testing.http.MockHttpTransport;
+import com.google.api.client.testing.http.MockLowLevelHttpRequest;
+import com.google.api.client.testing.http.MockLowLevelHttpResponse;
+import com.google.api.client.testing.util.TestableByteArrayInputStream;
+import com.google.common.collect.Iterables;
+import com.google.datastore.utils.DatastoreFactory;
+import com.google.datastore.utils.DatastoreOptions;
+import com.google.protobuf.Message;
+import com.google.rpc.Code;
+import com.google.rpc.Status;
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.util.List;
+
+/** Fake Datastore factory used for testing purposes when a true Datastore service is not needed. */
+public class MockDatastoreFactory extends DatastoreFactory {
+  private int nextStatus;
+  private Message nextResponse;
+  private Status nextError;
+  private IOException nextException;
+
+  private String lastPath;
+  private String lastMimeType;
+  private byte[] lastBody;
+  private List<String> lastCookies;
+  private String lastApiFormatHeaderValue;
+
+  public void setNextResponse(Message response) {
+    nextStatus = HttpStatusCodes.STATUS_CODE_OK;
+    nextResponse = response;
+    nextError = null;
+    nextException = null;
+  }
+
+  public void setNextError(int status, Code code, String message) {
+    nextStatus = status;
+    nextResponse = null;
+    nextError = makeErrorContent(message, code);
+    nextException = null;
+  }
+
+  public void setNextException(IOException exception) {
+    nextStatus = 0;
+    nextResponse = null;
+    nextError = null;
+    nextException = exception;
+  }
+
+  @Override
+  public HttpRequestFactory makeClient(DatastoreOptions options) {
+    HttpTransport transport =
+        new MockHttpTransport() {
+          @Override
+          public LowLevelHttpRequest buildRequest(String method, String url) {
+            return new MockLowLevelHttpRequest(url) {
+              @Override
+              public LowLevelHttpResponse execute() throws IOException {
+                lastPath = new GenericUrl(getUrl()).getRawPath();
+                lastMimeType = getContentType();
+                lastCookies = getHeaderValues("Cookie");
+                lastApiFormatHeaderValue =
+                    Iterables.getOnlyElement(getHeaderValues("X-Goog-Api-Format-Version"));
+                ByteArrayOutputStream out = new ByteArrayOutputStream();
+                getStreamingContent().writeTo(out);
+                lastBody = out.toByteArray();
+                if (nextException != null) {
+                  throw nextException;
+                }
+                MockLowLevelHttpResponse response =
+                    new MockLowLevelHttpResponse()
+                        .setStatusCode(nextStatus)
+                        .setContentType("application/x-protobuf");
+                if (nextError != null) {
+                  checkState(nextResponse == null);
+                  response.setContent(new TestableByteArrayInputStream(nextError.toByteArray()));
+                } else {
+                  response.setContent(new TestableByteArrayInputStream(nextResponse.toByteArray()));
+                }
+                return response;
+              }
+            };
+          }
+        };
+    Credential credential = options.getCredential();
+    return transport.createRequestFactory(credential);
+  }
+
+  public String getLastPath() {
+    return lastPath;
+  }
+
+  public String getLastMimeType() {
+    return lastMimeType;
+  }
+
+  public String getLastApiFormatHeaderValue() {
+    return lastApiFormatHeaderValue;
+  }
+
+  public byte[] getLastBody() {
+    return lastBody;
+  }
+
+  public List<String> getLastCookies() {
+    return lastCookies;
+  }
+
+  private static Status makeErrorContent(String message, Code code) {
+    return Status.newBuilder().setCode(code.getNumber()).setMessage(message).build();
+  }
+}
diff --git a/google-cloud-datastore-utils/src/test/java/com/google/datastore/utils/DatastoreClientTest.java b/google-cloud-datastore-utils/src/test/java/com/google/datastore/utils/DatastoreClientTest.java
new file mode 100644
index 00000000..31b0f644
--- /dev/null
+++ b/google-cloud-datastore-utils/src/test/java/com/google/datastore/utils/DatastoreClientTest.java
@@ -0,0 +1,407 @@
+/*
+ * Copyright 2024 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.utils;
+
+import static com.google.common.truth.Truth.assertThat;
+import static org.junit.Assert.assertArrayEquals;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertSame;
+import static org.junit.Assert.assertThrows;
+import static org.junit.Assert.fail;
+
+import com.google.api.client.http.HttpRequest;
+import com.google.api.client.http.HttpRequestInitializer;
+import com.google.datastore.utils.testing.MockCredential;
+import com.google.datastore.utils.testing.MockDatastoreFactory;
+import com.google.datastore.v1.AllocateIdsRequest;
+import com.google.datastore.v1.AllocateIdsResponse;
+import com.google.datastore.v1.BeginTransactionRequest;
+import com.google.datastore.v1.BeginTransactionResponse;
+import com.google.datastore.v1.CommitRequest;
+import com.google.datastore.v1.CommitResponse;
+import com.google.datastore.v1.EntityResult;
+import com.google.datastore.v1.LookupRequest;
+import com.google.datastore.v1.LookupResponse;
+import com.google.datastore.v1.QueryResultBatch;
+import com.google.datastore.v1.ReserveIdsRequest;
+import com.google.datastore.v1.ReserveIdsResponse;
+import com.google.datastore.v1.RollbackRequest;
+import com.google.datastore.v1.RollbackResponse;
+import com.google.datastore.v1.RunAggregationQueryRequest;
+import com.google.datastore.v1.RunAggregationQueryResponse;
+import com.google.datastore.v1.RunQueryRequest;
+import com.google.datastore.v1.RunQueryResponse;
+import com.google.protobuf.ByteString;
+import com.google.protobuf.Message;
+import com.google.rpc.Code;
+import java.io.IOException;
+import java.lang.reflect.InvocationTargetException;
+import java.lang.reflect.Method;
+import java.net.SocketTimeoutException;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+/** Tests for {@link DatastoreFactory} and {@link Datastore}. */
+@RunWith(JUnit4.class)
+public class DatastoreClientTest {
+  private static final String PROJECT_ID = "project-id";
+
+  private DatastoreFactory factory = new MockDatastoreFactory();
+  private DatastoreOptions.Builder options =
+      new DatastoreOptions.Builder().projectId(PROJECT_ID).credential(new MockCredential());
+
+  @Test
+  public void options_NoProjectIdOrProjectEndpoint() {
+    IllegalArgumentException exception =
+        assertThrows(
+            IllegalArgumentException.class,
+            () -> factory.create(new DatastoreOptions.Builder().build()));
+    assertThat(exception)
+        .hasMessageThat()
+        .contains("Either project ID or project endpoint must be provided");
+    factory.create(options.build());
+  }
+
+  @Test
+  public void options_ProjectIdAndProjectEndpoint() throws Exception {
+    IllegalArgumentException exception =
+        assertThrows(
+            IllegalArgumentException.class,
+            () ->
+                new DatastoreOptions.Builder()
+                    .projectId(PROJECT_ID)
+                    .projectEndpoint(
+                        "http://localhost:1234/datastore/v1beta42/projects/project-id"));
+    assertThat(exception)
+        .hasMessageThat()
+        .contains("Cannot set both project endpoint and project ID");
+  }
+
+  @Test
+  public void options_LocalHostAndProjectEndpoint() throws Exception {
+    IllegalArgumentException exception =
+        assertThrows(
+            IllegalArgumentException.class,
+            () ->
+                new DatastoreOptions.Builder()
+                    .localHost("localhost:8080")
+                    .projectEndpoint(
+                        "http://localhost:1234/datastore/v1beta42/projects/project-id"));
+    assertThat(exception)
+        .hasMessageThat()
+        .contains("Can set at most one of project endpoint, host, and local host");
+  }
+
+  @Test
+  public void options_HostAndProjectEndpoint() throws Exception {
+    IllegalArgumentException exception =
+        assertThrows(
+            IllegalArgumentException.class,
+            () ->
+                new DatastoreOptions.Builder()
+                    .host("foo-datastore.googleapis.com")
+                    .projectEndpoint(
+                        "http://localhost:1234/datastore/v1beta42/projects/project-id"));
+    assertThat(exception)
+        .hasMessageThat()
+        .contains("Can set at most one of project endpoint, host, and local host");
+  }
+
+  @Test
+  public void options_HostAndLocalHost() throws Exception {
+    IllegalArgumentException exception =
+        assertThrows(
+            IllegalArgumentException.class,
+            () ->
+                new DatastoreOptions.Builder()
+                    .host("foo-datastore.googleapis.com")
+                    .localHost("localhost:8080"));
+    assertThat(exception)
+        .hasMessageThat()
+        .contains("Can set at most one of project endpoint, host, and local host");
+  }
+
+  @Test
+  public void options_InvalidLocalHost() throws Exception {
+    IllegalArgumentException exception =
+        assertThrows(
+            IllegalArgumentException.class,
+            () ->
+                factory.create(
+                    new DatastoreOptions.Builder()
+                        .projectId(PROJECT_ID)
+                        .localHost("!not a valid url!")
+                        .build()));
+    assertThat(exception).hasMessageThat().contains("Illegal character");
+  }
+
+  @Test
+  public void options_SchemeInLocalHost() {
+    IllegalArgumentException exception =
+        assertThrows(
+            IllegalArgumentException.class,
+            () -> new DatastoreOptions.Builder().localHost("http://localhost:8080"));
+    assertThat(exception)
+        .hasMessageThat()
+        .contains("Local host \"http://localhost:8080\" must not include scheme");
+  }
+
+  @Test
+  public void options_InvalidHost() {
+    IllegalArgumentException exception =
+        assertThrows(
+            IllegalArgumentException.class,
+            () ->
+                factory.create(
+                    new DatastoreOptions.Builder()
+                        .projectId(PROJECT_ID)
+                        .host("!not a valid url!")
+                        .build()));
+    assertThat(exception).hasMessageThat().contains("Illegal character");
+  }
+
+  @Test
+  public void options_SchemeInHost() {
+    IllegalArgumentException exception =
+        assertThrows(
+            IllegalArgumentException.class,
+            () -> new DatastoreOptions.Builder().host("http://foo-datastore.googleapis.com"));
+
+    assertThat(exception)
+        .hasMessageThat()
+        .contains("Host \"http://foo-datastore.googleapis.com\" must not include scheme.");
+  }
+
+  @Test
+  public void create_NullOptions() throws Exception {
+    assertThrows(NullPointerException.class, () -> factory.create(null));
+  }
+
+  @Test
+  public void create_Host() {
+    Datastore datastore =
+        factory.create(
+            new DatastoreOptions.Builder()
+                .projectId(PROJECT_ID)
+                .host("foo-datastore.googleapis.com")
+                .build());
+    assertThat(datastore.remoteRpc.getUrl())
+        .isEqualTo("https://foo-datastore.googleapis.com/v1/projects/project-id");
+  }
+
+  @Test
+  public void create_LocalHost() {
+    Datastore datastore =
+        factory.create(
+            new DatastoreOptions.Builder()
+                .projectId(PROJECT_ID)
+                .localHost("localhost:8080")
+                .build());
+    assertThat(datastore.remoteRpc.getUrl())
+        .isEqualTo("http://localhost:8080/v1/projects/project-id");
+  }
+
+  @Test
+  public void create_LocalHostIp() {
+    Datastore datastore =
+        factory.create(
+            new DatastoreOptions.Builder()
+                .projectId(PROJECT_ID)
+                .localHost("127.0.0.1:8080")
+                .build());
+    assertThat(datastore.remoteRpc.getUrl())
+        .isEqualTo("http://127.0.0.1:8080/v1/projects/project-id");
+  }
+
+  @Test
+  public void create_DefaultHost() {
+    Datastore datastore =
+        factory.create(new DatastoreOptions.Builder().projectId(PROJECT_ID).build());
+    assertThat(datastore.remoteRpc.getUrl())
+        .isEqualTo("https://datastore.googleapis.com/v1/projects/project-id");
+  }
+
+  @Test
+  public void create_ProjectEndpoint() {
+    Datastore datastore =
+        factory.create(
+            new DatastoreOptions.Builder()
+                .projectEndpoint("http://prom-qa/datastore/v1beta42/projects/project-id")
+                .build());
+    assertThat(datastore.remoteRpc.getUrl())
+        .isEqualTo("http://prom-qa/datastore/v1beta42/projects/project-id");
+  }
+
+  @Test
+  public void create_ProjectEndpointNoScheme() {
+    IllegalArgumentException exception =
+        assertThrows(
+            IllegalArgumentException.class,
+            () ->
+                factory.create(
+                    new DatastoreOptions.Builder()
+                        .projectEndpoint("localhost:1234/datastore/v1beta42/projects/project-id")
+                        .build()));
+    assertThat(exception)
+        .hasMessageThat()
+        .contains(
+            "Project endpoint \"localhost:1234/datastore/v1beta42/projects/project-id\" must"
+                + " include scheme.");
+  }
+
+  @Test
+  public void initializer() throws Exception {
+    options.initializer(
+        new HttpRequestInitializer() {
+          @Override
+          public void initialize(HttpRequest request) {
+            request.getHeaders().setCookie("magic");
+          }
+        });
+    Datastore datastore = factory.create(options.build());
+    MockDatastoreFactory mockClient = (MockDatastoreFactory) factory;
+    AllocateIdsRequest request = AllocateIdsRequest.newBuilder().build();
+    AllocateIdsResponse response = AllocateIdsResponse.newBuilder().build();
+    mockClient.setNextResponse(response);
+    assertEquals(response, datastore.allocateIds(request));
+    assertEquals("magic", mockClient.getLastCookies().get(0));
+  }
+
+  @Test
+  public void allocateIds() throws Exception {
+    AllocateIdsRequest.Builder request = AllocateIdsRequest.newBuilder();
+    AllocateIdsResponse.Builder response = AllocateIdsResponse.newBuilder();
+    expectRpc("allocateIds", request.build(), response.build());
+  }
+
+  @Test
+  public void lookup() throws Exception {
+    LookupRequest.Builder request = LookupRequest.newBuilder();
+    LookupResponse.Builder response = LookupResponse.newBuilder();
+    expectRpc("lookup", request.build(), response.build());
+  }
+
+  @Test
+  public void beginTransaction() throws Exception {
+    BeginTransactionRequest.Builder request = BeginTransactionRequest.newBuilder();
+    BeginTransactionResponse.Builder response = BeginTransactionResponse.newBuilder();
+    response.setTransaction(ByteString.copyFromUtf8("project-id"));
+    expectRpc("beginTransaction", request.build(), response.build());
+  }
+
+  @Test
+  public void commit() throws Exception {
+    CommitRequest.Builder request = CommitRequest.newBuilder();
+    request.setTransaction(ByteString.copyFromUtf8("project-id"));
+    CommitResponse.Builder response = CommitResponse.newBuilder();
+    expectRpc("commit", request.build(), response.build());
+  }
+
+  @Test
+  public void reserveIds() throws Exception {
+    ReserveIdsRequest.Builder request = ReserveIdsRequest.newBuilder();
+    ReserveIdsResponse.Builder response = ReserveIdsResponse.newBuilder();
+    expectRpc("reserveIds", request.build(), response.build());
+  }
+
+  @Test
+  public void rollback() throws Exception {
+    RollbackRequest.Builder request = RollbackRequest.newBuilder();
+    request.setTransaction(ByteString.copyFromUtf8("project-id"));
+    RollbackResponse.Builder response = RollbackResponse.newBuilder();
+    expectRpc("rollback", request.build(), response.build());
+  }
+
+  @Test
+  public void runQuery() throws Exception {
+    RunQueryRequest.Builder request = RunQueryRequest.newBuilder();
+    request.getQueryBuilder();
+    RunQueryResponse.Builder response = RunQueryResponse.newBuilder();
+    response
+        .getBatchBuilder()
+        .setEntityResultType(EntityResult.ResultType.FULL)
+        .setMoreResults(QueryResultBatch.MoreResultsType.NOT_FINISHED);
+    expectRpc("runQuery", request.build(), response.build());
+  }
+
+  @Test
+  public void runAggregationQuery() throws Exception {
+    RunAggregationQueryRequest.Builder request = RunAggregationQueryRequest.newBuilder();
+    RunAggregationQueryResponse.Builder response = RunAggregationQueryResponse.newBuilder();
+    expectRpc("runAggregationQuery", request.build(), response.build());
+  }
+
+  private void expectRpc(String methodName, Message request, Message response) throws Exception {
+    Datastore datastore = factory.create(options.build());
+    MockDatastoreFactory mockClient = (MockDatastoreFactory) factory;
+
+    mockClient.setNextResponse(response);
+    @SuppressWarnings("rawtypes")
+    Class[] methodArgs = {request.getClass()};
+    Method call = Datastore.class.getMethod(methodName, methodArgs);
+    Object[] callArgs = {request};
+    assertEquals(response, call.invoke(datastore, callArgs));
+
+    assertEquals("/v1/projects/project-id:" + methodName, mockClient.getLastPath());
+    assertEquals("application/x-protobuf", mockClient.getLastMimeType());
+    assertEquals("2", mockClient.getLastApiFormatHeaderValue());
+    assertArrayEquals(request.toByteArray(), mockClient.getLastBody());
+    assertEquals(1, datastore.getRpcCount());
+
+    datastore.resetRpcCount();
+    assertEquals(0, datastore.getRpcCount());
+
+    mockClient.setNextError(400, Code.INVALID_ARGUMENT, "oops");
+    try {
+      call.invoke(datastore, callArgs);
+      fail();
+    } catch (InvocationTargetException targetException) {
+      DatastoreException exception = (DatastoreException) targetException.getCause();
+      assertEquals(Code.INVALID_ARGUMENT, exception.getCode());
+      assertEquals(methodName, exception.getMethodName());
+      assertEquals("oops", exception.getMessage());
+    }
+
+    SocketTimeoutException socketTimeoutException = new SocketTimeoutException("ste");
+    mockClient.setNextException(socketTimeoutException);
+    try {
+      call.invoke(datastore, callArgs);
+      fail();
+    } catch (InvocationTargetException targetException) {
+      DatastoreException exception = (DatastoreException) targetException.getCause();
+      assertEquals(Code.DEADLINE_EXCEEDED, exception.getCode());
+      assertEquals(methodName, exception.getMethodName());
+      assertEquals("Deadline exceeded", exception.getMessage());
+      assertSame(socketTimeoutException, exception.getCause());
+    }
+
+    IOException ioException = new IOException("ioe");
+    mockClient.setNextException(ioException);
+    try {
+      call.invoke(datastore, callArgs);
+      fail();
+    } catch (InvocationTargetException targetException) {
+      DatastoreException exception = (DatastoreException) targetException.getCause();
+      assertEquals(Code.UNAVAILABLE, exception.getCode());
+      assertEquals(methodName, exception.getMethodName());
+      assertEquals("I/O error", exception.getMessage());
+      assertSame(ioException, exception.getCause());
+    }
+
+    assertEquals(3, datastore.getRpcCount());
+  }
+}
diff --git a/google-cloud-datastore-utils/src/test/java/com/google/datastore/utils/DatastoreFactoryTest.java b/google-cloud-datastore-utils/src/test/java/com/google/datastore/utils/DatastoreFactoryTest.java
new file mode 100644
index 00000000..2a3d5a38
--- /dev/null
+++ b/google-cloud-datastore-utils/src/test/java/com/google/datastore/utils/DatastoreFactoryTest.java
@@ -0,0 +1,93 @@
+/*
+ * Copyright 2024 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.utils;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertNotSame;
+import static org.junit.Assert.assertTrue;
+
+import com.google.api.client.googleapis.auth.oauth2.GoogleCredential;
+import com.google.api.client.http.HttpRequestFactory;
+import com.google.api.client.http.javanet.NetHttpTransport;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+/** Test for {@link DatastoreFactory}. */
+@RunWith(JUnit4.class)
+public class DatastoreFactoryTest {
+  private static final String PROJECT_ID = "project-id";
+
+  private DatastoreFactory factory = DatastoreFactory.get();
+
+  /**
+   * Without specifying a credential or transport, the factory will create a default transport on
+   * its own.
+   */
+  @Test
+  public void makeClient_Default() {
+    DatastoreOptions options = new DatastoreOptions.Builder().projectId(PROJECT_ID).build();
+    HttpRequestFactory f = factory.makeClient(options);
+    assertNotNull(f.getTransport());
+    assertTrue(f.getTransport() instanceof NetHttpTransport);
+  }
+
+  /**
+   * Specifying a credential, but not a transport, the factory will use the transport from the
+   * credential.
+   */
+  @Test
+  public void makeClient_WithCredential() {
+    NetHttpTransport transport = new NetHttpTransport();
+    GoogleCredential credential = new GoogleCredential.Builder().setTransport(transport).build();
+    DatastoreOptions options =
+        new DatastoreOptions.Builder().projectId(PROJECT_ID).credential(credential).build();
+    HttpRequestFactory f = factory.makeClient(options);
+    assertEquals(transport, f.getTransport());
+  }
+
+  /** Specifying a transport, but not a credential, the factory will use the transport specified. */
+  @Test
+  public void makeClient_WithTransport() {
+    NetHttpTransport transport = new NetHttpTransport();
+    DatastoreOptions options =
+        new DatastoreOptions.Builder().projectId(PROJECT_ID).transport(transport).build();
+    HttpRequestFactory f = factory.makeClient(options);
+    assertEquals(transport, f.getTransport());
+  }
+
+  /**
+   * Specifying both credential and transport, the factory will use the transport specified and not
+   * the one in the credential.
+   */
+  @Test
+  public void makeClient_WithCredentialTransport() {
+    NetHttpTransport credTransport = new NetHttpTransport();
+    NetHttpTransport transport = new NetHttpTransport();
+    GoogleCredential credential =
+        new GoogleCredential.Builder().setTransport(credTransport).build();
+    DatastoreOptions options =
+        new DatastoreOptions.Builder()
+            .projectId(PROJECT_ID)
+            .credential(credential)
+            .transport(transport)
+            .build();
+    HttpRequestFactory f = factory.makeClient(options);
+    assertNotSame(credTransport, f.getTransport());
+    assertEquals(transport, f.getTransport());
+  }
+}
diff --git a/google-cloud-datastore-utils/src/test/java/com/google/datastore/utils/DatastoreHelperTest.java b/google-cloud-datastore-utils/src/test/java/com/google/datastore/utils/DatastoreHelperTest.java
new file mode 100644
index 00000000..24620244
--- /dev/null
+++ b/google-cloud-datastore-utils/src/test/java/com/google/datastore/utils/DatastoreHelperTest.java
@@ -0,0 +1,317 @@
+/*
+ * Copyright 2024 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.utils;
+
+import static com.google.datastore.utils.DatastoreHelper.*;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.fail;
+
+import com.google.datastore.v1.Key;
+import com.google.datastore.v1.PartitionId;
+import com.google.datastore.v1.Value;
+import com.google.datastore.v1.Value.ValueTypeCase;
+import com.google.protobuf.ByteString;
+import com.google.protobuf.Timestamp;
+import java.util.Date;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+/** Tests for {@link DatastoreHelper}. */
+@RunWith(JUnit4.class)
+public class DatastoreHelperTest {
+
+  private static final Key PARENT =
+      Key.newBuilder().addPath(Key.PathElement.newBuilder().setKind("Parent").setId(23L)).build();
+  private static final Key GRANDPARENT =
+      Key.newBuilder()
+          .addPath(Key.PathElement.newBuilder().setKind("Grandparent").setId(24L))
+          .build();
+  private static final Key CHILD =
+      Key.newBuilder().addPath(Key.PathElement.newBuilder().setKind("Child").setId(26L)).build();
+
+  @Test
+  public void testMakeKey_BadTypeForKind() {
+    try {
+      DatastoreHelper.makeKey(new Object());
+      fail("Expected IllegalArgumentException");
+    } catch (IllegalArgumentException expected) {
+    }
+  }
+
+  @Test
+  public void testMakeKey_BadTypeForNameId() {
+    try {
+      DatastoreHelper.makeKey("kind", new Object());
+      fail("Expected IllegalArgumentException");
+    } catch (IllegalArgumentException expected) {
+    }
+  }
+
+  @Test
+  public void testMakeKey_Empty() {
+    assertEquals(Key.newBuilder().build(), DatastoreHelper.makeKey().build());
+  }
+
+  @Test
+  public void testMakeKey_Incomplete() {
+    assertEquals(
+        Key.newBuilder().addPath(Key.PathElement.newBuilder().setKind("Foo")).build(),
+        makeKey("Foo").build());
+  }
+
+  @Test
+  public void testMakeKey_IdInt() {
+    assertEquals(
+        Key.newBuilder().addPath(Key.PathElement.newBuilder().setKind("Foo").setId(1)).build(),
+        makeKey("Foo", 1).build());
+  }
+
+  @Test
+  public void testMakeKey_IdLong() {
+    assertEquals(
+        Key.newBuilder().addPath(Key.PathElement.newBuilder().setKind("Foo").setId(1)).build(),
+        makeKey("Foo", 1L).build());
+  }
+
+  @Test
+  public void testMakeKey_IdShort() {
+    assertEquals(
+        Key.newBuilder().addPath(Key.PathElement.newBuilder().setKind("Foo").setId(1)).build(),
+        makeKey("Foo", (short) 1).build());
+  }
+
+  @Test
+  public void testMakeKey_Name() {
+    assertEquals(
+        Key.newBuilder().addPath(Key.PathElement.newBuilder().setKind("Foo").setName("hi")).build(),
+        makeKey("Foo", "hi").build());
+  }
+
+  @Test
+  public void testMakeKey_KindNameKind() {
+    assertEquals(
+        Key.newBuilder()
+            .addPath(Key.PathElement.newBuilder().setKind("Foo").setName("hi"))
+            .addPath(Key.PathElement.newBuilder().setKind("Bar"))
+            .build(),
+        makeKey("Foo", "hi", "Bar").build());
+  }
+
+  @Test
+  public void testMakeKey_KeyKind() {
+    // 1 key at the beginning of the series
+    assertEquals(
+        Key.newBuilder()
+            .addPath(PARENT.getPath(0))
+            .addPath(Key.PathElement.newBuilder().setKind("Child"))
+            .build(),
+        makeKey(PARENT, "Child").build());
+  }
+
+  @Test
+  public void testMakeKey_KindIdKeyKind() {
+    // 1 key in the middle of the series
+    assertEquals(
+        Key.newBuilder()
+            .addPath(Key.PathElement.newBuilder().setKind("Grandparent").setId(24L))
+            .addPath(PARENT.getPath(0))
+            .addPath(Key.PathElement.newBuilder().setKind("Child"))
+            .build(),
+        makeKey("Grandparent", 24L, PARENT, "Child").build());
+  }
+
+  @Test
+  public void testMakeKey_KindIdKey() {
+    // 1 key at the end of the series
+    assertEquals(
+        Key.newBuilder()
+            .addPath(Key.PathElement.newBuilder().setKind("Grandparent").setId(24L))
+            .addPath(PARENT.getPath(0))
+            .build(),
+        makeKey("Grandparent", 24L, PARENT).build());
+  }
+
+  @Test
+  public void testMakeKey_KeyKindIdKey() {
+    // 1 key at the beginning and 1 key at the end of the series
+    assertEquals(
+        Key.newBuilder()
+            .addPath(GRANDPARENT.getPath(0))
+            .addPath(Key.PathElement.newBuilder().setKind("Parent").setId(23L))
+            .addPath(CHILD.getPath(0))
+            .build(),
+        makeKey(GRANDPARENT, "Parent", 23, CHILD).build());
+  }
+
+  @Test
+  public void testMakeKey_Key() {
+    // Just 1 key
+    assertEquals(Key.newBuilder().addPath(CHILD.getPath(0)).build(), makeKey(CHILD).build());
+  }
+
+  @Test
+  public void testMakeKey_KeyKey() {
+    // Just 2 keys
+    assertEquals(
+        Key.newBuilder().addPath(PARENT.getPath(0)).addPath(CHILD.getPath(0)).build(),
+        makeKey(PARENT, CHILD).build());
+  }
+
+  @Test
+  public void testMakeKey_KeyKeyKey() {
+    // Just 3 keys
+    assertEquals(
+        Key.newBuilder()
+            .addPath(GRANDPARENT.getPath(0))
+            .addPath(PARENT.getPath(0))
+            .addPath(CHILD.getPath(0))
+            .build(),
+        makeKey(GRANDPARENT, PARENT, CHILD).build());
+  }
+
+  @Test
+  public void testMakeKey_KeyMultiLevelKey() {
+    // 1 key with 2 elements
+    assertEquals(
+        Key.newBuilder()
+            .addPath(GRANDPARENT.getPath(0))
+            .addPath(PARENT.getPath(0))
+            .addPath(CHILD.getPath(0))
+            .build(),
+        makeKey(GRANDPARENT, makeKey(PARENT, CHILD).build()).build());
+  }
+
+  @Test
+  public void testMakeKey_MultiLevelKeyKey() {
+    // 1 key with 2 elements
+    assertEquals(
+        Key.newBuilder()
+            .addPath(GRANDPARENT.getPath(0))
+            .addPath(PARENT.getPath(0))
+            .addPath(CHILD.getPath(0))
+            .build(),
+        makeKey(makeKey(GRANDPARENT, PARENT).build(), CHILD).build());
+  }
+
+  @Test
+  public void testMakeKey_MultiLevelKey() {
+    // 1 key with 3 elements
+    assertEquals(
+        Key.newBuilder()
+            .addPath(GRANDPARENT.getPath(0))
+            .addPath(PARENT.getPath(0))
+            .addPath(CHILD.getPath(0))
+            .build(),
+        makeKey(makeKey(GRANDPARENT, PARENT, CHILD).build()).build());
+  }
+
+  @Test
+  public void testMakeKey_PartitionId() {
+    PartitionId partitionId = PartitionId.newBuilder().setNamespaceId("namespace-id").build();
+    Key parent = PARENT.toBuilder().setPartitionId(partitionId).build();
+    assertEquals(
+        Key.newBuilder()
+            .setPartitionId(partitionId)
+            .addPath(PARENT.getPath(0))
+            .addPath(Key.PathElement.newBuilder().setKind("Child"))
+            .build(),
+        makeKey(parent, "Child").build());
+  }
+
+  @Test
+  public void testMakeKey_NonMatchingPartitionId2() {
+    PartitionId partitionId1 = PartitionId.newBuilder().setNamespaceId("namespace-id").build();
+    PartitionId partitionId2 =
+        PartitionId.newBuilder().setNamespaceId("another-namespace-id").build();
+    try {
+      makeKey(
+          PARENT.toBuilder().setPartitionId(partitionId1).build(),
+          CHILD.toBuilder().setPartitionId(partitionId2).build());
+      fail("expected IllegalArgumentException");
+    } catch (IllegalArgumentException expected) {
+    }
+  }
+
+  @Test
+  public void testMakeTimestampValue() throws Exception {
+    // Test cases with nanos == 0.
+    assertConversion(-50_000, -50, 0);
+    assertConversion(-1_000, -1, 0);
+    assertConversion(0, 0, 0);
+    assertConversion(1_000, 1, 0);
+    assertConversion(2_000, 2, 0);
+    assertConversion(100_000, 100, 0);
+
+    // Test cases with nanos % 1_000_000 == 0 (no loss of precision).
+    assertConversion(2, 0, 2_000_000);
+    assertConversion(1_003, 1, 3_000_000);
+    assertConversion(2_005, 2, 5_000_000);
+
+    // Timestamp specification requires that nanos >= 0 even if the timestamp
+    // is before the epoch.
+    assertConversion(0, 0, 0);
+    assertConversion(-250, -1, 750_000_000); // 1/4 second before epoch
+    assertConversion(-500, -1, 500_000_000); // 1/2 second before epoch
+    assertConversion(-750, -1, 250_000_000); // 3/4 second before epoch
+
+    // If nanos % 1_000_000 != 0, precision is lost (via truncation) when
+    // converting to milliseconds.
+    assertTimestampToMilliseconds(3_100, 3, 100_000_999);
+    assertMillisecondsToTimestamp(3_100, 3, 100_000_000);
+    assertTimestampToMilliseconds(5_999, 5, 999_999_999);
+    assertMillisecondsToTimestamp(5_999, 5, 999_000_000);
+    assertTimestampToMilliseconds(7_100, 7, 100_000_001);
+    assertMillisecondsToTimestamp(7_100, 7, 100_000_000);
+  }
+
+  private void assertConversion(long millis, long seconds, int nanos) {
+    assertMillisecondsToTimestamp(millis, seconds, nanos);
+    assertTimestampToMilliseconds(millis, seconds, nanos);
+  }
+
+  private void assertMillisecondsToTimestamp(long millis, long seconds, long nanos) {
+    Value timestampValue = makeValue(new Date(millis)).build();
+    assertEquals(ValueTypeCase.TIMESTAMP_VALUE, timestampValue.getValueTypeCase());
+    assertEquals(seconds, timestampValue.getTimestampValue().getSeconds());
+    assertEquals(nanos, timestampValue.getTimestampValue().getNanos());
+  }
+
+  private void assertTimestampToMilliseconds(long millis, long seconds, int nanos) {
+    Value.Builder value =
+        Value.newBuilder()
+            .setTimestampValue(Timestamp.newBuilder().setSeconds(seconds).setNanos(nanos));
+    assertEquals(millis, DatastoreHelper.toDate(value.build()).getTime());
+  }
+
+  @Test
+  public void testProjectionHandling() {
+    assertEquals(
+        ByteString.copyFromUtf8("hi"), getByteString(makeValue("hi").setMeaning(18).build()));
+    try {
+      getByteString(makeValue("hi").build());
+      fail("Expected IllegalArgumentException");
+    } catch (IllegalArgumentException expected) {
+    }
+
+    assertEquals(new Date(1), toDate(makeValue(1000).setMeaning(18).build()));
+    try {
+      toDate(makeValue(1000).build());
+      fail("Expected IllegalArgumentException");
+    } catch (IllegalArgumentException expected) {
+    }
+  }
+}
diff --git a/google-cloud-datastore-utils/src/test/java/com/google/datastore/utils/QuerySplitterTest.java b/google-cloud-datastore-utils/src/test/java/com/google/datastore/utils/QuerySplitterTest.java
new file mode 100644
index 00000000..483eca82
--- /dev/null
+++ b/google-cloud-datastore-utils/src/test/java/com/google/datastore/utils/QuerySplitterTest.java
@@ -0,0 +1,378 @@
+/*
+ * Copyright 2024 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.utils;
+
+import static com.google.common.truth.Truth.assertThat;
+import static com.google.datastore.utils.DatastoreHelper.*;
+import static org.junit.Assert.assertArrayEquals;
+import static org.junit.Assert.assertThrows;
+
+import com.google.datastore.utils.testing.MockCredential;
+import com.google.datastore.utils.testing.MockDatastoreFactory;
+import com.google.datastore.v1.*;
+import com.google.datastore.v1.EntityResult.ResultType;
+import com.google.datastore.v1.PropertyFilter.Operator;
+import com.google.datastore.v1.PropertyOrder.Direction;
+import com.google.datastore.v1.QueryResultBatch.MoreResultsType;
+import com.google.protobuf.Int32Value;
+import com.google.protobuf.Timestamp;
+import java.util.List;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+/** Tests for {@link com.google.datastore.utils.QuerySplitterImpl}. */
+@RunWith(JUnit4.class)
+public class QuerySplitterTest {
+  private static final String PROJECT_ID = "project-id";
+  private static final PartitionId PARTITION =
+      PartitionId.newBuilder().setProjectId(PROJECT_ID).build();
+  private static final String KIND = "test-kind";
+
+  private DatastoreFactory factory = new MockDatastoreFactory();
+  private com.google.datastore.utils.DatastoreOptions.Builder options =
+      new DatastoreOptions.Builder().projectId(PROJECT_ID).credential(new MockCredential());
+
+  private Filter propertyFilter = makeFilter("foo", Operator.EQUAL, makeValue("value")).build();
+
+  private Query query =
+      Query.newBuilder()
+          .addKind(KindExpression.newBuilder().setName(KIND).build())
+          .setFilter(propertyFilter)
+          .build();
+
+  private Query splitQuery =
+      Query.newBuilder()
+          .addKind(KindExpression.newBuilder().setName(KIND).build())
+          .addOrder(makeOrder("__scatter__", Direction.ASCENDING))
+          .addProjection(Projection.newBuilder().setProperty(makePropertyReference("__key__")))
+          .build();
+
+  private Key splitKey0 = makeKey(KIND, String.format("%05d", 1)).setPartitionId(PARTITION).build();
+  private Key splitKey1 =
+      makeKey(KIND, String.format("%05d", 101)).setPartitionId(PARTITION).build();
+  private Key splitKey2 =
+      makeKey(KIND, String.format("%05d", 201)).setPartitionId(PARTITION).build();
+  private Key splitKey3 =
+      makeKey(KIND, String.format("%05d", 301)).setPartitionId(PARTITION).build();
+
+  @Test
+  public void disallowsSortOrder() {
+    com.google.datastore.utils.Datastore datastore = factory.create(options.build());
+    Query queryWithOrder =
+        query.toBuilder().addOrder(makeOrder("bar", Direction.ASCENDING)).build();
+    IllegalArgumentException exception =
+        assertThrows(
+            IllegalArgumentException.class,
+            () ->
+                com.google.datastore.utils.QuerySplitterImpl.INSTANCE.getSplits(
+                    queryWithOrder, PARTITION, 2, datastore));
+    assertThat(exception).hasMessageThat().contains("Query cannot have any sort orders.");
+  }
+
+  @Test
+  public void disallowsMultipleKinds() {
+    com.google.datastore.utils.Datastore datastore = factory.create(options.build());
+    Query queryWithMultipleKinds =
+        query
+            .toBuilder()
+            .addKind(KindExpression.newBuilder().setName("another-kind").build())
+            .build();
+    IllegalArgumentException exception =
+        assertThrows(
+            IllegalArgumentException.class,
+            () ->
+                com.google.datastore.utils.QuerySplitterImpl.INSTANCE.getSplits(
+                    queryWithMultipleKinds, PARTITION, 2, datastore));
+    assertThat(exception).hasMessageThat().contains("Query must have exactly one kind.");
+  }
+
+  @Test
+  public void disallowsKindlessQuery() {
+    com.google.datastore.utils.Datastore datastore = factory.create(options.build());
+    Query kindlessQuery = query.toBuilder().clearKind().build();
+    IllegalArgumentException exception =
+        assertThrows(
+            IllegalArgumentException.class,
+            () ->
+                com.google.datastore.utils.QuerySplitterImpl.INSTANCE.getSplits(
+                    kindlessQuery, PARTITION, 2, datastore));
+    assertThat(exception).hasMessageThat().contains("Query must have exactly one kind.");
+  }
+
+  @Test
+  public void disallowsInequalityFilter() {
+    com.google.datastore.utils.Datastore datastore = factory.create(options.build());
+    Query queryWithInequality =
+        query
+            .toBuilder()
+            .setFilter(makeFilter("foo", Operator.GREATER_THAN, makeValue("value")))
+            .build();
+    IllegalArgumentException exception =
+        assertThrows(
+            IllegalArgumentException.class,
+            () ->
+                com.google.datastore.utils.QuerySplitterImpl.INSTANCE.getSplits(
+                    queryWithInequality, PARTITION, 2, datastore));
+    assertThat(exception).hasMessageThat().contains("Query cannot have any inequality filters.");
+  }
+
+  @Test
+  public void splitsMustBePositive() {
+    com.google.datastore.utils.Datastore datastore = factory.create(options.build());
+    IllegalArgumentException exception =
+        assertThrows(
+            IllegalArgumentException.class,
+            () ->
+                com.google.datastore.utils.QuerySplitterImpl.INSTANCE.getSplits(
+                    query, PARTITION, 0, datastore));
+    assertThat(exception).hasMessageThat().contains("The number of splits must be greater than 0.");
+  }
+
+  @Test
+  public void getSplits() throws Exception {
+    com.google.datastore.utils.Datastore datastore = factory.create(options.build());
+    MockDatastoreFactory mockClient = (MockDatastoreFactory) factory;
+
+    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> splittedQueries =
+        com.google.datastore.utils.QuerySplitterImpl.INSTANCE.getSplits(
+            query, PARTITION, 3, datastore);
+
+    assertThat(splittedQueries)
+        .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)
+            .setQuery(
+                splitQuery.toBuilder().setLimit(Int32Value.newBuilder().setValue(2 * 32).build()))
+            .build();
+
+    assertArrayEquals(expectedSplitQueryRequest.toByteArray(), mockClient.getLastBody());
+  }
+
+  @Test
+  public void getSplitsWithDatabaseId() throws Exception {
+    com.google.datastore.utils.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 =
+        com.google.datastore.utils.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();
+
+    assertArrayEquals(expectedSplitQueryRequest.toByteArray(), mockClient.getLastBody());
+  }
+
+  @Test
+  public void notEnoughSplits() throws Exception {
+    com.google.datastore.utils.Datastore datastore = factory.create(options.build());
+    MockDatastoreFactory mockClient = (MockDatastoreFactory) factory;
+
+    RunQueryResponse splitQueryResponse =
+        RunQueryResponse.newBuilder()
+            .setQuery(splitQuery)
+            .setBatch(
+                QueryResultBatch.newBuilder()
+                    .setEntityResultType(ResultType.KEY_ONLY)
+                    .setMoreResults(MoreResultsType.NO_MORE_RESULTS)
+                    .addEntityResults(makeKeyOnlyEntity(splitKey0))
+                    .build())
+            .build();
+
+    mockClient.setNextResponse(splitQueryResponse);
+
+    List<Query> splittedQueries =
+        com.google.datastore.utils.QuerySplitterImpl.INSTANCE.getSplits(
+            query, PARTITION, 100, datastore);
+
+    assertThat(splittedQueries)
+        .containsExactly(
+            query
+                .toBuilder()
+                .setFilter(makeFilterWithKeyRange(propertyFilter, null, splitKey0))
+                .build(),
+            query
+                .toBuilder()
+                .setFilter(makeFilterWithKeyRange(propertyFilter, splitKey0, null))
+                .build());
+
+    RunQueryRequest expectedSplitQueryRequest =
+        RunQueryRequest.newBuilder()
+            .setPartitionId(PARTITION)
+            .setProjectId(PROJECT_ID)
+            .setQuery(
+                splitQuery.toBuilder().setLimit(Int32Value.newBuilder().setValue(99 * 32).build()))
+            .build();
+
+    assertArrayEquals(expectedSplitQueryRequest.toByteArray(), mockClient.getLastBody());
+  }
+
+  @Test
+  public void getSplits_withReadTime() throws Exception {
+    Datastore datastore = factory.create(options.build());
+    MockDatastoreFactory mockClient = (MockDatastoreFactory) factory;
+
+    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);
+
+    Timestamp readTime = Timestamp.newBuilder().setSeconds(1654651341L).build();
+
+    List<Query> splittedQueries =
+        com.google.datastore.utils.QuerySplitterImpl.INSTANCE.getSplits(
+            query, PARTITION, 3, datastore, readTime);
+
+    assertThat(splittedQueries)
+        .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)
+            .setQuery(
+                splitQuery.toBuilder().setLimit(Int32Value.newBuilder().setValue(2 * 32).build()))
+            .setReadOptions(ReadOptions.newBuilder().setReadTime(readTime))
+            .build();
+
+    assertArrayEquals(expectedSplitQueryRequest.toByteArray(), mockClient.getLastBody());
+  }
+
+  private static EntityResult makeKeyOnlyEntity(Key key) {
+    return EntityResult.newBuilder().setEntity(Entity.newBuilder().setKey(key).build()).build();
+  }
+
+  private static Filter makeFilterWithKeyRange(Filter originalFilter, Key startKey, Key endKey) {
+    Filter startKeyFilter =
+        startKey == null
+            ? null
+            : makeFilter("__key__", Operator.GREATER_THAN_OR_EQUAL, makeValue(startKey)).build();
+
+    Filter endKeyFilter =
+        endKey == null
+            ? null
+            : makeFilter("__key__", Operator.LESS_THAN, makeValue(endKey)).build();
+
+    if (startKeyFilter == null && endKeyFilter == null) {
+      throw new IllegalArgumentException();
+    }
+
+    if (startKeyFilter != null && endKeyFilter == null) {
+      return makeAndFilter(originalFilter, startKeyFilter).build();
+    }
+
+    if (startKeyFilter == null && endKeyFilter != null) {
+      return makeAndFilter(originalFilter, endKeyFilter).build();
+    }
+
+    return makeAndFilter(originalFilter, startKeyFilter, endKeyFilter).build();
+  }
+}
diff --git a/google-cloud-datastore-utils/src/test/java/com/google/datastore/utils/RemoteRpcTest.java b/google-cloud-datastore-utils/src/test/java/com/google/datastore/utils/RemoteRpcTest.java
new file mode 100644
index 00000000..ae4d7a23
--- /dev/null
+++ b/google-cloud-datastore-utils/src/test/java/com/google/datastore/utils/RemoteRpcTest.java
@@ -0,0 +1,341 @@
+/*
+ * Copyright 2024 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.utils;
+
+import static org.junit.Assert.*;
+
+import com.google.api.client.http.HttpRequest;
+import com.google.api.client.http.HttpTransport;
+import com.google.api.client.http.LowLevelHttpRequest;
+import com.google.api.client.http.LowLevelHttpResponse;
+import com.google.api.client.http.protobuf.ProtoHttpContent;
+import com.google.api.client.util.Charsets;
+import com.google.datastore.v1.BeginTransactionResponse;
+import com.google.datastore.v1.RollbackRequest;
+import com.google.protobuf.ByteString;
+import com.google.protobuf.MessageLite;
+import com.google.rpc.Code;
+import com.google.rpc.Status;
+import java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.io.InputStream;
+import java.util.zip.GZIPOutputStream;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+/** Test for {@link RemoteRpc}. */
+@RunWith(JUnit4.class)
+public class RemoteRpcTest {
+
+  private static final String METHOD_NAME = "methodName";
+
+  @Test
+  public void testException() {
+    Status statusProto =
+        Status.newBuilder()
+            .setCode(Code.UNAUTHENTICATED_VALUE)
+            .setMessage("The request does not have valid authentication credentials.")
+            .build();
+    DatastoreException exception =
+        RemoteRpc.makeException(
+            "url",
+            METHOD_NAME,
+            new ByteArrayInputStream(statusProto.toByteArray()),
+            "application/x-protobuf",
+            Charsets.UTF_8,
+            new RuntimeException(),
+            401);
+    assertEquals(Code.UNAUTHENTICATED, exception.getCode());
+    assertEquals(
+        "The request does not have valid authentication credentials.", exception.getMessage());
+    assertEquals(METHOD_NAME, exception.getMethodName());
+  }
+
+  @Test
+  public void testInvalidProtoException() {
+    DatastoreException exception =
+        RemoteRpc.makeException(
+            "url",
+            METHOD_NAME,
+            new ByteArrayInputStream("<invalid proto>".getBytes()),
+            "application/x-protobuf",
+            Charsets.UTF_8,
+            new RuntimeException(),
+            401);
+    assertEquals(Code.INTERNAL, exception.getCode());
+    assertEquals(
+        "Unable to parse Status protocol buffer: HTTP status code was 401.",
+        exception.getMessage());
+    assertEquals(METHOD_NAME, exception.getMethodName());
+  }
+
+  @Test
+  public void testEmptyProtoException() {
+    Status statusProto = Status.newBuilder().build();
+    DatastoreException exception =
+        RemoteRpc.makeException(
+            "url",
+            METHOD_NAME,
+            new ByteArrayInputStream(statusProto.toByteArray()),
+            "application/x-protobuf",
+            Charsets.UTF_8,
+            new RuntimeException(),
+            404);
+    assertEquals(Code.INTERNAL, exception.getCode());
+    assertEquals(
+        "Unexpected OK error code with HTTP status code of 404. Message: .",
+        exception.getMessage());
+    assertEquals(METHOD_NAME, exception.getMethodName());
+  }
+
+  @Test
+  public void testEmptyProtoExceptionUnauthenticated() {
+    Status statusProto = Status.newBuilder().build();
+    DatastoreException exception =
+        RemoteRpc.makeException(
+            "url",
+            METHOD_NAME,
+            new ByteArrayInputStream(statusProto.toByteArray()),
+            "application/x-protobuf",
+            Charsets.UTF_8,
+            new RuntimeException(),
+            401);
+    assertEquals(Code.UNAUTHENTICATED, exception.getCode());
+    assertEquals("Unauthenticated.", exception.getMessage());
+    assertEquals(METHOD_NAME, exception.getMethodName());
+  }
+
+  @Test
+  public void testPlainTextException() {
+    DatastoreException exception =
+        RemoteRpc.makeException(
+            "url",
+            METHOD_NAME,
+            new ByteArrayInputStream("Text Error".getBytes()),
+            "text/plain",
+            Charsets.UTF_8,
+            new RuntimeException(),
+            401);
+    assertEquals(Code.INTERNAL, exception.getCode());
+    assertEquals(
+        "Non-protobuf error: Text Error. HTTP status code was 401.", exception.getMessage());
+    assertEquals(METHOD_NAME, exception.getMethodName());
+  }
+
+  @Test
+  public void testGzip() throws IOException, DatastoreException {
+    BeginTransactionResponse response = newBeginTransactionResponse();
+    InjectedTestValues injectedTestValues =
+        new InjectedTestValues(gzip(response), new byte[1], true);
+    RemoteRpc rpc = newRemoteRpc(injectedTestValues);
+
+    InputStream is =
+        rpc.call("beginTransaction", BeginTransactionResponse.getDefaultInstance(), "", "");
+    BeginTransactionResponse parsedResponse = BeginTransactionResponse.parseFrom(is);
+    is.close();
+
+    assertEquals(response, parsedResponse);
+    // Check that the underlying stream is exhausted.
+    assertEquals(-1, injectedTestValues.inputStream.read());
+  }
+
+  @Test
+  public void testHttpHeaders_apiFormat() throws IOException {
+    String projectId = "project-id";
+    MessageLite request =
+        RollbackRequest.newBuilder().setTransaction(ByteString.copyFromUtf8(projectId)).build();
+    RemoteRpc rpc =
+        newRemoteRpc(
+            new InjectedTestValues(gzip(newBeginTransactionResponse()), new byte[1], true));
+    HttpRequest httpRequest =
+        rpc.getClient().buildPostRequest(rpc.resolveURL("blah"), new ProtoHttpContent(request));
+    rpc.setHeaders(request, httpRequest, projectId, "");
+    assertNotNull(
+        httpRequest.getHeaders().getFirstHeaderStringValue(RemoteRpc.API_FORMAT_VERSION_HEADER));
+  }
+
+  @Test
+  public void testHttpHeaders_prefixHeader() throws IOException {
+    String projectId = "my-project";
+    String databaseId = "my-db";
+    MessageLite request =
+        RollbackRequest.newBuilder()
+            .setTransaction(ByteString.copyFromUtf8(projectId))
+            .setDatabaseId(databaseId)
+            .build();
+    RemoteRpc rpc =
+        newRemoteRpc(
+            new InjectedTestValues(gzip(newBeginTransactionResponse()), new byte[1], true));
+    HttpRequest httpRequest =
+        rpc.getClient().buildPostRequest(rpc.resolveURL("blah"), new ProtoHttpContent(request));
+    rpc.setHeaders(request, httpRequest, projectId, databaseId);
+    assertEquals(
+        "project_id=my-project&database_id=my-db",
+        httpRequest.getHeaders().get(RemoteRpc.X_GOOG_REQUEST_PARAMS_HEADER));
+
+    MessageLite request2 =
+        RollbackRequest.newBuilder().setTransaction(ByteString.copyFromUtf8(projectId)).build();
+    RemoteRpc rpc2 =
+        newRemoteRpc(
+            new InjectedTestValues(gzip(newBeginTransactionResponse()), new byte[1], true));
+    HttpRequest httpRequest2 =
+        rpc2.getClient().buildPostRequest(rpc2.resolveURL("blah"), new ProtoHttpContent(request2));
+    rpc2.setHeaders(request, httpRequest2, projectId, "");
+    assertEquals(
+        "project_id=my-project",
+        httpRequest2.getHeaders().get(RemoteRpc.X_GOOG_REQUEST_PARAMS_HEADER));
+  }
+
+  private static BeginTransactionResponse newBeginTransactionResponse() {
+    return BeginTransactionResponse.newBuilder()
+        .setTransaction(ByteString.copyFromUtf8("blah-blah-blah"))
+        .build();
+  }
+
+  private static RemoteRpc newRemoteRpc(InjectedTestValues injectedTestValues) {
+    return new RemoteRpc(
+        new MyHttpTransport(injectedTestValues).createRequestFactory(),
+        null,
+        "https://www.example.com/v1/projects/p");
+  }
+
+  private byte[] gzip(BeginTransactionResponse response) throws IOException {
+    ByteArrayOutputStream bytesOut = new ByteArrayOutputStream();
+    try (GZIPOutputStream gzipOut = new GZIPOutputStream(bytesOut)) {
+      response.writeTo(gzipOut);
+    }
+    return bytesOut.toByteArray();
+  }
+
+  private static class InjectedTestValues {
+    private final InputStream inputStream;
+    private final int contentLength;
+    private final boolean isGzip;
+
+    public InjectedTestValues(byte[] messageBytes, byte[] additionalBytes, boolean isGzip) {
+      byte[] allBytes = concat(messageBytes, additionalBytes);
+      this.inputStream = new ByteArrayInputStream(allBytes);
+      this.contentLength = allBytes.length;
+      this.isGzip = isGzip;
+    }
+
+    private static byte[] concat(byte[] a, byte[] b) {
+      byte[] c = new byte[a.length + b.length];
+      System.arraycopy(a, 0, c, 0, a.length);
+      System.arraycopy(b, 0, c, a.length, b.length);
+      return c;
+    }
+  }
+
+  /** {@link HttpTransport} that allows injection of the returned {@link LowLevelHttpRequest}. */
+  private static class MyHttpTransport extends HttpTransport {
+
+    private final InjectedTestValues injectedTestValues;
+
+    public MyHttpTransport(InjectedTestValues injectedTestValues) {
+      this.injectedTestValues = injectedTestValues;
+    }
+
+    @Override
+    protected LowLevelHttpRequest buildRequest(String method, String url) throws IOException {
+      return new MyLowLevelHttpRequest(injectedTestValues);
+    }
+  }
+
+  /**
+   * {@link LowLevelHttpRequest} that allows injection of the returned {@link LowLevelHttpResponse}.
+   */
+  private static class MyLowLevelHttpRequest extends LowLevelHttpRequest {
+
+    private final InjectedTestValues injectedTestValues;
+
+    public MyLowLevelHttpRequest(InjectedTestValues injectedTestValues) {
+      this.injectedTestValues = injectedTestValues;
+    }
+
+    @Override
+    public void addHeader(String name, String value) throws IOException {
+      // Do nothing.
+    }
+
+    @Override
+    public LowLevelHttpResponse execute() throws IOException {
+      return new MyLowLevelHttpResponse(injectedTestValues);
+    }
+  }
+
+  /** {@link LowLevelHttpResponse} that allows injected properties. */
+  private static class MyLowLevelHttpResponse extends LowLevelHttpResponse {
+
+    private final InjectedTestValues injectedTestValues;
+
+    public MyLowLevelHttpResponse(InjectedTestValues injectedTestValues) {
+      this.injectedTestValues = injectedTestValues;
+    }
+
+    @Override
+    public InputStream getContent() throws IOException {
+      return injectedTestValues.inputStream;
+    }
+
+    @Override
+    public String getContentEncoding() throws IOException {
+      return injectedTestValues.isGzip ? "gzip" : "";
+    }
+
+    @Override
+    public long getContentLength() throws IOException {
+      return injectedTestValues.contentLength;
+    }
+
+    @Override
+    public String getContentType() throws IOException {
+      return "application/x-protobuf";
+    }
+
+    @Override
+    public String getStatusLine() throws IOException {
+      return null;
+    }
+
+    @Override
+    public int getStatusCode() throws IOException {
+      return 200;
+    }
+
+    @Override
+    public String getReasonPhrase() throws IOException {
+      return null;
+    }
+
+    @Override
+    public int getHeaderCount() throws IOException {
+      return 0;
+    }
+
+    @Override
+    public String getHeaderName(int index) throws IOException {
+      return null;
+    }
+
+    @Override
+    public String getHeaderValue(int index) throws IOException {
+      return null;
+    }
+  }
+}
diff --git a/google-cloud-datastore-utils/src/test/java/com/google/datastore/utils/it/ITDatastoreProtoClientTest.java b/google-cloud-datastore-utils/src/test/java/com/google/datastore/utils/it/ITDatastoreProtoClientTest.java
new file mode 100644
index 00000000..d30c1cbd
--- /dev/null
+++ b/google-cloud-datastore-utils/src/test/java/com/google/datastore/utils/it/ITDatastoreProtoClientTest.java
@@ -0,0 +1,90 @@
+/*
+ * Copyright 2024 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.utils.it;
+
+import static com.google.datastore.utils.DatastoreHelper.makeFilter;
+import static com.google.datastore.utils.DatastoreHelper.makeValue;
+
+import com.google.common.truth.Truth;
+import com.google.datastore.utils.Datastore;
+import com.google.datastore.utils.DatastoreException;
+import com.google.datastore.utils.DatastoreHelper;
+import com.google.datastore.v1.*;
+import java.io.IOException;
+import java.security.GeneralSecurityException;
+import java.util.List;
+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);
+    Truth.assertThat(splits).isNotEmpty();
+    splits.forEach(
+        split -> {
+          Truth.assertThat(split.getKind(0).getName()).isEqualTo(KIND);
+          Truth.assertThat(split.getFilter()).isEqualTo(propertyFilter);
+        });
+  }
+
+  @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);
+
+    Truth.assertThat(splits).isNotEmpty();
+    splits.forEach(
+        split -> {
+          Truth.assertThat(split.getKind(0).getName()).isEqualTo(KIND);
+          Truth.assertThat(split.getFilter()).isEqualTo(propertyFilter);
+        });
+  }
+}