feat: expose reserveIds API (#56)

googleapis · Mar 2, 2020 · c7f795b · c7f795b
1 parent 2af29c5
commit c7f795b
Show file tree

Hide file tree

Showing 7 changed files with 125 additions and 0 deletions.
diff --git a/google-cloud-datastore/clirr-ignored-differences.xml 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
@@ -168,6 +168,23 @@ interface TransactionCallable<T> {
    */
   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
@@ -31,6 +31,7 @@
 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 @@ public com.google.datastore.v1.LookupResponse call() throws DatastoreException {
     }
   }
 
+  @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
@@ -26,6 +26,8 @@
 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 @@ BeginTransactionResponse beginTransaction(BeginTransactionRequest request)
    */
   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
@@ -30,6 +30,8 @@
 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 LookupResponse lookup(LookupRequest request) {
     }
   }
 
+  @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/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.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.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 void testAllocateIdArray() {
     }
   }
 
+  @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
@@ -597,6 +597,15 @@ public void testAllocateId() {
     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);