allocateId(IncompleteKey... keys);
+ /**
+ * Reserve one or more keys, preventing them from being automatically allocated by Datastore.
+ *
+ * Example of reserving multiple ids in a single batch.
+ *
+ *
{@code
+ * KeyFactory keyFactory = datastore.newKeyFactory().setKind("MyKind");
+ * Key key1 = keyFactory.newKey(10);
+ * Key key2 = keyFactory.newKey("name");
+ * List keys = datastore.reserveIds(key1, key2);
+ *
+ * }
+ *
+ * @throws DatastoreException upon failure
+ */
+ List 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 84102c75b..be31594e7 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.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 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 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() {
+ @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 d39f3a300..5e64c9255 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.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
index 27ff8cec0..7d3434108 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.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
index bef15e409..8049e0957 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.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 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 af1848ad1..798b37e0a 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 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 keyList = DATASTORE.reserveIds(key1, key2);
+ assertEquals(2, keyList.size());
+ }
+
@Test
public void testAllocateIdArray() {
KeyFactory keyFactory = DATASTORE.newKeyFactory().setKind(KIND1);