refactor: refactor grpc query splitter logic

google-cloud-datastore-utils/pom.xml

@@ -0,0 +1,110 @@
+<?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.20.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.20.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.codehaus.mojo</groupId>
+                <artifactId>flatten-maven-plugin</artifactId>
+            </plugin>
+        </plugins>
+    </build>
+    <profiles>
+        <profile>
+            <id>native</id>
+            <properties>
+                <!--
+                Override the default match pattern for native image tests since we have a custom
+                test specifically for native image validation.
+                -->
+                <test>com.google.cloud.datastore.it.ITDatastoreConceptsTest</test>
+            </properties>
+        </profile>
+    </profiles>
+</project>

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

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