LevelDB schema changes for CollectionGroup queries.

Author: Michael Lehenbauer

Firestore/Example/Tests/Local/FSTLevelDBMigrationsTests.mm

@@ -16,6 +16,7 @@
 
 #import <XCTest/XCTest.h>
 
+#include <map>
 #include <memory>
 #include <string>
 #include <vector>
@@ -37,6 +38,7 @@
 NS_ASSUME_NONNULL_BEGIN
 
 using firebase::firestore::FirestoreErrorCode;
+using firebase::firestore::local::LevelDbCollectionParentKey;
 using firebase::firestore::local::LevelDbDocumentMutationKey;
 using firebase::firestore::local::LevelDbDocumentTargetKey;
 using firebase::firestore::local::LevelDbMigrations;
@@ -334,7 +336,6 @@ - (void)testRemovesMutationBatches {
     // Verify
     std::string buffer;
     LevelDbTransaction transaction(_db.get(), "Verify");
-    auto it = transaction.NewIterator();
     // verify that we deleted the correct batches
     XCTAssertTrue(transaction.Get(LevelDbMutationKey::Key("foo", 1), &buffer).IsNotFound());
     XCTAssertTrue(transaction.Get(LevelDbMutationKey::Key("foo", 2), &buffer).IsNotFound());
@@ -360,6 +361,60 @@ - (void)testRemovesMutationBatches {
   }
 }
 
+- (void)testCreateCollectionParentsIndex {
+  // This test creates a database with schema version 5 that has a few
+  // mutations and a few remote documents and then ensures that appropriate
+  // entries are written to the collectionParentIndex.
+  std::vector<std::string> write_paths{"cg1/x", "cg1/y", "cg1/x/cg1/x", "cg2/x", "cg1/x/cg2/x"};
+  std::vector<std::string> remote_doc_paths{"cg1/z", "cg1/y/cg1/x", "cg2/x/cg3/x",
+                                            "blah/x/blah/x/cg3/x"};
+  std::map<std::string, std::vector<std::string>> expected_parents{
+      {"cg1", {"", "cg1/x", "cg1/y"}}, {"cg2", {"", "cg1/x"}}, {"cg3", {"blah/x/blah/x", "cg2/x"}}};
+
+  std::string empty_buffer;
+  LevelDbMigrations::RunMigrations(_db.get(), 5);
+  {
+    LevelDbTransaction transaction(_db.get(), "Write Mutations and Remote Documents");
+    // Write mutations.
+    for (auto write_path : write_paths) {
+      // We "cheat" and only write the DbDocumentMutation index entries, since
+      // that's all the migration uses.
+      DocumentKey key = DocumentKey::FromPathString(write_path);
+      transaction.Put(LevelDbDocumentMutationKey::Key("dummy-uid", key, /*dummy batchId=*/123),
+                      empty_buffer);
+    }
+
+    // Write remote document entries.
+    for (auto remote_doc_path : remote_doc_paths) {
+      DocumentKey key = DocumentKey::FromPathString(remote_doc_path);
+      transaction.Put(LevelDbRemoteDocumentKey::Key(key), empty_buffer);
+    }
+
+    transaction.Commit();
+  }
+
+  // Migrate to v6 and verify index entries.
+  LevelDbMigrations::RunMigrations(_db.get(), 6);
+  {
+    LevelDbTransaction transaction(_db.get(), "Verify");
+
+    std::map<std::string, std::vector<std::string>> actual_parents;
+    auto index_iterator = transaction.NewIterator();
+    std::string index_prefix = LevelDbCollectionParentKey::KeyPrefix();
+    LevelDbCollectionParentKey row_key;
+    for (index_iterator->Seek(index_prefix); index_iterator->Valid(); index_iterator->Next()) {
+      if (!absl::StartsWith(index_iterator->key(), index_prefix) ||
+          !row_key.Decode(index_iterator->key()))
+        break;
+
+      std::vector<std::string> &parents = actual_parents[row_key.collection_id()];
+      parents.push_back(row_key.parent().CanonicalString());
+    }
+
+    XCTAssertEqual(actual_parents, expected_parents);
+  }
+}
+
 - (void)testCanDowngrade {
   // First, run all of the migrations
   LevelDbMigrations::RunMigrations(_db.get());

Firestore/core/src/firebase/firestore/local/leveldb_key.cc

@@ -45,6 +45,7 @@ const char* kQueryTargetsTable = "query_target";
 const char* kTargetDocumentsTable = "target_document";
 const char* kDocumentTargetsTable = "document_target";
 const char* kRemoteDocumentsTable = "remote_document";
+const char* kCollectionParentsTable = "collection_parent";
 
 /**
  * Labels for the components of keys. These serve to make keys self-describing.
@@ -89,6 +90,12 @@ enum ComponentLabel {
   /** A component containing a user Id. */
   UserId = 13,
 
+  /**
+   * A component containing a standalone collection ID (e.g. as used by the
+   * collection_parent table, but not for collection IDs within paths).
+   */
+  CollectionId = 14,
+
   /**
    * A path segment describes just a single segment in a resource path. Path
    * segments that occur sequentially in a key represent successive segments in
@@ -159,6 +166,23 @@ class Reader {
     return ReadLabeledString(ComponentLabel::UserId);
   }
 
+  std::string ReadCollectionId() {
+    return ReadLabeledString(ComponentLabel::CollectionId);
+  }
+
+  /**
+   * Reads component labels and strings from the key until it finds a component
+   * label other than ComponentLabel::PathSegment (or the key is exhausted).
+   * All matched path segments are assembled into a ResourcePath.
+   *
+   * If the read is unsuccessful, returns a default ResourcePath and fails the
+   * Reader.
+   *
+   * Otherwise returns the decoded ResourcePath and the Reader advances to the
+   * next unread byte.
+   */
+  ResourcePath ReadResourcePath();
+
   /**
    * Reads component labels and strings from the key until it finds a component
    * label other than ComponentLabel::PathSegment (or the key is exhausted).
@@ -358,7 +382,7 @@ class Reader {
   bool ok_;
 };
 
-DocumentKey Reader::ReadDocumentKey() {
+ResourcePath Reader::ReadResourcePath() {
   std::vector<std::string> path_segments;
   while (!empty()) {
     // Advance a temporary slice to avoid advancing contents into the next key
@@ -375,7 +399,16 @@ DocumentKey Reader::ReadDocumentKey() {
     path_segments.push_back(std::move(segment));
   }
 
-  ResourcePath path{std::move(path_segments)};
+  if (ok_) {
+    return ResourcePath{std::move(path_segments)};
+  } else {
+    return ResourcePath{};
+  }
+}
+
+DocumentKey Reader::ReadDocumentKey() {
+  ResourcePath path = ReadResourcePath();
+
   if (ok_ && !path.empty() && DocumentKey::IsDocumentKey(path)) {
     return DocumentKey{std::move(path)};
   }
@@ -419,10 +452,10 @@ std::string Reader::Describe() {
     src_ = saved_source;
 
     if (label == ComponentLabel::PathSegment) {
-      DocumentKey document_key = ReadDocumentKey();
+      ResourcePath resource_path = ReadResourcePath();
       if (ok_) {
         absl::StrAppend(&description,
-                        " key=", document_key.path().CanonicalString());
+                        " path=", resource_path.CanonicalString());
       }
 
     } else if (label == ComponentLabel::TableName) {
@@ -455,6 +488,12 @@ std::string Reader::Describe() {
         absl::StrAppend(&description, " user_id=", user_id);
       }
 
+    } else if (label == ComponentLabel::CollectionId) {
+      std::string collection_id = ReadCollectionId();
+      if (ok_) {
+        absl::StrAppend(&description, " collection_id=", collection_id);
+      }
+
     } else {
       absl::StrAppend(&description, " unknown label=", static_cast<int>(label));
       Fail();
@@ -502,6 +541,10 @@ class Writer {
     WriteLabeledString(ComponentLabel::UserId, user_id);
   }
 
+  void WriteCollectionId(absl::string_view collection_id) {
+    WriteLabeledString(ComponentLabel::CollectionId, collection_id);
+  }
+
   /**
    * For each segment in the given resource path writes a
    * ComponentLabel::PathSegment component label and a string containing the
@@ -848,6 +891,39 @@ bool LevelDbRemoteDocumentKey::Decode(absl::string_view key) {
   return reader.ok();
 }
 
+std::string LevelDbCollectionParentKey::KeyPrefix() {
+  Writer writer;
+  writer.WriteTableName(kCollectionParentsTable);
+  return writer.result();
+}
+
+std::string LevelDbCollectionParentKey::KeyPrefix(
+    absl::string_view collection_id) {
+  Writer writer;
+  writer.WriteTableName(kCollectionParentsTable);
+  writer.WriteCollectionId(collection_id);
+  return writer.result();
+}
+
+std::string LevelDbCollectionParentKey::Key(absl::string_view collection_id,
+                                            const ResourcePath& parent) {
+  Writer writer;
+  writer.WriteTableName(kCollectionParentsTable);
+  writer.WriteCollectionId(collection_id);
+  writer.WriteResourcePath(parent);
+  writer.WriteTerminator();
+  return writer.result();
+}
+
+bool LevelDbCollectionParentKey::Decode(absl::string_view key) {
+  Reader reader{key};
+  reader.ReadTableNameMatching(kCollectionParentsTable);
+  collection_id_ = reader.ReadCollectionId();
+  parent_ = reader.ReadResourcePath();
+  reader.ReadTerminator();
+  return reader.ok();
+}
+
 }  // namespace local
 }  // namespace firestore
 }  // namespace firebase

Firestore/core/src/firebase/firestore/local/leveldb_key.h

@@ -77,6 +77,11 @@ namespace local {
 // remote_documents:
 //   - table_name: string = "remote_document"
 //   - path: ResourcePath
+//
+// collection_parents:
+//   - table_name: string = "collection_parent"
+//   - collectionId: string
+//   - parent: ResourcePath
 
 /**
  * Parses the given key and returns a human readable description of its
@@ -520,6 +525,58 @@ class LevelDbRemoteDocumentKey {
   model::DocumentKey document_key_;
 };
 
+/**
+ * A key in the collection parents index, which stores an association between a
+ * Collection ID (e.g. 'messages') to a parent path (e.g. '/chats/123') that
+ * contains it as a (sub)collection. This is used to efficiently find all
+ * collections to query when performing a Collection Group query.
+ */
+class LevelDbCollectionParentKey {
+ public:
+  /**
+   * Creates a key prefix that points just before the first key in the table.
+   */
+  static std::string KeyPrefix();
+
+  /**
+   * Creates a key prefix that points just before the first key for the given
+   * collection_id.
+   */
+  static std::string KeyPrefix(absl::string_view collection_id);
+
+  /**
+   * Creates a complete key that points to a specific collection_id and parent.
+   */
+  static std::string Key(absl::string_view collection_id,
+                         const model::ResourcePath& parent);
+
+  /**
+   * Decodes the given complete key, storing the decoded values in this
+   * instance.
+   *
+   * @return true if the key successfully decoded, false otherwise. If false is
+   * returned, this instance is in an undefined state until the next call to
+   * `Decode()`.
+   */
+  ABSL_MUST_USE_RESULT
+  bool Decode(absl::string_view key);
+
+  /** The collection_id, as encoded in the key. */
+  const std::string& collection_id() const {
+    return collection_id_;
+  }
+
+  /** The parent path, as encoded in the key. */
+  const model::ResourcePath& parent() const {
+    return parent_;
+  }
+
+ private:
+  // Deliberately uninitialized: will be assigned in Decode
+  std::string collection_id_;
+  model::ResourcePath parent_;
+};
+
 }  // namespace local
 }  // namespace firestore
 }  // namespace firebase