Update Icing from upstream.

Descriptions:
========================================================================
Handle version changes in the schema store.
========================================================================

Bug: 280697513
Change-Id: I6da14f236cd8e4cf2af75c1fb2c115d1df512114

8 files changed, 602 insertions, 17 deletions

diff --git a/icing/file/version-util.cc b/icing/file/version-util.cc
index 625f7f5..468bde5 100644
--- a/icing/file/version-util.cc
+++ b/icing/file/version-util.cc
@@ -83,7 +83,7 @@ StateChange GetVersionStateChange(const VersionInfo& existing_version_info,
   }
 
   if (existing_version_info.version == 0) {
-    return (existing_version_info.max_version < curr_version)
+    return (existing_version_info.max_version == existing_version_info.version)
                ? StateChange::kVersionZeroUpgrade
                : StateChange::kVersionZeroRollForward;
   }
@@ -93,7 +93,7 @@ StateChange GetVersionStateChange(const VersionInfo& existing_version_info,
   } else if (existing_version_info.version > curr_version) {
     return StateChange::kRollBack;
   } else {  // existing_version_info.version < curr_version
-    return (existing_version_info.max_version < curr_version)
+    return (existing_version_info.max_version == existing_version_info.version)
                ? StateChange::kUpgrade
                : StateChange::kRollForward;
   }
diff --git a/icing/file/version-util.h b/icing/file/version-util.h
index a57b67c..7fa7fbd 100644
--- a/icing/file/version-util.h
+++ b/icing/file/version-util.h
@@ -30,6 +30,7 @@ namespace version_util {
 // - Version 0: Android T. Can be identified only by flash index magic.
 // - Version 1: mainline release 2023-06.
 inline static constexpr int32_t kVersion = 1;
+inline static constexpr int32_t kVersionOne = 1;
 
 inline static constexpr int kVersionZeroFlashIndexMagic = 0x6dfba6ae;
 
diff --git a/icing/file/version-util_test.cc b/icing/file/version-util_test.cc
index e2d7e96..78cdb7d 100644
--- a/icing/file/version-util_test.cc
+++ b/icing/file/version-util_test.cc
@@ -281,7 +281,7 @@ INSTANTIATE_TEST_SUITE_P(
         VersionUtilStateChangeTestParam(
             /*existing_version_info_in=*/VersionInfo(0, 1),
             /*curr_version_in=*/2,
-            /*expected_state_change_in=*/StateChange::kVersionZeroUpgrade),
+            /*expected_state_change_in=*/StateChange::kVersionZeroRollForward),
 
         // - version 0, max_version 2
         // - Current version = 2
diff --git a/icing/icing-search-engine.cc b/icing/icing-search-engine.cc
index 0347c84..e7b6ae9 100644
--- a/icing/icing-search-engine.cc
+++ b/icing/icing-search-engine.cc
@@ -581,7 +581,7 @@ libtextclassifier3::Status IcingSearchEngine::InitializeMembers(
     // Step 1: migrate schema according to the version state change.
     ICING_RETURN_IF_ERROR(SchemaStore::MigrateSchema(
         filesystem_.get(), MakeSchemaDirectoryPath(options_.base_dir()),
-        version_state_change));
+        version_state_change, version_util::kVersion));
 
     // Step 2: discard all derived data
     ICING_RETURN_IF_ERROR(DiscardDerivedFiles());
diff --git a/icing/schema/schema-store.cc b/icing/schema/schema-store.cc
index de26573..f262206 100644
--- a/icing/schema/schema-store.cc
+++ b/icing/schema/schema-store.cc
@@ -16,6 +16,7 @@
 
 #include <algorithm>
 #include <cstdint>
+#include <limits>
 #include <memory>
 #include <string>
 #include <string_view>
@@ -151,7 +152,11 @@ SchemaStore::Header::Read(const Filesystem* filesystem,
           absl_ports::StrCat("Invalid header kMagic for file: ", path));
     }
   } else {
-    return absl_ports::InternalError("");
+    int legacy_header_size = sizeof(LegacyHeader);
+    int header_size = sizeof(Header);
+    return absl_ports::InternalError(IcingStringUtil::StringPrintf(
+        "Unexpected header size %" PRId64 ". Expected %d or %d", file_size,
+        legacy_header_size, header_size));
   }
   return header;
 }
@@ -201,9 +206,27 @@ libtextclassifier3::StatusOr<std::unique_ptr<SchemaStore>> SchemaStore::Create(
   return schema_store;
 }
 
+/* static */ libtextclassifier3::Status SchemaStore::DiscardOverlaySchema(
+    const Filesystem* filesystem, const std::string& base_dir, Header& header) {
+  std::string header_filename = MakeHeaderFilename(base_dir);
+  if (header.overlay_created()) {
+    header.SetOverlayInfo(
+        /*overlay_created=*/false,
+        /*min_overlay_version_compatibility=*/ std::numeric_limits<
+            int32_t>::max());
+    ICING_RETURN_IF_ERROR(header.Write(filesystem, header_filename));
+  }
+  std::string schema_overlay_filename = MakeOverlaySchemaFilename(base_dir);
+  if (!filesystem->DeleteFile(schema_overlay_filename.c_str())) {
+    return absl_ports::InternalError(
+        "Unable to delete stale schema overlay file.");
+  }
+  return libtextclassifier3::Status::OK;
+}
+
 /* static */ libtextclassifier3::Status SchemaStore::MigrateSchema(
     const Filesystem* filesystem, const std::string& base_dir,
-    version_util::StateChange version_state_change) {
+    version_util::StateChange version_state_change, int32_t new_version) {
   if (!filesystem->DirectoryExists(base_dir.c_str())) {
     // Situations when schema store directory doesn't exist:
     // - Initializing new Icing instance: don't have to do anything now. The
@@ -215,8 +238,64 @@ libtextclassifier3::StatusOr<std::unique_ptr<SchemaStore>> SchemaStore::Create(
     return libtextclassifier3::Status::OK;
   }
 
-  // TODO(b/280697513): implement schema migration (backup and new) according to
-  // version_state_change.
+  std::string overlay_schema_filename = MakeOverlaySchemaFilename(base_dir);
+  if (!filesystem->FileExists(overlay_schema_filename.c_str())) {
+    // The overlay doesn't exist. So there should be nothing particularly
+    // interesting to worry about.
+    return libtextclassifier3::Status::OK;
+  }
+
+  std::string header_filename = MakeHeaderFilename(base_dir);
+  libtextclassifier3::StatusOr<Header> header_or;
+  switch (version_state_change) {
+    // No necessary actions for normal upgrades or no version change. The data
+    // that was produced by the previous version is fully compatible with this
+    // version and there's no stale data for us to clean up.
+    // The same is true for a normal rollforward. A normal rollforward implies
+    // that the previous version was one that understood the concept of the
+    // overlay schema and would have already discarded it if it was unusable.
+    case version_util::StateChange::kVersionZeroUpgrade:
+      // fallthrough
+    case version_util::StateChange::kUpgrade:
+      // fallthrough
+    case version_util::StateChange::kRollForward:
+      // fallthrough
+    case version_util::StateChange::kCompatible:
+      return libtextclassifier3::Status::OK;
+    case version_util::StateChange::kVersionZeroRollForward:
+      // We've rolled forward. The schema overlay file, if it exists, is
+      // possibly stale. We must throw it out.
+      header_or = Header::Read(filesystem, header_filename);
+      if (!header_or.ok()) {
+        return header_or.status();
+      }
+      return SchemaStore::DiscardOverlaySchema(filesystem, base_dir,
+                                               header_or.ValueOrDie());
+    case version_util::StateChange::kRollBack:
+      header_or = Header::Read(filesystem, header_filename);
+      if (!header_or.ok()) {
+        return header_or.status();
+      }
+      if (header_or.ValueOrDie().min_overlay_version_compatibility() <=
+          new_version) {
+        // We've been rolled back, but the overlay schema claims that it
+        // supports this version. So we can safely return.
+        return libtextclassifier3::Status::OK;
+      }
+      // We've been rolled back to a version that the overlay schema doesn't
+      // support. We must throw it out.
+      return SchemaStore::DiscardOverlaySchema(filesystem, base_dir,
+                                               header_or.ValueOrDie());
+    case version_util::StateChange::kUndetermined:
+      // It's not clear what version we're on, but the base schema should always
+      // be safe to use. Throw out the overlay.
+      header_or = Header::Read(filesystem, header_filename);
+      if (!header_or.ok()) {
+        return header_or.status();
+      }
+      return SchemaStore::DiscardOverlaySchema(filesystem, base_dir,
+                                               header_or.ValueOrDie());
+  }
   return libtextclassifier3::Status::OK;
 }
 
@@ -320,7 +399,13 @@ libtextclassifier3::Status SchemaStore::LoadSchema() {
 
   // Something has gone wrong. We've lost part of the schema ground truth.
   // Return an error.
-  return absl_ports::InternalError("");
+  bool overlay_created = header_->overlay_created();
+  bool base_schema_exists = base_schema_state.ok();
+  return absl_ports::InternalError(IcingStringUtil::StringPrintf(
+      "Unable to properly load schema. Header {exists:%d, overlay_created:%d}, "
+      "base schema exists: %d, overlay_schema_exists: %d",
+      header_exists, overlay_created, base_schema_exists,
+      overlay_schema_file_exists));
 }
 
 libtextclassifier3::Status SchemaStore::InitializeInternal(
@@ -399,8 +484,9 @@ libtextclassifier3::Status SchemaStore::RegenerateDerivedFiles(
           std::make_unique<SchemaProto>(std::move(base_schema));
       ICING_RETURN_IF_ERROR(schema_file_->Write(std::move(base_schema_ptr)));
 
-      header_->set_overlay_created(true);
-
+      header_->SetOverlayInfo(
+          /*overlay_created=*/true,
+          /*min_overlay_version_compatibility=*/version_util::kVersionOne);
       // Rebuild in memory data - references to the old schema will be invalid
       // now.
       BuildInMemoryCache();
diff --git a/icing/schema/schema-store.h b/icing/schema/schema-store.h
index abbdd63..73d7848 100644
--- a/icing/schema/schema-store.h
+++ b/icing/schema/schema-store.h
@@ -17,6 +17,7 @@
 
 #include <cstdint>
 #include <cstring>
+#include <limits>
 #include <memory>
 #include <string>
 #include <string_view>
@@ -67,10 +68,19 @@ class SchemaStore {
    public:
     static constexpr int32_t kMagic = 0x72650d0a;
 
-    explicit Header() : magic_(kMagic), checksum_(0), overlay_created_(false) {
+    explicit Header()
+        : magic_(kMagic),
+          checksum_(0),
+          overlay_created_(false),
+          min_overlay_version_compatibility_(
+              std::numeric_limits<int32_t>::max()) {
       memset(padding, 0, kPaddingSize);
     }
 
+    // RETURNS:
+    //   - On success, a valid Header instance
+    //   - NOT_FOUND if header file doesn't exist
+    //   - INTERNAL if unable to read header
     static libtextclassifier3::StatusOr<Header> Read(
         const Filesystem* filesystem, const std::string& path);
 
@@ -83,8 +93,15 @@ class SchemaStore {
     void set_checksum(uint32_t checksum) { checksum_ = checksum; }
 
     bool overlay_created() const { return overlay_created_; }
-    void set_overlay_created(bool overlay_created) {
+
+    int32_t min_overlay_version_compatibility() const {
+      return min_overlay_version_compatibility_;
+    }
+
+    void SetOverlayInfo(bool overlay_created,
+                        int32_t min_overlay_version_compatibility) {
       overlay_created_ = overlay_created;
+      min_overlay_version_compatibility_ = min_overlay_version_compatibility;
     }
 
    private:
@@ -96,8 +113,10 @@ class SchemaStore {
 
     bool overlay_created_;
 
-    static constexpr int kPaddingSize =
-        1024 - sizeof(magic_) - sizeof(checksum_) - sizeof(overlay_created_);
+    int32_t min_overlay_version_compatibility_;
+
+    static constexpr int kPaddingSize = 1008;
+    // Padding exists just to reserve space for additional values.
     uint8_t padding[kPaddingSize];
   };
   static_assert(sizeof(Header) == 1024);
@@ -193,7 +212,7 @@ class SchemaStore {
   //   OK on success or nothing to migrate
   static libtextclassifier3::Status MigrateSchema(
       const Filesystem* filesystem, const std::string& base_dir,
-      version_util::StateChange version_state_change);
+      version_util::StateChange version_state_change, int32_t new_version);
 
   // Discards all derived data in the schema store.
   //
@@ -403,6 +422,15 @@ class SchemaStore {
   explicit SchemaStore(const Filesystem* filesystem, std::string base_dir,
                        const Clock* clock);
 
+  // Deletes the overlay schema and ensures that the Header is correctly set.
+  //
+  // RETURNS:
+  //   OK on success
+  //   INTERNAL_ERROR on any IO errors
+  static libtextclassifier3::Status DiscardOverlaySchema(
+      const Filesystem* filesystem, const std::string& base_dir,
+      Header& header);
+
   // Verifies that there is no error retrieving a previously set schema. Then
   // initializes like normal.
   //
@@ -483,6 +511,13 @@ class SchemaStore {
                : absl_ports::FailedPreconditionError("Schema not set yet.");
   }
 
+  // Correctly loads the Header, schema_file_ and (if present) the
+  // overlay_schema_file_.
+  // RETURNS:
+  //   - OK on success
+  //   - INTERNAL if an IO error is encountered when reading the Header or
+  //   schemas.
+  //     Or an invalid schema configuration is present.
   libtextclassifier3::Status LoadSchema();
 
   const Filesystem* filesystem_;
diff --git a/icing/schema/schema-store_test.cc b/icing/schema/schema-store_test.cc
index 9a14ac3..3298b75 100644
--- a/icing/schema/schema-store_test.cc
+++ b/icing/schema/schema-store_test.cc
@@ -25,6 +25,7 @@
 #include "icing/document-builder.h"
 #include "icing/file/filesystem.h"
 #include "icing/file/mock-filesystem.h"
+#include "icing/file/version-util.h"
 #include "icing/portable/equals-proto.h"
 #include "icing/proto/debug.pb.h"
 #include "icing/proto/document.pb.h"
@@ -2594,6 +2595,468 @@ TEST_F(SchemaStoreTest, LoadSchemaNoOverlayHeaderMissing) {
   }
 }
 
+TEST_F(SchemaStoreTest, MigrateSchemaCompatibleNoChange) {
+  // Create a schema that is rollback incompatible and will trigger us to create
+  // a backup schema.
+  SchemaTypeConfigProto type_a =
+      SchemaTypeConfigBuilder()
+          .SetType("type_a")
+          .AddProperty(
+              PropertyConfigBuilder()
+                  .SetName("propRfc")
+                  .SetCardinality(CARDINALITY_OPTIONAL)
+                  .SetDataTypeString(TERM_MATCH_PREFIX, TOKENIZER_RFC822))
+          .Build();
+  SchemaProto schema = SchemaBuilder().AddType(type_a).Build();
+
+  {
+    // Create an instance of the schema store and set the schema.
+    ICING_ASSERT_OK_AND_ASSIGN(
+        std::unique_ptr<SchemaStore> schema_store,
+        SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
+    ICING_ASSERT_OK(schema_store->SetSchema(
+        schema, /*ignore_errors_and_delete_documents=*/false,
+        /*allow_circular_schema_definitions=*/false));
+
+    EXPECT_THAT(schema_store->GetSchema(),
+                IsOkAndHolds(Pointee(EqualsProto(schema))));
+  }
+
+  ICING_EXPECT_OK(SchemaStore::MigrateSchema(
+      &filesystem_, schema_store_dir_, version_util::StateChange::kCompatible,
+      version_util::kVersion));
+
+  {
+    // Create a new of the schema store and check that the same schema is
+    // present.
+    ICING_ASSERT_OK_AND_ASSIGN(
+        std::unique_ptr<SchemaStore> schema_store,
+        SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
+    EXPECT_THAT(schema_store->GetSchema(),
+                IsOkAndHolds(Pointee(EqualsProto(schema))));
+  }
+}
+
+TEST_F(SchemaStoreTest, MigrateSchemaUpgradeNoChange) {
+  // Create a schema that is rollback incompatible and will trigger us to create
+  // a backup schema.
+  SchemaTypeConfigProto type_a =
+      SchemaTypeConfigBuilder()
+          .SetType("type_a")
+          .AddProperty(
+              PropertyConfigBuilder()
+                  .SetName("propRfc")
+                  .SetCardinality(CARDINALITY_OPTIONAL)
+                  .SetDataTypeString(TERM_MATCH_PREFIX, TOKENIZER_RFC822))
+          .Build();
+  SchemaProto schema = SchemaBuilder().AddType(type_a).Build();
+
+  {
+    // Create an instance of the schema store and set the schema.
+    ICING_ASSERT_OK_AND_ASSIGN(
+        std::unique_ptr<SchemaStore> schema_store,
+        SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
+    ICING_ASSERT_OK(schema_store->SetSchema(
+        schema, /*ignore_errors_and_delete_documents=*/false,
+        /*allow_circular_schema_definitions=*/false));
+
+    EXPECT_THAT(schema_store->GetSchema(),
+                IsOkAndHolds(Pointee(EqualsProto(schema))));
+  }
+
+  ICING_EXPECT_OK(SchemaStore::MigrateSchema(
+      &filesystem_, schema_store_dir_, version_util::StateChange::kUpgrade,
+      version_util::kVersion + 1));
+
+  {
+    // Create a new of the schema store and check that the same schema is
+    // present.
+    ICING_ASSERT_OK_AND_ASSIGN(
+        std::unique_ptr<SchemaStore> schema_store,
+        SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
+    EXPECT_THAT(schema_store->GetSchema(),
+                IsOkAndHolds(Pointee(EqualsProto(schema))));
+  }
+}
+
+TEST_F(SchemaStoreTest, MigrateSchemaVersionZeroUpgradeNoChange) {
+  // Because we are upgrading from version zero, the schema must be compatible
+  // with version zero.
+  SchemaTypeConfigProto type_a =
+      SchemaTypeConfigBuilder()
+          .SetType("type_a")
+          .AddProperty(
+              PropertyConfigBuilder()
+                  .SetName("propRfc")
+                  .SetCardinality(CARDINALITY_OPTIONAL)
+                  .SetDataTypeString(TERM_MATCH_PREFIX, TOKENIZER_PLAIN))
+          .Build();
+  SchemaProto schema = SchemaBuilder().AddType(type_a).Build();
+
+  {
+    // Create an instance of the schema store and set the schema.
+    ICING_ASSERT_OK_AND_ASSIGN(
+        std::unique_ptr<SchemaStore> schema_store,
+        SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
+    ICING_ASSERT_OK(schema_store->SetSchema(
+        schema, /*ignore_errors_and_delete_documents=*/false,
+        /*allow_circular_schema_definitions=*/false));
+
+    EXPECT_THAT(schema_store->GetSchema(),
+                IsOkAndHolds(Pointee(EqualsProto(schema))));
+  }
+
+  ICING_EXPECT_OK(
+      SchemaStore::MigrateSchema(&filesystem_, schema_store_dir_,
+                                 version_util::StateChange::kVersionZeroUpgrade,
+                                 version_util::kVersion + 1));
+
+  {
+    // Create a new of the schema store and check that the same schema is
+    // present.
+    ICING_ASSERT_OK_AND_ASSIGN(
+        std::unique_ptr<SchemaStore> schema_store,
+        SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
+    EXPECT_THAT(schema_store->GetSchema(),
+                IsOkAndHolds(Pointee(EqualsProto(schema))));
+  }
+}
+
+TEST_F(SchemaStoreTest, MigrateSchemaRollbackDiscardsOverlaySchema) {
+  // Because we are upgrading from version zero, the schema must be compatible
+  // with version zero.
+  SchemaTypeConfigProto type_a =
+      SchemaTypeConfigBuilder()
+          .SetType("type_a")
+          .AddProperty(
+              PropertyConfigBuilder()
+                  .SetName("propRfc")
+                  .SetCardinality(CARDINALITY_OPTIONAL)
+                  .SetDataTypeString(TERM_MATCH_PREFIX, TOKENIZER_RFC822))
+          .Build();
+  SchemaProto schema = SchemaBuilder().AddType(type_a).Build();
+
+  {
+    // Create an instance of the schema store and set the schema.
+    ICING_ASSERT_OK_AND_ASSIGN(
+        std::unique_ptr<SchemaStore> schema_store,
+        SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
+    ICING_ASSERT_OK(schema_store->SetSchema(
+        schema, /*ignore_errors_and_delete_documents=*/false,
+        /*allow_circular_schema_definitions=*/false));
+
+    EXPECT_THAT(schema_store->GetSchema(),
+                IsOkAndHolds(Pointee(EqualsProto(schema))));
+  }
+
+  // Rollback to a version before kVersion. The schema header will declare that
+  // the overlay is compatible with any version starting with kVersion. So
+  // kVersion - 1 is incompatible and will throw out the schema.
+  ICING_EXPECT_OK(SchemaStore::MigrateSchema(
+      &filesystem_, schema_store_dir_, version_util::StateChange::kRollBack,
+      version_util::kVersion - 1));
+
+  {
+    // Create a new of the schema store and check that we fell back to the
+    // base schema.
+    ICING_ASSERT_OK_AND_ASSIGN(
+        std::unique_ptr<SchemaStore> schema_store,
+        SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
+
+    SchemaTypeConfigProto other_type_a =
+        SchemaTypeConfigBuilder()
+            .SetType("type_a")
+            .AddProperty(PropertyConfigBuilder()
+                             .SetName("propRfc")
+                             .SetCardinality(CARDINALITY_OPTIONAL)
+                             .SetDataType(TYPE_STRING))
+            .Build();
+    SchemaProto base_schema = SchemaBuilder().AddType(other_type_a).Build();
+    EXPECT_THAT(schema_store->GetSchema(),
+                IsOkAndHolds(Pointee(EqualsProto(base_schema))));
+  }
+}
+
+TEST_F(SchemaStoreTest, MigrateSchemaCompatibleRollbackKeepsOverlaySchema) {
+  // Because we are upgrading from version zero, the schema must be compatible
+  // with version zero.
+  SchemaTypeConfigProto type_a =
+      SchemaTypeConfigBuilder()
+          .SetType("type_a")
+          .AddProperty(
+              PropertyConfigBuilder()
+                  .SetName("propRfc")
+                  .SetCardinality(CARDINALITY_OPTIONAL)
+                  .SetDataTypeString(TERM_MATCH_PREFIX, TOKENIZER_RFC822))
+          .Build();
+  SchemaProto schema = SchemaBuilder().AddType(type_a).Build();
+
+  {
+    // Create an instance of the schema store and set the schema.
+    ICING_ASSERT_OK_AND_ASSIGN(
+        std::unique_ptr<SchemaStore> schema_store,
+        SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
+    ICING_ASSERT_OK(schema_store->SetSchema(
+        schema, /*ignore_errors_and_delete_documents=*/false,
+        /*allow_circular_schema_definitions=*/false));
+
+    EXPECT_THAT(schema_store->GetSchema(),
+                IsOkAndHolds(Pointee(EqualsProto(schema))));
+  }
+
+  // Rollback to kVersion. The schema header will declare that the overlay is
+  // compatible with any version starting with kVersion. So we will be
+  // compatible and retain the overlay schema.
+  ICING_EXPECT_OK(SchemaStore::MigrateSchema(
+      &filesystem_, schema_store_dir_, version_util::StateChange::kRollBack,
+      version_util::kVersion));
+
+  {
+    // Create a new of the schema store and check that the same schema is
+    // present.
+    ICING_ASSERT_OK_AND_ASSIGN(
+        std::unique_ptr<SchemaStore> schema_store,
+        SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
+    EXPECT_THAT(schema_store->GetSchema(),
+                IsOkAndHolds(Pointee(EqualsProto(schema))));
+  }
+}
+
+TEST_F(SchemaStoreTest, MigrateSchemaRollforwardRetainsBaseSchema) {
+  SchemaTypeConfigProto type_a =
+      SchemaTypeConfigBuilder()
+          .SetType("type_a")
+          .AddProperty(
+              PropertyConfigBuilder()
+                  .SetName("propRfc")
+                  .SetCardinality(CARDINALITY_OPTIONAL)
+                  .SetDataTypeString(TERM_MATCH_PREFIX, TOKENIZER_RFC822))
+          .Build();
+  SchemaProto schema = SchemaBuilder().AddType(type_a).Build();
+  {
+    // Create an instance of the schema store and set the schema.
+    ICING_ASSERT_OK_AND_ASSIGN(
+        std::unique_ptr<SchemaStore> schema_store,
+        SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
+    ICING_ASSERT_OK(schema_store->SetSchema(
+        schema, /*ignore_errors_and_delete_documents=*/false,
+        /*allow_circular_schema_definitions=*/false));
+
+    EXPECT_THAT(schema_store->GetSchema(),
+                IsOkAndHolds(Pointee(EqualsProto(schema))));
+  }
+
+  // Rollback to a version before kVersion. The schema header will declare that
+  // the overlay is compatible with any version starting with kVersion. So
+  // kVersion - 1 is incompatible and will throw out the schema.
+  ICING_EXPECT_OK(SchemaStore::MigrateSchema(
+      &filesystem_, schema_store_dir_, version_util::StateChange::kRollBack,
+      version_util::kVersion - 1));
+
+  SchemaTypeConfigProto other_type_a =
+      SchemaTypeConfigBuilder()
+          .SetType("type_a")
+          .AddProperty(PropertyConfigBuilder()
+                           .SetName("propRfc")
+                           .SetCardinality(CARDINALITY_OPTIONAL)
+                           .SetDataType(TYPE_STRING))
+          .Build();
+  SchemaProto base_schema = SchemaBuilder().AddType(other_type_a).Build();
+
+  {
+    // Create a new of the schema store and check that we fell back to the
+    // base schema.
+    ICING_ASSERT_OK_AND_ASSIGN(
+        std::unique_ptr<SchemaStore> schema_store,
+        SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
+
+    EXPECT_THAT(schema_store->GetSchema(),
+                IsOkAndHolds(Pointee(EqualsProto(base_schema))));
+  }
+
+  // Now rollforward to a new version. This should accept whatever schema is
+  // present (currently base schema)
+  ICING_EXPECT_OK(SchemaStore::MigrateSchema(
+      &filesystem_, schema_store_dir_, version_util::StateChange::kRollForward,
+      version_util::kVersion));
+  {
+    // Create a new of the schema store and check that we fell back to the
+    // base schema.
+    ICING_ASSERT_OK_AND_ASSIGN(
+        std::unique_ptr<SchemaStore> schema_store,
+        SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
+
+    EXPECT_THAT(schema_store->GetSchema(),
+                IsOkAndHolds(Pointee(EqualsProto(base_schema))));
+  }
+}
+
+TEST_F(SchemaStoreTest, MigrateSchemaRollforwardRetainsOverlaySchema) {
+  SchemaTypeConfigProto type_a =
+      SchemaTypeConfigBuilder()
+          .SetType("type_a")
+          .AddProperty(
+              PropertyConfigBuilder()
+                  .SetName("propRfc")
+                  .SetCardinality(CARDINALITY_OPTIONAL)
+                  .SetDataTypeString(TERM_MATCH_PREFIX, TOKENIZER_RFC822))
+          .Build();
+  SchemaProto schema = SchemaBuilder().AddType(type_a).Build();
+  {
+    // Create an instance of the schema store and set the schema.
+    ICING_ASSERT_OK_AND_ASSIGN(
+        std::unique_ptr<SchemaStore> schema_store,
+        SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
+    ICING_ASSERT_OK(schema_store->SetSchema(
+        schema, /*ignore_errors_and_delete_documents=*/false,
+        /*allow_circular_schema_definitions=*/false));
+
+    EXPECT_THAT(schema_store->GetSchema(),
+                IsOkAndHolds(Pointee(EqualsProto(schema))));
+  }
+
+  // Rollback to kVersion. The schema header will declare that the overlay is
+  // compatible with any version starting with kVersion. So we will be
+  // compatible and retain the overlay schema.
+  ICING_EXPECT_OK(SchemaStore::MigrateSchema(
+      &filesystem_, schema_store_dir_, version_util::StateChange::kRollBack,
+      version_util::kVersion));
+
+  {
+    // Create a new of the schema store and check that the same schema is
+    // present.
+    ICING_ASSERT_OK_AND_ASSIGN(
+        std::unique_ptr<SchemaStore> schema_store,
+        SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
+
+    EXPECT_THAT(schema_store->GetSchema(),
+                IsOkAndHolds(Pointee(EqualsProto(schema))));
+  }
+
+  // Now rollforward to a new version. This should accept whatever schema is
+  // present (currently overlay schema)
+  ICING_EXPECT_OK(SchemaStore::MigrateSchema(
+      &filesystem_, schema_store_dir_, version_util::StateChange::kRollForward,
+      version_util::kVersion));
+  {
+    // Create a new of the schema store and check that the same schema is
+    // present.
+    ICING_ASSERT_OK_AND_ASSIGN(
+        std::unique_ptr<SchemaStore> schema_store,
+        SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
+
+    EXPECT_THAT(schema_store->GetSchema(),
+                IsOkAndHolds(Pointee(EqualsProto(schema))));
+  }
+}
+
+TEST_F(SchemaStoreTest,
+       MigrateSchemaVersionZeroRollforwardDiscardsOverlaySchema) {
+  SchemaTypeConfigProto type_a =
+      SchemaTypeConfigBuilder()
+          .SetType("type_a")
+          .AddProperty(
+              PropertyConfigBuilder()
+                  .SetName("propRfc")
+                  .SetCardinality(CARDINALITY_OPTIONAL)
+                  .SetDataTypeString(TERM_MATCH_PREFIX, TOKENIZER_RFC822))
+          .Build();
+  SchemaProto schema = SchemaBuilder().AddType(type_a).Build();
+  {
+    // Create an instance of the schema store and set the schema.
+    ICING_ASSERT_OK_AND_ASSIGN(
+        std::unique_ptr<SchemaStore> schema_store,
+        SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
+    ICING_ASSERT_OK(schema_store->SetSchema(
+        schema, /*ignore_errors_and_delete_documents=*/false,
+        /*allow_circular_schema_definitions=*/false));
+
+    EXPECT_THAT(schema_store->GetSchema(),
+                IsOkAndHolds(Pointee(EqualsProto(schema))));
+  }
+
+  // A VersionZeroRollforward will always discard the overlay schema because it
+  // could be stale.
+  ICING_EXPECT_OK(SchemaStore::MigrateSchema(
+      &filesystem_, schema_store_dir_,
+      version_util::StateChange::kVersionZeroRollForward,
+      version_util::kVersion));
+
+  SchemaTypeConfigProto other_type_a =
+      SchemaTypeConfigBuilder()
+          .SetType("type_a")
+          .AddProperty(PropertyConfigBuilder()
+                           .SetName("propRfc")
+                           .SetCardinality(CARDINALITY_OPTIONAL)
+                           .SetDataType(TYPE_STRING))
+          .Build();
+  SchemaProto base_schema = SchemaBuilder().AddType(other_type_a).Build();
+
+  {
+    // Create a new of the schema store and check that we fell back to the
+    // base schema.
+    ICING_ASSERT_OK_AND_ASSIGN(
+        std::unique_ptr<SchemaStore> schema_store,
+        SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
+
+    EXPECT_THAT(schema_store->GetSchema(),
+                IsOkAndHolds(Pointee(EqualsProto(base_schema))));
+  }
+}
+
+TEST_F(SchemaStoreTest, MigrateSchemaVersionUndeterminedDiscardsOverlaySchema) {
+  SchemaTypeConfigProto type_a =
+      SchemaTypeConfigBuilder()
+          .SetType("type_a")
+          .AddProperty(
+              PropertyConfigBuilder()
+                  .SetName("propRfc")
+                  .SetCardinality(CARDINALITY_OPTIONAL)
+                  .SetDataTypeString(TERM_MATCH_PREFIX, TOKENIZER_RFC822))
+          .Build();
+  SchemaProto schema = SchemaBuilder().AddType(type_a).Build();
+  {
+    // Create an instance of the schema store and set the schema.
+    ICING_ASSERT_OK_AND_ASSIGN(
+        std::unique_ptr<SchemaStore> schema_store,
+        SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
+    ICING_ASSERT_OK(schema_store->SetSchema(
+        schema, /*ignore_errors_and_delete_documents=*/false,
+        /*allow_circular_schema_definitions=*/false));
+
+    EXPECT_THAT(schema_store->GetSchema(),
+                IsOkAndHolds(Pointee(EqualsProto(schema))));
+  }
+
+  // An Undetermined will always discard the overlay schema because it doesn't
+  // know which state we're in and so it fallback to the base schema because
+  // it should always be valid.
+  ICING_EXPECT_OK(SchemaStore::MigrateSchema(
+      &filesystem_, schema_store_dir_, version_util::StateChange::kUndetermined,
+      version_util::kVersion));
+
+  SchemaTypeConfigProto other_type_a =
+      SchemaTypeConfigBuilder()
+          .SetType("type_a")
+          .AddProperty(PropertyConfigBuilder()
+                           .SetName("propRfc")
+                           .SetCardinality(CARDINALITY_OPTIONAL)
+                           .SetDataType(TYPE_STRING))
+          .Build();
+  SchemaProto base_schema = SchemaBuilder().AddType(other_type_a).Build();
+
+  {
+    // Create a new of the schema store and check that we fell back to the
+    // base schema.
+    ICING_ASSERT_OK_AND_ASSIGN(
+        std::unique_ptr<SchemaStore> schema_store,
+        SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
+
+    EXPECT_THAT(schema_store->GetSchema(),
+                IsOkAndHolds(Pointee(EqualsProto(base_schema))));
+  }
+}
+
 }  // namespace
 
 }  // namespace lib
diff --git a/synced_AOSP_CL_number.txt b/synced_AOSP_CL_number.txt
index 615a911..ae59ff7 100644
--- a/synced_AOSP_CL_number.txt
+++ b/synced_AOSP_CL_number.txt
@@ -1 +1 @@
-set(synced_AOSP_CL_number=531106228)
+set(synced_AOSP_CL_number=531296607)