DocumentOverlays schema and scaffolding (#3007)

Author: wu-hui

firebase-firestore/firebase-firestore.gradle

@@ -106,11 +106,19 @@ android.libraryVariants.all { variant ->
     } catch (FileNotFoundException e) {
     }
 
+    // TODO(Indexing): Delete below once indexing is shipped.
     if (localProps['firestoreEnableIndexing']) {
         variant.buildConfigField("boolean", "ENABLE_INDEXING", "true")
     } else {
         variant.buildConfigField("boolean", "ENABLE_INDEXING", "false")
     }
+
+    // TODO(Overlay): Delete below once overlay is shipped.
+    if (localProps['firestoreEnableOverlay']) {
+        variant.buildConfigField("boolean", "ENABLE_OVERLAY", "true")
+    } else {
+        variant.buildConfigField("boolean", "ENABLE_OVERLAY", "false")
+    }
 }
 
 configurations.all {

firebase-firestore/src/androidTest/java/com/google/firebase/firestore/testutil/IntegrationTestUtil.java

@@ -255,7 +255,8 @@ public static FirebaseFirestore testFirestore(
     // This unfortunately is a global setting that affects existing Firestore clients.
     Logger.setLogLevel(logLevel);
 
-    // TODO: Remove this once this is ready to ship.
+    // TODO(Overlay): Remove below once this is ready to ship.
+    Persistence.OVERLAY_SUPPORT_ENABLED = true;
     Persistence.INDEXING_SUPPORT_ENABLED = true;
 
     Context context = ApplicationProvider.getApplicationContext();

firebase-firestore/src/main/java/com/google/firebase/firestore/local/DocumentOverlay.java

@@ -0,0 +1,40 @@
+// Copyright 2021 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.local;
+
+import androidx.annotation.Nullable;
+import com.google.firebase.firestore.model.DocumentKey;
+import com.google.firebase.firestore.model.mutation.Mutation;
+
+/**
+ * Provides methods to read and write document overlays.
+ *
+ * <p>An overlay is a saved {@link Mutation}, that gives a local view of a document when applied to
+ * the remote version of the document.
+ */
+public interface DocumentOverlay {
+  /**
+   * Gets the saved overlay mutation for the given document key. Returns null if there is no overlay
+   * for that key.
+   */
+  @Nullable
+  Mutation getOverlay(DocumentKey key);
+
+  /** Saves the given mutation as overlay for the given document key. */
+  void saveOverlay(DocumentKey key, Mutation mutation);
+
+  /** Removes the overlay associated for the given document key. */
+  void removeOverlay(DocumentKey key);
+}

firebase-firestore/src/main/java/com/google/firebase/firestore/local/LocalSerializer.java

@@ -327,4 +327,12 @@ public FieldIndex decodeFieldIndex(
         fieldIndex.withVersion(new SnapshotVersion(new Timestamp(updateSeconds, updateNanos)));
     return fieldIndex;
   }
+
+  public Mutation decodeMutation(Write mutation) {
+    return rpcSerializer.decodeMutation(mutation);
+  }
+
+  public Write encodeMutation(Mutation mutation) {
+    return rpcSerializer.encodeMutation(mutation);
+  }
 }

firebase-firestore/src/main/java/com/google/firebase/firestore/local/MemoryDocumentOverlay.java

@@ -0,0 +1,41 @@
+// Copyright 2021 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.local;
+
+import androidx.annotation.Nullable;
+import com.google.firebase.firestore.model.DocumentKey;
+import com.google.firebase.firestore.model.mutation.Mutation;
+import java.util.HashMap;
+import java.util.Map;
+
+public class MemoryDocumentOverlay implements DocumentOverlay {
+  private Map<DocumentKey, Mutation> overlays = new HashMap<>();
+
+  @Nullable
+  @Override
+  public Mutation getOverlay(DocumentKey key) {
+    return overlays.get(key);
+  }
+
+  @Override
+  public void saveOverlay(DocumentKey key, Mutation mutation) {
+    overlays.put(key, mutation);
+  }
+
+  @Override
+  public void removeOverlay(DocumentKey key) {
+    overlays.remove(key);
+  }
+}

firebase-firestore/src/main/java/com/google/firebase/firestore/local/MemoryPersistence.java

@@ -32,6 +32,7 @@ public final class MemoryPersistence extends Persistence {
   // LocalStore wrapping this Persistence instance and this will make the in-memory persistence
   // layer behave as if it were actually persisting values.
   private final Map<User, MemoryMutationQueue> mutationQueues;
+  private final Map<User, MemoryDocumentOverlay> overlays;
   private final MemoryIndexManager indexManager;
   private final MemoryTargetCache targetCache;
   private final MemoryBundleCache bundleCache;
@@ -61,6 +62,7 @@ private MemoryPersistence() {
     targetCache = new MemoryTargetCache(this);
     bundleCache = new MemoryBundleCache();
     remoteDocumentCache = new MemoryRemoteDocumentCache(this);
+    overlays = new HashMap<>();
   }
 
   @Override
@@ -125,6 +127,16 @@ BundleCache getBundleCache() {
     return bundleCache;
   }
 
+  @Override
+  DocumentOverlay getDocumentOverlay(User user) {
+    MemoryDocumentOverlay overlay = overlays.get(user);
+    if (overlay == null) {
+      overlay = new MemoryDocumentOverlay();
+      overlays.put(user, overlay);
+    }
+    return overlay;
+  }
+
   @Override
   void runTransaction(String action, Runnable operation) {
     referenceDelegate.onTransactionStarted();

firebase-firestore/src/main/java/com/google/firebase/firestore/local/Persistence.java

@@ -50,8 +50,12 @@
 public abstract class Persistence {
   static final String TAG = Persistence.class.getSimpleName();
 
+  /** Temporary setting for enabling document overlays. */
+  // TODO(Overlay): Remove this.
+  public static boolean OVERLAY_SUPPORT_ENABLED = BuildConfig.ENABLE_OVERLAY;
+
   /** Temporary setting for enabling indexing-specific code paths while in development. */
-  // TODO: Remove this.
+  // TODO(Indexing): Remove this.
   public static boolean INDEXING_SUPPORT_ENABLED = BuildConfig.ENABLE_INDEXING;
 
   // Local subclasses only, please.
@@ -93,6 +97,9 @@ public abstract class Persistence {
   /** Returns a BundleCache representing the persisted cache of loaded bundles. */
   abstract BundleCache getBundleCache();
 
+  /** Returns a DocumentOverlay representing the documents that are mutated locally. */
+  abstract DocumentOverlay getDocumentOverlay(User user);
+
   /**
    * Performs an operation inside a persistence transaction. Any reads or writes against persistence
    * must be performed within a transaction. Writes will be committed atomically once the

firebase-firestore/src/main/java/com/google/firebase/firestore/local/SQLiteDocumentOverlay.java

@@ -0,0 +1,75 @@
+// Copyright 2021 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.local;
+
+import static com.google.firebase.firestore.util.Assert.fail;
+
+import androidx.annotation.Nullable;
+import com.google.firebase.firestore.auth.User;
+import com.google.firebase.firestore.model.DocumentKey;
+import com.google.firebase.firestore.model.mutation.Mutation;
+import com.google.firestore.v1.Write;
+import com.google.protobuf.InvalidProtocolBufferException;
+
+public class SQLiteDocumentOverlay implements DocumentOverlay {
+  private final SQLitePersistence db;
+  private final LocalSerializer serializer;
+  private final String uid;
+
+  public SQLiteDocumentOverlay(SQLitePersistence db, LocalSerializer serializer, User user) {
+    this.db = db;
+    this.serializer = serializer;
+    this.uid = user.isAuthenticated() ? user.getUid() : "";
+  }
+
+  @Nullable
+  @Override
+  public Mutation getOverlay(DocumentKey key) {
+    String path = EncodedPath.encode(key.getPath());
+    return db.query("SELECT overlay_mutation FROM document_overlays WHERE uid = ? AND path = ?")
+        .binding(uid, path)
+        .firstValue(
+            row -> {
+              if (row != null) {
+                try {
+                  Write mutation = Write.parseFrom(row.getBlob(0));
+                  return serializer.decodeMutation(mutation);
+                } catch (InvalidProtocolBufferException e) {
+                  throw fail("Overlay failed to parse: %s", e);
+                }
+              }
+
+              return null;
+            });
+  }
+
+  @Override
+  public void saveOverlay(DocumentKey key, Mutation mutation) {
+    db.execute(
+        "INSERT OR REPLACE INTO document_overlays "
+            + "(uid, path, overlay_mutation) VALUES (?, ?, ?)",
+        uid,
+        EncodedPath.encode(key.getPath()),
+        serializer.encodeMutation(mutation).toByteArray());
+  }
+
+  @Override
+  public void removeOverlay(DocumentKey key) {
+    db.execute(
+        "DELETE FROM document_overlays WHERE uid = ? AND path = ?",
+        uid,
+        EncodedPath.encode(key.getPath()));
+  }
+}

firebase-firestore/src/main/java/com/google/firebase/firestore/local/SQLitePersistence.java

@@ -192,6 +192,11 @@ BundleCache getBundleCache() {
     return bundleCache;
   }
 
+  @Override
+  DocumentOverlay getDocumentOverlay(User user) {
+    return new SQLiteDocumentOverlay(this, this.serializer, user);
+  }
+
   @Override
   RemoteDocumentCache getRemoteDocumentCache() {
     return remoteDocumentCache;

firebase-firestore/src/main/java/com/google/firebase/firestore/local/SQLiteSchema.java

@@ -51,8 +51,10 @@ class SQLiteSchema {
    */
   static final int VERSION = 12;
 
+  static final int OVERLAY_SUPPORT_VERSION = VERSION + 1;
+
   // TODO(indexing): Remove this constant and increment VERSION to enable indexing support
-  static final int INDEXING_SUPPORT_VERSION = VERSION + 1;
+  static final int INDEXING_SUPPORT_VERSION = OVERLAY_SUPPORT_VERSION + 1;
 
   /**
    * The batch size for the sequence number migration in `ensureSequenceNumbers()`.
@@ -77,8 +79,14 @@ void runMigrations() {
   }
 
   void runMigrations(int fromVersion) {
-    runMigrations(
-        fromVersion, Persistence.INDEXING_SUPPORT_ENABLED ? INDEXING_SUPPORT_VERSION : VERSION);
+    int toVersion = VERSION;
+    if (Persistence.OVERLAY_SUPPORT_ENABLED) {
+      toVersion = OVERLAY_SUPPORT_VERSION;
+    }
+    if (Persistence.INDEXING_SUPPORT_ENABLED) {
+      toVersion = INDEXING_SUPPORT_VERSION;
+    }
+    runMigrations(fromVersion, toVersion);
   }
 
   /**
@@ -174,6 +182,11 @@ void runMigrations(int fromVersion, int toVersion) {
      *    maintained invariants from later versions, so migrations that update values cannot assume
      *    that existing values have been properly maintained. Calculate them again, if applicable.
      */
+    if (fromVersion < OVERLAY_SUPPORT_VERSION && toVersion >= OVERLAY_SUPPORT_VERSION) {
+      Preconditions.checkState(
+          Persistence.OVERLAY_SUPPORT_ENABLED || Persistence.INDEXING_SUPPORT_ENABLED);
+      createOverlays();
+    }
 
     if (fromVersion < INDEXING_SUPPORT_VERSION && toVersion >= INDEXING_SUPPORT_VERSION) {
       Preconditions.checkState(Persistence.INDEXING_SUPPORT_ENABLED);
@@ -600,6 +613,19 @@ private void createBundleCache() {
         });
   }
 
+  private void createOverlays() {
+    ifTablesDontExist(
+        new String[] {"document_overlays"},
+        () -> {
+          db.execSQL(
+              "CREATE TABLE document_overlays ("
+                  + "uid TEXT, "
+                  + "path TEXT, "
+                  + "overlay_mutation BLOB, "
+                  + "PRIMARY KEY (uid, path))");
+        });
+  }
+
   private boolean tableExists(String table) {
     return !new SQLitePersistence.Query(db, "SELECT 1=1 FROM sqlite_master WHERE tbl_name = ?")
         .binding(table)