1 files changed, 39 insertions, 3 deletions
diff --git a/Firestore/Source/Local/FSTLevelDBMigrations.mm b/Firestore/Source/Local/FSTLevelDBMigrations.mm
index 49af893..7595c53 100644
--- a/Firestore/Source/Local/FSTLevelDBMigrations.mm
+++ b/Firestore/Source/Local/FSTLevelDBMigrations.mm
@@ -16,21 +16,22 @@
 
 #include "Firestore/Source/Local/FSTLevelDBMigrations.h"
 
-#include <leveldb/db.h>
-#include <leveldb/write_batch.h>
+#include "leveldb/write_batch.h"
 
 #import "Firestore/Protos/objc/firestore/local/Target.pbobjc.h"
 #import "Firestore/Source/Local/FSTLevelDB.h"
 #import "Firestore/Source/Local/FSTLevelDBKey.h"
 #import "Firestore/Source/Local/FSTLevelDBQueryCache.h"
 #import "Firestore/Source/Local/FSTWriteGroup.h"
+#import "Firestore/Source/Util/FSTAssert.h"
 
 NS_ASSUME_NONNULL_BEGIN
 
 // Current version of the schema defined in this file.
-static FSTLevelDBSchemaVersion kSchemaVersion = 1;
+static FSTLevelDBSchemaVersion kSchemaVersion = 2;
 
 using leveldb::DB;
+using leveldb::Iterator;
 using leveldb::Status;
 using leveldb::Slice;
 using leveldb::WriteOptions;
@@ -57,6 +58,31 @@ static void SaveVersion(FSTLevelDBSchemaVersion version, FSTWriteGroup *group) {
   [group setData:version_string forKey:key];
 }
 
+/**
+ * This function counts the number of targets that currently exist in the given db. It
+ * then reads the target global row, adds the count to the metadata from that row, and writes
+ * the metadata back.
+ *
+ * It assumes the metadata has already been written and is able to be read in this transaction.
+ */
+static void AddTargetCount(std::shared_ptr<DB> db, FSTWriteGroup *group) {
+  std::unique_ptr<Iterator> it(db->NewIterator([FSTLevelDB standardReadOptions]));
+  Slice start_key = [FSTLevelDBTargetKey keyPrefix];
+  it->Seek(start_key);
+
+  int32_t count = 0;
+  while (it->Valid() && it->key().starts_with(start_key)) {
+    count++;
+    it->Next();
+  }
+
+  FSTPBTargetGlobal *targetGlobal = [FSTLevelDBQueryCache readTargetMetadataFromDB:db];
+  FSTCAssert(targetGlobal != nil,
+             @"We should have a metadata row as it was added in an earlier migration");
+  targetGlobal.targetCount = count;
+  [group setMessage:targetGlobal forKey:[FSTLevelDBTargetGlobalKey key]];
+}
+
 @implementation FSTLevelDBMigrations
 
 + (FSTLevelDBSchemaVersion)schemaVersionForDB:(std::shared_ptr<DB>)db {
@@ -80,6 +106,16 @@ static void SaveVersion(FSTLevelDBSchemaVersion version, FSTWriteGroup *group) {
     case 0:
       EnsureTargetGlobal(db, group);
       // Fallthrough
+    case 1:
+      // We need to make sure we have metadata, since we're going to read and modify it
+      // in this migration. Commit the current transaction and start a new one. Since we're
+      // committing, we need to save a version. It's safe to save this one, if we crash
+      // after saving we'll resume from this step when we try to migrate.
+      SaveVersion(1, group);
+      [group writeToDB:db];
+      group = [FSTWriteGroup groupWithAction:@"Migrations"];
+      AddTargetCount(db, group);
+      // Fallthrough
     default:
       if (currentVersion < kSchemaVersion) {
         SaveVersion(kSchemaVersion, group);