Skip to content

Commit 0412836

Browse files
aozarovaozarov
committed
Provide an option for an add with auto-id allocation.
Former-commit-id: cc001fe
1 parent 14ab215 commit 0412836

14 files changed

Lines changed: 264 additions & 104 deletions

‎src/main/java/com/google/gcloud/datastore/BatchWriter.java‎

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -28,10 +28,20 @@ public interface BatchWriter extends DatastoreWriter {
2828
@Override
2929
void add(Entity... entity);
3030

31+
/**
32+
* Datastore add operation.
33+
* This method will automatically allocate id for any entity with incomplete key.
34+
*
35+
* @throws IllegalArgumentException if any of the given entities is missing a key
36+
* @throws DatastoreServiceException if a given entity with a complete key was already added to
37+
* this batch or if batch is no longer active
38+
*/
39+
void add(PartialEntity... entity);
40+
3141
/**
3242
* {@inheritDoc}
3343
* This operation will be converted to {@link #put} operation for entities that were already
34-
* added or put in this batch.
44+
* added or put in this batch
3545
* @throws DatastoreServiceException if an entity is marked for deletion in this batch or if
3646
* batch is no longer active
3747
*/
@@ -41,7 +51,7 @@ public interface BatchWriter extends DatastoreWriter {
4151
/**
4252
* {@inheritDoc}
4353
* This operation will also remove from this batch any prior writes for entities with the same
44-
* keys.
54+
* keys
4555
* @throws DatastoreServiceException if batch is no longer active
4656
*/
4757
@Override

‎src/main/java/com/google/gcloud/datastore/BatchWriterImpl.java‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,12 +7,15 @@
77

88
import java.util.LinkedHashMap;
99
import java.util.LinkedHashSet;
10+
import java.util.LinkedList;
11+
import java.util.List;
1012
import java.util.Map;
1113
import java.util.Set;
1214

1315
class BatchWriterImpl implements BatchWriter {
1416

1517
private final Map<Key, Entity> toAdd = new LinkedHashMap<>();
18+
private final List<PartialEntity> toAddAutoId = new LinkedList<>();
1619
private final Map<Key, Entity> toUpdate = new LinkedHashMap<>();
1720
private final Map<Key, Entity> toPut = new LinkedHashMap<>();
1821
private final Set<Key> toDelete = new LinkedHashSet<>();
@@ -59,6 +62,18 @@ public void add(Entity... entities) {
5962
}
6063
}
6164

65+
@Override
66+
public void add(PartialEntity... entities) {
67+
validateActive();
68+
for (PartialEntity entity : entities) {
69+
if (entity instanceof Entity) {
70+
add((Entity) entity);
71+
} else {
72+
toAddAutoId.add(entity);
73+
}
74+
}
75+
}
76+
6277
@Override
6378
public void update(Entity... entities) {
6479
validateActive();
@@ -103,6 +118,9 @@ public void delete(Key... keys) {
103118
public void submit() {
104119
validateActive();
105120
DatastoreV1.Mutation.Builder mutationPb = DatastoreV1.Mutation.newBuilder();
121+
for (PartialEntity entity : toAddAutoId) {
122+
mutationPb.addInsertAutoId(entity.toPb());
123+
}
106124
for (Entity entity : toAdd.values()) {
107125
mutationPb.addInsert(entity.toPb());
108126
}

‎src/main/java/com/google/gcloud/datastore/DatastoreHelper.java‎

Lines changed: 21 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,6 @@
1010
/**
1111
* Adds some functionality to DatastoreService that should
1212
* be provided statically to the interface (Java 8).
13-
*
1413
*/
1514
public class DatastoreHelper implements DatastoreService {
1615

@@ -26,8 +25,8 @@ public Entity get(Key key) {
2625
}
2726

2827
@Override
29-
public Iterator<Entity> get(Key key, Key... others) {
30-
return delegate.get(key, others);
28+
public Iterator<Entity> get(Key... key) {
29+
return delegate.get(key);
3130
}
3231

3332
@Override
@@ -56,8 +55,18 @@ public Key allocateId(PartialKey key) {
5655
}
5756

5857
@Override
59-
public Iterator<Key> allocateId(PartialKey key, PartialKey... others) {
60-
return delegate.allocateId(key, others);
58+
public List<Key> allocateId(PartialKey... key) {
59+
return delegate.allocateId(key);
60+
}
61+
62+
@Override
63+
public Entity add(PartialEntity entity) {
64+
return delegate.add(entity);
65+
}
66+
67+
@Override
68+
public List<Entity> add(PartialEntity... entity) {
69+
return delegate.add(entity);
6170
}
6271

6372
@Override
@@ -91,17 +100,17 @@ public KeyFactory newKeyFactory() {
91100
* Returns a list with a value for each given key (ordered by input).
92101
* A {@code null} would be returned for non-existing keys.
93102
*/
94-
public List<Entity> fetch(Key key, Key... others) {
95-
Iterator<Entity> entities = delegate.get(key, others);
96-
Map<Key, Entity> map = Maps.newHashMapWithExpectedSize(1 + others.length);
103+
public List<Entity> fetch(Key... keys) {
104+
Iterator<Entity> entities = delegate.get(keys);
105+
Map<Key, Entity> map = Maps.newHashMapWithExpectedSize(keys.length);
97106
while (entities.hasNext()) {
98107
Entity entity = entities.next();
99108
map.put(entity.key(), entity);
100109
}
101-
List<Entity> list = new ArrayList<>(1 + others.length);
102-
list.add(map.get(key));
103-
for (Key other : others) {
104-
list.add(map.get(other));
110+
List<Entity> list = new ArrayList<>(keys.length);
111+
for (Key key : keys) {
112+
// this will include nulls for non-existing keys
113+
list.add(map.get(key));
105114
}
106115
return list;
107116
}

‎src/main/java/com/google/gcloud/datastore/DatastoreReader.java‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,8 +22,9 @@ public interface DatastoreReader {
2222
* {@link Iterator#next next} methods.
2323
*
2424
* @throws DatastoreServiceException upon failure.
25+
* @see #get(Key)
2526
*/
26-
Iterator<Entity> get(Key key, Key... others);
27+
Iterator<Entity> get(Key... key);
2728

2829
/**
2930
* Submit a {@link Query} and returns its result.

‎src/main/java/com/google/gcloud/datastore/DatastoreService.java‎

Lines changed: 26 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
package com.google.gcloud.datastore;
22

3-
import java.util.Iterator;
3+
import java.util.List;
44

55
/**
66
* An interface for Google Cloud Datastore dataset.
@@ -37,10 +37,34 @@ public interface DatastoreService extends DatastoreReaderWriter {
3737
/**
3838
* Returns a list of keys using the allocated ids ordered by the input.
3939
*
40+
* @throws DatastoreServiceException upon failure
4041
* @see #allocateId(PartialKey)
42+
*/
43+
List<Key> allocateId(PartialKey... key);
44+
45+
/**
46+
* Datastore add operation.
47+
* This method will automatically allocate an id if necessary.
48+
*
49+
* @param entity the entity to add
50+
* @return an {@code Entity} with the same properties and a key that is either newly allocated
51+
* or the same one if was already complete
52+
* @throws DatastoreServiceException upon failure
53+
* @throws IllegalArgumentException if the given entity is missing a key
54+
*/
55+
Entity add(PartialEntity entity);
56+
57+
/**
58+
* Datastore add operation.
59+
* This method will automatically allocate id for any entity with incomplete key.
60+
*
61+
* @return a list of {@code Entity} ordered by input with the same properties and a key that is
62+
* either newly allocated or the same one if was already complete
4163
* @throws DatastoreServiceException upon failure
64+
* @throws IllegalArgumentException if any of the given entities is missing a key
65+
* @see #add(PartialKey)
4266
*/
43-
Iterator<Key> allocateId(PartialKey key, PartialKey... others);
67+
List<Entity> add(PartialEntity... entity);
4468

4569
/**
4670
* {@inheritDoc}

‎src/main/java/com/google/gcloud/datastore/DatastoreServiceException.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -114,7 +114,7 @@ static DatastoreServiceException translateAndThrow(DatastoreException exception)
114114
String reason = "";
115115
if (message != null) {
116116
try {
117-
JSONObject json = new JSONObject(new JSONTokener(exception.getMessage()));
117+
JSONObject json = new JSONObject(new JSONTokener(message));
118118
JSONObject error = json.getJSONObject("error").getJSONArray("errors").getJSONObject(0);
119119
reason = error.getString("reason");
120120
message = error.getString("message");

0 commit comments

Comments
 (0)