Repository: googleapis/java-datastore (java)

Issue:
[Datastore] expose 'reserveIds' API
Expose the  `'reserveIds'` API  to Prevents the supplied keys' IDs from being auto-allocated by Cloud Datastore.

- [ ] [reserveIds](https://cloud.google.com/datastore/docs/reference/data/rest/v1/projects/reserveIds)

/cc: @BenWhitehead 

Pull request:
feat: expose reserveIds API
Fixes #53 

Proposed Design [Doc](https://docs.google.com/document/d/1FYfg3-R6-AoWbv5JWlYoJYtG-L-JpQCpISy3IVDBpK0/edit#).


~/cc: @BenWhitehead 

Patch:
diff --git a/google-cloud-datastore/clirr-ignored-differences.xml b/google-cloud-datastore/clirr-ignored-differences.xml
new file mode 100644
index 00000000..110f22f7
--- /dev/null
+++ b/google-cloud-datastore/clirr-ignored-differences.xml
@@ -0,0 +1,14 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<!-- see http://mojo.codehaus.org/clirr-maven-plugin/examples/ignored-differences.html -->
+<differences>
+    <difference>
+        <className>com/google/cloud/datastore/Datastore</className>
+        <method>java.util.List reserveIds(com.google.cloud.datastore.Key[])</method>
+        <differenceType>7012</differenceType>
+    </difference>
+    <difference>
+        <className>com/google/cloud/datastore/spi/v1/DatastoreRpc</className>
+        <method>com.google.datastore.v1.ReserveIdsResponse reserveIds(com.google.datastore.v1.ReserveIdsRequest)</method>
+        <differenceType>7012</differenceType>
+    </difference>
+</differences>
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Datastore.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Datastore.java
index 0edbadc1..bb115995 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Datastore.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Datastore.java
@@ -168,6 +168,23 @@ public interface Datastore extends Service<DatastoreOptions>, DatastoreReaderWri
    */
   List<Key> allocateId(IncompleteKey... keys);
 
+  /**
+   * Reserve one or more keys, preventing them from being automatically allocated by Datastore.
+   *
+   * <p>Example of reserving multiple ids in a single batch.
+   *
+   * <pre>{@code
+   * KeyFactory keyFactory = datastore.newKeyFactory().setKind("MyKind");
+   * Key key1 = keyFactory.newKey(10);
+   * Key key2 = keyFactory.newKey("name");
+   * List<Key> keys = datastore.reserveIds(key1, key2);
+   *
+   * }</pre>
+   *
+   * @throws DatastoreException upon failure
+   */
+  List<Key> reserveIds(Key... keys);
+
   /**
    * {@inheritDoc}
    *
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java
index 84102c75..be31594e 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java
@@ -31,6 +31,7 @@ import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Iterables;
 import com.google.common.collect.Sets;
 import com.google.datastore.v1.ReadOptions.ReadConsistency;
+import com.google.datastore.v1.ReserveIdsRequest;
 import com.google.datastore.v1.TransactionOptions;
 import com.google.protobuf.ByteString;
 import java.util.ArrayList;
@@ -400,6 +401,40 @@ final class DatastoreImpl extends BaseService<DatastoreOptions> implements Datas
     }
   }
 
+  @Override
+  public List<Key> reserveIds(Key... keys) {
+    ReserveIdsRequest.Builder requestPb = ReserveIdsRequest.newBuilder();
+    for (Key key : keys) {
+      requestPb.addKeys(key.toPb());
+    }
+    com.google.datastore.v1.ReserveIdsResponse responsePb = reserveIds(requestPb.build());
+    ImmutableList.Builder<Key> keyList = ImmutableList.builder();
+    if (responsePb.isInitialized()) {
+      for (Key key : keys) {
+        keyList.add(key);
+      }
+    }
+    return keyList.build();
+  }
+
+  com.google.datastore.v1.ReserveIdsResponse reserveIds(
+      final com.google.datastore.v1.ReserveIdsRequest requestPb) {
+    try {
+      return RetryHelper.runWithRetries(
+          new Callable<com.google.datastore.v1.ReserveIdsResponse>() {
+            @Override
+            public com.google.datastore.v1.ReserveIdsResponse call() throws DatastoreException {
+              return datastoreRpc.reserveIds(requestPb);
+            }
+          },
+          retrySettings,
+          EXCEPTION_HANDLER,
+          getOptions().getClock());
+    } catch (RetryHelperException e) {
+      throw DatastoreException.translateAndThrow(e);
+    }
+  }
+
   @Override
   public void update(Entity... entities) {
     if (entities.length > 0) {
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/spi/v1/DatastoreRpc.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/spi/v1/DatastoreRpc.java
index d39f3a30..5e64c925 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/spi/v1/DatastoreRpc.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/spi/v1/DatastoreRpc.java
@@ -26,6 +26,8 @@ import com.google.datastore.v1.CommitRequest;
 import com.google.datastore.v1.CommitResponse;
 import com.google.datastore.v1.LookupRequest;
 import com.google.datastore.v1.LookupResponse;
+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.RunQueryRequest;
@@ -63,6 +65,13 @@ public interface DatastoreRpc extends ServiceRpc {
    */
   LookupResponse lookup(LookupRequest request);
 
+  /**
+   * Sends a reserveIds request.
+   *
+   * @throws DatastoreException upon failure
+   */
+  ReserveIdsResponse reserveIds(ReserveIdsRequest request);
+
   /**
    * Sends a rollback request.
    *
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/spi/v1/HttpDatastoreRpc.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/spi/v1/HttpDatastoreRpc.java
index 27ff8cec..7d343410 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/spi/v1/HttpDatastoreRpc.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/spi/v1/HttpDatastoreRpc.java
@@ -30,6 +30,8 @@ import com.google.datastore.v1.CommitRequest;
 import com.google.datastore.v1.CommitResponse;
 import com.google.datastore.v1.LookupRequest;
 import com.google.datastore.v1.LookupResponse;
+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.RunQueryRequest;
@@ -164,6 +166,15 @@ public class HttpDatastoreRpc implements DatastoreRpc {
     }
   }
 
+  @Override
+  public ReserveIdsResponse reserveIds(ReserveIdsRequest request) {
+    try {
+      return client.reserveIds(request);
+    } catch (com.google.datastore.v1.client.DatastoreException ex) {
+      throw translate(ex);
+    }
+  }
+
   @Override
   public RollbackResponse rollback(RollbackRequest request) {
     try {
diff --git a/pom.xml b/pom.xml
index db8e0d97..793b3e8a 100644
--- a/pom.xml
+++ b/pom.xml
@@ -152,11 +152,11 @@
     <github.global.server>github</github.global.server>
     <site.installationModule>google-cloud-datastore-parent</site.installationModule>
     <project.javadoc.protobufBaseURL>https://googleapis.dev/java/google-api-grpc/latest</project.javadoc.protobufBaseURL>
-    <google.core.version>1.92.5</google.core.version>
+    <google.core.version>1.93.0</google.core.version>
     <google.api-common.version>1.8.1</google.api-common.version>
     <google.common-protos.version>1.17.0</google.common-protos.version>
-    <gax.version>1.53.1</gax.version>
-    <grpc.version>1.27.1</grpc.version>
+    <gax.version>1.54.0</gax.version>
+    <grpc.version>1.27.2</grpc.version>
     <protobuf.version>3.11.4</protobuf.version>
     <junit.version>4.13</junit.version>
     <guava.version>28.2-android</guava.version>
@@ -164,7 +164,7 @@
     <opencensus.version>0.24.0</opencensus.version>
     <mockito.version>1.10.19</mockito.version>
     <javax.annotations.version>1.3.2</javax.annotations.version>
-    <truth.version>1.0</truth.version>
+    <truth.version>1.0.1</truth.version>
   </properties>
 
   <dependencyManagement>


Test patch:
diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/ITDatastoreTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/ITDatastoreTest.java
index bef15e40..8049e095 100644
--- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/ITDatastoreTest.java
+++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/ITDatastoreTest.java
@@ -16,6 +16,9 @@
 
 package com.google.cloud.datastore;
 
+import static org.easymock.EasyMock.createStrictMock;
+import static org.easymock.EasyMock.replay;
+import static org.easymock.EasyMock.verify;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotEquals;
@@ -47,6 +50,8 @@ import com.google.datastore.v1.PartitionId;
 import com.google.datastore.v1.QueryResultBatch;
 import com.google.datastore.v1.ReadOptions;
 import com.google.datastore.v1.ReadOptions.ReadConsistency;
+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.RunQueryRequest;
@@ -55,6 +60,7 @@ import com.google.datastore.v1.TransactionOptions;
 import com.google.protobuf.ByteString;
 import java.io.IOException;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collections;
 import java.util.HashSet;
 import java.util.Iterator;
@@ -840,6 +846,30 @@ public class ITDatastoreTest {
     }
   }
 
+  @Test
+  public void testReserveIds() {
+    ReserveIdsRequest reserveIdsRequest =
+        ReserveIdsRequest.newBuilder().addKeys(KEY1.toPb()).build();
+    EasyMock.expect(rpcMock.reserveIds(reserveIdsRequest))
+        .andReturn(ReserveIdsResponse.newBuilder().build())
+        .times(1);
+    EasyMock.replay(rpcFactoryMock, rpcMock);
+    Datastore datastore = rpcMockOptions.getService();
+    datastore.reserveIds(KEY1);
+    EasyMock.verify(rpcFactoryMock, rpcMock);
+  }
+
+  @Test
+  public void testReserveIdsWithKeys() {
+    Datastore datastore = createStrictMock(Datastore.class);
+    EasyMock.expect(datastore.reserveIds(KEY1, KEY2)).andReturn(Arrays.asList(KEY1, KEY2));
+    replay(datastore);
+    List<Key> result = datastore.reserveIds(KEY1, KEY2);
+    assertEquals(KEY1, result.get(0));
+    assertEquals(KEY2, result.get(1));
+    verify(datastore);
+  }
+
   @Test
   public void testGet() {
     Entity entity = datastore.get(KEY3);
diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITDatastoreTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITDatastoreTest.java
index af1848ad..798b37e0 100644
--- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITDatastoreTest.java
+++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITDatastoreTest.java
@@ -597,6 +597,15 @@ public class ITDatastoreTest {
     assertEquals(Key.newBuilder(pk1, key2.getId()).build(), key2);
   }
 
+  @Test
+  public void testReserveIds() {
+    KeyFactory keyFactory = DATASTORE.newKeyFactory().setKind("MyKind");
+    Key key1 = keyFactory.newKey(10);
+    Key key2 = keyFactory.newKey("name");
+    List<Key> keyList = DATASTORE.reserveIds(key1, key2);
+    assertEquals(2, keyList.size());
+  }
+
   @Test
   public void testAllocateIdArray() {
     KeyFactory keyFactory = DATASTORE.newKeyFactory().setKind(KIND1);