Fix dangling pointer from failure in SchemaStore::RegenerateDerivedFiles.

Essentially, SchemaStore::SetSchema would work in three steps:
1. Write the new schema
2. Throw out derived data structures
3. Regenerate the derived date structures

A failure in the middle of step 3 could cause SectionManager to maintain a a reference to a now-deleted instance of KeyMapper which would then crash on the next access. This resolves that issue via shadowing to ensure that changes are only applied if everything is successfully generated.

You can find the full description of the bug at b/231416401#comment10.

This is a cherrypick of cl/449537037 and cl/450024848. The changes in cl/449849170 (and replaced with direct calls to delete the temp directories) were excluded to keep the size of cherrypick smaller.

Test: Builds

Change-Id: Ic018ade45b5aab4b797a0a654eaea6ac424469ab

10 files changed, 471 insertions, 124 deletions

diff --git a/icing/file/file-backed-proto.h b/icing/file/file-backed-proto.h
index 15a1953..d7d9bad 100644
--- a/icing/file/file-backed-proto.h
+++ b/icing/file/file-backed-proto.h
@@ -63,6 +63,17 @@ class FileBackedProto {
   // file_path : Must be a path within in a directory that already exists.
   FileBackedProto(const Filesystem& filesystem, std::string_view file_path);
 
+  // Reset the internal file_path for the file backed proto.
+  // Example use:
+  //   auto file_backed_proto1 = *FileBackedProto<Proto>::Create(...);
+  //   auto file_backed_proto2 = *FileBackedProto<Proto>::Create(...);
+  //   filesystem.SwapFiles(file1, file2);
+  //   file_backed_proto1.SetSwappedFilepath(file2);
+  //   file_backed_proto2.SetSwappedFilepath(file1);
+  void SetSwappedFilepath(std::string_view swapped_to_file_path) {
+    file_path_ = swapped_to_file_path;
+  }
+
   // Returns a reference to the proto read from the file. It
   // internally caches the read proto so that future calls are fast.
   //
@@ -99,7 +110,7 @@ class FileBackedProto {
   mutable absl_ports::shared_mutex mutex_;
 
   const Filesystem* const filesystem_;
-  const std::string file_path_;
+  std::string file_path_;
 
   mutable std::unique_ptr<ProtoT> cached_proto_ ICING_GUARDED_BY(mutex_);
 };
diff --git a/icing/icing-search-engine_test.cc b/icing/icing-search-engine_test.cc
index 5244f4c..13e77b8 100644
--- a/icing/icing-search-engine_test.cc
+++ b/icing/icing-search-engine_test.cc
@@ -765,8 +765,7 @@ TEST_F(IcingSearchEngineTest, FailToWriteSchema) {
 
   auto mock_filesystem = std::make_unique<MockFilesystem>();
   // This fails FileBackedProto::Write()
-  ON_CALL(*mock_filesystem,
-          OpenForWrite(Eq(icing_options.base_dir() + "/schema_dir/schema.pb")))
+  ON_CALL(*mock_filesystem, OpenForWrite(HasSubstr("schema.pb")))
       .WillByDefault(Return(-1));
 
   TestIcingSearchEngine icing(icing_options, std::move(mock_filesystem),
@@ -2993,13 +2992,17 @@ TEST_F(IcingSearchEngineTest, OptimizationFailureUninitializesIcing) {
   };
   ON_CALL(*mock_filesystem, CreateDirectoryRecursively)
       .WillByDefault(create_dir_lambda);
+
   auto swap_lambda = [&just_swapped_files](const char* first_dir,
                                            const char* second_dir) {
     just_swapped_files = true;
     return false;
   };
-  ON_CALL(*mock_filesystem, SwapFiles).WillByDefault(swap_lambda);
-  TestIcingSearchEngine icing(GetDefaultIcingOptions(),
+  IcingSearchEngineOptions options = GetDefaultIcingOptions();
+  ON_CALL(*mock_filesystem, SwapFiles(HasSubstr("document_dir_optimize_tmp"),
+                                      HasSubstr("document_dir")))
+      .WillByDefault(swap_lambda);
+  TestIcingSearchEngine icing(options, std::move(mock_filesystem),
                               std::move(mock_filesystem),
                               std::make_unique<IcingFilesystem>(),
                               std::make_unique<FakeClock>(), GetTestJniCache());
@@ -3752,7 +3755,8 @@ TEST_F(IcingSearchEngineTest, IcingShouldWorkFineIfOptimizationIsAborted) {
   // fails. This will fail IcingSearchEngine::OptimizeDocumentStore() and makes
   // it return ABORTED_ERROR.
   auto mock_filesystem = std::make_unique<MockFilesystem>();
-  ON_CALL(*mock_filesystem, DeleteDirectoryRecursively)
+  ON_CALL(*mock_filesystem,
+          DeleteDirectoryRecursively(HasSubstr("_optimize_tmp")))
       .WillByDefault(Return(false));
 
   TestIcingSearchEngine icing(GetDefaultIcingOptions(),
@@ -3799,7 +3803,8 @@ TEST_F(IcingSearchEngineTest,
   // Creates a mock filesystem in which SwapFiles() always fails and deletes the
   // directories. This will fail IcingSearchEngine::OptimizeDocumentStore().
   auto mock_filesystem = std::make_unique<MockFilesystem>();
-  ON_CALL(*mock_filesystem, SwapFiles)
+  ON_CALL(*mock_filesystem, SwapFiles(HasSubstr("document_dir_optimize_tmp"),
+                                      HasSubstr("document_dir")))
       .WillByDefault([this](const char* one, const char* two) {
         filesystem()->DeleteDirectoryRecursively(one);
         filesystem()->DeleteDirectoryRecursively(two);
@@ -3870,7 +3875,8 @@ TEST_F(IcingSearchEngineTest, OptimizationShouldRecoverIfDataFilesAreMissing) {
   // Creates a mock filesystem in which SwapFiles() always fails and empties the
   // directories. This will fail IcingSearchEngine::OptimizeDocumentStore().
   auto mock_filesystem = std::make_unique<MockFilesystem>();
-  ON_CALL(*mock_filesystem, SwapFiles)
+  ON_CALL(*mock_filesystem, SwapFiles(HasSubstr("document_dir_optimize_tmp"),
+                                      HasSubstr("document_dir")))
       .WillByDefault([this](const char* one, const char* two) {
         filesystem()->DeleteDirectoryRecursively(one);
         filesystem()->CreateDirectoryRecursively(one);
diff --git a/icing/index/index-processor_test.cc b/icing/index/index-processor_test.cc
index bd310de..7746688 100644
--- a/icing/index/index-processor_test.cc
+++ b/icing/index/index-processor_test.cc
@@ -153,9 +153,12 @@ class IndexProcessorTest : public Test {
         normalizer_factory::Create(
             /*max_term_byte_size=*/std::numeric_limits<int32_t>::max()));
 
+    std::string schema_store_dir = GetTestTempDir() + "/schema_store";
+    ASSERT_TRUE(
+        filesystem_.CreateDirectoryRecursively(schema_store_dir.c_str()));
     ICING_ASSERT_OK_AND_ASSIGN(
         schema_store_,
-        SchemaStore::Create(&filesystem_, GetTestTempDir(), &fake_clock_));
+        SchemaStore::Create(&filesystem_, schema_store_dir, &fake_clock_));
     SchemaProto schema =
         SchemaBuilder()
             .AddType(
diff --git a/icing/query/query-processor_test.cc b/icing/query/query-processor_test.cc
index 950f739..eaa0efc 100644
--- a/icing/query/query-processor_test.cc
+++ b/icing/query/query-processor_test.cc
@@ -77,12 +77,14 @@ class QueryProcessorTest : public Test {
   QueryProcessorTest()
       : test_dir_(GetTestTempDir() + "/icing"),
         store_dir_(test_dir_ + "/store"),
+        schema_store_dir_(test_dir_ + "/schema_store"),
         index_dir_(test_dir_ + "/index") {}
 
   void SetUp() override {
     filesystem_.DeleteDirectoryRecursively(test_dir_.c_str());
     filesystem_.CreateDirectoryRecursively(index_dir_.c_str());
     filesystem_.CreateDirectoryRecursively(store_dir_.c_str());
+    filesystem_.CreateDirectoryRecursively(schema_store_dir_.c_str());
 
     if (!IsCfStringTokenization() && !IsReverseJniTokenization()) {
       // If we've specified using the reverse-JNI method for segmentation (i.e.
@@ -129,6 +131,7 @@ class QueryProcessorTest : public Test {
   Filesystem filesystem_;
   const std::string test_dir_;
   const std::string store_dir_;
+  const std::string schema_store_dir_;
   std::unique_ptr<Index> index_;
   std::unique_ptr<LanguageSegmenter> language_segmenter_;
   std::unique_ptr<Normalizer> normalizer_;
@@ -176,7 +179,7 @@ TEST_F(QueryProcessorTest, EmptyGroupMatchAllDocuments) {
 
   ICING_ASSERT_OK_AND_ASSIGN(
       schema_store_,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
   ASSERT_THAT(schema_store_->SetSchema(schema), IsOk());
 
   ICING_ASSERT_OK_AND_ASSIGN(
@@ -227,7 +230,7 @@ TEST_F(QueryProcessorTest, EmptyQueryMatchAllDocuments) {
 
   ICING_ASSERT_OK_AND_ASSIGN(
       schema_store_,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
   ASSERT_THAT(schema_store_->SetSchema(schema), IsOk());
 
   ICING_ASSERT_OK_AND_ASSIGN(
@@ -278,7 +281,7 @@ TEST_F(QueryProcessorTest, QueryTermNormalized) {
 
   ICING_ASSERT_OK_AND_ASSIGN(
       schema_store_,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
   ASSERT_THAT(schema_store_->SetSchema(schema), IsOk());
 
   ICING_ASSERT_OK_AND_ASSIGN(
@@ -354,7 +357,7 @@ TEST_F(QueryProcessorTest, OneTermPrefixMatch) {
 
   ICING_ASSERT_OK_AND_ASSIGN(
       schema_store_,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
   ASSERT_THAT(schema_store_->SetSchema(schema), IsOk());
 
   ICING_ASSERT_OK_AND_ASSIGN(
@@ -422,7 +425,7 @@ TEST_F(QueryProcessorTest, OneTermExactMatch) {
 
   ICING_ASSERT_OK_AND_ASSIGN(
       schema_store_,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
   ASSERT_THAT(schema_store_->SetSchema(schema), IsOk());
 
   ICING_ASSERT_OK_AND_ASSIGN(
@@ -490,7 +493,7 @@ TEST_F(QueryProcessorTest, AndSameTermExactMatch) {
 
   ICING_ASSERT_OK_AND_ASSIGN(
       schema_store_,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
   ASSERT_THAT(schema_store_->SetSchema(schema), IsOk());
 
   ICING_ASSERT_OK_AND_ASSIGN(
@@ -560,7 +563,7 @@ TEST_F(QueryProcessorTest, AndTwoTermExactMatch) {
 
   ICING_ASSERT_OK_AND_ASSIGN(
       schema_store_,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
   ASSERT_THAT(schema_store_->SetSchema(schema), IsOk());
 
   ICING_ASSERT_OK_AND_ASSIGN(
@@ -635,7 +638,7 @@ TEST_F(QueryProcessorTest, AndSameTermPrefixMatch) {
 
   ICING_ASSERT_OK_AND_ASSIGN(
       schema_store_,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
   ASSERT_THAT(schema_store_->SetSchema(schema), IsOk());
 
   ICING_ASSERT_OK_AND_ASSIGN(
@@ -705,7 +708,7 @@ TEST_F(QueryProcessorTest, AndTwoTermPrefixMatch) {
 
   ICING_ASSERT_OK_AND_ASSIGN(
       schema_store_,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
   ASSERT_THAT(schema_store_->SetSchema(schema), IsOk());
 
   ICING_ASSERT_OK_AND_ASSIGN(
@@ -781,7 +784,7 @@ TEST_F(QueryProcessorTest, AndTwoTermPrefixAndExactMatch) {
 
   ICING_ASSERT_OK_AND_ASSIGN(
       schema_store_,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
   ASSERT_THAT(schema_store_->SetSchema(schema), IsOk());
 
   ICING_ASSERT_OK_AND_ASSIGN(
@@ -857,7 +860,7 @@ TEST_F(QueryProcessorTest, OrTwoTermExactMatch) {
 
   ICING_ASSERT_OK_AND_ASSIGN(
       schema_store_,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
   ASSERT_THAT(schema_store_->SetSchema(schema), IsOk());
 
   ICING_ASSERT_OK_AND_ASSIGN(
@@ -946,7 +949,7 @@ TEST_F(QueryProcessorTest, OrTwoTermPrefixMatch) {
 
   ICING_ASSERT_OK_AND_ASSIGN(
       schema_store_,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
   ASSERT_THAT(schema_store_->SetSchema(schema), IsOk());
 
   ICING_ASSERT_OK_AND_ASSIGN(
@@ -1034,7 +1037,7 @@ TEST_F(QueryProcessorTest, OrTwoTermPrefixAndExactMatch) {
 
   ICING_ASSERT_OK_AND_ASSIGN(
       schema_store_,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
   ASSERT_THAT(schema_store_->SetSchema(schema), IsOk());
 
   ICING_ASSERT_OK_AND_ASSIGN(
@@ -1121,7 +1124,7 @@ TEST_F(QueryProcessorTest, CombinedAndOrTerms) {
 
   ICING_ASSERT_OK_AND_ASSIGN(
       schema_store_,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
   ASSERT_THAT(schema_store_->SetSchema(schema), IsOk());
 
   ICING_ASSERT_OK_AND_ASSIGN(
@@ -1307,7 +1310,7 @@ TEST_F(QueryProcessorTest, OneGroup) {
 
   ICING_ASSERT_OK_AND_ASSIGN(
       schema_store_,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
   ASSERT_THAT(schema_store_->SetSchema(schema), IsOk());
 
   ICING_ASSERT_OK_AND_ASSIGN(
@@ -1383,7 +1386,7 @@ TEST_F(QueryProcessorTest, TwoGroups) {
 
   ICING_ASSERT_OK_AND_ASSIGN(
       schema_store_,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
   ASSERT_THAT(schema_store_->SetSchema(schema), IsOk());
 
   ICING_ASSERT_OK_AND_ASSIGN(
@@ -1461,7 +1464,7 @@ TEST_F(QueryProcessorTest, ManyLevelNestedGrouping) {
 
   ICING_ASSERT_OK_AND_ASSIGN(
       schema_store_,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
   ASSERT_THAT(schema_store_->SetSchema(schema), IsOk());
 
   ICING_ASSERT_OK_AND_ASSIGN(
@@ -1537,7 +1540,7 @@ TEST_F(QueryProcessorTest, OneLevelNestedGrouping) {
 
   ICING_ASSERT_OK_AND_ASSIGN(
       schema_store_,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
   ASSERT_THAT(schema_store_->SetSchema(schema), IsOk());
 
   ICING_ASSERT_OK_AND_ASSIGN(
@@ -1614,7 +1617,7 @@ TEST_F(QueryProcessorTest, ExcludeTerm) {
 
   ICING_ASSERT_OK_AND_ASSIGN(
       schema_store_,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
   ASSERT_THAT(schema_store_->SetSchema(schema), IsOk());
 
   ICING_ASSERT_OK_AND_ASSIGN(
@@ -1679,7 +1682,7 @@ TEST_F(QueryProcessorTest, ExcludeNonexistentTerm) {
 
   ICING_ASSERT_OK_AND_ASSIGN(
       schema_store_,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
   ASSERT_THAT(schema_store_->SetSchema(schema), IsOk());
 
   ICING_ASSERT_OK_AND_ASSIGN(
@@ -1742,7 +1745,7 @@ TEST_F(QueryProcessorTest, ExcludeAnd) {
 
   ICING_ASSERT_OK_AND_ASSIGN(
       schema_store_,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
   ASSERT_THAT(schema_store_->SetSchema(schema), IsOk());
 
   ICING_ASSERT_OK_AND_ASSIGN(
@@ -1832,7 +1835,7 @@ TEST_F(QueryProcessorTest, ExcludeOr) {
 
   ICING_ASSERT_OK_AND_ASSIGN(
       schema_store_,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
   ASSERT_THAT(schema_store_->SetSchema(schema), IsOk());
 
   ICING_ASSERT_OK_AND_ASSIGN(
@@ -1928,7 +1931,7 @@ TEST_F(QueryProcessorTest, DeletedFilter) {
 
   ICING_ASSERT_OK_AND_ASSIGN(
       schema_store_,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
   ASSERT_THAT(schema_store_->SetSchema(schema), IsOk());
 
   ICING_ASSERT_OK_AND_ASSIGN(
@@ -2002,7 +2005,7 @@ TEST_F(QueryProcessorTest, NamespaceFilter) {
 
   ICING_ASSERT_OK_AND_ASSIGN(
       schema_store_,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
   ASSERT_THAT(schema_store_->SetSchema(schema), IsOk());
 
   ICING_ASSERT_OK_AND_ASSIGN(
@@ -2078,7 +2081,7 @@ TEST_F(QueryProcessorTest, SchemaTypeFilter) {
 
   ICING_ASSERT_OK_AND_ASSIGN(
       schema_store_,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
   ASSERT_THAT(schema_store_->SetSchema(schema), IsOk());
 
   ICING_ASSERT_OK_AND_ASSIGN(
@@ -2155,7 +2158,7 @@ TEST_F(QueryProcessorTest, SectionFilterForOneDocument) {
 
   ICING_ASSERT_OK_AND_ASSIGN(
       schema_store_,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
   ASSERT_THAT(schema_store_->SetSchema(schema), IsOk());
 
   ICING_ASSERT_OK_AND_ASSIGN(
@@ -2237,7 +2240,7 @@ TEST_F(QueryProcessorTest, SectionFilterAcrossSchemaTypes) {
 
   ICING_ASSERT_OK_AND_ASSIGN(
       schema_store_,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
   ASSERT_THAT(schema_store_->SetSchema(schema), IsOk());
 
   ICING_ASSERT_OK_AND_ASSIGN(
@@ -2320,7 +2323,7 @@ TEST_F(QueryProcessorTest, SectionFilterWithinSchemaType) {
 
   ICING_ASSERT_OK_AND_ASSIGN(
       schema_store_,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
   ASSERT_THAT(schema_store_->SetSchema(schema), IsOk());
 
   ICING_ASSERT_OK_AND_ASSIGN(
@@ -2404,7 +2407,7 @@ TEST_F(QueryProcessorTest, SectionFilterRespectsDifferentSectionIds) {
 
   ICING_ASSERT_OK_AND_ASSIGN(
       schema_store_,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
   ASSERT_THAT(schema_store_->SetSchema(schema), IsOk());
 
   ICING_ASSERT_OK_AND_ASSIGN(
@@ -2477,7 +2480,7 @@ TEST_F(QueryProcessorTest, NonexistentSectionFilterReturnsEmptyResults) {
 
   ICING_ASSERT_OK_AND_ASSIGN(
       schema_store_,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
   ASSERT_THAT(schema_store_->SetSchema(schema), IsOk());
 
   ICING_ASSERT_OK_AND_ASSIGN(
@@ -2544,7 +2547,7 @@ TEST_F(QueryProcessorTest, UnindexedSectionFilterReturnsEmptyResults) {
 
   ICING_ASSERT_OK_AND_ASSIGN(
       schema_store_,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
   ASSERT_THAT(schema_store_->SetSchema(schema), IsOk());
 
   ICING_ASSERT_OK_AND_ASSIGN(
@@ -2614,7 +2617,7 @@ TEST_F(QueryProcessorTest, SectionFilterTermAndUnrestrictedTerm) {
 
   ICING_ASSERT_OK_AND_ASSIGN(
       schema_store_,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
   ASSERT_THAT(schema_store_->SetSchema(schema), IsOk());
 
   ICING_ASSERT_OK_AND_ASSIGN(
@@ -2689,7 +2692,7 @@ TEST_F(QueryProcessorTest, DocumentBeforeTtlNotFilteredOut) {
 
   ICING_ASSERT_OK_AND_ASSIGN(
       schema_store_,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
   ASSERT_THAT(schema_store_->SetSchema(schema), IsOk());
 
   // Arbitrary value, just has to be less than the document's creation
@@ -2748,7 +2751,7 @@ TEST_F(QueryProcessorTest, DocumentPastTtlFilteredOut) {
 
   ICING_ASSERT_OK_AND_ASSIGN(
       schema_store_,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
   ASSERT_THAT(schema_store_->SetSchema(schema), IsOk());
 
   // Arbitrary value, just has to be greater than the document's creation
diff --git a/icing/schema/schema-store.cc b/icing/schema/schema-store.cc
index acc5030..fc50ea6 100644
--- a/icing/schema/schema-store.cc
+++ b/icing/schema/schema-store.cc
@@ -108,27 +108,60 @@ libtextclassifier3::StatusOr<std::unique_ptr<SchemaStore>> SchemaStore::Create(
   ICING_RETURN_ERROR_IF_NULL(filesystem);
   ICING_RETURN_ERROR_IF_NULL(clock);
 
+  if (!filesystem->DirectoryExists(base_dir.c_str())) {
+    return absl_ports::FailedPreconditionError(
+        "Schema store base directory does not exist!");
+  }
   std::unique_ptr<SchemaStore> schema_store = std::unique_ptr<SchemaStore>(
       new SchemaStore(filesystem, base_dir, clock));
   ICING_RETURN_IF_ERROR(schema_store->Initialize(initialize_stats));
   return schema_store;
 }
 
+libtextclassifier3::StatusOr<std::unique_ptr<SchemaStore>> SchemaStore::Create(
+    const Filesystem* filesystem, const std::string& base_dir,
+    const Clock* clock, SchemaProto schema) {
+  ICING_RETURN_ERROR_IF_NULL(filesystem);
+  ICING_RETURN_ERROR_IF_NULL(clock);
+
+  if (!filesystem->DirectoryExists(base_dir.c_str())) {
+    return absl_ports::FailedPreconditionError(
+        "Schema store base directory does not exist!");
+  }
+  std::unique_ptr<SchemaStore> schema_store = std::unique_ptr<SchemaStore>(
+      new SchemaStore(filesystem, base_dir, clock));
+  ICING_RETURN_IF_ERROR(schema_store->Initialize(std::move(schema)));
+  return schema_store;
+}
+
 SchemaStore::SchemaStore(const Filesystem* filesystem, std::string base_dir,
                          const Clock* clock)
-    : filesystem_(*filesystem),
+    : filesystem_(filesystem),
       base_dir_(std::move(base_dir)),
-      clock_(*clock),
-      schema_file_(*filesystem, MakeSchemaFilename(base_dir_)) {}
+      clock_(clock),
+      schema_file_(std::make_unique<FileBackedProto<SchemaProto>>(
+          *filesystem, MakeSchemaFilename(base_dir_))) {}
 
 SchemaStore::~SchemaStore() {
-  if (has_schema_successfully_set_) {
+  if (has_schema_successfully_set_ && schema_file_ != nullptr &&
+      schema_type_mapper_ != nullptr && section_manager_ != nullptr) {
     if (!PersistToDisk().ok()) {
       ICING_LOG(ERROR) << "Error persisting to disk in SchemaStore destructor";
     }
   }
 }
 
+libtextclassifier3::Status SchemaStore::Initialize(SchemaProto new_schema) {
+  if (!absl_ports::IsNotFound(GetSchema().status())) {
+    return absl_ports::FailedPreconditionError(
+        "Incorrectly tried to initialize schema store with a new schema, when "
+        "one is already set!");
+  }
+  ICING_RETURN_IF_ERROR(schema_file_->Write(
+      std::make_unique<SchemaProto>(std::move(new_schema))));
+  return InitializeInternal(/*initialize_stats=*/nullptr);
+}
+
 libtextclassifier3::Status SchemaStore::Initialize(
     InitializeStatsProto* initialize_stats) {
   auto schema_proto_or = GetSchema();
@@ -139,13 +172,16 @@ libtextclassifier3::Status SchemaStore::Initialize(
     // Real error when trying to read the existing schema
     return schema_proto_or.status();
   }
-  has_schema_successfully_set_ = true;
+  return InitializeInternal(initialize_stats);
+}
 
+libtextclassifier3::Status SchemaStore::InitializeInternal(
+    InitializeStatsProto* initialize_stats) {
   if (!InitializeDerivedFiles().ok()) {
     ICING_VLOG(3)
         << "Couldn't find derived files or failed to initialize them, "
            "regenerating derived files for SchemaStore.";
-    std::unique_ptr<Timer> regenerate_timer = clock_.GetNewTimer();
+    std::unique_ptr<Timer> regenerate_timer = clock_->GetNewTimer();
     if (initialize_stats != nullptr) {
       initialize_stats->set_schema_store_recovery_cause(
           InitializeStatsProto::IO_ERROR);
@@ -161,6 +197,7 @@ libtextclassifier3::Status SchemaStore::Initialize(
     initialize_stats->set_num_schema_types(type_config_map_.size());
   }
 
+  has_schema_successfully_set_ = true;
   return libtextclassifier3::Status::OK;
 }
 
@@ -172,8 +209,8 @@ libtextclassifier3::Status SchemaStore::InitializeDerivedFiles() {
   }
 
   SchemaStore::Header header;
-  if (!filesystem_.Read(MakeHeaderFilename(base_dir_).c_str(), &header,
-                        sizeof(header))) {
+  if (!filesystem_->Read(MakeHeaderFilename(base_dir_).c_str(), &header,
+                         sizeof(header))) {
     return absl_ports::InternalError(
         absl_ports::StrCat("Couldn't read: ", MakeHeaderFilename(base_dir_)));
   }
@@ -185,7 +222,7 @@ libtextclassifier3::Status SchemaStore::InitializeDerivedFiles() {
 
   ICING_ASSIGN_OR_RETURN(
       schema_type_mapper_,
-      KeyMapper<SchemaTypeId>::Create(filesystem_,
+      KeyMapper<SchemaTypeId>::Create(*filesystem_,
                                       MakeSchemaTypeMapperFilename(base_dir_),
                                       kSchemaTypeMapperMaxSize));
 
@@ -236,12 +273,12 @@ libtextclassifier3::Status SchemaStore::RegenerateDerivedFiles() {
 }
 
 bool SchemaStore::HeaderExists() {
-  if (!filesystem_.FileExists(MakeHeaderFilename(base_dir_).c_str())) {
+  if (!filesystem_->FileExists(MakeHeaderFilename(base_dir_).c_str())) {
     return false;
   }
 
   int64_t file_size =
-      filesystem_.GetFileSize(MakeHeaderFilename(base_dir_).c_str());
+      filesystem_->GetFileSize(MakeHeaderFilename(base_dir_).c_str());
 
   // If it's been truncated to size 0 before, we consider it to be a new file
   return file_size != 0 && file_size != Filesystem::kBadFileSize;
@@ -254,11 +291,11 @@ libtextclassifier3::Status SchemaStore::UpdateHeader(const Crc32& checksum) {
   header.checksum = checksum.Get();
 
   ScopedFd scoped_fd(
-      filesystem_.OpenForWrite(MakeHeaderFilename(base_dir_).c_str()));
+      filesystem_->OpenForWrite(MakeHeaderFilename(base_dir_).c_str()));
   // This should overwrite the header.
   if (!scoped_fd.is_valid() ||
-      !filesystem_.Write(scoped_fd.get(), &header, sizeof(header)) ||
-      !filesystem_.DataSync(scoped_fd.get())) {
+      !filesystem_->Write(scoped_fd.get(), &header, sizeof(header)) ||
+      !filesystem_->DataSync(scoped_fd.get())) {
     return absl_ports::InternalError(absl_ports::StrCat(
         "Failed to write SchemaStore header: ", MakeHeaderFilename(base_dir_)));
   }
@@ -271,7 +308,7 @@ libtextclassifier3::Status SchemaStore::ResetSchemaTypeMapper() {
   // TODO(b/216487496): Implement a more robust version of TC_RETURN_IF_ERROR
   // that can support error logging.
   libtextclassifier3::Status status = KeyMapper<SchemaTypeId>::Delete(
-      filesystem_, MakeSchemaTypeMapperFilename(base_dir_));
+      *filesystem_, MakeSchemaTypeMapperFilename(base_dir_));
   if (!status.ok()) {
     ICING_LOG(ERROR) << status.error_message()
                      << "Failed to delete old schema_type mapper";
@@ -279,7 +316,7 @@ libtextclassifier3::Status SchemaStore::ResetSchemaTypeMapper() {
   }
   ICING_ASSIGN_OR_RETURN(
       schema_type_mapper_,
-      KeyMapper<SchemaTypeId>::Create(filesystem_,
+      KeyMapper<SchemaTypeId>::Create(*filesystem_,
                                       MakeSchemaTypeMapperFilename(base_dir_),
                                       kSchemaTypeMapperMaxSize));
 
@@ -287,17 +324,17 @@ libtextclassifier3::Status SchemaStore::ResetSchemaTypeMapper() {
 }
 
 libtextclassifier3::StatusOr<Crc32> SchemaStore::ComputeChecksum() const {
-  Crc32 total_checksum;
-  if (!has_schema_successfully_set_) {
-    // Nothing to checksum
-    return total_checksum;
+  auto schema_proto_or = GetSchema();
+  if (absl_ports::IsNotFound(schema_proto_or.status())) {
+    return Crc32();
   }
-  ICING_ASSIGN_OR_RETURN(const SchemaProto* schema_proto, GetSchema());
+  ICING_ASSIGN_OR_RETURN(const SchemaProto* schema_proto, schema_proto_or);
   Crc32 schema_checksum;
   schema_checksum.Append(schema_proto->SerializeAsString());
 
   Crc32 schema_type_mapper_checksum = schema_type_mapper_->ComputeChecksum();
 
+  Crc32 total_checksum;
   total_checksum.Append(std::to_string(schema_checksum.Get()));
   total_checksum.Append(std::to_string(schema_type_mapper_checksum.Get()));
 
@@ -306,7 +343,7 @@ libtextclassifier3::StatusOr<Crc32> SchemaStore::ComputeChecksum() const {
 
 libtextclassifier3::StatusOr<const SchemaProto*> SchemaStore::GetSchema()
     const {
-  return schema_file_.Read();
+  return schema_file_->Read();
 }
 
 // TODO(cassiewang): Consider removing this definition of SetSchema if it's not
@@ -393,17 +430,80 @@ SchemaStore::SetSchema(SchemaProto&& new_schema,
   result.success = result.success || ignore_errors_and_delete_documents;
 
   if (result.success) {
-    // Write the schema (and potentially overwrite a previous schema)
-    ICING_RETURN_IF_ERROR(
-        schema_file_.Write(std::make_unique<SchemaProto>(new_schema)));
+    ICING_RETURN_IF_ERROR(ApplySchemaChange(std::move(new_schema)));
     has_schema_successfully_set_ = true;
-
-    ICING_RETURN_IF_ERROR(RegenerateDerivedFiles());
   }
 
   return result;
 }
 
+libtextclassifier3::Status SchemaStore::ApplySchemaChange(
+    SchemaProto new_schema) {
+  // We need to ensure that we either 1) successfully set the schema and
+  // update all derived data structures or 2) fail and leave the schema store
+  // unchanged.
+  // So, first, we create an empty temporary directory to build a new schema
+  // store in.
+  std::string temp_schema_store_dir_path = base_dir_ + "_temp";
+  if (!filesystem_->DeleteDirectoryRecursively(
+          temp_schema_store_dir_path.c_str())) {
+    ICING_LOG(WARNING) << "Failed to recursively delete "
+                     << temp_schema_store_dir_path.c_str();
+    return absl_ports::InternalError(
+        "Unable to delete temp directory to prepare to build new schema "
+        "store.");
+  }
+
+  if (!filesystem_->CreateDirectoryRecursively(
+      temp_schema_store_dir_path.c_str())) {
+    return absl_ports::InternalError(
+        "Unable to create temp directory to build new schema store.");
+  }
+
+  // Then we create our new schema store with the new schema.
+  auto new_schema_store_or =
+      SchemaStore::Create(filesystem_, temp_schema_store_dir_path, clock_,
+                          std::move(new_schema));
+  if (!new_schema_store_or.ok()) {
+    // Attempt to clean up the temp directory.
+    if (!filesystem_->DeleteDirectoryRecursively(
+            temp_schema_store_dir_path.c_str())) {
+      // Nothing to do here. Just log an error.
+      ICING_LOG(WARNING) << "Failed to recursively delete "
+                       << temp_schema_store_dir_path.c_str();
+    }
+    return new_schema_store_or.status();
+  }
+  std::unique_ptr<SchemaStore> new_schema_store =
+      std::move(new_schema_store_or).ValueOrDie();
+
+  // Then we swap the new schema file + new derived files with the old files.
+  if (!filesystem_->SwapFiles(base_dir_.c_str(),
+                              temp_schema_store_dir_path.c_str())) {
+    // Attempt to clean up the temp directory.
+    if (!filesystem_->DeleteDirectoryRecursively(
+            temp_schema_store_dir_path.c_str())) {
+      // Nothing to do here. Just log an error.
+      ICING_LOG(WARNING) << "Failed to recursively delete "
+                       << temp_schema_store_dir_path.c_str();
+    }
+    return absl_ports::InternalError(
+        "Unable to apply new schema due to failed swap!");
+  }
+
+  std::string old_base_dir = std::move(base_dir_);
+  *this = std::move(*new_schema_store);
+
+  // After the std::move, the filepaths saved in this instance and in the
+  // schema_file_ instance will still be the one from temp_schema_store_dir
+  // even though they now point to files that are within old_base_dir.
+  // Manually set them to the correct paths.
+  base_dir_ = std::move(old_base_dir);
+  schema_file_->SetSwappedFilepath(MakeSchemaFilename(base_dir_));
+
+  return libtextclassifier3::Status::OK;
+}
+
 libtextclassifier3::StatusOr<const SchemaTypeConfigProto*>
 SchemaStore::GetSchemaTypeConfig(std::string_view schema_type) const {
   ICING_RETURN_IF_ERROR(CheckSchemaSet());
@@ -463,7 +563,7 @@ libtextclassifier3::Status SchemaStore::PersistToDisk() {
 
 SchemaStoreStorageInfoProto SchemaStore::GetStorageInfo() const {
   SchemaStoreStorageInfoProto storage_info;
-  int64_t directory_size = filesystem_.GetDiskUsage(base_dir_.c_str());
+  int64_t directory_size = filesystem_->GetDiskUsage(base_dir_.c_str());
   storage_info.set_schema_store_size(
       Filesystem::SanitizeFileSize(directory_size));
   ICING_ASSIGN_OR_RETURN(const SchemaProto* schema, GetSchema(), storage_info);
diff --git a/icing/schema/schema-store.h b/icing/schema/schema-store.h
index 2d3aca7..58e5477 100644
--- a/icing/schema/schema-store.h
+++ b/icing/schema/schema-store.h
@@ -130,8 +130,10 @@ class SchemaStore {
   static libtextclassifier3::StatusOr<std::unique_ptr<SchemaStore>> Create(
       const Filesystem* filesystem, const std::string& base_dir,
       const Clock* clock, InitializeStatsProto* initialize_stats = nullptr);
+  
+  SchemaStore(SchemaStore&&) = default;
+  SchemaStore& operator=(SchemaStore&&) = default;
 
-  // Not copyable
   SchemaStore(const SchemaStore&) = delete;
   SchemaStore& operator=(const SchemaStore&) = delete;
 
@@ -265,17 +267,51 @@ class SchemaStore {
   libtextclassifier3::StatusOr<SchemaDebugInfoProto> GetDebugInfo() const;
 
  private:
+  // Factory function to create a SchemaStore and set its schema. The created
+  // instance does not take ownership of any input components and all pointers
+  // must refer to valid objects that outlive the created SchemaStore instance.
+  // The base_dir must already exist. No schema must have set in base_dir prior
+  // to this.
+  //
+  // Returns:
+  //   A SchemaStore on success
+  //   FAILED_PRECONDITION on any null pointer input or if there has already
+  //       been a schema set for this path.
+  //   INTERNAL_ERROR on any IO errors
+  static libtextclassifier3::StatusOr<std::unique_ptr<SchemaStore>> Create(
+      const Filesystem* filesystem, const std::string& base_dir,
+      const Clock* clock, SchemaProto schema);
+
+
   // Use SchemaStore::Create instead.
   explicit SchemaStore(const Filesystem* filesystem, std::string base_dir,
                        const Clock* clock);
 
-  // Handles initializing the SchemaStore and regenerating any data if needed.
+  // Verifies that there is no error retrieving a previously set schema. Then
+  // initializes like normal.
   //
   // Returns:
   //   OK on success
   //   INTERNAL_ERROR on IO error
   libtextclassifier3::Status Initialize(InitializeStatsProto* initialize_stats);
 
+  // First, blindly writes new_schema to the schema_file. Then initializes like
+  // normal.
+  //
+  // Returns:
+  //   OK on success
+  //   INTERNAL_ERROR on IO error
+  //   FAILED_PRECONDITION if there is already a schema set for the schema_file.
+  libtextclassifier3::Status Initialize(SchemaProto new_schema);
+
+  // Handles initializing the SchemaStore and regenerating any data if needed.
+  //
+  // Returns:
+  //   OK on success
+  //   INTERNAL_ERROR on IO error
+  libtextclassifier3::Status InitializeInternal(
+      InitializeStatsProto* initialize_stats);
+
   // Creates sub-components and verifies the integrity of each sub-component.
   //
   // Returns:
@@ -310,15 +346,25 @@ class SchemaStore {
   // Returns any IO errors.
   libtextclassifier3::Status ResetSchemaTypeMapper();
 
+  // Creates a new schema store with new_schema and then swaps that new schema
+  // store with the existing one. This function guarantees that either: this
+  // instance will be fully updated to the new schema or no changes will take
+  // effect.
+  //
+  // Returns:
+  //   OK on success
+  //   INTERNAL on I/O error.
+  libtextclassifier3::Status ApplySchemaChange(SchemaProto new_schema);
+
   libtextclassifier3::Status CheckSchemaSet() const {
     return has_schema_successfully_set_
                ? libtextclassifier3::Status::OK
                : absl_ports::FailedPreconditionError("Schema not set yet.");
   }
 
-  const Filesystem& filesystem_;
-  const std::string base_dir_;
-  const Clock& clock_;
+  const Filesystem* filesystem_;
+  std::string base_dir_;
+  const Clock* clock_;
 
   // Used internally to indicate whether the class has been successfully
   // initialized with a valid schema. Will be false if Initialize failed or no
@@ -326,7 +372,7 @@ class SchemaStore {
   bool has_schema_successfully_set_ = false;
 
   // Cached schema
-  FileBackedProto<SchemaProto> schema_file_;
+  std::unique_ptr<FileBackedProto<SchemaProto>> schema_file_;
 
   // A hash map of (type config name -> type config), allows faster lookup of
   // type config in schema. The O(1) type config access makes schema-related and
diff --git a/icing/schema/schema-store_test.cc b/icing/schema/schema-store_test.cc
index 541918f..3fd41c4 100644
--- a/icing/schema/schema-store_test.cc
+++ b/icing/schema/schema-store_test.cc
@@ -21,7 +21,9 @@
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
 #include "icing/absl_ports/str_cat.h"
+#include "icing/document-builder.h"
 #include "icing/file/filesystem.h"
+#include "icing/file/mock-filesystem.h"
 #include "icing/portable/equals-proto.h"
 #include "icing/proto/document.pb.h"
 #include "icing/proto/schema.pb.h"
@@ -33,6 +35,7 @@
 #include "icing/testing/common-matchers.h"
 #include "icing/testing/fake-clock.h"
 #include "icing/testing/tmp-directory.h"
+#include "icing/text_classifier/lib3/utils/base/status.h"
 #include "icing/util/crc32.h"
 
 namespace icing {
@@ -45,8 +48,11 @@ using ::testing::ElementsAre;
 using ::testing::Eq;
 using ::testing::Ge;
 using ::testing::Gt;
+using ::testing::HasSubstr;
 using ::testing::Not;
 using ::testing::Pointee;
+using ::testing::Return;
+using ::testing::SizeIs;
 
 constexpr PropertyConfigProto::Cardinality::Code CARDINALITY_OPTIONAL =
     PropertyConfigProto::Cardinality::OPTIONAL;
@@ -66,8 +72,10 @@ constexpr PropertyConfigProto::DataType::Code TYPE_DOUBLE =
 
 class SchemaStoreTest : public ::testing::Test {
  protected:
-  SchemaStoreTest() : test_dir_(GetTestTempDir() + "/icing") {
-    filesystem_.CreateDirectoryRecursively(test_dir_.c_str());
+  void SetUp() override {
+    temp_dir_ = GetTestTempDir() + "/icing";
+    schema_store_dir_ = temp_dir_ + "/schema_store";
+    filesystem_.CreateDirectoryRecursively(schema_store_dir_.c_str());
 
     schema_ =
         SchemaBuilder()
@@ -81,26 +89,112 @@ class SchemaStoreTest : public ::testing::Test {
   }
 
   void TearDown() override {
-    filesystem_.DeleteDirectoryRecursively(test_dir_.c_str());
+    // Check that the schema store directory is the *only* directory in the
+    // schema_store_dir_. IOW, ensure that all temporary directories have been
+    // properly cleaned up.
+    std::vector<std::string> sub_dirs;
+    ASSERT_TRUE(filesystem_.ListDirectory(temp_dir_.c_str(), &sub_dirs));
+    ASSERT_THAT(sub_dirs, ElementsAre("schema_store"));
+
+    // Finally, clean everything up.
+    ASSERT_TRUE(filesystem_.DeleteDirectoryRecursively(temp_dir_.c_str()));
   }
 
-  const Filesystem filesystem_;
-  const std::string test_dir_;
+  Filesystem filesystem_;
+  std::string temp_dir_;
+  std::string schema_store_dir_;
   SchemaProto schema_;
-  const FakeClock fake_clock_;
+  FakeClock fake_clock_;
 };
 
 TEST_F(SchemaStoreTest, CreationWithNullPointerShouldFail) {
   EXPECT_THAT(
-      SchemaStore::Create(/*filesystem=*/nullptr, test_dir_, &fake_clock_),
+      SchemaStore::Create(/*filesystem=*/nullptr, schema_store_dir_, &fake_clock_),
       StatusIs(libtextclassifier3::StatusCode::FAILED_PRECONDITION));
 }
 
+TEST_F(SchemaStoreTest, SchemaStoreMoveConstructible) {
+  // Create an instance of SchemaStore.
+  SchemaProto schema =
+      SchemaBuilder()
+          .AddType(SchemaTypeConfigBuilder().SetType("TypeA").AddProperty(
+              PropertyConfigBuilder()
+                  .SetName("prop1")
+                  .SetDataTypeString(MATCH_EXACT, TOKENIZER_PLAIN)
+                  .SetCardinality(CARDINALITY_OPTIONAL)))
+          .Build();
+
+  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));
+  ICING_ASSERT_OK_AND_ASSIGN(Crc32 expected_checksum,
+                             schema_store->ComputeChecksum());
+
+  // Move construct an instance of SchemaStore
+  SchemaStore move_constructed_schema_store(std::move(*schema_store));
+  EXPECT_THAT(move_constructed_schema_store.GetSchema(),
+              IsOkAndHolds(Pointee(EqualsProto(schema))));
+  EXPECT_THAT(move_constructed_schema_store.ComputeChecksum(),
+              IsOkAndHolds(Eq(expected_checksum)));
+  SectionMetadata expected_metadata(/*id_in=*/0, MATCH_EXACT, TOKENIZER_PLAIN,
+                                    "prop1");
+  EXPECT_THAT(move_constructed_schema_store.GetSectionMetadata("TypeA"),
+              IsOkAndHolds(Pointee(ElementsAre(expected_metadata))));
+}
+
+TEST_F(SchemaStoreTest, SchemaStoreMoveAssignment) {
+  // Create an instance of SchemaStore.
+  SchemaProto schema1 =
+      SchemaBuilder()
+          .AddType(SchemaTypeConfigBuilder().SetType("TypeA").AddProperty(
+              PropertyConfigBuilder()
+                  .SetName("prop1")
+                  .SetDataTypeString(MATCH_EXACT, TOKENIZER_PLAIN)
+                  .SetCardinality(CARDINALITY_OPTIONAL)))
+          .Build();
+
+  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(schema1));
+  ICING_ASSERT_OK_AND_ASSIGN(Crc32 expected_checksum,
+                             schema_store->ComputeChecksum());
+
+  // Construct another instance of SchemaStore
+  SchemaProto schema2 =
+      SchemaBuilder()
+          .AddType(SchemaTypeConfigBuilder().SetType("TypeB").AddProperty(
+              PropertyConfigBuilder()
+                  .SetName("prop2")
+                  .SetDataTypeString(MATCH_EXACT, TOKENIZER_PLAIN)
+                  .SetCardinality(CARDINALITY_OPTIONAL)))
+          .Build();
+
+  ICING_ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<SchemaStore> move_assigned_schema_store,
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
+  ICING_ASSERT_OK(schema_store->SetSchema(schema2));
+
+  // Move assign the first instance into the second one.
+  *move_assigned_schema_store = std::move(*schema_store);
+  EXPECT_THAT(move_assigned_schema_store->GetSchema(),
+              IsOkAndHolds(Pointee(EqualsProto(schema1))));
+  EXPECT_THAT(move_assigned_schema_store->ComputeChecksum(),
+              IsOkAndHolds(Eq(expected_checksum)));
+  SectionMetadata expected_metadata(/*id_in=*/0, MATCH_EXACT, TOKENIZER_PLAIN,
+                                    "prop1");
+  EXPECT_THAT(move_assigned_schema_store->GetSectionMetadata("TypeA"),
+              IsOkAndHolds(Pointee(ElementsAre(expected_metadata))));
+}
+
 TEST_F(SchemaStoreTest, CorruptSchemaError) {
   {
     ICING_ASSERT_OK_AND_ASSIGN(
         std::unique_ptr<SchemaStore> schema_store,
-        SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+        SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
 
     // Set it for the first time
     SchemaStore::SetSchemaResult result;
@@ -121,14 +215,14 @@ TEST_F(SchemaStoreTest, CorruptSchemaError) {
           .AddType(SchemaTypeConfigBuilder().SetType("corrupted"))
           .Build();
 
-  const std::string schema_file = absl_ports::StrCat(test_dir_, "/schema.pb");
+  const std::string schema_file = absl_ports::StrCat(schema_store_dir_, "/schema.pb");
   const std::string serialized_schema = corrupt_schema.SerializeAsString();
 
   filesystem_.Write(schema_file.c_str(), serialized_schema.data(),
                     serialized_schema.size());
 
   // If ground truth was corrupted, we won't know what to do
-  EXPECT_THAT(SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_),
+  EXPECT_THAT(SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_),
               StatusIs(libtextclassifier3::StatusCode::INTERNAL));
 }
 
@@ -136,7 +230,7 @@ TEST_F(SchemaStoreTest, RecoverCorruptDerivedFileOk) {
   {
     ICING_ASSERT_OK_AND_ASSIGN(
         std::unique_ptr<SchemaStore> schema_store,
-        SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+        SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
 
     // Set it for the first time
     SchemaStore::SetSchemaResult result;
@@ -156,12 +250,12 @@ TEST_F(SchemaStoreTest, RecoverCorruptDerivedFileOk) {
   // regenerated from ground truth
 
   const std::string schema_type_mapper_dir =
-      absl_ports::StrCat(test_dir_, "/schema_type_mapper");
+      absl_ports::StrCat(schema_store_dir_, "/schema_type_mapper");
   filesystem_.DeleteDirectoryRecursively(schema_type_mapper_dir.c_str());
 
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<SchemaStore> schema_store,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
 
   // Everything looks fine, ground truth and derived data
   ICING_ASSERT_OK_AND_ASSIGN(const SchemaProto* actual_schema,
@@ -174,7 +268,7 @@ TEST_F(SchemaStoreTest, RecoverBadChecksumOk) {
   {
     ICING_ASSERT_OK_AND_ASSIGN(
         std::unique_ptr<SchemaStore> schema_store,
-        SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+        SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
 
     // Set it for the first time
     SchemaStore::SetSchemaResult result;
@@ -193,7 +287,7 @@ TEST_F(SchemaStoreTest, RecoverBadChecksumOk) {
   // the recalculated checksum on initialization. This will force a regeneration
   // of derived files from ground truth.
   const std::string header_file =
-      absl_ports::StrCat(test_dir_, "/schema_store_header");
+      absl_ports::StrCat(schema_store_dir_, "/schema_store_header");
   SchemaStore::Header header;
   header.magic = SchemaStore::Header::kMagic;
   header.checksum = 10;  // Arbitrary garbage checksum
@@ -202,7 +296,7 @@ TEST_F(SchemaStoreTest, RecoverBadChecksumOk) {
 
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<SchemaStore> schema_store,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
 
   // Everything looks fine, ground truth and derived data
   ICING_ASSERT_OK_AND_ASSIGN(const SchemaProto* actual_schema,
@@ -214,7 +308,7 @@ TEST_F(SchemaStoreTest, RecoverBadChecksumOk) {
 TEST_F(SchemaStoreTest, CreateNoPreviousSchemaOk) {
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<SchemaStore> store,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
 
   // The apis to retrieve information about the schema should fail gracefully.
   EXPECT_THAT(store->GetSchema(),
@@ -247,7 +341,7 @@ TEST_F(SchemaStoreTest, CreateNoPreviousSchemaOk) {
 TEST_F(SchemaStoreTest, CreateWithPreviousSchemaOk) {
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<SchemaStore> schema_store,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
 
   SchemaStore::SetSchemaResult result;
   result.success = true;
@@ -256,7 +350,7 @@ TEST_F(SchemaStoreTest, CreateWithPreviousSchemaOk) {
               IsOkAndHolds(EqualsSetSchemaResult(result)));
 
   schema_store.reset();
-  EXPECT_THAT(SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_),
+  EXPECT_THAT(SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_),
               IsOk());
 }
 
@@ -269,7 +363,7 @@ TEST_F(SchemaStoreTest, MultipleCreateOk) {
 
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<SchemaStore> schema_store,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
 
   SchemaStore::SetSchemaResult result;
   result.success = true;
@@ -289,7 +383,7 @@ TEST_F(SchemaStoreTest, MultipleCreateOk) {
 
   schema_store.reset();
   ICING_ASSERT_OK_AND_ASSIGN(
-      schema_store, SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      schema_store, SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
 
   // Verify that our in-memory structures are ok
   EXPECT_THAT(schema_store->GetSchemaTypeConfig("email"),
@@ -305,7 +399,7 @@ TEST_F(SchemaStoreTest, MultipleCreateOk) {
 TEST_F(SchemaStoreTest, SetNewSchemaOk) {
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<SchemaStore> schema_store,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
 
   // Set it for the first time
   SchemaStore::SetSchemaResult result;
@@ -321,7 +415,7 @@ TEST_F(SchemaStoreTest, SetNewSchemaOk) {
 TEST_F(SchemaStoreTest, SetSameSchemaOk) {
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<SchemaStore> schema_store,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
 
   // Set it for the first time
   SchemaStore::SetSchemaResult result;
@@ -345,7 +439,7 @@ TEST_F(SchemaStoreTest, SetSameSchemaOk) {
 TEST_F(SchemaStoreTest, SetIncompatibleSchemaOk) {
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<SchemaStore> schema_store,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
 
   // Set it for the first time
   SchemaStore::SetSchemaResult result;
@@ -372,7 +466,7 @@ TEST_F(SchemaStoreTest, SetIncompatibleSchemaOk) {
 TEST_F(SchemaStoreTest, SetSchemaWithAddedTypeOk) {
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<SchemaStore> schema_store,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
 
   SchemaProto schema = SchemaBuilder()
                            .AddType(SchemaTypeConfigBuilder().SetType("email"))
@@ -406,7 +500,7 @@ TEST_F(SchemaStoreTest, SetSchemaWithAddedTypeOk) {
 TEST_F(SchemaStoreTest, SetSchemaWithDeletedTypeOk) {
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<SchemaStore> schema_store,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
 
   SchemaProto schema =
       SchemaBuilder()
@@ -464,7 +558,7 @@ TEST_F(SchemaStoreTest, SetSchemaWithDeletedTypeOk) {
 TEST_F(SchemaStoreTest, SetSchemaWithReorderedTypesOk) {
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<SchemaStore> schema_store,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
 
   SchemaProto schema =
       SchemaBuilder()
@@ -507,7 +601,7 @@ TEST_F(SchemaStoreTest, SetSchemaWithReorderedTypesOk) {
 TEST_F(SchemaStoreTest, IndexedPropertyChangeRequiresReindexingOk) {
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<SchemaStore> schema_store,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
 
   SchemaProto schema =
       SchemaBuilder()
@@ -551,7 +645,7 @@ TEST_F(SchemaStoreTest, IndexedPropertyChangeRequiresReindexingOk) {
 TEST_F(SchemaStoreTest, IndexNestedDocumentsChangeRequiresReindexingOk) {
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<SchemaStore> schema_store,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
 
   // Make two schemas. One that sets index_nested_properties to false and one
   // that sets it to true.
@@ -620,7 +714,7 @@ TEST_F(SchemaStoreTest, IndexNestedDocumentsChangeRequiresReindexingOk) {
 TEST_F(SchemaStoreTest, SetSchemaWithIncompatibleTypesOk) {
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<SchemaStore> schema_store,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
 
   SchemaProto schema =
       SchemaBuilder()
@@ -682,7 +776,7 @@ TEST_F(SchemaStoreTest, SetSchemaWithIncompatibleTypesOk) {
 TEST_F(SchemaStoreTest, SetSchemaWithIncompatibleNestedTypesOk) {
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<SchemaStore> schema_store,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
 
   // 1. Create a ContactPoint type with a repeated property and set that schema
   SchemaTypeConfigBuilder contact_point_repeated_label =
@@ -749,7 +843,7 @@ TEST_F(SchemaStoreTest, SetSchemaWithIncompatibleNestedTypesOk) {
 TEST_F(SchemaStoreTest, SetSchemaWithIndexIncompatibleNestedTypesOk) {
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<SchemaStore> schema_store,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
 
   // 1. Create a ContactPoint type with label that matches prefix and set that
   // schema
@@ -804,7 +898,7 @@ TEST_F(SchemaStoreTest, SetSchemaWithIndexIncompatibleNestedTypesOk) {
 TEST_F(SchemaStoreTest, SetSchemaWithCompatibleNestedTypesOk) {
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<SchemaStore> schema_store,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
 
   // 1. Create a ContactPoint type with a optional property and set that schema
   SchemaTypeConfigBuilder contact_point_optional_label =
@@ -857,7 +951,7 @@ TEST_F(SchemaStoreTest, SetSchemaWithCompatibleNestedTypesOk) {
 TEST_F(SchemaStoreTest, GetSchemaTypeId) {
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<SchemaStore> schema_store,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
 
   schema_.clear_types();
 
@@ -885,7 +979,7 @@ TEST_F(SchemaStoreTest, GetSchemaTypeId) {
 TEST_F(SchemaStoreTest, ComputeChecksumDefaultOnEmptySchemaStore) {
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<SchemaStore> schema_store,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
 
   Crc32 default_checksum;
   EXPECT_THAT(schema_store->ComputeChecksum(), IsOkAndHolds(default_checksum));
@@ -894,7 +988,7 @@ TEST_F(SchemaStoreTest, ComputeChecksumDefaultOnEmptySchemaStore) {
 TEST_F(SchemaStoreTest, ComputeChecksumSameBetweenCalls) {
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<SchemaStore> schema_store,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
 
   SchemaProto foo_schema =
       SchemaBuilder().AddType(SchemaTypeConfigBuilder().SetType("foo")).Build();
@@ -910,7 +1004,7 @@ TEST_F(SchemaStoreTest, ComputeChecksumSameBetweenCalls) {
 TEST_F(SchemaStoreTest, ComputeChecksumSameAcrossInstances) {
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<SchemaStore> schema_store,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
 
   SchemaProto foo_schema =
       SchemaBuilder().AddType(SchemaTypeConfigBuilder().SetType("foo")).Build();
@@ -923,14 +1017,14 @@ TEST_F(SchemaStoreTest, ComputeChecksumSameAcrossInstances) {
   schema_store.reset();
 
   ICING_ASSERT_OK_AND_ASSIGN(
-      schema_store, SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      schema_store, SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
   EXPECT_THAT(schema_store->ComputeChecksum(), IsOkAndHolds(checksum));
 }
 
 TEST_F(SchemaStoreTest, ComputeChecksumChangesOnModification) {
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<SchemaStore> schema_store,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
 
   SchemaProto foo_schema =
       SchemaBuilder().AddType(SchemaTypeConfigBuilder().SetType("foo")).Build();
@@ -954,7 +1048,7 @@ TEST_F(SchemaStoreTest, ComputeChecksumChangesOnModification) {
 TEST_F(SchemaStoreTest, PersistToDiskFineForEmptySchemaStore) {
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<SchemaStore> schema_store,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
 
   // Persisting is fine and shouldn't affect anything
   ICING_EXPECT_OK(schema_store->PersistToDisk());
@@ -963,7 +1057,7 @@ TEST_F(SchemaStoreTest, PersistToDiskFineForEmptySchemaStore) {
 TEST_F(SchemaStoreTest, PersistToDiskPreservesAcrossInstances) {
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<SchemaStore> schema_store,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
 
   SchemaProto schema =
       SchemaBuilder().AddType(SchemaTypeConfigBuilder().SetType("foo")).Build();
@@ -988,7 +1082,7 @@ TEST_F(SchemaStoreTest, PersistToDiskPreservesAcrossInstances) {
 
   // And we get the same schema back on reinitialization
   ICING_ASSERT_OK_AND_ASSIGN(
-      schema_store, SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      schema_store, SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
   ICING_ASSERT_OK_AND_ASSIGN(actual_schema, schema_store->GetSchema());
   EXPECT_THAT(*actual_schema, EqualsProto(schema));
 }
@@ -996,7 +1090,7 @@ TEST_F(SchemaStoreTest, PersistToDiskPreservesAcrossInstances) {
 TEST_F(SchemaStoreTest, SchemaStoreStorageInfoProto) {
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<SchemaStore> schema_store,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
 
   // Create a schema with two types: one simple type and one type that uses all
   // 16 sections.
@@ -1048,7 +1142,7 @@ TEST_F(SchemaStoreTest, SchemaStoreStorageInfoProto) {
 TEST_F(SchemaStoreTest, GetDebugInfo) {
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<SchemaStore> schema_store,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
 
   // Set schema
   ASSERT_THAT(
@@ -1067,7 +1161,7 @@ TEST_F(SchemaStoreTest, GetDebugInfo) {
 TEST_F(SchemaStoreTest, GetDebugInfoForEmptySchemaStore) {
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<SchemaStore> schema_store,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
 
   // Check debug info before setting a schema
   ICING_ASSERT_OK_AND_ASSIGN(SchemaDebugInfoProto out,
@@ -1077,6 +1171,82 @@ TEST_F(SchemaStoreTest, GetDebugInfoForEmptySchemaStore) {
   EXPECT_THAT(out, EqualsProto(expected_out));
 }
 
+TEST_F(SchemaStoreTest, InitializeRegenerateDerivedFilesFailure) {
+  // This test covers the first point that RegenerateDerivedFiles could fail.
+  // This should simply result in SetSchema::Create returning an INTERNAL error.
+
+  {
+    ICING_ASSERT_OK_AND_ASSIGN(
+        std::unique_ptr<SchemaStore> schema_store,
+        SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
+    SchemaProto schema = SchemaBuilder()
+                             .AddType(SchemaTypeConfigBuilder().SetType("Type"))
+                             .Build();
+    ICING_ASSERT_OK(schema_store->SetSchema(std::move(schema)));
+  }
+
+  auto mock_filesystem = std::make_unique<MockFilesystem>();
+  ON_CALL(*mock_filesystem,
+          CreateDirectoryRecursively(HasSubstr("key_mapper_dir")))
+      .WillByDefault(Return(false));
+  {
+    EXPECT_THAT(SchemaStore::Create(mock_filesystem.get(), schema_store_dir_,
+                                    &fake_clock_),
+                StatusIs(libtextclassifier3::StatusCode::INTERNAL));
+  }
+}
+
+TEST_F(SchemaStoreTest, SetSchemaRegenerateDerivedFilesFailure) {
+  // This test covers the second point that RegenerateDerivedFiles could fail.
+  // If handled correctly, the schema store and section manager should still be
+  // in the original, valid state.
+  SchemaTypeConfigProto type =
+      SchemaTypeConfigBuilder()
+          .SetType("Type")
+          .AddProperty(PropertyConfigBuilder()
+                           .SetName("prop1")
+                           .SetDataTypeString(MATCH_EXACT, TOKENIZER_PLAIN)
+                           .SetCardinality(CARDINALITY_OPTIONAL))
+          .Build();
+  {
+    ICING_ASSERT_OK_AND_ASSIGN(
+        std::unique_ptr<SchemaStore> schema_store,
+        SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
+    SchemaProto schema = SchemaBuilder().AddType(type).Build();
+    ICING_ASSERT_OK(schema_store->SetSchema(std::move(schema)));
+  }
+
+  {
+    auto mock_filesystem = std::make_unique<MockFilesystem>();
+    ICING_ASSERT_OK_AND_ASSIGN(
+        std::unique_ptr<SchemaStore> schema_store,
+        SchemaStore::Create(mock_filesystem.get(), schema_store_dir_,
+                            &fake_clock_));
+
+    ON_CALL(*mock_filesystem,
+            CreateDirectoryRecursively(HasSubstr("key_mapper_dir")))
+        .WillByDefault(Return(false));
+    SchemaProto schema =
+        SchemaBuilder()
+            .AddType(type)
+            .AddType(SchemaTypeConfigBuilder().SetType("Type2"))
+            .Build();
+    EXPECT_THAT(schema_store->SetSchema(std::move(schema)),
+                StatusIs(libtextclassifier3::StatusCode::INTERNAL));
+    DocumentProto document = DocumentBuilder()
+                                 .SetSchema("Type")
+                                 .AddStringProperty("prop1", "foo bar baz")
+                                 .Build();
+    SectionMetadata expected_metadata(/*id_in=*/0, MATCH_EXACT, TOKENIZER_PLAIN,
+                                      "prop1");
+    ICING_ASSERT_OK_AND_ASSIGN(std::vector<Section> sections,
+                               schema_store->ExtractSections(document));
+    ASSERT_THAT(sections, SizeIs(1));
+    EXPECT_THAT(sections.at(0).metadata, Eq(expected_metadata));
+    EXPECT_THAT(sections.at(0).content, ElementsAre("foo bar baz"));
+  }
+}
+
 }  // namespace
 
 }  // namespace lib
diff --git a/icing/schema/section.h b/icing/schema/section.h
index 40e623a..8b2ba55 100644
--- a/icing/schema/section.h
+++ b/icing/schema/section.h
@@ -77,6 +77,11 @@ struct SectionMetadata {
         id(id_in),
         tokenizer(tokenizer),
         term_match_type(term_match_type_in) {}
+
+  bool operator==(const SectionMetadata& rhs) const {
+    return path == rhs.path && id == rhs.id && tokenizer == rhs.tokenizer &&
+           term_match_type == rhs.term_match_type;
+  }
 };
 
 // Section is an icing internal concept similar to document property but with
diff --git a/icing/scoring/scoring-processor_test.cc b/icing/scoring/scoring-processor_test.cc
index f169039..b42ba31 100644
--- a/icing/scoring/scoring-processor_test.cc
+++ b/icing/scoring/scoring-processor_test.cc
@@ -60,7 +60,7 @@ class ScoringProcessorTest : public testing::Test {
 
     ICING_ASSERT_OK_AND_ASSIGN(
         schema_store_,
-        SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+        SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
 
     ICING_ASSERT_OK_AND_ASSIGN(
         DocumentStore::CreateResult create_result,
diff --git a/icing/util/document-validator_test.cc b/icing/util/document-validator_test.cc
index 2261c37..45c23e0 100644
--- a/icing/util/document-validator_test.cc
+++ b/icing/util/document-validator_test.cc
@@ -93,9 +93,11 @@ class DocumentValidatorTest : public ::testing::Test {
                             .SetCardinality(CARDINALITY_REPEATED)))
             .Build();
 
+    schema_dir_ = GetTestTempDir() + "/schema_store";
+    ASSERT_TRUE(filesystem_.CreateDirectory(schema_dir_.c_str()));
     ICING_ASSERT_OK_AND_ASSIGN(
         schema_store_,
-        SchemaStore::Create(&filesystem_, GetTestTempDir(), &fake_clock_));
+        SchemaStore::Create(&filesystem_, schema_dir_, &fake_clock_));
     ASSERT_THAT(schema_store_->SetSchema(schema), IsOk());
 
     document_validator_ =
@@ -122,6 +124,7 @@ class DocumentValidatorTest : public ::testing::Test {
                              SimpleEmailBuilder().Build());
   }
 
+  std::string schema_dir_;
   std::unique_ptr<DocumentValidator> document_validator_;
   std::unique_ptr<SchemaStore> schema_store_;
   Filesystem filesystem_;