Skip to content

Commit c7f795b

Browse files
authored
feat: expose reserveIds API (#56)
1 parent 2af29c5 commit c7f795b

File tree

7 files changed

+125
-0
lines changed

7 files changed

+125
-0
lines changed
Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,14 @@
1+
<?xml version="1.0" encoding="UTF-8"?>
2+
<!-- see http://mojo.codehaus.org/clirr-maven-plugin/examples/ignored-differences.html -->
3+
<differences>
4+
<difference>
5+
<className>com/google/cloud/datastore/Datastore</className>
6+
<method>java.util.List reserveIds(com.google.cloud.datastore.Key[])</method>
7+
<differenceType>7012</differenceType>
8+
</difference>
9+
<difference>
10+
<className>com/google/cloud/datastore/spi/v1/DatastoreRpc</className>
11+
<method>com.google.datastore.v1.ReserveIdsResponse reserveIds(com.google.datastore.v1.ReserveIdsRequest)</method>
12+
<differenceType>7012</differenceType>
13+
</difference>
14+
</differences>

google-cloud-datastore/src/main/java/com/google/cloud/datastore/Datastore.java

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -168,6 +168,23 @@ interface TransactionCallable<T> {
168168
*/
169169
List<Key> allocateId(IncompleteKey... keys);
170170

171+
/**
172+
* Reserve one or more keys, preventing them from being automatically allocated by Datastore.
173+
*
174+
* <p>Example of reserving multiple ids in a single batch.
175+
*
176+
* <pre>{@code
177+
* KeyFactory keyFactory = datastore.newKeyFactory().setKind("MyKind");
178+
* Key key1 = keyFactory.newKey(10);
179+
* Key key2 = keyFactory.newKey("name");
180+
* List<Key> keys = datastore.reserveIds(key1, key2);
181+
*
182+
* }</pre>
183+
*
184+
* @throws DatastoreException upon failure
185+
*/
186+
List<Key> reserveIds(Key... keys);
187+
171188
/**
172189
* {@inheritDoc}
173190
*

google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@
3131
import com.google.common.collect.Iterables;
3232
import com.google.common.collect.Sets;
3333
import com.google.datastore.v1.ReadOptions.ReadConsistency;
34+
import com.google.datastore.v1.ReserveIdsRequest;
3435
import com.google.datastore.v1.TransactionOptions;
3536
import com.google.protobuf.ByteString;
3637
import java.util.ArrayList;
@@ -400,6 +401,40 @@ public com.google.datastore.v1.LookupResponse call() throws DatastoreException {
400401
}
401402
}
402403

404+
@Override
405+
public List<Key> reserveIds(Key... keys) {
406+
ReserveIdsRequest.Builder requestPb = ReserveIdsRequest.newBuilder();
407+
for (Key key : keys) {
408+
requestPb.addKeys(key.toPb());
409+
}
410+
com.google.datastore.v1.ReserveIdsResponse responsePb = reserveIds(requestPb.build());
411+
ImmutableList.Builder<Key> keyList = ImmutableList.builder();
412+
if (responsePb.isInitialized()) {
413+
for (Key key : keys) {
414+
keyList.add(key);
415+
}
416+
}
417+
return keyList.build();
418+
}
419+
420+
com.google.datastore.v1.ReserveIdsResponse reserveIds(
421+
final com.google.datastore.v1.ReserveIdsRequest requestPb) {
422+
try {
423+
return RetryHelper.runWithRetries(
424+
new Callable<com.google.datastore.v1.ReserveIdsResponse>() {
425+
@Override
426+
public com.google.datastore.v1.ReserveIdsResponse call() throws DatastoreException {
427+
return datastoreRpc.reserveIds(requestPb);
428+
}
429+
},
430+
retrySettings,
431+
EXCEPTION_HANDLER,
432+
getOptions().getClock());
433+
} catch (RetryHelperException e) {
434+
throw DatastoreException.translateAndThrow(e);
435+
}
436+
}
437+
403438
@Override
404439
public void update(Entity... entities) {
405440
if (entities.length > 0) {

google-cloud-datastore/src/main/java/com/google/cloud/datastore/spi/v1/DatastoreRpc.java

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,8 @@
2626
import com.google.datastore.v1.CommitResponse;
2727
import com.google.datastore.v1.LookupRequest;
2828
import com.google.datastore.v1.LookupResponse;
29+
import com.google.datastore.v1.ReserveIdsRequest;
30+
import com.google.datastore.v1.ReserveIdsResponse;
2931
import com.google.datastore.v1.RollbackRequest;
3032
import com.google.datastore.v1.RollbackResponse;
3133
import com.google.datastore.v1.RunQueryRequest;
@@ -63,6 +65,13 @@ BeginTransactionResponse beginTransaction(BeginTransactionRequest request)
6365
*/
6466
LookupResponse lookup(LookupRequest request);
6567

68+
/**
69+
* Sends a reserveIds request.
70+
*
71+
* @throws DatastoreException upon failure
72+
*/
73+
ReserveIdsResponse reserveIds(ReserveIdsRequest request);
74+
6675
/**
6776
* Sends a rollback request.
6877
*

google-cloud-datastore/src/main/java/com/google/cloud/datastore/spi/v1/HttpDatastoreRpc.java

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,8 @@
3030
import com.google.datastore.v1.CommitResponse;
3131
import com.google.datastore.v1.LookupRequest;
3232
import com.google.datastore.v1.LookupResponse;
33+
import com.google.datastore.v1.ReserveIdsRequest;
34+
import com.google.datastore.v1.ReserveIdsResponse;
3335
import com.google.datastore.v1.RollbackRequest;
3436
import com.google.datastore.v1.RollbackResponse;
3537
import com.google.datastore.v1.RunQueryRequest;
@@ -164,6 +166,15 @@ public LookupResponse lookup(LookupRequest request) {
164166
}
165167
}
166168

169+
@Override
170+
public ReserveIdsResponse reserveIds(ReserveIdsRequest request) {
171+
try {
172+
return client.reserveIds(request);
173+
} catch (com.google.datastore.v1.client.DatastoreException ex) {
174+
throw translate(ex);
175+
}
176+
}
177+
167178
@Override
168179
public RollbackResponse rollback(RollbackRequest request) {
169180
try {

google-cloud-datastore/src/test/java/com/google/cloud/datastore/ITDatastoreTest.java

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,9 @@
1616

1717
package com.google.cloud.datastore;
1818

19+
import static org.easymock.EasyMock.createStrictMock;
20+
import static org.easymock.EasyMock.replay;
21+
import static org.easymock.EasyMock.verify;
1922
import static org.junit.Assert.assertEquals;
2023
import static org.junit.Assert.assertFalse;
2124
import static org.junit.Assert.assertNotEquals;
@@ -47,6 +50,8 @@
4750
import com.google.datastore.v1.QueryResultBatch;
4851
import com.google.datastore.v1.ReadOptions;
4952
import com.google.datastore.v1.ReadOptions.ReadConsistency;
53+
import com.google.datastore.v1.ReserveIdsRequest;
54+
import com.google.datastore.v1.ReserveIdsResponse;
5055
import com.google.datastore.v1.RollbackRequest;
5156
import com.google.datastore.v1.RollbackResponse;
5257
import com.google.datastore.v1.RunQueryRequest;
@@ -55,6 +60,7 @@
5560
import com.google.protobuf.ByteString;
5661
import java.io.IOException;
5762
import java.util.ArrayList;
63+
import java.util.Arrays;
5864
import java.util.Collections;
5965
import java.util.HashSet;
6066
import java.util.Iterator;
@@ -840,6 +846,30 @@ public void testAllocateIdArray() {
840846
}
841847
}
842848

849+
@Test
850+
public void testReserveIds() {
851+
ReserveIdsRequest reserveIdsRequest =
852+
ReserveIdsRequest.newBuilder().addKeys(KEY1.toPb()).build();
853+
EasyMock.expect(rpcMock.reserveIds(reserveIdsRequest))
854+
.andReturn(ReserveIdsResponse.newBuilder().build())
855+
.times(1);
856+
EasyMock.replay(rpcFactoryMock, rpcMock);
857+
Datastore datastore = rpcMockOptions.getService();
858+
datastore.reserveIds(KEY1);
859+
EasyMock.verify(rpcFactoryMock, rpcMock);
860+
}
861+
862+
@Test
863+
public void testReserveIdsWithKeys() {
864+
Datastore datastore = createStrictMock(Datastore.class);
865+
EasyMock.expect(datastore.reserveIds(KEY1, KEY2)).andReturn(Arrays.asList(KEY1, KEY2));
866+
replay(datastore);
867+
List<Key> result = datastore.reserveIds(KEY1, KEY2);
868+
assertEquals(KEY1, result.get(0));
869+
assertEquals(KEY2, result.get(1));
870+
verify(datastore);
871+
}
872+
843873
@Test
844874
public void testGet() {
845875
Entity entity = datastore.get(KEY3);

google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITDatastoreTest.java

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -597,6 +597,15 @@ public void testAllocateId() {
597597
assertEquals(Key.newBuilder(pk1, key2.getId()).build(), key2);
598598
}
599599

600+
@Test
601+
public void testReserveIds() {
602+
KeyFactory keyFactory = DATASTORE.newKeyFactory().setKind("MyKind");
603+
Key key1 = keyFactory.newKey(10);
604+
Key key2 = keyFactory.newKey("name");
605+
List<Key> keyList = DATASTORE.reserveIds(key1, key2);
606+
assertEquals(2, keyList.size());
607+
}
608+
600609
@Test
601610
public void testAllocateIdArray() {
602611
KeyFactory keyFactory = DATASTORE.newKeyFactory().setKind(KIND1);

0 commit comments

Comments
 (0)