Updating datastore write-related functions to v1beta3

gcloud-java-datastore/src/main/java/com/google/gcloud/datastore/BaseDatastoreBatchWriter.java

@@ -16,11 +16,11 @@
 
 package com.google.gcloud.datastore;
 
-import com.google.api.services.datastore.DatastoreV1;
 import com.google.common.base.Preconditions;
 import com.google.common.collect.Iterables;
 import com.google.common.collect.Lists;
 
+import java.util.ArrayList;
 import java.util.Collections;
 import java.util.Iterator;
 import java.util.LinkedHashMap;
@@ -88,9 +88,7 @@ public final List<Entity> add(FullEntity<?>... entities) {
     for (FullEntity<?> entity : entities) {
       IncompleteKey key = entity.key();
       Preconditions.checkArgument(key != null, "Entity must have a key");
-      if (key instanceof Key) {
-        addInternal((FullEntity<Key>) entity);
-      } else {
+      if (!(key instanceof Key)) {
         incompleteKeys.add(key);
       }
     }
@@ -104,6 +102,7 @@ public final List<Entity> add(FullEntity<?>... entities) {
     List<Entity> answer = Lists.newArrayListWithExpectedSize(entities.length);
     for (FullEntity<?> entity : entities) {
       if (entity.key() instanceof Key) {
+        addInternal((FullEntity<Key>) entity);
         answer.add(Entity.convert((FullEntity<Key>) entity));
       } else {
         Entity entityWithAllocatedId = Entity.builder(allocated.next(), entity).build();
@@ -184,6 +183,10 @@ protected Map<Key, FullEntity<Key>> toPut() {
   protected Set<Key> toDelete() {
     return toDelete;
   }
+
+  protected int numAutoAllocatedIds() {
+    return toAddAutoId.size();
+  }
 
   protected void deactivate() {
     active = false;
@@ -199,25 +202,30 @@ protected DatastoreException newInvalidRequest(String msg, Object... params) {
     return DatastoreException.throwInvalidRequest(String.format(msg, params));
   }
 
-  protected DatastoreV1.Mutation.Builder toMutationPb() {
-    DatastoreV1.Mutation.Builder mutationPb = DatastoreV1.Mutation.newBuilder();
+  protected List<com.google.datastore.v1beta3.Mutation> toMutationPbList() {
+    List<com.google.datastore.v1beta3.Mutation> mutationsPb = 
+        new ArrayList<>();
     for (FullEntity<IncompleteKey> entity : toAddAutoId()) {
-      mutationPb.addInsertAutoId(entity.toPb());
+      mutationsPb.add(
+          com.google.datastore.v1beta3.Mutation.newBuilder().setInsert(entity.toPb()).build());
     }
     for (FullEntity<Key> entity : toAdd().values()) {
-      mutationPb.addInsert(entity.toPb());
+      mutationsPb.add(
+          com.google.datastore.v1beta3.Mutation.newBuilder().setInsert(entity.toPb()).build());
     }
     for (FullEntity<Key> entity : toUpdate().values()) {
-      mutationPb.addUpdate(entity.toPb());
+      mutationsPb.add(
+          com.google.datastore.v1beta3.Mutation.newBuilder().setUpdate(entity.toPb()).build());
     }
     for (FullEntity<Key> entity : toPut().values()) {
-      mutationPb.addUpsert(entity.toPb());
+      mutationsPb.add(
+          com.google.datastore.v1beta3.Mutation.newBuilder().setUpsert(entity.toPb()).build());
     }
     for (Key key : toDelete()) {
-      // TODO(ajaykannan): uncomment when possible in datastore v1beta3 transition
-      //mutationPb.addDelete(key.toPb());
+      mutationsPb.add(
+          com.google.datastore.v1beta3.Mutation.newBuilder().setDelete(key.toPb()).build());
     }
-    return mutationPb;
+    return mutationsPb;
   }
 
   protected abstract Datastore datastore();

gcloud-java-datastore/src/main/java/com/google/gcloud/datastore/BaseEntity.java

@@ -27,7 +27,6 @@
 import static com.google.gcloud.datastore.NullValue.of;
 import static com.google.gcloud.datastore.StringValue.of;
 
-import com.google.api.services.datastore.DatastoreV1;
 import com.google.common.collect.ImmutableSortedMap;
 import com.google.common.collect.Maps;
 import com.google.protobuf.InvalidProtocolBufferException;
@@ -48,7 +47,8 @@
  * @see <a href="https://cloud.google.com/datastore/docs/concepts/entities">Google Cloud Datastore
  *     Entities, Properties, and Keys</a>
  */
-public abstract class BaseEntity<K extends IncompleteKey> extends Serializable<DatastoreV1.Entity> {
+public abstract class BaseEntity<K extends IncompleteKey> 
+    extends Serializable<com.google.datastore.v1beta3.Entity> {
 
   private static final long serialVersionUID = 8175618724683792766L;
 
@@ -90,16 +90,15 @@ private B self() {
     }
 
     @SuppressWarnings("unchecked")
-    protected B fill(DatastoreV1.Entity entityPb) {
+    protected B fill(com.google.datastore.v1beta3.Entity entityPb) {
       Map<String, Value<?>> copiedProperties = Maps.newHashMap();
-      for (DatastoreV1.Property property : entityPb.getPropertyList()) {
-        // TODO(ajaykannan): Uncomment when possible in datastore v1beta3 transition
-        //copiedProperties.put(property.getName(), Value.fromPb(property.getValue()));
+      for (Map.Entry<String, com.google.datastore.v1beta3.Value> entry : 
+           entityPb.getProperties().entrySet()) {
+        copiedProperties.put(entry.getKey(), Value.fromPb(entry.getValue()));
       }
       properties(copiedProperties);
       if (entityPb.hasKey()) {
-        // TODO(ajaykannan): Uncomment when possible in datastore v1beta3 transition
-        //key((K) IncompleteKey.fromPb(entityPb.getKey()));
+        key((K) IncompleteKey.fromPb(entityPb.getKey()));
       }
       return self();
     }
@@ -379,25 +378,21 @@ ImmutableSortedMap<String, Value<?>> properties() {
   @Override
   protected Object fromPb(byte[] bytesPb) throws InvalidProtocolBufferException {
     Builder<?, ?> builder = emptyBuilder();
-    builder.fill(DatastoreV1.Entity.parseFrom(bytesPb));
+    builder.fill(com.google.datastore.v1beta3.Entity.parseFrom(bytesPb));
     return builder.build();
   }
 
   protected abstract Builder<?, ?> emptyBuilder();
 
   @Override
-  protected final DatastoreV1.Entity toPb() {
-    DatastoreV1.Entity.Builder entityPb = DatastoreV1.Entity.newBuilder();
+  protected final com.google.datastore.v1beta3.Entity toPb() {
+    com.google.datastore.v1beta3.Entity.Builder entityPb = com.google.datastore.v1beta3.Entity.newBuilder();
+    Map<String, com.google.datastore.v1beta3.Value> propertiesPb = entityPb.getMutableProperties();
     for (Map.Entry<String, Value<?>> entry : properties.entrySet()) {
-      DatastoreV1.Property.Builder propertyPb = DatastoreV1.Property.newBuilder();
-      propertyPb.setName(entry.getKey());
-      // TODO(ajaykannan): Uncomment when possible in datastore v1beta3 transition
-      //propertyPb.setValue(entry.getValue().toPb());
-      entityPb.addProperty(propertyPb.build());
+      propertiesPb.put(entry.getKey(), entry.getValue().toPb());
     }
     if (key != null) {
-      // TODO(ajaykannan): uncomment when possible in datastore v1beta3 transition
-      //entityPb.setKey(key.toPb());
+      entityPb.setKey(key.toPb());
     }
     return entityPb.build();
   }

gcloud-java-datastore/src/main/java/com/google/gcloud/datastore/BatchImpl.java

@@ -16,65 +16,53 @@
 
 package com.google.gcloud.datastore;
 
-import com.google.api.services.datastore.DatastoreV1;
-import com.google.common.base.Function;
-import com.google.common.collect.Lists;
-import com.google.gcloud.datastore.BatchOption.ForceWrites;
-
+import java.util.Iterator;
+import java.util.LinkedList;
 import java.util.List;
-import java.util.Map;
 
 
 class BatchImpl extends BaseDatastoreBatchWriter implements Batch {
 
   private final DatastoreImpl datastore;
-  private final boolean force;
 
   static class ResponseImpl implements Batch.Response {
 
-    private final DatastoreV1.CommitResponse response;
+    private final com.google.datastore.v1beta3.CommitResponse response;
+    private final int numAutoAllocatedIds;
 
-    ResponseImpl(DatastoreV1.CommitResponse response) {
+    ResponseImpl(com.google.datastore.v1beta3.CommitResponse response, int numAutoAllocatedIds) {
       this.response = response;
+      this.numAutoAllocatedIds = numAutoAllocatedIds;
     }
 
     @Override
     public List<Key> generatedKeys() {
-      return Lists.transform(response.getMutationResult().getInsertAutoIdKeyList(),
-          new Function<DatastoreV1.Key, Key>() {
-            @Override public Key apply(DatastoreV1.Key keyPb) {
-              // TODO(ajaykannan): uncomment when possible in datastore v1beta3 transition
-              //return Key.fromPb(keyPb);
-              return Key.builder(null).build(); // TODO(ajaykannan): remove this line when possible
-            }
-          });
+      Iterator<com.google.datastore.v1beta3.MutationResult> results = 
+          response.getMutationResultsList().iterator();
+      List<Key> generated = new LinkedList<Key>();
+      for (int i = 0; i < numAutoAllocatedIds; i++) {
+        generated.add(Key.fromPb(results.next().getKey()));
+      }
+      return generated;
     }
   }
 
-  BatchImpl(DatastoreImpl datastore, BatchOption... options) {
+  BatchImpl(DatastoreImpl datastore) {
     super("batch");
     this.datastore = datastore;
-    Map<Class<? extends BatchOption>, BatchOption> optionsMap = BatchOption.asImmutableMap(options);
-    if (optionsMap.containsKey(ForceWrites.class)) {
-      force = ((ForceWrites) optionsMap.get(ForceWrites.class)).force();
-    } else {
-      force = datastore.options().force();
-    }
   }
 
   @Override
   public Batch.Response submit() {
     validateActive();
-    DatastoreV1.Mutation.Builder mutationPb = toMutationPb();
-    if (force) {
-      mutationPb.setForce(force);
-    }
-    DatastoreV1.CommitRequest.Builder requestPb = DatastoreV1.CommitRequest.newBuilder();
-    requestPb.setMode(DatastoreV1.CommitRequest.Mode.NON_TRANSACTIONAL);
-    requestPb.setMutation(mutationPb);
-    DatastoreV1.CommitResponse responsePb = datastore.commit(requestPb.build());
+    List<com.google.datastore.v1beta3.Mutation> mutationsPb = toMutationPbList();
+    com.google.datastore.v1beta3.CommitRequest.Builder requestPb = 
+        com.google.datastore.v1beta3.CommitRequest.newBuilder();
+    requestPb.setMode(com.google.datastore.v1beta3.CommitRequest.Mode.NON_TRANSACTIONAL);
+    requestPb.addAllMutations(mutationsPb);
+    com.google.datastore.v1beta3.CommitResponse responsePb = datastore.commit(requestPb.build());
     deactivate();
-    return new ResponseImpl(responsePb);
+    return new ResponseImpl(responsePb, numAutoAllocatedIds());
   }
 
   @Override

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

@@ -30,7 +30,7 @@ public interface Datastore extends Service<DatastoreOptions>, DatastoreReaderWri
    *
    * @throws DatastoreException upon failure
    */
-  Transaction newTransaction(TransactionOption... options);
+  Transaction newTransaction();
 
 
   /**
@@ -57,12 +57,12 @@ interface TransactionCallable<T> {
    * @param options the options for the created transaction
    * @throws DatastoreException upon failure
    */
-  <T> T runInTransaction(TransactionCallable<T> callable, TransactionOption... options);
+  <T> T runInTransaction(TransactionCallable<T> callable);
 
   /**
    * Returns a new Batch for processing multiple write operations in one request.
    */
-  Batch newBatch(BatchOption... options);
+  Batch newBatch();
 
   /**
    * Allocate a unique id for the given key.

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

@@ -70,8 +70,8 @@ static List<Entity> fetch(DatastoreReader reader, Key... keys) {
   }
 
   static <T> T runInTransaction(Datastore datastore,
-      Datastore.TransactionCallable<T> callable, TransactionOption... options) {
-    Transaction transaction = datastore.newTransaction(options);
+      Datastore.TransactionCallable<T> callable) {
+    Transaction transaction = datastore.newTransaction();
     try {
       T value = callable.run(transaction);
       transaction.commit();