Add verify support (#1107)

Author: Brian Chen

firebase-firestore/src/androidTest/java/com/google/firebase/firestore/TransactionTest.java

@@ -34,7 +34,6 @@
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
-import java.util.Map;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.atomic.AtomicInteger;
 import org.junit.After;
@@ -203,19 +202,6 @@ public void tearDown() {
     IntegrationTestUtil.tearDown();
   }
 
-  @Test
-  public void testGetDocuments() {
-    FirebaseFirestore firestore = testFirestore();
-    DocumentReference doc = firestore.collection("spaces").document();
-    Map<String, Object> value = map("foo", 1, "desc", "Stuff", "owner", "Jonny");
-    waitFor(doc.set(value));
-    Exception e = waitForException(firestore.runTransaction(transaction -> transaction.get(doc)));
-    // We currently require every document read to also be written.
-    // TODO: Fix this check once we drop that requirement.
-    assertEquals(Code.INVALID_ARGUMENT, ((FirebaseFirestoreException) e).getCode());
-    assertEquals("Every document read in a transaction must also be written.", e.getMessage());
-  }
-
   @Test
   public void testRunsTransactionsAfterGettingExistingDoc() {
     FirebaseFirestore firestore = testFirestore();
@@ -513,38 +499,32 @@ public void testUpdatePOJOTransactionally() {
   }
 
   @Test
-  public void testHandleReadingOneDocAndWritingAnother() {
+  public void testRetriesWhenDocumentThatWasReadWithoutBeingWrittenChanges() {
     FirebaseFirestore firestore = testFirestore();
     DocumentReference doc1 = firestore.collection("counters").document();
     DocumentReference doc2 = firestore.collection("counters").document();
     CountDownLatch latch = new CountDownLatch(2);
     waitFor(doc1.set(map("count", 15)));
-    Exception e =
-        waitForException(
-            firestore.runTransaction(
-                transaction -> {
-                  latch.countDown();
-                  // Get the first doc.
-                  transaction.get(doc1);
-                  // Do a write outside of the transaction. The first time the
-                  // transaction is tried, this will bump the version, which
-                  // will cause the write to doc2 to fail. The second time, it
-                  // will be a no-op and not bump the version.
-                  waitFor(doc1.set(map("count", 1234)));
-                  // Now try to update the other doc from within the transaction.
-                  // This should fail once, because we read 15 earlier.
-                  transaction.set(doc2, map("count", 16));
-                  return null;
-                }));
-    // We currently require every document read to also be written.
-    // TODO: Add this check back once we drop that.
-    // Document snapshot = waitFor(doc1.get());
-    // assertEquals(0, tries.getCount());
-    // assertEquals(1234, snapshot.getDouble("count"));
-    // snapshot = waitFor(doc2.get());
-    // assertEquals(16, snapshot.getDouble("count"));
-    assertEquals(Code.INVALID_ARGUMENT, ((FirebaseFirestoreException) e).getCode());
-    assertEquals("Every document read in a transaction must also be written.", e.getMessage());
+    waitFor(
+        firestore.runTransaction(
+            transaction -> {
+              latch.countDown();
+              // Get the first doc.
+              transaction.get(doc1);
+              // Do a write outside of the transaction. The first time the
+              // transaction is tried, this will bump the version, which
+              // will cause the write to doc2 to fail. The second time, it
+              // will be a no-op and not bump the version.
+              waitFor(doc1.set(map("count", 1234)));
+              // Now try to update the other doc from within the transaction.
+              // This should fail once, because we read 15 earlier.
+              transaction.set(doc2, map("count", 16));
+              return null;
+            }));
+    DocumentSnapshot snapshot = waitFor(doc1.get());
+    assertEquals(0, latch.getCount());
+    assertTrue(snapshot.exists());
+    assertEquals(1234L, snapshot.getData().get("count"));
   }
 
   @Test
@@ -616,17 +596,20 @@ public void testReadAndUpdateNonExistentDocumentWithExternalWrite() {
   }
 
   @Test
-  public void testCannotHaveAGetWithoutMutations() {
+  public void testCanHaveGetsWithoutMutations() {
     FirebaseFirestore firestore = testFirestore();
     DocumentReference doc = firestore.collection("foo").document();
+    DocumentReference doc2 = firestore.collection("foo").document();
     waitFor(doc.set(map("foo", "bar")));
-
-    Exception e = waitForException(firestore.runTransaction(transaction -> transaction.get(doc)));
-    // We currently require every document read to also be written.
-    // TODO: Add this check back once we drop that.
-    // assertEquals("bar", snapshot.getString("foo"));
-    assertEquals(Code.INVALID_ARGUMENT, ((FirebaseFirestoreException) e).getCode());
-    assertEquals("Every document read in a transaction must also be written.", e.getMessage());
+    waitFor(
+        firestore.runTransaction(
+            transaction -> {
+              transaction.get(doc2);
+              return transaction.get(doc);
+            }));
+    DocumentSnapshot snapshot = waitFor(doc.get());
+    assertTrue(snapshot.exists());
+    assertEquals(map("foo", "bar"), snapshot.getData());
   }
 
   @Test

firebase-firestore/src/main/java/com/google/firebase/firestore/core/Transaction.java

@@ -32,6 +32,7 @@
 import com.google.firebase.firestore.model.mutation.DeleteMutation;
 import com.google.firebase.firestore.model.mutation.Mutation;
 import com.google.firebase.firestore.model.mutation.Precondition;
+import com.google.firebase.firestore.model.mutation.VerifyMutation;
 import com.google.firebase.firestore.remote.Datastore;
 import com.google.firebase.firestore.util.Executors;
 import java.util.ArrayList;
@@ -136,10 +137,9 @@ public Task<Void> commit() {
     for (Mutation mutation : mutations) {
       unwritten.remove(mutation.getKey());
     }
-    if (unwritten.size() > 0) {
-      return Tasks.forException(
-          new FirebaseFirestoreException(
-              "Every document read in a transaction must also be written.", Code.INVALID_ARGUMENT));
+    // For each document that was read but not written to, we want to perform a `verify` operation.
+    for (DocumentKey key : unwritten) {
+      mutations.add(new VerifyMutation(key, precondition(key)));
     }
     committed = true;
     return datastore

firebase-firestore/src/main/java/com/google/firebase/firestore/model/mutation/VerifyMutation.java

@@ -0,0 +1,78 @@
+// Copyright 2020 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package com.google.firebase.firestore.model.mutation;
+
+import androidx.annotation.Nullable;
+import com.google.firebase.Timestamp;
+import com.google.firebase.firestore.model.DocumentKey;
+import com.google.firebase.firestore.model.MaybeDocument;
+import com.google.firebase.firestore.model.value.ObjectValue;
+import com.google.firebase.firestore.util.Assert;
+
+/**
+ * A mutation that verifies the existence of the document at the given key with the provided
+ * precondition.
+ *
+ * <p>The `verify` operation is only used in Transactions, and this class serves primarily to
+ * facilitate serialization into protos.
+ */
+public final class VerifyMutation extends Mutation {
+
+  public VerifyMutation(DocumentKey key, Precondition precondition) {
+    super(key, precondition);
+  }
+
+  @Override
+  public boolean equals(Object o) {
+    if (this == o) {
+      return true;
+    }
+    if (o == null || getClass() != o.getClass()) {
+      return false;
+    }
+
+    VerifyMutation that = (VerifyMutation) o;
+    return hasSameKeyAndPrecondition(that);
+  }
+
+  @Override
+  public int hashCode() {
+    return keyAndPreconditionHashCode();
+  }
+
+  @Override
+  public String toString() {
+    return "VerifyMutation{" + keyAndPreconditionToString() + "}";
+  }
+
+  @Override
+  public MaybeDocument applyToRemoteDocument(
+      @Nullable MaybeDocument maybeDoc, MutationResult mutationResult) {
+    throw Assert.fail("VerifyMutation should only be used in Transactions.");
+  }
+
+  @Nullable
+  @Override
+  public MaybeDocument applyToLocalView(
+      @Nullable MaybeDocument maybeDoc, @Nullable MaybeDocument baseDoc, Timestamp localWriteTime) {
+    throw Assert.fail("VerifyMutation should only be used in Transactions.");
+  }
+
+  @Nullable
+  @Override
+  public ObjectValue extractBaseValue(@Nullable MaybeDocument maybeDoc) {
+    throw Assert.fail("VerifyMutation should only be used in Transactions.");
+  }
+}

firebase-firestore/src/main/java/com/google/firebase/firestore/remote/RemoteSerializer.java

@@ -51,6 +51,7 @@
 import com.google.firebase.firestore.model.mutation.SetMutation;
 import com.google.firebase.firestore.model.mutation.TransformMutation;
 import com.google.firebase.firestore.model.mutation.TransformOperation;
+import com.google.firebase.firestore.model.mutation.VerifyMutation;
 import com.google.firebase.firestore.model.value.ArrayValue;
 import com.google.firebase.firestore.model.value.BlobValue;
 import com.google.firebase.firestore.model.value.BooleanValue;
@@ -443,6 +444,8 @@ public com.google.firestore.v1.Write encodeMutation(Mutation mutation) {
       builder.setTransform(transformBuilder);
     } else if (mutation instanceof DeleteMutation) {
       builder.setDelete(encodeKey(mutation.getKey()));
+    } else if (mutation instanceof VerifyMutation) {
+      builder.setVerify(encodeKey(mutation.getKey()));
     } else {
       throw fail("unknown mutation type %s", mutation.getClass());
     }
@@ -489,6 +492,9 @@ public Mutation decodeMutation(com.google.firestore.v1.Write mutation) {
         return new TransformMutation(
             decodeKey(mutation.getTransform().getDocument()), fieldTransforms);
 
+      case VERIFY:
+        return new VerifyMutation(decodeKey(mutation.getVerify()), precondition);
+
       default:
         throw fail("Unknown mutation operation: %d", mutation.getOperationCase());
     }

firebase-firestore/src/proto/google/firestore/v1/write.proto

@@ -42,6 +42,11 @@ message Write {
     // `projects/{project_id}/databases/{database_id}/documents/{document_path}`.
     string delete = 2;
 
+    // The name of a document on which to verify the `current_document`
+    // precondition.
+    // This only requires read access to the document.
+    string verify = 5;
+
     // Applies a tranformation to a document.
     // At most one `transform` per document is allowed in a given request.
     // An `update` cannot follow a `transform` on the same document in a given

firebase-firestore/src/test/java/com/google/firebase/firestore/remote/RemoteSerializerTest.java

@@ -27,6 +27,7 @@
 import static com.google.firebase.firestore.testutil.TestUtil.ref;
 import static com.google.firebase.firestore.testutil.TestUtil.setMutation;
 import static com.google.firebase.firestore.testutil.TestUtil.transformMutation;
+import static com.google.firebase.firestore.testutil.TestUtil.verifyMutation;
 import static com.google.firebase.firestore.testutil.TestUtil.wrap;
 import static java.util.Arrays.asList;
 import static org.junit.Assert.assertEquals;
@@ -317,6 +318,21 @@ public void testEncodeDeleteMutation() {
     assertRoundTripForMutation(mutation, expected);
   }
 
+  @Test
+  public void testEncodeVerifyMutation() {
+    Mutation mutation = verifyMutation("docs/1", 4);
+
+    com.google.firestore.v1.Write expected =
+        Write.newBuilder()
+            .setVerify("projects/p/databases/d/documents/docs/1")
+            .setCurrentDocument(
+                Precondition.newBuilder()
+                    .setUpdateTime(Timestamp.newBuilder().setNanos(4000).build())
+                    .build())
+            .build();
+    assertRoundTripForMutation(mutation, expected);
+  }
+
   @Test
   public void testEncodeSetMutation() {
     Mutation mutation = setMutation("docs/1", map("key", "value"));

firebase-firestore/src/testUtil/java/com/google/firebase/firestore/testutil/TestUtil.java

@@ -62,6 +62,7 @@
 import com.google.firebase.firestore.model.mutation.Precondition;
 import com.google.firebase.firestore.model.mutation.SetMutation;
 import com.google.firebase.firestore.model.mutation.TransformMutation;
+import com.google.firebase.firestore.model.mutation.VerifyMutation;
 import com.google.firebase.firestore.model.value.FieldValue;
 import com.google.firebase.firestore.model.value.ObjectValue;
 import com.google.firebase.firestore.remote.RemoteEvent;
@@ -506,6 +507,10 @@ public static DeleteMutation deleteMutation(String path) {
     return new DeleteMutation(key(path), Precondition.NONE);
   }
 
+  public static VerifyMutation verifyMutation(String path, int micros) {
+    return new VerifyMutation(key(path), Precondition.updateTime(version(micros)));
+  }
+
   /**
    * Creates a TransformMutation by parsing any FieldValue sentinels in the provided data. The data
    * is expected to use dotted-notation for nested fields (i.e. { "foo.bar": FieldValue.foo() } and