Snap for 13188152 from d8984eee1609af239dbf57f5826d7d9c9648d37a to androidx-javascriptengine-release

Change-Id: I40822f112b30b4f6534d089a85f199fec5fbfc43
diff --git a/icing/file/version-util.cc b/icing/file/version-util.cc
index 22116fb..c44ac69 100644
--- a/icing/file/version-util.cc
+++ b/icing/file/version-util.cc
@@ -307,6 +307,7 @@
   bool should_rebuild = false;
   int32_t existing_version = existing_version_info.version;
   while (existing_version < curr_version) {
+    // LINT.IfChange(should_rebuild_derived_files_upgrade_check)
     switch (existing_version) {
       case 1: {
         // version 1 -> version 2 upgrade, no need to rebuild
@@ -324,10 +325,15 @@
         // version 4 -> version 5 upgrade, no need to rebuild
         break;
       }
+      case 5: {
+        // version 5 -> version 6 upgrade, no need to rebuild
+        break;
+      }
       default:
         // This should not happen. Rebuild anyway if unsure.
         should_rebuild |= true;
     }
+    // LINT.ThenChange(//depot/google3/icing/file/version-util.h:kVersion)
     ++existing_version;
   }
   return should_rebuild;
diff --git a/icing/icing-search-engine.cc b/icing/icing-search-engine.cc
index 2e58eaf..0f85590 100644
--- a/icing/icing-search-engine.cc
+++ b/icing/icing-search-engine.cc
@@ -1166,6 +1166,16 @@
 
 SetSchemaResultProto IcingSearchEngine::SetSchema(
     SchemaProto&& new_schema, bool ignore_errors_and_delete_documents) {
+  SetSchemaRequestProto set_schema_request;
+  *set_schema_request.mutable_schema() = std::move(new_schema);
+  set_schema_request.set_ignore_errors_and_delete_documents(
+      ignore_errors_and_delete_documents);
+
+  return SetSchema(std::move(set_schema_request));
+}
+
+SetSchemaResultProto IcingSearchEngine::SetSchema(
+    SetSchemaRequestProto&& set_schema_request) {
   ICING_VLOG(1) << "Setting new Schema";
 
   SetSchemaResultProto result_proto;
@@ -1207,8 +1217,8 @@
   std::unique_ptr<MarkerFile> marker_file =
       std::move(marker_file_or).ValueOrDie();
 
-  auto set_schema_result_or = schema_store_->SetSchema(
-      std::move(new_schema), ignore_errors_and_delete_documents);
+  auto set_schema_result_or =
+      schema_store_->SetSchema(std::move(set_schema_request));
   if (!set_schema_result_or.ok()) {
     TransformStatus(set_schema_result_or.status(), result_status);
     return result_proto;
@@ -1400,6 +1410,26 @@
   return result_proto;
 }
 
+BatchPutResultProto IcingSearchEngine::BatchPut(
+    PutDocumentRequest&& put_document_request) {
+  BatchPutResultProto batch_put_result_proto;
+
+  // TODO(b/394875109): right now we lock in the Put(DocumentProto&&) for each
+  // document. We should considering just locking once for the whole batch here.
+  for (DocumentProto& document_proto :
+       *(put_document_request.mutable_documents())) {
+    batch_put_result_proto.mutable_put_result_protos()->Add(
+        Put(std::move(document_proto)));
+  }
+
+  if (put_document_request.persist_type() != PersistType::UNKNOWN) {
+    *batch_put_result_proto.mutable_persist_to_disk_result_proto() =
+        PersistToDisk(put_document_request.persist_type());
+  }
+
+  return batch_put_result_proto;
+}
+
 PutResultProto IcingSearchEngine::Put(const DocumentProto& document) {
   return Put(DocumentProto(document));
 }
diff --git a/icing/icing-search-engine.h b/icing/icing-search-engine.h
index 8008ce1..1fc8e23 100644
--- a/icing/icing-search-engine.h
+++ b/icing/icing-search-engine.h
@@ -106,9 +106,22 @@
   //   NOT_FOUND if missing some internal data
   InitializeResultProto Initialize() ICING_LOCKS_EXCLUDED(mutex_);
 
+  // TODO: b/337913932 - Remove this method once all callers are migrated to the
+  // new SetSchema method.
+  //
+  // This method is deprecated. Please use
+  // `IcingSearchEngine::SetSchema(SetSchemaRequestProto)` instead.
+  //
   // Specifies the schema to be applied on all Documents that are already
-  // stored as well as future documents. A schema can be 'invalid' and/or
-  // 'incompatible'. These are two independent concepts.
+  // stored as well as future documents.
+  //
+  // This SetSchema call only allows setting schemas in the default empty
+  // database. Any non-empty database field in `new_schema.types` invalidates
+  // this SetSchema request. To set a schema for a non-empty database, please
+  // use `IcingSearchEngine::SetSchema(SetSchemaRequestProto)`.
+  //
+  // A schema can be 'invalid' and/or 'incompatible'. These are two independent
+  // concepts.
   //
   // An 'invalid' schema is one that is not constructed properly. For example,
   // a PropertyConfigProto is missing the property name field. A schema can be
@@ -169,6 +182,63 @@
                                  bool ignore_errors_and_delete_documents =
                                      false) ICING_LOCKS_EXCLUDED(mutex_);
 
+  // Specifies the schema to be applied on all Documents that are already
+  // stored as well as future documents.
+  //
+  // This operation sets the schema for the single database specified in
+  // `set_schema_request.database()`. If unset, the default empty
+  // database is assumed.
+  //
+  // A schema can be 'invalid' and/or 'incompatible'. These are two independent
+  // concepts.
+  // - An 'invalid' schema is one that is not constructed properly.
+  //   - For example, a PropertyConfigProto is missing the property name field.
+  //   - A schema can be 'invalid' even if there is no previously existing
+  //     schema.
+  // - An 'incompatible' schema is one that is incompatible with a previously
+  //   existing schema.
+  //   - If there is no previously existing schema, then a new schema cannot be
+  //     incompatible. An incompatible schema is one that invalidates
+  //     pre-existing data.
+  //   - For example, a previously OPTIONAL field is now REQUIRED in the new
+  //     schema, and pre-existing data is considered invalid against the new
+  //     schema now.
+  //
+  // Default behavior will not allow a new schema to be set if it is invalid or
+  // incompatible.
+  // - `set_schema_request.ignore_errors_and_delete_documents' can be set to
+  //   true to force set an incompatible schema.
+  //   - In that case, documents that are invalidated by the new schema would be
+  //     deleted from Icing.
+  //   - This cannot be used to force set an invalid schema.
+  //
+  // This schema is persisted to disk and used across multiple instances.
+  // So, callers should only have to call this if the schema changed.
+  // However, calling it multiple times with the same schema is a no-op.
+  //
+  // On some errors, Icing will keep using the older schema, but on
+  // INTERNAL_ERROR, it is undefined to continue using Icing.
+  //
+  // Returns:
+  // - OK on success
+  // - ALREADY_EXISTS if 'set_schema_request.schema' contains multiple
+  //     definitions of the same type or contains a type that has multiple
+  //     properties with the same name.
+  // - INVALID_ARGUMENT if 'set_schema_request.schema' is invalid, or if
+  //     `set_schema_request.database` does not match the database fields of
+  //     `set_schema_request.schema.types`.
+  // - FAILED_PRECONDITION if 'set_schema_request.schema' is incompatible, or
+  //     IcingSearchEngine has not been initialized yet.
+  // - INTERNAL_ERROR if Icing failed to store the new schema or upgrade
+  //     existing data based on the new schema. Using Icing beyond this error is
+  //     undefined and may cause crashes.
+  // - DATA_LOSS_ERROR if 'set_schema_request.schema' requires the index to be
+  //     rebuilt and an IO error leads to some documents being excluded from the
+  //     index. These documents will still be retrievable via Get, but won't
+  //     match queries.
+  SetSchemaResultProto SetSchema(SetSchemaRequestProto&& set_schema_request)
+      ICING_LOCKS_EXCLUDED(mutex_);
+
   // Get Icing's current copy of the schema.
   //
   // Returns:
@@ -205,6 +275,16 @@
   GetSchemaTypeResultProto GetSchemaType(std::string_view schema_type)
       ICING_LOCKS_EXCLUDED(mutex_);
 
+  // Batch puts the documents into icing search engine so that they're stored
+  // and indexed. Documents are automatically written to disk, callers can also
+  // call PersistToDisk() to flush changes immediately.
+  //
+  // Returns: BatchPutResultProto with a list of PutResultProtos for each
+  // document, and a PersistToDiskResultProto if persist_type is set in the
+  // request.
+  BatchPutResultProto BatchPut(PutDocumentRequest&& put_document_request)
+      ICING_LOCKS_EXCLUDED(mutex_);
+
   // Puts the document into icing search engine so that it's stored and
   // indexed. Documents are automatically written to disk, callers can also
   // call PersistToDisk() to flush changes immediately.
diff --git a/icing/icing-search-engine_initialization_test.cc b/icing/icing-search-engine_initialization_test.cc
index a1bc9fe..5fc862b 100644
--- a/icing/icing-search-engine_initialization_test.cc
+++ b/icing/icing-search-engine_initialization_test.cc
@@ -330,6 +330,19 @@
   return scoring_spec;
 }
 
+SetSchemaRequestProto CreateSetSchemaRequestProto(
+    SchemaProto schema, std::string database,
+    bool ignore_errors_and_delete_documents) {
+  SetSchemaRequestProto set_schema_request;
+
+  *set_schema_request.mutable_schema() = std::move(schema);
+  set_schema_request.set_database(std::move(database));
+  set_schema_request.set_ignore_errors_and_delete_documents(
+      ignore_errors_and_delete_documents);
+
+  return set_schema_request;
+}
+
 // TODO(b/272145329): create SearchSpecBuilder, JoinSpecBuilder,
 // SearchResultProtoBuilder and ResultProtoBuilder for unit tests and build all
 // instances by them.
@@ -6991,27 +7004,49 @@
                            .SetDataTypeInt64(NUMERIC_MATCH_RANGE)
                            .SetCardinality(CARDINALITY_OPTIONAL))
           .Build();
+  SchemaTypeConfigProto db1_email_type_with_db =
+      SchemaTypeConfigBuilder(db1_email_type).SetDatabase("db1").Build();
+  SchemaTypeConfigProto db2_email_type_with_db =
+      SchemaTypeConfigBuilder(db2_email_type).SetDatabase("db2").Build();
 
+  SchemaProto previous_version_db1_schema;
+  SchemaProto previous_version_db2_schema;
+  SchemaBuilder previous_version_full_schema_builder;
+  SetSchemaRequestProto set_schema_request;
   if (previous_version_has_schema_database_enabled) {
-    // Populate the database field for the db1/email type.
-    db1_email_type =
-        SchemaTypeConfigBuilder(db1_email_type).SetDatabase("db1").Build();
+    // Set db1/email schema to always populate the database field.
+    previous_version_full_schema_builder.AddType(db1_email_type_with_db);
+    previous_version_db1_schema =
+        SchemaBuilder().AddType(db1_email_type_with_db).Build();
+
     if (existing_version >= version_util::kSchemaDatabaseVersion) {
       // Populate the database field for the db2/email type only if previous
       // version is a post schema-database version.
-      db2_email_type =
-          SchemaTypeConfigBuilder(db2_email_type).SetDatabase("db2").Build();
+      previous_version_full_schema_builder.AddType(db2_email_type_with_db);
+      previous_version_db2_schema =
+          SchemaBuilder().AddType(db2_email_type_with_db).Build();
+    } else {
+      // Otherwise, the database field is not populated for db2/email type. This
+      // is to simulate the following situation:
+      // 1. Icing is initialized on a version>kSchemaDatabaseVersion with schema
+      //    database enabled, and db1/email is set with the database field
+      //    populated.
+      // 2. Icing gets rolled back to pre-schema database version, db2/email is
+      //    set during this time so the database field is not populated.
+      previous_version_full_schema_builder.AddType(db2_email_type);
+      previous_version_db2_schema =
+          SchemaBuilder().AddType(db2_email_type).Build();
     }
-    // Otherwise, the database field is not populated for db2/email type. This
-    // is to simulate the following situation:
-    // 1. Icing is initialized on a version>kSchemaDatabaseVersion with schema
-    //    database enabled, and db1/email is set with the database field
-    //    populated.
-    // 2. Icing gets rolled back to pre-schema database version, db2/email is
-    //    set during this time so the database field is not populated.
+  } else {
+    previous_version_full_schema_builder.AddType(db1_email_type)
+        .AddType(db2_email_type);
+    previous_version_db1_schema =
+        SchemaBuilder().AddType(db1_email_type).Build();
+    previous_version_db2_schema =
+        SchemaBuilder().AddType(db2_email_type).Build();
   }
   SchemaProto previous_version_schema =
-      SchemaBuilder().AddType(db1_email_type).AddType(db2_email_type).Build();
+      previous_version_full_schema_builder.Build();
 
   DocumentProto db1_email_doc =
       DocumentBuilder()
@@ -7040,15 +7075,21 @@
     EXPECT_THAT(icing.Initialize().status(), ProtoIsOk());
     // 1. Set schema.
     if (options.enable_schema_database()) {
-      // Need to set schemas with a single database field at a time.
-      ASSERT_THAT(
-          icing.SetSchema(SchemaBuilder().AddType(db1_email_type).Build())
-              .status(),
-          ProtoIsOk());
-      ASSERT_THAT(
-          icing.SetSchema(SchemaBuilder().AddType(db2_email_type).Build())
-              .status(),
-          ProtoIsOk());
+      // Can only set schema for a single database at a time.
+      ASSERT_THAT(icing
+                      .SetSchema(CreateSetSchemaRequestProto(
+                          previous_version_db1_schema,
+                          previous_version_db1_schema.types(0).database(),
+                          /*ignore_errors_and_delete_documents=*/false))
+                      .status(),
+                  ProtoIsOk());
+      ASSERT_THAT(icing
+                      .SetSchema(CreateSetSchemaRequestProto(
+                          previous_version_db2_schema,
+                          previous_version_db2_schema.types(0).database(),
+                          /*ignore_errors_and_delete_documents=*/false))
+                      .status(),
+                  ProtoIsOk());
     } else {
       ASSERT_THAT(icing.SetSchema(previous_version_schema).status(),
                   ProtoIsOk());
diff --git a/icing/icing-search-engine_schema_test.cc b/icing/icing-search-engine_schema_test.cc
index dbab065..269daaa 100644
--- a/icing/icing-search-engine_schema_test.cc
+++ b/icing/icing-search-engine_schema_test.cc
@@ -169,6 +169,19 @@
   return scoring_spec;
 }
 
+SetSchemaRequestProto CreateSetSchemaRequestProto(
+    SchemaProto schema, std::string database,
+    bool ignore_errors_and_delete_documents) {
+  SetSchemaRequestProto set_schema_request;
+
+  *set_schema_request.mutable_schema() = std::move(schema);
+  set_schema_request.set_database(std::move(database));
+  set_schema_request.set_ignore_errors_and_delete_documents(
+      ignore_errors_and_delete_documents);
+
+  return set_schema_request;
+}
+
 // TODO(b/272145329): create SearchSpecBuilder, JoinSpecBuilder,
 // SearchResultProtoBuilder and ResultProtoBuilder for unit tests and build all
 // instances by them.
@@ -865,7 +878,9 @@
                            .SetCardinality(CARDINALITY_REQUIRED))
           .Build();
   SchemaProto db1_schema = SchemaBuilder().AddType(db1_type).Build();
-  SetSchemaResultProto set_schema_result = icing.SetSchema(db1_schema);
+  SetSchemaResultProto set_schema_result =
+      icing.SetSchema(CreateSetSchemaRequestProto(
+          db1_schema, "db1", /*ignore_errors_and_delete_documents=*/false));
   // Ignore latency numbers. They're covered elsewhere.
   set_schema_result.clear_latency_ms();
   SetSchemaResultProto expected_set_schema_result;
@@ -928,7 +943,8 @@
                            .SetCardinality(CARDINALITY_REQUIRED))
           .Build();
   SchemaProto db2_schema = SchemaBuilder().AddType(db2_type).Build();
-  set_schema_result = icing.SetSchema(db2_schema);
+  set_schema_result = icing.SetSchema(CreateSetSchemaRequestProto(
+      db2_schema, "db2", /*ignore_errors_and_delete_documents=*/false));
   // Ignore latency numbers. They're covered elsewhere.
   set_schema_result.clear_latency_ms();
   expected_set_schema_result = SetSchemaResultProto();
@@ -1016,7 +1032,9 @@
                            .SetCardinality(CARDINALITY_REQUIRED))
           .Build();
   SchemaProto db1_schema = SchemaBuilder().AddType(db1_type).Build();
-  SetSchemaResultProto set_schema_result = icing.SetSchema(db1_schema);
+  SetSchemaResultProto set_schema_result =
+      icing.SetSchema(CreateSetSchemaRequestProto(
+          db1_schema, "db1", /*ignore_errors_and_delete_documents=*/false));
   // Ignore latency numbers. They're covered elsewhere.
   set_schema_result.clear_latency_ms();
   SetSchemaResultProto expected_set_schema_result;
@@ -1042,7 +1060,8 @@
                            .SetCardinality(CARDINALITY_REQUIRED))
           .Build();
   SchemaProto db2_schema = SchemaBuilder().AddType(db2_type).Build();
-  set_schema_result = icing.SetSchema(db2_schema);
+  set_schema_result = icing.SetSchema(CreateSetSchemaRequestProto(
+      db2_schema, "db2", /*ignore_errors_and_delete_documents=*/false));
   // Ignore latency numbers. They're covered elsewhere.
   set_schema_result.clear_latency_ms();
   expected_set_schema_result = SetSchemaResultProto();
@@ -1118,7 +1137,8 @@
           .SetCardinality(CARDINALITY_OPTIONAL)
           .Build());
   db1_schema = SchemaBuilder().AddType(db1_type).Build();
-  set_schema_result = icing.SetSchema(db1_schema);
+  set_schema_result = icing.SetSchema(CreateSetSchemaRequestProto(
+      db1_schema, "db1", /*ignore_errors_and_delete_documents=*/false));
   // Ignore latency numbers. They're covered elsewhere.
   set_schema_result.clear_latency_ms();
   expected_set_schema_result = SetSchemaResultProto();
@@ -1206,6 +1226,197 @@
               EqualsProto(expected_get_schema_result_proto_db2_full));
 }
 
+TEST_F(IcingSearchEngineSchemaTest, SetSchemaEmptySchemaClearsDatabase) {
+  IcingSearchEngine icing(GetDefaultIcingOptions(), GetTestJniCache());
+  ASSERT_THAT(icing.Initialize().status(), ProtoIsOk());
+
+  // Create and set schema in db1 with 2 properties:
+  // - 'b': string type, indexed.
+  // - 'c': int64 type, indexed.
+  SchemaTypeConfigProto db1_type =
+      SchemaTypeConfigBuilder()
+          .SetType("db1_type")
+          .SetDatabase("db1")
+          .AddProperty(PropertyConfigBuilder()
+                           .SetName("b")
+                           .SetDataTypeString(TERM_MATCH_EXACT, TOKENIZER_PLAIN)
+                           .SetCardinality(CARDINALITY_REQUIRED))
+          .AddProperty(PropertyConfigBuilder()
+                           .SetName("c")
+                           .SetDataTypeInt64(NUMERIC_MATCH_RANGE)
+                           .SetCardinality(CARDINALITY_REQUIRED))
+          .Build();
+  SchemaProto db1_schema = SchemaBuilder().AddType(db1_type).Build();
+  SetSchemaResultProto set_schema_result =
+      icing.SetSchema(CreateSetSchemaRequestProto(
+          db1_schema, "db1", /*ignore_errors_and_delete_documents=*/true));
+  // Ignore latency numbers. They're covered elsewhere.
+  set_schema_result.clear_latency_ms();
+  SetSchemaResultProto expected_set_schema_result;
+  expected_set_schema_result.mutable_status()->set_code(StatusProto::OK);
+  expected_set_schema_result.mutable_new_schema_types()->Add("db1_type");
+  EXPECT_THAT(set_schema_result, EqualsProto(expected_set_schema_result));
+
+  // Add a schema for db2:
+  // - 'b': string type, indexed.
+  // - 'd': int64 type, indexed.
+  SchemaTypeConfigProto db2_type =
+      SchemaTypeConfigBuilder()
+          .SetType("db2_type")
+          .SetDatabase("db2")
+          .AddProperty(
+              PropertyConfigBuilder()
+                  .SetName("b")
+                  .SetDataTypeString(TERM_MATCH_PREFIX, TOKENIZER_PLAIN)
+                  .SetCardinality(CARDINALITY_REQUIRED))
+          .AddProperty(PropertyConfigBuilder()
+                           .SetName("d")
+                           .SetDataTypeInt64(NUMERIC_MATCH_RANGE)
+                           .SetCardinality(CARDINALITY_REQUIRED))
+          .Build();
+  SchemaProto db2_schema = SchemaBuilder().AddType(db2_type).Build();
+  set_schema_result = icing.SetSchema(CreateSetSchemaRequestProto(
+      db2_schema, "db2", /*ignore_errors_and_delete_documents=*/true));
+  // Ignore latency numbers. They're covered elsewhere.
+  set_schema_result.clear_latency_ms();
+  expected_set_schema_result = SetSchemaResultProto();
+  expected_set_schema_result.mutable_status()->set_code(StatusProto::OK);
+  expected_set_schema_result.mutable_new_schema_types()->Add("db2_type");
+  EXPECT_THAT(set_schema_result, EqualsProto(expected_set_schema_result));
+
+  // Add documents
+  DocumentProto db1_document1 =
+      DocumentBuilder()
+          .SetKey("namespace", "uri1")
+          .SetSchema("db1_type")
+          .AddStringProperty("b", "message body")
+          .AddInt64Property("c", 123)
+          .SetCreationTimestampMs(kDefaultCreationTimestampMs)
+          .Build();
+  DocumentProto db2_document =
+      DocumentBuilder()
+          .SetKey("namespace", "uri2")
+          .SetSchema("db2_type")
+          .AddStringProperty("b", "message body")
+          .AddInt64Property("d", 123)
+          .SetCreationTimestampMs(kDefaultCreationTimestampMs)
+          .Build();
+  EXPECT_THAT(icing.Put(db1_document1).status(), ProtoIsOk());
+  EXPECT_THAT(icing.Put(db2_document).status(), ProtoIsOk());
+
+  // Verify term search. Should match both docs.
+  SearchSpecProto search_spec1;
+  search_spec1.set_query("b:message");
+  search_spec1.set_term_match_type(TermMatchType::EXACT_ONLY);
+
+  SearchResultProto expected_result_db1doc1;
+  expected_result_db1doc1.mutable_status()->set_code(StatusProto::OK);
+  *expected_result_db1doc1.mutable_results()->Add()->mutable_document() =
+      db1_document1;
+
+  SearchResultProto expected_result_db1doc1_db2;
+  expected_result_db1doc1_db2.mutable_status()->set_code(StatusProto::OK);
+  *expected_result_db1doc1_db2.mutable_results()->Add()->mutable_document() =
+      db2_document;
+  *expected_result_db1doc1_db2.mutable_results()->Add()->mutable_document() =
+      db1_document1;
+
+  SearchResultProto actual_results =
+      icing.Search(search_spec1, GetDefaultScoringSpec(),
+                   ResultSpecProto::default_instance());
+  EXPECT_THAT(actual_results, EqualsSearchResultIgnoreStatsAndScores(
+                                  expected_result_db1doc1_db2));
+
+  // Verify numeric (integer) search. Should only match db1_document1.
+  SearchSpecProto search_spec2;
+  search_spec2.set_query("c == 123");
+  search_spec2.add_enabled_features(std::string(kNumericSearchFeature));
+
+  actual_results = icing.Search(search_spec2, GetDefaultScoringSpec(),
+                                ResultSpecProto::default_instance());
+  EXPECT_THAT(actual_results,
+              EqualsSearchResultIgnoreStatsAndScores(expected_result_db1doc1));
+
+  // Clear db1 by setting an empty schema:
+  // - db1_schema (deleted)
+  // - db2_schema:
+  //   - 'b': string type, indexed.
+  //   - 'd': int64 type, indexed.
+  db1_schema = SchemaProto();
+  set_schema_result = icing.SetSchema(CreateSetSchemaRequestProto(
+      db1_schema, "db1", /*ignore_errors_and_delete_documents=*/true));
+  // Ignore latency numbers. They're covered elsewhere.
+  set_schema_result.clear_latency_ms();
+  expected_set_schema_result = SetSchemaResultProto();
+  expected_set_schema_result.mutable_status()->set_code(StatusProto::OK);
+  expected_set_schema_result.mutable_deleted_schema_types()->Add("db1_type");
+  EXPECT_THAT(set_schema_result, EqualsProto(expected_set_schema_result));
+
+  // Adding new document fails because db1_type is deleted.
+  DocumentProto db1_document2 =
+      DocumentBuilder()
+          .SetKey("namespace", "uri3")
+          .SetSchema("db1_type")
+          .AddStringProperty("a", "message body")
+          .AddStringProperty("b", "string value")
+          .AddInt64Property("c", 123)
+          .SetCreationTimestampMs(kDefaultCreationTimestampMs)
+          .Build();
+  EXPECT_THAT(icing.Put(db1_document2).status(),
+              ProtoStatusIs(StatusProto::NOT_FOUND));
+
+  // Check that original db1 docs are deleted, and that db2 docs are still
+  // searchable.
+  // Verify term search. "b:message" only matches db2_document.
+  actual_results = icing.Search(search_spec1, GetDefaultScoringSpec(),
+                                ResultSpecProto::default_instance());
+  SearchResultProto expected_result_db2doc;
+  expected_result_db2doc.mutable_status()->set_code(StatusProto::OK);
+  *expected_result_db2doc.mutable_results()->Add()->mutable_document() =
+      db2_document;
+  EXPECT_THAT(actual_results,
+              EqualsSearchResultIgnoreStatsAndScores(expected_result_db2doc));
+
+  // Verify numeric (integer) search. "c == 123" should not match anything.
+  actual_results = icing.Search(search_spec2, GetDefaultScoringSpec(),
+                                ResultSpecProto::default_instance());
+  SearchResultProto expected_result_empty;
+  expected_result_empty.mutable_status()->set_code(StatusProto::OK);
+  EXPECT_THAT(actual_results,
+              EqualsSearchResultIgnoreStatsAndScores(expected_result_empty));
+
+  // "message" should only match db2_document.
+  SearchSpecProto search_spec3;
+  search_spec3.set_query("message");
+  search_spec3.set_term_match_type(TermMatchType::PREFIX);
+
+  actual_results = icing.Search(search_spec3, GetDefaultScoringSpec(),
+                                ResultSpecProto::default_instance());
+  EXPECT_THAT(actual_results,
+              EqualsSearchResultIgnoreStatsAndScores(expected_result_db2doc));
+
+  // Get full schema
+  GetSchemaResultProto expected_get_schema_result_proto_full;
+  expected_get_schema_result_proto_full.mutable_status()->set_code(
+      StatusProto::OK);
+  *expected_get_schema_result_proto_full.mutable_schema() =
+      SchemaBuilder().AddType(db2_type).Build();
+  EXPECT_THAT(icing.GetSchema(),
+              EqualsProto(expected_get_schema_result_proto_full));
+
+  // Get db1 schema should return NOT_FOUND.
+  EXPECT_THAT(icing.GetSchema("db1").status(),
+              ProtoStatusIs(StatusProto::NOT_FOUND));
+
+  // Get db2 schema
+  GetSchemaResultProto expected_get_schema_result_proto_db2_full;
+  expected_get_schema_result_proto_db2_full.mutable_status()->set_code(
+      StatusProto::OK);
+  *expected_get_schema_result_proto_db2_full.mutable_schema() = db2_schema;
+  EXPECT_THAT(icing.GetSchema("db2"),
+              EqualsProto(expected_get_schema_result_proto_db2_full));
+}
+
 TEST_F(IcingSearchEngineSchemaTest,
        SetSchemaNewIndexedStringPropertyTriggersIndexRestorationAndReturnsOk) {
   IcingSearchEngine icing(GetDefaultIcingOptions(), GetTestJniCache());
@@ -3325,7 +3536,9 @@
           .Build();
 
   SchemaProto db1_schema = SchemaBuilder().AddType(db1_type).Build();
-  SetSchemaResultProto set_schema_result = icing.SetSchema(db1_schema);
+  SetSchemaResultProto set_schema_result =
+      icing.SetSchema(CreateSetSchemaRequestProto(
+          db1_schema, "db1", /*ignore_errors_and_delete_documents=*/false));
   // Ignore latency numbers. They're covered elsewhere.
   set_schema_result.clear_latency_ms();
   SetSchemaResultProto expected_set_schema_result;
@@ -3350,7 +3563,8 @@
           .Build();
 
   SchemaProto db2_schema = SchemaBuilder().AddType(db2_type).Build();
-  set_schema_result = icing.SetSchema(db2_schema);
+  set_schema_result = icing.SetSchema(CreateSetSchemaRequestProto(
+      db2_schema, "db2", /*ignore_errors_and_delete_documents=*/false));
   // Ignore latency numbers. They're covered elsewhere.
   set_schema_result.clear_latency_ms();
   expected_set_schema_result = SetSchemaResultProto();
@@ -3409,7 +3623,9 @@
                                         .SetDataTypeInt64(NUMERIC_MATCH_RANGE)
                                         .SetCardinality(CARDINALITY_REQUIRED)))
           .Build();
-  SetSchemaResultProto set_schema_result = icing.SetSchema(db1_schema);
+  SetSchemaResultProto set_schema_result =
+      icing.SetSchema(CreateSetSchemaRequestProto(
+          db1_schema, "db1", /*ignore_errors_and_delete_documents=*/false));
   // Ignore latency numbers. They're covered elsewhere.
   set_schema_result.clear_latency_ms();
   SetSchemaResultProto expected_set_schema_result;
diff --git a/icing/jni/icing-search-engine-jni.cc b/icing/jni/icing-search-engine-jni.cc
index 6ed42ee..c17ea0b 100644
--- a/icing/jni/icing-search-engine-jni.cc
+++ b/icing/jni/icing-search-engine-jni.cc
@@ -126,6 +126,28 @@
 
   return SerializeProtoToJniByteArray(env, set_schema_result_proto);
 }
+// TODO : b/337913932 - pre-register this API once Jetpack build is dropped back
+// into g3
+JNIEXPORT jbyteArray JNICALL
+Java_com_google_android_icing_IcingSearchEngineImpl_nativeSetSchemaWithRequestProto(
+    JNIEnv* env, jclass clazz, jobject object,
+    jbyteArray set_schema_request_bytes) {
+  icing::lib::IcingSearchEngine* icing =
+      GetIcingSearchEnginePointer(env, object);
+
+  icing::lib::SetSchemaRequestProto set_schema_request;
+  if (!ParseProtoFromJniByteArray(env, set_schema_request_bytes,
+                                  &set_schema_request)) {
+    ICING_LOG(icing::lib::ERROR)
+        << "Failed to parse SetSchemaRequestProto in nativeSetSchema";
+    return nullptr;
+  }
+
+  icing::lib::SetSchemaResultProto set_schema_result_proto =
+      icing->SetSchema(std::move(set_schema_request));
+
+  return SerializeProtoToJniByteArray(env, set_schema_result_proto);
+}
 
 jbyteArray nativeGetSchema(JNIEnv* env, jclass clazz, jobject object) {
   icing::lib::IcingSearchEngine* icing =
@@ -136,8 +158,8 @@
   return SerializeProtoToJniByteArray(env, get_schema_result_proto);
 }
 
-jbyteArray nativeGetSchemaForDatabase(
-    JNIEnv* env, jclass clazz, jobject object, jstring database) {
+jbyteArray nativeGetSchemaForDatabase(JNIEnv* env, jclass clazz, jobject object,
+                                      jstring database) {
   icing::lib::IcingSearchEngine* icing =
       GetIcingSearchEnginePointer(env, object);
 
@@ -195,12 +217,8 @@
     return nullptr;
   }
 
-  icing::lib::BatchPutResultProto batch_put_result_proto;
-  for (icing::lib::DocumentProto& document_proto :
-       *(put_document_request.mutable_documents())) {
-    batch_put_result_proto.mutable_put_result_protos()->Add(
-        icing->Put(std::move(document_proto)));
-  }
+  icing::lib::BatchPutResultProto batch_put_result_proto =
+      icing->BatchPut(std::move(put_document_request));
 
   return SerializeProtoToJniByteArray(env, batch_put_result_proto);
 }
diff --git a/icing/schema-builder.h b/icing/schema-builder.h
index ab14d7e..0c46159 100644
--- a/icing/schema-builder.h
+++ b/icing/schema-builder.h
@@ -15,7 +15,6 @@
 #ifndef ICING_SCHEMA_BUILDER_H_
 #define ICING_SCHEMA_BUILDER_H_
 
-#include <cstdint>
 #include <initializer_list>
 #include <string>
 #include <string_view>
diff --git a/icing/schema/schema-store.cc b/icing/schema/schema-store.cc
index 308c217..dda7746 100644
--- a/icing/schema/schema-store.cc
+++ b/icing/schema/schema-store.cc
@@ -856,37 +856,45 @@
   return schema_proto;
 }
 
-// TODO(cassiewang): Consider removing this definition of SetSchema if it's not
-// needed by production code. It's currently being used by our tests, but maybe
-// it's trivial to change our test code to also use the
-// SetSchema(SchemaProto&& new_schema)
+// TODO - b/337913932 - Remove this method once all callers are migrated to
+// SetSchema(SetSchemaRequestProto&& set_schema_request). This should just be
+// used in our tests.
 libtextclassifier3::StatusOr<SchemaStore::SetSchemaResult>
-SchemaStore::SetSchema(const SchemaProto& new_schema,
+SchemaStore::SetSchema(SchemaProto new_schema,
                        bool ignore_errors_and_delete_documents) {
-  return SetSchema(SchemaProto(new_schema), ignore_errors_and_delete_documents);
+  SetSchemaRequestProto set_schema_request;
+  *set_schema_request.mutable_schema() = std::move(new_schema);
+  set_schema_request.set_ignore_errors_and_delete_documents(
+      ignore_errors_and_delete_documents);
+
+  return SetSchema(std::move(set_schema_request));
 }
 
 libtextclassifier3::StatusOr<SchemaStore::SetSchemaResult>
-SchemaStore::SetSchema(SchemaProto&& new_schema,
-                       bool ignore_errors_and_delete_documents) {
+SchemaStore::SetSchema(SetSchemaRequestProto&& set_schema_request) {
+  bool ignore_errors_and_delete_documents =
+      set_schema_request.ignore_errors_and_delete_documents();
+
   if (feature_flags_->enable_schema_database()) {
     // Step 1: (Only required if schema database is enabled)
     // Do some preliminary checks on the new schema before formal validation and
     // delta computation. This checks that:
-    // - The new schema only contains types from a single database.
+    // - The database field in the new schema's types match the provided
+    //   database.
     // - The new schema's type names are not already in use from other
-    // databases.
-    ICING_ASSIGN_OR_RETURN(std::string database,
-                           ValidateAndGetDatabase(new_schema));
+    //   databases.
+    ICING_RETURN_IF_ERROR(ValidateSchemaDatabase(
+        set_schema_request.schema(), set_schema_request.database()));
 
     // Step 2: Schema validation and delta computation -- try to get the
     // existing schema for the database to compare to the new schema.
     libtextclassifier3::StatusOr<SchemaProto> schema_proto =
-        GetSchema(database);
+        GetSchema(set_schema_request.database());
     if (absl_ports::IsNotFound(schema_proto.status())) {
       // Case 1: No preexisting schema for this database.
-      return SetInitialSchemaForDatabase(std::move(new_schema),
-                                         ignore_errors_and_delete_documents);
+      return SetInitialSchemaForDatabase(
+          std::move(*set_schema_request.mutable_schema()),
+          set_schema_request.database(), ignore_errors_and_delete_documents);
     }
 
     if (!schema_proto.ok()) {
@@ -897,16 +905,18 @@
     // Case 3: At this point, we're guaranteed that we have an existing schema
     // for this database.
     const SchemaProto& old_schema = schema_proto.ValueOrDie();
-    return SetSchemaWithDatabaseOverride(std::move(new_schema), old_schema,
-                                         ignore_errors_and_delete_documents);
+    return SetSchemaWithDatabaseOverride(
+        std::move(*set_schema_request.mutable_schema()), old_schema,
+        set_schema_request.database(), ignore_errors_and_delete_documents);
   }
 
   // Get the full schema if schema database is disabled.
   libtextclassifier3::StatusOr<const SchemaProto*> schema_proto = GetSchema();
   if (absl_ports::IsNotFound(schema_proto.status())) {
     // Case 1: No preexisting schema
-    return SetInitialSchemaForDatabase(std::move(new_schema),
-                                       ignore_errors_and_delete_documents);
+    return SetInitialSchemaForDatabase(
+        std::move(*set_schema_request.mutable_schema()),
+        set_schema_request.database(), ignore_errors_and_delete_documents);
   }
 
   if (!schema_proto.ok()) {
@@ -916,13 +926,15 @@
 
   // Case 3: At this point, we're guaranteed that we have an existing schema
   const SchemaProto& old_schema = *schema_proto.ValueOrDie();
-  return SetSchemaWithDatabaseOverride(std::move(new_schema), old_schema,
-                                       ignore_errors_and_delete_documents);
+  return SetSchemaWithDatabaseOverride(
+      std::move(*set_schema_request.mutable_schema()), old_schema,
+      set_schema_request.database(), ignore_errors_and_delete_documents);
 }
 
 libtextclassifier3::StatusOr<SchemaStore::SetSchemaResult>
 SchemaStore::SetInitialSchemaForDatabase(
-    SchemaProto new_schema, bool ignore_errors_and_delete_documents) {
+    SchemaProto new_schema, const std::string& database,
+    bool ignore_errors_and_delete_documents) {
   SetSchemaResult result;
 
   ICING_RETURN_IF_ERROR(SchemaUtil::Validate(new_schema, *feature_flags_));
@@ -936,7 +948,7 @@
   // schema file.
   ICING_ASSIGN_OR_RETURN(
       SchemaProto full_new_schema,
-      GetFullSchemaProtoWithUpdatedDb(std::move(new_schema)));
+      GetFullSchemaProtoWithUpdatedDb(std::move(new_schema), database));
   ICING_RETURN_IF_ERROR(ApplySchemaChange(std::move(full_new_schema)));
   has_schema_successfully_set_ = true;
 
@@ -946,12 +958,25 @@
 libtextclassifier3::StatusOr<SchemaStore::SetSchemaResult>
 SchemaStore::SetSchemaWithDatabaseOverride(
     SchemaProto new_schema, const SchemaProto& old_schema,
-    bool ignore_errors_and_delete_documents) {
+    const std::string& database, bool ignore_errors_and_delete_documents) {
   // Assume we can set the schema unless proven otherwise.
   SetSchemaResult result;
   result.success = true;
 
   if (feature_flags_->enable_schema_database()) {
+    // Sanity check to make sure that we're comparing schemas from the same
+    // database.
+    // The new code path ensures that old_schema contains types from exactly one
+    // database since it's obtained using GetSchema(database), which is
+    // guaranteed to only return types from the single provided database.
+    libtextclassifier3::Status validate_old_schema_database =
+        ValidateSchemaDatabase(old_schema, database);
+    if (!validate_old_schema_database.ok()) {
+      return absl_ports::InvalidArgumentError(
+          "Schema database mismatch between new and old schemas. This should "
+          "never happen");
+    }
+
     // Check if the schema types are the same between the new and old schema,
     // ignoring order.
     if (AreSchemaTypesEqual(new_schema, old_schema)) {
@@ -1015,7 +1040,7 @@
   // for writing the full proto to the schema file.
   ICING_ASSIGN_OR_RETURN(
       SchemaProto full_new_schema,
-      GetFullSchemaProtoWithUpdatedDb(std::move(new_schema)));
+      GetFullSchemaProtoWithUpdatedDb(std::move(new_schema), database));
 
   // We still need to update old_schema_type_ids_changed. We need to retrieve
   // the entire old schema for this, as type ids are assigned for the entire
@@ -1390,15 +1415,12 @@
   return blob_property_map;
 }
 
-libtextclassifier3::StatusOr<std::string> SchemaStore::ValidateAndGetDatabase(
-    const SchemaProto& new_schema) const {
-  std::string database;
-
+libtextclassifier3::Status SchemaStore::ValidateSchemaDatabase(
+    const SchemaProto& new_schema, const std::string& database) const {
   if (!feature_flags_->enable_schema_database() || new_schema.types().empty()) {
-    return database;
+    return libtextclassifier3::Status::OK;
   }
 
-  database = new_schema.types(0).database();
   // Loop through new_schema's types and validate it. The input SchemaProto
   // contains a list of SchemaTypeConfigProtos without deduplication. We need to
   // check that:
@@ -1409,10 +1431,10 @@
   for (const SchemaTypeConfigProto& type_config : new_schema.types()) {
     // Check database consistency.
     if (database != type_config.database()) {
-      return absl_ports::InvalidArgumentError(
-          "SetSchema only accepts a SchemaProto with types from a single "
-          "database at a time. Please make separate calls for each database if "
-          "you need to set the schema for multiple databases.");
+      return absl_ports::InvalidArgumentError(absl_ports::StrCat(
+          "Mismatch between the set schema request's database and the new "
+          "schema types' database. Expected '",
+          database, "' but got '", type_config.database(), "'."));
     }
 
     // Check type name uniqueness. This is only necessary if there is a
@@ -1427,12 +1449,13 @@
       }
     }
   }
-  return database;
+  return libtextclassifier3::Status::OK;
 }
 
 libtextclassifier3::StatusOr<SchemaProto>
 SchemaStore::GetFullSchemaProtoWithUpdatedDb(
-    SchemaProto input_database_schema) const {
+    SchemaProto input_database_schema,
+    const std::string& database_to_update) const {
   if (!feature_flags_->enable_schema_database()) {
     // If the schema database is not enabled, the input schema is already the
     // full schema, so we don't need to do any merges.
@@ -1458,13 +1481,8 @@
 
   // At this point, we have a pre-existing schema -- we need to merge the
   // updated database with the existing schema.
-  if (input_database_schema.types().empty()) {
-    return *schema_proto.ValueOrDie();
-  }
-
-  std::string input_database = input_database_schema.types(0).database();
   if (database_type_map_.size() == 1 &&
-      database_type_map_.find(input_database) != database_type_map_.end()) {
+      database_type_map_.find(database_to_update) != database_type_map_.end()) {
     // No other databases in the schema -- we can return the input database
     // schema.
     return input_database_schema;
@@ -1475,18 +1493,18 @@
 
   // 1. Add types from the existing schema, replacing existing types with the
   // input types if the database is the one being updated by the input schema.
-  // - For the input_database, we replace the existing types with the input
-  //   types. An exisiting type is deleted if it's not included in
-  //   input_database.
-  // - If there are more input types than existing types for the input_database,
+  // - For database_to_update, we replace the existing types with the input
+  //   types. Any existing type not included in input_database_schema is
+  //   deleted.
+  // - If there are more input types than existing types for database_to_update,
   //   the rest of the input types are appended to the end of the full_schema.
-  // - If there are fewer input types than existing types for the
-  //   input_database, we shift all existing that come after input_database
-  //   forward.
-  // - For existing types from other databases, we add the types in their
-  //   original order to full_schema. Note that the type-ids of existing types
-  //   might still change if some types deleted in input_database as this will
-  //   cause all subsequent types ids to shift forward.
+  // - If there are fewer input types than existing types for
+  //   database_to_update, we shift forward all existing types that appear after
+  //   the last input type.
+  // - For existing types from other databases, we preserve the existing order
+  //   after adding to full_schema. Note that the type-ids of existing types
+  //   might still change if some types are deleted in the database_to_update as
+  //   this will cause all subsequent types ids to shift forward.
   int input_schema_index = 0, existing_schema_index = 0;
   while (input_schema_index < input_database_schema.types().size() &&
          existing_schema_index < existing_schema->types().size()) {
@@ -1495,12 +1513,7 @@
     SchemaTypeConfigProto& input_type_config =
         *input_database_schema.mutable_types(input_schema_index);
 
-    if (input_type_config.database() != input_database) {
-      return absl_ports::InvalidArgumentError(
-          "Can only update a single database at a time.");
-    }
-
-    if (existing_type_config.database() == input_database) {
+    if (existing_type_config.database() == database_to_update) {
       // If the database is the one being updated by the input schema, replace
       // the existing type with a type from the input schema.
       *full_schema.add_types() = std::move(input_type_config);
@@ -1529,7 +1542,7 @@
     // that are from input_database, since existing types from input_database
     // are replaced with input_database_schema.
     if (existing_schema->types(existing_schema_index).database() !=
-        input_database) {
+        database_to_update) {
       *full_schema.add_types() = existing_schema->types(existing_schema_index);
     }
   }
diff --git a/icing/schema/schema-store.h b/icing/schema/schema-store.h
index 8567a2a..b5d89ce 100644
--- a/icing/schema/schema-store.h
+++ b/icing/schema/schema-store.h
@@ -253,6 +253,8 @@
 
   static constexpr std::string_view kSchemaTypeWildcard = "*";
 
+  static constexpr std::string_view kDefaultEmptySchemaDatabase = "";
+
   // Factory function to create a SchemaStore which 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.
@@ -323,9 +325,22 @@
   // schema or schema with types from multiple databases. Compatibility rules
   // defined by SchemaUtil::ComputeCompatibilityDelta.
   //
-  // The schema types in the new schema proto must all be from a single
-  // database. Does not support setting schema types across multiple databases
-  // at once.
+  // NOTE: This method is deprecated. Please use
+  // `SetSchema(SetSchemaRequestProto&& set_schema_request)` instead.
+  //
+  // TODO: b/337913932 - Remove this method once all callers (currently only
+  // used in tests) are migrated to the new SetSchema method.
+  libtextclassifier3::StatusOr<SetSchemaResult> SetSchema(
+      SchemaProto new_schema, bool ignore_errors_and_delete_documents);
+
+  // Update our current schema if it's compatible. Does not accept incompatible
+  // schema or schema with types from multiple databases. Compatibility rules
+  // defined by SchemaUtil::ComputeCompatibilityDelta.
+  //
+  // Does not support setting the schema across multiple databases if
+  // `feature_flags_->enable_schema_database()` is true. This means that:
+  // - All types within the new schema must have their `database` field matching
+  //  `set_schema_request.database()`.
   //
   // If ignore_errors_and_delete_documents is set to true, then incompatible
   // schema are allowed and we'll force set the schema, meaning
@@ -337,12 +352,12 @@
   //   - INTERNAL_ERROR on any IO errors
   //   - ALREADY_EXISTS_ERROR if type names in the new schema are already in use
   //     by a different database.
-  //   - INVALID_ARGUMENT_ERROR if the schema is invalid, or if the schema types
-  //     are from multiple databases (once schema database is enabled).
+  //   - INVALID_ARGUMENT_ERROR if the schema is invalid. This can happen if
+  //     the schema is malformed, if the new schema contains types where the
+  //     database field does not match the database field in the
+  //     set_schema_request.
   libtextclassifier3::StatusOr<SetSchemaResult> SetSchema(
-      const SchemaProto& new_schema, bool ignore_errors_and_delete_documents);
-  libtextclassifier3::StatusOr<SetSchemaResult> SetSchema(
-      SchemaProto&& new_schema, bool ignore_errors_and_delete_documents);
+      SetSchemaRequestProto&& set_schema_request);
 
   // Get the SchemaTypeConfigProto of schema_type name.
   //
@@ -725,19 +740,18 @@
   // Sets the schema for a database for the first time.
   //
   // Note that when schema database is disabled, this function sets the entire
-  // schema, with all under the default empty database.
+  // schema, with all types under the default empty database.
   //
   // Requires:
-  //   - All types in new_schema are from the same database.
-  //   - new_schema does not contain type names that are already in use by a
-  //     different database.
+  //   - `new_schema` is valid according to `ValidateSchemaDatabase'
   //
   // Returns:
   //   - SetSchemaResult that indicates if the new schema can be set.
   //   - INTERNAL_ERROR on any IO errors.
-  //   - INVALID_ARGUMENT_ERROR if the schema is invalid.
+  //   - INVALID_ARGUMENT_ERROR if the new schema is invalid.
   libtextclassifier3::StatusOr<SchemaStore::SetSchemaResult>
   SetInitialSchemaForDatabase(SchemaProto new_schema,
+                              const std::string& database,
                               bool ignore_errors_and_delete_documents);
 
   // Sets the schema for a database, overriding any existing schema for that
@@ -747,18 +761,24 @@
   // overrides the entire schema.
   //
   // Requires:
-  //   - All types in new_schema are from the same database.
-  //   - new_schema does not contain type names that are already in use by a
-  //     different database.
+  //   - `new_schema` and `database` are valid according to
+  //     `ValidateSchemaDatabase(new_schema, database)`
+  //   - Types in `new_schema` and `old_schema` all belong to the provided
+  //     database.
+  //     - The old schema is guaranteed to contain types from exactly one
+  //       database when schema database is enabled, because it was obtained
+  //       using `GetSchema(database)`.
   //
   // Returns:
   //   - SetSchemaResult that encapsulates the differences between the old and
   //     new schema, as well as if the new schema can be set.
   //   - INTERNAL_ERROR on any IO errors.
-  //   - INVALID_ARGUMENT_ERROR if the schema is invalid.
+  //   - INVALID_ARGUMENT_ERROR if the schema is invalid, or if there are
+  //     mismatches between the schema databases.
   libtextclassifier3::StatusOr<SchemaStore::SetSchemaResult>
   SetSchemaWithDatabaseOverride(SchemaProto new_schema,
                                 const SchemaProto& old_schema,
+                                const std::string& database,
                                 bool ignore_errors_and_delete_documents);
 
   // Initial validation on the SchemaProto for SetSchema. This is intended as a
@@ -769,22 +789,25 @@
   // an empty string is returned as the database.
   //
   // Checks that:
-  // - The new schema only contains types from a single database.
+  // - The new schema only contains types from a single database, which matches
+  //   the provided database.
   // - The schema's type names are not already in use in other databases. This
   //   is done outside of `SchemaUtil::Validate` because we need to know all
   //   existing type names, which is stored in the SchemaStore and not known to
   //   SchemaUtil.
   //
   // Returns:
-  //   - new_schema's database on success
-  //   - INVALID_ARGUMENT_ERROR if new_schema contains types from multiple
-  //     databases
+  //   - OK on success
+  //   - INVALID_ARGUMENT_ERROR if new_schema.types's databases do not match the
+  //     provided database.
   //   - ALREADY_EXISTS_ERROR if new_schema's types names are not unique
-  libtextclassifier3::StatusOr<std::string> ValidateAndGetDatabase(
-      const SchemaProto& new_schema) const;
+  libtextclassifier3::Status ValidateSchemaDatabase(
+      const SchemaProto& new_schema, const std::string& database) const;
 
   // Returns a SchemaProto representing the full schema, which is a combination
-  // of the existing schema and the input database schema.
+  // of the existing schema and the input database schema. Deletes all types
+  // belonging to the specified database if input_database_schema is an empty
+  // proto.
   //
   // For the database being updated by the input database schema:
   // - If the existing schema does not contain the database, the input types
@@ -800,7 +823,8 @@
   //     existing types from unaffected databases.
   //
   // Requires:
-  //   - input_database_schema must not contain types from multiple databases.
+  //   - input_database_schema is valid according to `ValidateSchemaDatabase`
+  //     and `SchemaUtil::Validate`.
   //
   // Returns:
   //   - SchemaProto on success
@@ -809,7 +833,8 @@
   //   - INVALID_ARGUMENT_ERROR if the input schema contains types from multiple
   //     databases.
   libtextclassifier3::StatusOr<SchemaProto> GetFullSchemaProtoWithUpdatedDb(
-      SchemaProto input_database_schema) const;
+      SchemaProto input_database_schema,
+      const std::string& database_to_update) const;
 
   const Filesystem* filesystem_;
   std::string base_dir_;
diff --git a/icing/schema/schema-store_test.cc b/icing/schema/schema-store_test.cc
index 7f74a8a..4993408 100644
--- a/icing/schema/schema-store_test.cc
+++ b/icing/schema/schema-store_test.cc
@@ -18,6 +18,7 @@
 #include <memory>
 #include <optional>
 #include <string>
+#include <string_view>
 #include <utility>
 #include <vector>
 
@@ -116,6 +117,19 @@
   FakeClock fake_clock_;
 };
 
+SetSchemaRequestProto CreateSetSchemaRequestProto(
+    SchemaProto schema, std::string database,
+    bool ignore_errors_and_delete_documents) {
+  SetSchemaRequestProto set_schema_request;
+
+  *set_schema_request.mutable_schema() = std::move(schema);
+  set_schema_request.set_database(std::move(database));
+  set_schema_request.set_ignore_errors_and_delete_documents(
+      ignore_errors_and_delete_documents);
+
+  return set_schema_request;
+}
+
 TEST_F(SchemaStoreTest, CreationWithFilesystemNullPointerShouldFail) {
   EXPECT_THAT(SchemaStore::Create(/*filesystem=*/nullptr, schema_store_dir_,
                                   &fake_clock_, feature_flags_.get()),
@@ -618,8 +632,9 @@
   result.success = true;
   result.schema_types_new_by_name.insert("db1_email");
   result.schema_types_new_by_name.insert("db1_message");
-  EXPECT_THAT(schema_store->SetSchema(
-                  db1_schema, /*ignore_errors_and_delete_documents=*/false),
+  EXPECT_THAT(schema_store->SetSchema(CreateSetSchemaRequestProto(
+                  db1_schema, /*database=*/"db1",
+                  /*ignore_errors_and_delete_documents=*/false)),
               IsOkAndHolds(EqualsSetSchemaResult(result)));
   EXPECT_THAT(schema_store->GetSchema(),
               IsOkAndHolds(Pointee(EqualsProto(db1_schema))));
@@ -639,8 +654,9 @@
   result.success = true;
   result.schema_types_new_by_name.insert("db2_email");
   result.schema_types_new_by_name.insert("db2_message");
-  EXPECT_THAT(schema_store->SetSchema(
-                  db2_schema, /*ignore_errors_and_delete_documents=*/false),
+  EXPECT_THAT(schema_store->SetSchema(CreateSetSchemaRequestProto(
+                  db2_schema, /*database=*/"db2",
+                  /*ignore_errors_and_delete_documents=*/false)),
               IsOkAndHolds(EqualsSetSchemaResult(result)));
 
   // Check the full schema. Databases that are updated last are appended to the
@@ -666,6 +682,56 @@
               IsOkAndHolds(EqualsProto(db2_schema)));
 }
 
+TEST_F(SchemaStoreTest, SetEmptyDatabaseSchemaOk) {
+  ICING_ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<SchemaStore> schema_store,
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_,
+                          feature_flags_.get()));
+
+  SchemaProto schema =
+      SchemaBuilder()
+          .AddType(SchemaTypeConfigBuilder().SetType("email"))
+          .AddType(SchemaTypeConfigBuilder().SetType("message"))
+          .Build();
+  SchemaStore::SetSchemaResult result;
+  result.success = true;
+  result.schema_types_new_by_name.insert("email");
+  result.schema_types_new_by_name.insert("message");
+  EXPECT_THAT(schema_store->SetSchema(CreateSetSchemaRequestProto(
+                  schema, /*database=*/"",
+                  /*ignore_errors_and_delete_documents=*/false)),
+              IsOkAndHolds(EqualsSetSchemaResult(result)));
+  EXPECT_THAT(schema_store->GetSchema(),
+              IsOkAndHolds(Pointee(EqualsProto(schema))));
+  EXPECT_THAT(schema_store->GetSchema(""), IsOkAndHolds(EqualsProto(schema)));
+
+  // Reset the schema. This should still reset the empty schema, and replace
+  // the existing 2 types.
+  schema =
+      SchemaBuilder()
+          .AddType(
+              SchemaTypeConfigBuilder().SetType("email_v2").SetDatabase(""))
+          .AddType(
+              SchemaTypeConfigBuilder().SetType("message_v2").SetDatabase(""))
+          .Build();
+  result = SchemaStore::SetSchemaResult();
+  result.success = true;
+  result.schema_types_new_by_name.insert("email_v2");
+  result.schema_types_new_by_name.insert("message_v2");
+  result.schema_types_deleted_by_name.insert("email");
+  result.schema_types_deleted_by_name.insert("message");
+  result.schema_types_deleted_by_id.insert(0);
+  result.schema_types_deleted_by_id.insert(1);
+
+  EXPECT_THAT(schema_store->SetSchema(CreateSetSchemaRequestProto(
+                  schema, /*database=*/"",
+                  /*ignore_errors_and_delete_documents=*/true)),
+              IsOkAndHolds(EqualsSetSchemaResult(result)));
+  EXPECT_THAT(schema_store->GetSchema(),
+              IsOkAndHolds(Pointee(EqualsProto(schema))));
+  EXPECT_THAT(schema_store->GetSchema(""), IsOkAndHolds(EqualsProto(schema)));
+}
+
 TEST_F(SchemaStoreTest, SetSameSchemaOk) {
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<SchemaStore> schema_store,
@@ -734,15 +800,17 @@
   result.success = true;
   result.schema_types_new_by_name.insert("db1_email");
   result.schema_types_new_by_name.insert("db1_message");
-  EXPECT_THAT(schema_store->SetSchema(
-                  db1_schema, /*ignore_errors_and_delete_documents=*/false),
+  EXPECT_THAT(schema_store->SetSchema(CreateSetSchemaRequestProto(
+                  db1_schema, /*database=*/"db1",
+                  /*ignore_errors_and_delete_documents=*/false)),
               IsOkAndHolds(EqualsSetSchemaResult(result)));
   result = SchemaStore::SetSchemaResult();
   result.success = true;
   result.schema_types_new_by_name.insert("db2_email");
   result.schema_types_new_by_name.insert("db2_message");
-  EXPECT_THAT(schema_store->SetSchema(
-                  db2_schema, /*ignore_errors_and_delete_documents=*/false),
+  EXPECT_THAT(schema_store->SetSchema(CreateSetSchemaRequestProto(
+                  db2_schema, /*database=*/"db2",
+                  /*ignore_errors_and_delete_documents=*/false)),
               IsOkAndHolds(EqualsSetSchemaResult(result)));
   ICING_ASSERT_OK_AND_ASSIGN(const SchemaProto* actual_full_schema,
                              schema_store->GetSchema());
@@ -751,10 +819,10 @@
   // Reset db1 with the same SchemaProto. The schema should be exactly the same.
   result = SchemaStore::SetSchemaResult();
   result.success = true;
-  EXPECT_THAT(
-      schema_store->SetSchema(db1_schema,
-                              /*ignore_errors_and_delete_documents=*/false),
-      IsOkAndHolds(EqualsSetSchemaResult(result)));
+  EXPECT_THAT(schema_store->SetSchema(CreateSetSchemaRequestProto(
+                  db1_schema, /*database=*/"db1",
+                  /*ignore_errors_and_delete_documents=*/false)),
+              IsOkAndHolds(EqualsSetSchemaResult(result)));
 
   // Check the schema, this should not have changed
   EXPECT_THAT(schema_store->GetSchema(),
@@ -821,24 +889,27 @@
   result.success = true;
   result.schema_types_new_by_name.insert("db1_email");
   result.schema_types_new_by_name.insert("db1_message");
-  EXPECT_THAT(schema_store->SetSchema(
-                  db1_schema, /*ignore_errors_and_delete_documents=*/false),
+  EXPECT_THAT(schema_store->SetSchema(CreateSetSchemaRequestProto(
+                  db1_schema, /*database=*/"db1",
+                  /*ignore_errors_and_delete_documents=*/false)),
               IsOkAndHolds(EqualsSetSchemaResult(result)));
   // Set schema for db2
   result = SchemaStore::SetSchemaResult();
   result.success = true;
   result.schema_types_new_by_name.insert("db2_email");
   result.schema_types_new_by_name.insert("db2_message");
-  EXPECT_THAT(schema_store->SetSchema(
-                  db2_schema, /*ignore_errors_and_delete_documents=*/false),
+  EXPECT_THAT(schema_store->SetSchema(CreateSetSchemaRequestProto(
+                  db2_schema, /*database=*/"db2",
+                  /*ignore_errors_and_delete_documents=*/false)),
               IsOkAndHolds(EqualsSetSchemaResult(result)));
   // Set schema for db3
   result = SchemaStore::SetSchemaResult();
   result.success = true;
   result.schema_types_new_by_name.insert("db3_email");
   result.schema_types_new_by_name.insert("db3_message");
-  EXPECT_THAT(schema_store->SetSchema(
-                  db3_schema, /*ignore_errors_and_delete_documents=*/false),
+  EXPECT_THAT(schema_store->SetSchema(CreateSetSchemaRequestProto(
+                  db3_schema, /*database=*/"db3",
+                  /*ignore_errors_and_delete_documents=*/false)),
               IsOkAndHolds(EqualsSetSchemaResult(result)));
   // Verify schema.
   ICING_ASSERT_OK_AND_ASSIGN(const SchemaProto* actual_full_schema,
@@ -859,8 +930,9 @@
   result.success = true;
 
   libtextclassifier3::StatusOr<SchemaStore::SetSchemaResult> actual_result =
-      schema_store->SetSchema(reordered_db2_schema,
-                              /*ignore_errors_and_delete_documents=*/false);
+      schema_store->SetSchema(CreateSetSchemaRequestProto(
+          reordered_db2_schema, /*database=*/"db2",
+          /*ignore_errors_and_delete_documents=*/false));
   EXPECT_THAT(actual_result, IsOkAndHolds(EqualsSetSchemaResult(result)));
   EXPECT_THAT(actual_result.ValueOrDie().old_schema_type_ids_changed,
               IsEmpty());
@@ -937,24 +1009,27 @@
   result.success = true;
   result.schema_types_new_by_name.insert("db1_email");
   result.schema_types_new_by_name.insert("db1_message");
-  EXPECT_THAT(schema_store->SetSchema(
-                  db1_schema, /*ignore_errors_and_delete_documents=*/false),
+  EXPECT_THAT(schema_store->SetSchema(CreateSetSchemaRequestProto(
+                  db1_schema, /*database=*/"db1",
+                  /*ignore_errors_and_delete_documents=*/false)),
               IsOkAndHolds(EqualsSetSchemaResult(result)));
   // Set schema for db2
   result = SchemaStore::SetSchemaResult();
   result.success = true;
   result.schema_types_new_by_name.insert("db2_email");
   result.schema_types_new_by_name.insert("db2_message");
-  EXPECT_THAT(schema_store->SetSchema(
-                  db2_schema, /*ignore_errors_and_delete_documents=*/false),
+  EXPECT_THAT(schema_store->SetSchema(CreateSetSchemaRequestProto(
+                  db2_schema, /*database=*/"db2",
+                  /*ignore_errors_and_delete_documents=*/false)),
               IsOkAndHolds(EqualsSetSchemaResult(result)));
   // Set schema for db3
   result = SchemaStore::SetSchemaResult();
   result.success = true;
   result.schema_types_new_by_name.insert("db3_email");
   result.schema_types_new_by_name.insert("db3_message");
-  EXPECT_THAT(schema_store->SetSchema(
-                  db3_schema, /*ignore_errors_and_delete_documents=*/false),
+  EXPECT_THAT(schema_store->SetSchema(CreateSetSchemaRequestProto(
+                  db3_schema, /*database=*/"db3",
+                  /*ignore_errors_and_delete_documents=*/false)),
               IsOkAndHolds(EqualsSetSchemaResult(result)));
   // Verify schema.
   ICING_ASSERT_OK_AND_ASSIGN(const SchemaProto* actual_full_schema,
@@ -1006,10 +1081,10 @@
   result = SchemaStore::SetSchemaResult();
   result.success = true;
   result.schema_types_new_by_name.insert("db2_recipient");
-  EXPECT_THAT(
-      schema_store->SetSchema(db2_schema,
-                              /*ignore_errors_and_delete_documents=*/false),
-      IsOkAndHolds(EqualsSetSchemaResult(result)));
+  EXPECT_THAT(schema_store->SetSchema(CreateSetSchemaRequestProto(
+                  db2_schema, /*database=*/"db2",
+                  /*ignore_errors_and_delete_documents=*/false)),
+              IsOkAndHolds(EqualsSetSchemaResult(result)));
 
   // Check the schema
   EXPECT_THAT(schema_store->GetSchema(),
@@ -1060,24 +1135,27 @@
   result.success = true;
   result.schema_types_new_by_name.insert("db1_email");
   result.schema_types_new_by_name.insert("db1_message");
-  EXPECT_THAT(schema_store->SetSchema(
-                  db1_schema, /*ignore_errors_and_delete_documents=*/false),
+  EXPECT_THAT(schema_store->SetSchema(CreateSetSchemaRequestProto(
+                  db1_schema, /*database=*/"db1",
+                  /*ignore_errors_and_delete_documents=*/false)),
               IsOkAndHolds(EqualsSetSchemaResult(result)));
   // Set schema for db2
   result = SchemaStore::SetSchemaResult();
   result.success = true;
   result.schema_types_new_by_name.insert("db2_email");
   result.schema_types_new_by_name.insert("db2_message");
-  EXPECT_THAT(schema_store->SetSchema(
-                  db2_schema, /*ignore_errors_and_delete_documents=*/false),
+  EXPECT_THAT(schema_store->SetSchema(CreateSetSchemaRequestProto(
+                  db2_schema, /*database=*/"db2",
+                  /*ignore_errors_and_delete_documents=*/false)),
               IsOkAndHolds(EqualsSetSchemaResult(result)));
   // Set schema for db3
   result = SchemaStore::SetSchemaResult();
   result.success = true;
   result.schema_types_new_by_name.insert("db3_email");
   result.schema_types_new_by_name.insert("db3_message");
-  EXPECT_THAT(schema_store->SetSchema(
-                  db3_schema, /*ignore_errors_and_delete_documents=*/false),
+  EXPECT_THAT(schema_store->SetSchema(CreateSetSchemaRequestProto(
+                  db3_schema, /*database=*/"db3",
+                  /*ignore_errors_and_delete_documents=*/false)),
               IsOkAndHolds(EqualsSetSchemaResult(result)));
   // Set schema again for db2 and add a type. The added type should be appended
   // to the end of the SchemaProto.
@@ -1095,8 +1173,9 @@
   result = SchemaStore::SetSchemaResult();
   result.success = true;
   result.schema_types_new_by_name.insert("db2_recipient");
-  EXPECT_THAT(schema_store->SetSchema(
-                  db2_schema, /*ignore_errors_and_delete_documents=*/false),
+  EXPECT_THAT(schema_store->SetSchema(CreateSetSchemaRequestProto(
+                  db2_schema, /*database=*/"db2",
+                  /*ignore_errors_and_delete_documents=*/false)),
               IsOkAndHolds(EqualsSetSchemaResult(result)));
   SchemaProto expected_full_schema =
       SchemaBuilder()
@@ -1161,10 +1240,10 @@
   result.old_schema_type_ids_changed.insert(3);  // db2_message
   result.old_schema_type_ids_changed.insert(4);  // db3_email
   result.old_schema_type_ids_changed.insert(5);  // db3_message
-  EXPECT_THAT(
-      schema_store->SetSchema(db2_schema,
-                              /*ignore_errors_and_delete_documents=*/true),
-      IsOkAndHolds(EqualsSetSchemaResult(result)));
+  EXPECT_THAT(schema_store->SetSchema(CreateSetSchemaRequestProto(
+                  db2_schema, /*database=*/"db2",
+                  /*ignore_errors_and_delete_documents=*/true)),
+              IsOkAndHolds(EqualsSetSchemaResult(result)));
 
   // Check the schema
   EXPECT_THAT(schema_store->GetSchema(),
@@ -1177,13 +1256,14 @@
               IsOkAndHolds(EqualsProto(db3_schema)));
 }
 
-TEST_F(SchemaStoreTest, SetEmptySchemaInDifferentDatabaseOk) {
+TEST_F(SchemaStoreTest, SetEmptySchemaClearsDatabase) {
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<SchemaStore> schema_store,
       SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_,
                           feature_flags_.get(),
                           /*initialize_stats=*/nullptr));
 
+  // Set schema for the first time
   SchemaProto db1_schema =
       SchemaBuilder()
           .AddType(
@@ -1192,35 +1272,117 @@
                        .SetType("db1_message")
                        .SetDatabase("db1"))
           .Build();
+  SchemaProto db2_schema =
+      SchemaBuilder()
+          .AddType(
+              SchemaTypeConfigBuilder().SetType("db2_email").SetDatabase("db2"))
+          .AddType(SchemaTypeConfigBuilder()
+                       .SetType("db2_message")
+                       .SetDatabase("db2"))
+          .Build();
+  SchemaProto db3_schema =
+      SchemaBuilder()
+          .AddType(
+              SchemaTypeConfigBuilder().SetType("db3_email").SetDatabase("db3"))
+          .AddType(SchemaTypeConfigBuilder()
+                       .SetType("db3_message")
+                       .SetDatabase("db3"))
+          .Build();
+
+  // Set schema for db1
   SchemaStore::SetSchemaResult result;
   result.success = true;
   result.schema_types_new_by_name.insert("db1_email");
   result.schema_types_new_by_name.insert("db1_message");
-  EXPECT_THAT(schema_store->SetSchema(
-                  db1_schema, /*ignore_errors_and_delete_documents=*/false),
+  EXPECT_THAT(schema_store->SetSchema(CreateSetSchemaRequestProto(
+                  db1_schema, /*database=*/"db1",
+                  /*ignore_errors_and_delete_documents=*/false)),
               IsOkAndHolds(EqualsSetSchemaResult(result)));
-  EXPECT_THAT(schema_store->GetSchema(),
-              IsOkAndHolds(Pointee(EqualsProto(db1_schema))));
-  EXPECT_THAT(schema_store->GetSchema("db1"),
-              IsOkAndHolds(EqualsProto(db1_schema)));
-
-  // Set an empty schema in a different database
-  SchemaProto db2_schema;
+  // Set schema for db2
   result = SchemaStore::SetSchemaResult();
   result.success = true;
-  EXPECT_THAT(schema_store->SetSchema(
-                  db2_schema, /*ignore_errors_and_delete_documents=*/false),
+  result.schema_types_new_by_name.insert("db2_email");
+  result.schema_types_new_by_name.insert("db2_message");
+  EXPECT_THAT(schema_store->SetSchema(CreateSetSchemaRequestProto(
+                  db2_schema, /*database=*/"db2",
+                  /*ignore_errors_and_delete_documents=*/false)),
+              IsOkAndHolds(EqualsSetSchemaResult(result)));
+  // Set schema for db3
+  result = SchemaStore::SetSchemaResult();
+  result.success = true;
+  result.schema_types_new_by_name.insert("db3_email");
+  result.schema_types_new_by_name.insert("db3_message");
+  EXPECT_THAT(schema_store->SetSchema(CreateSetSchemaRequestProto(
+                  db3_schema, /*database=*/"db3",
+                  /*ignore_errors_and_delete_documents=*/false)),
+              IsOkAndHolds(EqualsSetSchemaResult(result)));
+  // Verify schema.
+  SchemaProto expected_full_schema =
+      SchemaBuilder()
+          .AddType(SchemaTypeConfigBuilder()
+                       .SetType("db1_email")  // SchemaTypeId 0
+                       .SetDatabase("db1"))
+          .AddType(SchemaTypeConfigBuilder()
+                       .SetType("db1_message")  // SchemaTypeId 1
+                       .SetDatabase("db1"))
+          .AddType(SchemaTypeConfigBuilder()
+                       .SetType("db2_email")  // SchemaTypeId 2
+                       .SetDatabase("db2"))
+          .AddType(SchemaTypeConfigBuilder()
+                       .SetType("db2_message")  // SchemaTypeId 3
+                       .SetDatabase("db2"))
+          .AddType(SchemaTypeConfigBuilder()
+                       .SetType("db3_email")  // SchemaTypeId 4
+                       .SetDatabase("db3"))
+          .AddType(SchemaTypeConfigBuilder()
+                       .SetType("db3_message")  // SchemaTypeId 5
+                       .SetDatabase("db3"))
+          .Build();
+  EXPECT_THAT(schema_store->GetSchema(),
+              IsOkAndHolds(Pointee(EqualsProto(expected_full_schema))));
+
+  // Set an empty schema for db2. This deletes all types from db2, and changes
+  // the type ids of types from db3 because they appear after db2 in the
+  // original schema.
+  db2_schema = SchemaProto();
+  result = SchemaStore::SetSchemaResult();
+  result.success = true;
+  result.schema_types_deleted_by_name.insert("db2_email");
+  result.schema_types_deleted_by_name.insert("db2_message");
+  result.schema_types_deleted_by_id.insert(2);   // db2_email
+  result.schema_types_deleted_by_id.insert(3);   // db2_message
+  result.old_schema_type_ids_changed.insert(4);  // db3_email
+  result.old_schema_type_ids_changed.insert(5);  // db3_message
+  EXPECT_THAT(schema_store->SetSchema(CreateSetSchemaRequestProto(
+                  db2_schema, /*database=*/"db2",
+                  /*ignore_errors_and_delete_documents=*/true)),
               IsOkAndHolds(EqualsSetSchemaResult(result)));
 
-  // Check the schema, this should not have changed
-  EXPECT_THAT(schema_store->GetSchema(),
-              IsOkAndHolds(Pointee(EqualsProto(db1_schema))));
+  // Check the schema. Schemas for db1 and db3 should be unchanged.
   EXPECT_THAT(schema_store->GetSchema("db1"),
               IsOkAndHolds(EqualsProto(db1_schema)));
+  EXPECT_THAT(schema_store->GetSchema("db3"),
+              IsOkAndHolds(EqualsProto(db3_schema)));
 
-  // GetSchema for an empty database should return NotFoundError
+  // GetSchema for db2 should return NotFoundError
   EXPECT_THAT(schema_store->GetSchema("db2"),
               StatusIs(libtextclassifier3::StatusCode::NOT_FOUND));
+
+  expected_full_schema =
+      SchemaBuilder()
+          .AddType(
+              SchemaTypeConfigBuilder().SetType("db1_email").SetDatabase("db1"))
+          .AddType(SchemaTypeConfigBuilder()
+                       .SetType("db1_message")
+                       .SetDatabase("db1"))
+          .AddType(
+              SchemaTypeConfigBuilder().SetType("db3_email").SetDatabase("db3"))
+          .AddType(SchemaTypeConfigBuilder()
+                       .SetType("db3_message")
+                       .SetDatabase("db3"))
+          .Build();
+  EXPECT_THAT(schema_store->GetSchema(),
+              IsOkAndHolds(Pointee(EqualsProto(expected_full_schema))));
 }
 
 TEST_F(SchemaStoreTest, SetIncompatibleSchemaOk) {
@@ -1294,15 +1456,17 @@
   result.success = true;
   result.schema_types_new_by_name.insert("db1_email");
   result.schema_types_new_by_name.insert("db1_message");
-  EXPECT_THAT(schema_store->SetSchema(
-                  db1_schema, /*ignore_errors_and_delete_documents=*/false),
+  EXPECT_THAT(schema_store->SetSchema(CreateSetSchemaRequestProto(
+                  db1_schema, /*database=*/"db1",
+                  /*ignore_errors_and_delete_documents=*/false)),
               IsOkAndHolds(EqualsSetSchemaResult(result)));
   result = SchemaStore::SetSchemaResult();
   result.success = true;
   result.schema_types_new_by_name.insert("db2_email");
   result.schema_types_new_by_name.insert("db2_message");
-  EXPECT_THAT(schema_store->SetSchema(
-                  db2_schema, /*ignore_errors_and_delete_documents=*/false),
+  EXPECT_THAT(schema_store->SetSchema(CreateSetSchemaRequestProto(
+                  db2_schema, /*database=*/"db2",
+                  /*ignore_errors_and_delete_documents=*/false)),
               IsOkAndHolds(EqualsSetSchemaResult(result)));
   ICING_ASSERT_OK_AND_ASSIGN(const SchemaProto* actual_full_schema,
                              schema_store->GetSchema());
@@ -1322,10 +1486,10 @@
   result.schema_types_deleted_by_name.insert("db2_message");
   result.schema_types_new_by_name.insert("db2_recipient");
   result.schema_types_deleted_by_id.insert(3);  // db2_message
-  EXPECT_THAT(
-      schema_store->SetSchema(db2_schema_incompatible,
-                              /*ignore_errors_and_delete_documents=*/false),
-      IsOkAndHolds(EqualsSetSchemaResult(result)));
+  EXPECT_THAT(schema_store->SetSchema(CreateSetSchemaRequestProto(
+                  db2_schema_incompatible, /*database=*/"db2",
+                  /*ignore_errors_and_delete_documents=*/false)),
+              IsOkAndHolds(EqualsSetSchemaResult(result)));
 
   // Check the schema, this should not have changed
   EXPECT_THAT(schema_store->GetSchema(),
@@ -1377,15 +1541,17 @@
   result.success = true;
   result.schema_types_new_by_name.insert("db1_email");
   result.schema_types_new_by_name.insert("db1_message");
-  EXPECT_THAT(schema_store->SetSchema(
-                  db1_schema, /*ignore_errors_and_delete_documents=*/false),
+  EXPECT_THAT(schema_store->SetSchema(CreateSetSchemaRequestProto(
+                  db1_schema, /*database=*/"db1",
+                  /*ignore_errors_and_delete_documents=*/false)),
               IsOkAndHolds(EqualsSetSchemaResult(result)));
   result = SchemaStore::SetSchemaResult();
   result.success = true;
   result.schema_types_new_by_name.insert("db2_email");
   result.schema_types_new_by_name.insert("db2_message");
-  EXPECT_THAT(schema_store->SetSchema(
-                  db2_schema, /*ignore_errors_and_delete_documents=*/false),
+  EXPECT_THAT(schema_store->SetSchema(CreateSetSchemaRequestProto(
+                  db2_schema, /*database=*/"db2",
+                  /*ignore_errors_and_delete_documents=*/false)),
               IsOkAndHolds(EqualsSetSchemaResult(result)));
   ICING_ASSERT_OK_AND_ASSIGN(const SchemaProto* actual_full_schema,
                              schema_store->GetSchema());
@@ -1405,10 +1571,11 @@
                                                          .AddProperty(prop)
                                                          .AddProperty(prop))
                                             .Build();
-  EXPECT_THAT(
-      schema_store->SetSchema(db2_schema_incompatible,
-                              /*ignore_errors_and_delete_documents=*/false),
-      StatusIs(libtextclassifier3::StatusCode::ALREADY_EXISTS));
+  EXPECT_THAT(schema_store->SetSchema(CreateSetSchemaRequestProto(
+                  db2_schema_incompatible,
+                  /*database=*/"db2",
+                  /*ignore_errors_and_delete_documents=*/false)),
+              StatusIs(libtextclassifier3::StatusCode::ALREADY_EXISTS));
 
   // Check the schema, this should not have changed
   EXPECT_THAT(schema_store->GetSchema(),
@@ -1426,7 +1593,6 @@
                           feature_flags_.get(),
                           /*initialize_stats=*/nullptr));
 
-  // Set schema for the first time
   SchemaProto combined_schema =
       SchemaBuilder()
           .AddType(
@@ -1440,10 +1606,64 @@
                        .SetType("db1_message")
                        .SetDatabase("db1"))
           .Build();
-  EXPECT_THAT(
-      schema_store->SetSchema(combined_schema,
-                              /*ignore_errors_and_delete_documents=*/false),
-      StatusIs(libtextclassifier3::StatusCode::INVALID_ARGUMENT));
+  EXPECT_THAT(schema_store->SetSchema(CreateSetSchemaRequestProto(
+                  combined_schema, /*database=*/"db1",
+                  /*ignore_errors_and_delete_documents=*/false)),
+              StatusIs(libtextclassifier3::StatusCode::INVALID_ARGUMENT));
+}
+
+TEST_F(SchemaStoreTest, SetSchemaWithMismatchedDbFails) {
+  ICING_ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<SchemaStore> schema_store,
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_,
+                          feature_flags_.get(),
+                          /*initialize_stats=*/nullptr));
+
+  SchemaProto schema =
+      SchemaBuilder()
+          // This type does not explicitly set its database, so it defaults to
+          // the empty database.
+          .AddType(SchemaTypeConfigBuilder().SetType("db1_email"))
+          .AddType(SchemaTypeConfigBuilder()
+                       .SetType("db1_message")
+                       .SetDatabase("db1"))
+          .Build();
+
+  EXPECT_THAT(schema_store->SetSchema(CreateSetSchemaRequestProto(
+                  schema, /*database=*/"db1",
+                  /*ignore_errors_and_delete_documents=*/false)),
+              StatusIs(libtextclassifier3::StatusCode::INVALID_ARGUMENT));
+
+  schema =
+      SchemaBuilder()
+          .AddType(
+              SchemaTypeConfigBuilder().SetType("db1_email").SetDatabase("db1"))
+          .AddType(SchemaTypeConfigBuilder()
+                       .SetType("db1_message")
+                       .SetDatabase("db1"))
+          .Build();
+  EXPECT_THAT(schema_store->SetSchema(CreateSetSchemaRequestProto(
+                  schema, /*database=*/"db_mismatch",
+                  /*ignore_errors_and_delete_documents=*/false)),
+              StatusIs(libtextclassifier3::StatusCode::INVALID_ARGUMENT));
+
+  schema =
+      SchemaBuilder()
+          .AddType(
+              SchemaTypeConfigBuilder().SetType("db1_email").SetDatabase("db1"))
+          .AddType(SchemaTypeConfigBuilder()
+                       .SetType("db1_message")
+                       .SetDatabase("db1"))
+          .AddType(
+              SchemaTypeConfigBuilder().SetType("db2_email").SetDatabase("db2"))
+          .AddType(SchemaTypeConfigBuilder()
+                       .SetType("db2_message")
+                       .SetDatabase("db2"))
+          .Build();
+  EXPECT_THAT(schema_store->SetSchema(CreateSetSchemaRequestProto(
+                  schema, /*database=*/"",
+                  /*ignore_errors_and_delete_documents=*/false)),
+              StatusIs(libtextclassifier3::StatusCode::INVALID_ARGUMENT));
 }
 
 TEST_F(SchemaStoreTest, SetSchemaWithDuplicateTypeNameAcrossDifferentDbFails) {
@@ -1466,8 +1686,9 @@
   result.success = true;
   result.schema_types_new_by_name.insert("email");
   result.schema_types_new_by_name.insert("db1_message");
-  EXPECT_THAT(schema_store->SetSchema(
-                  db1_schema, /*ignore_errors_and_delete_documents=*/false),
+  EXPECT_THAT(schema_store->SetSchema(CreateSetSchemaRequestProto(
+                  db1_schema, /*database=*/"db1",
+                  /*ignore_errors_and_delete_documents=*/false)),
               IsOkAndHolds(EqualsSetSchemaResult(result)));
   EXPECT_THAT(schema_store->GetSchema(),
               IsOkAndHolds(Pointee(EqualsProto(db1_schema))));
@@ -1483,8 +1704,9 @@
                        .SetType("db2_message")
                        .SetDatabase("db2"))
           .Build();
-  EXPECT_THAT(schema_store->SetSchema(
-                  db2_schema, /*ignore_errors_and_delete_documents=*/false),
+  EXPECT_THAT(schema_store->SetSchema(CreateSetSchemaRequestProto(
+                  db2_schema, /*database=*/"db2",
+                  /*ignore_errors_and_delete_documents=*/false)),
               StatusIs(libtextclassifier3::StatusCode::ALREADY_EXISTS));
 
   // Check schema, this should not have changed
diff --git a/icing/tokenization/combined-tokenizer_test.cc b/icing/tokenization/combined-tokenizer_test.cc
index acbe306..520ea47 100644
--- a/icing/tokenization/combined-tokenizer_test.cc
+++ b/icing/tokenization/combined-tokenizer_test.cc
@@ -282,8 +282,9 @@
       tokenizer_factory::CreateIndexingTokenizer(
           StringIndexingConfig::TokenizerType::PLAIN, lang_segmenter_.get()));
 
-  if (GetIcuTokenizationVersion() >= 72) {
-    // In ICU 72+ and above, ':' are no longer considered word connectors.
+  int icu_version = GetIcuTokenizationVersion();
+  if (icu_version >= 72 && icu_version < 77) {
+    // In ICU 72+ and before 77, ':' are not considered word connectors.
     constexpr std::string_view kText = "foo:bar";
     ICING_ASSERT_OK_AND_ASSIGN(std::vector<Token> indexing_tokens,
                                indexing_tokenizer->TokenizeAll(kText));
diff --git a/icing/tokenization/icu/icu-language-segmenter_test.cc b/icing/tokenization/icu/icu-language-segmenter_test.cc
index 664ccce..cde62f2 100644
--- a/icing/tokenization/icu/icu-language-segmenter_test.cc
+++ b/icing/tokenization/icu/icu-language-segmenter_test.cc
@@ -229,12 +229,19 @@
   //   2. '@' became a word connector
   //   3. <numeric><word-connector><numeric> such as "3'14" is now considered as
   //      a single token.
-  if (GetIcuTokenizationVersion() >= 72) {
-    EXPECT_THAT(
-        language_segmenter->GetAllTerms("com:google:android"),
-        IsOkAndHolds(ElementsAre("com", ":", "google", ":", "android")));
+  int icu_version = GetIcuTokenizationVersion();
+  if (icu_version >= 72) {
+    // In ICU 77, the rules for ':' were reverted.
+    if (icu_version >= 77) {
+      EXPECT_THAT(language_segmenter->GetAllTerms("com:google:android"),
+                  IsOkAndHolds(ElementsAre("com:google:android")));
+    } else {
+      EXPECT_THAT(
+          language_segmenter->GetAllTerms("com:google:android"),
+          IsOkAndHolds(ElementsAre("com", ":", "google", ":", "android")));
+    }
     // In ICU 74, the rules for '@' were reverted.
-    if (GetIcuTokenizationVersion() >= 74) {
+    if (icu_version >= 74) {
       EXPECT_THAT(
           language_segmenter->GetAllTerms("com@google@android"),
           IsOkAndHolds(ElementsAre("com", "@", "google", "@", "android")));
diff --git a/java/src/com/google/android/icing/IcingSearchEngine.java b/java/src/com/google/android/icing/IcingSearchEngine.java
index 930dc33..9856e09 100644
--- a/java/src/com/google/android/icing/IcingSearchEngine.java
+++ b/java/src/com/google/android/icing/IcingSearchEngine.java
@@ -45,6 +45,7 @@
 import com.google.android.icing.proto.ScoringSpecProto;
 import com.google.android.icing.proto.SearchResultProto;
 import com.google.android.icing.proto.SearchSpecProto;
+import com.google.android.icing.proto.SetSchemaRequestProto;
 import com.google.android.icing.proto.SetSchemaResultProto;
 import com.google.android.icing.proto.StorageInfoResultProto;
 import com.google.android.icing.proto.SuggestionResponse;
@@ -89,7 +90,7 @@
 
   @Override
   public @NonNull SetSchemaResultProto setSchema(@NonNull SchemaProto schema) {
-    return setSchema(schema, /*ignoreErrorsAndDeleteDocuments=*/ false);
+    return setSchema(schema, /* ignoreErrorsAndDeleteDocuments= */ false);
   }
 
   @Override
@@ -100,6 +101,13 @@
   }
 
   @Override
+  public @NonNull SetSchemaResultProto setSchemaWithRequestProto(
+      @NonNull SetSchemaRequestProto setSchemaRequest) {
+    return IcingSearchEngineUtils.byteArrayToSetSchemaResultProto(
+        icingSearchEngineImpl.setSchemaWithRequestProto(setSchemaRequest.toByteArray()));
+  }
+
+  @Override
   public @NonNull GetSchemaResultProto getSchema() {
     return IcingSearchEngineUtils.byteArrayToGetSchemaResultProto(
         icingSearchEngineImpl.getSchema());
@@ -222,7 +230,7 @@
 
   @Override
   public @NonNull DeleteByQueryResultProto deleteByQuery(@NonNull SearchSpecProto searchSpec) {
-    return deleteByQuery(searchSpec, /*returnDeletedDocumentInfo=*/ false);
+    return deleteByQuery(searchSpec, /* returnDeletedDocumentInfo= */ false);
   }
 
   @Override
diff --git a/java/src/com/google/android/icing/IcingSearchEngineImpl.java b/java/src/com/google/android/icing/IcingSearchEngineImpl.java
index 1596fc8..ad5a5aa 100644
--- a/java/src/com/google/android/icing/IcingSearchEngineImpl.java
+++ b/java/src/com/google/android/icing/IcingSearchEngineImpl.java
@@ -89,6 +89,12 @@
   }
 
   @Nullable
+  public byte[] setSchemaWithRequestProto(@NonNull byte[] setSchemaRequestBytes) {
+    throwIfClosed();
+    return nativeSetSchemaWithRequestProto(this, setSchemaRequestBytes);
+  }
+
+  @Nullable
   public byte[] getSchema() {
     throwIfClosed();
     return nativeGetSchema(this);
@@ -296,6 +302,9 @@
   private static native byte[] nativeSetSchema(
       IcingSearchEngineImpl instance, byte[] schemaBytes, boolean ignoreErrorsAndDeleteDocuments);
 
+  private static native byte[] nativeSetSchemaWithRequestProto(
+      IcingSearchEngineImpl instance, byte[] setSchemaRequestBytes);
+
   private static native byte[] nativeGetSchema(IcingSearchEngineImpl instance);
 
   private static native byte[] nativeGetSchemaForDatabase(
diff --git a/java/src/com/google/android/icing/IcingSearchEngineInterface.java b/java/src/com/google/android/icing/IcingSearchEngineInterface.java
index 71f2e80..c381d7c 100644
--- a/java/src/com/google/android/icing/IcingSearchEngineInterface.java
+++ b/java/src/com/google/android/icing/IcingSearchEngineInterface.java
@@ -29,6 +29,7 @@
 import com.google.android.icing.proto.ScoringSpecProto;
 import com.google.android.icing.proto.SearchResultProto;
 import com.google.android.icing.proto.SearchSpecProto;
+import com.google.android.icing.proto.SetSchemaRequestProto;
 import com.google.android.icing.proto.SetSchemaResultProto;
 import com.google.android.icing.proto.StorageInfoResultProto;
 import com.google.android.icing.proto.SuggestionResponse;
@@ -45,17 +46,32 @@
    */
   InitializeResultProto initialize();
 
-  /** Sets the schema for the icing instance. */
+  /**
+   * Sets the schema for the icing instance.
+   *
+   * <p>Note: This method is deprecated. Please use {@link
+   * #setSchemaWithRequestProto(SetSchemaRequestProto)} instead.
+   */
   SetSchemaResultProto setSchema(SchemaProto schema);
 
   /**
    * Sets the schema for the icing instance.
    *
+   * <p>Note: This method is deprecated. Please use {@link
+   * #setSchemaWithRequestProto(SetSchemaRequestProto)} instead.
+   *
    * @param ignoreErrorsAndDeleteDocuments force to set the schema and delete documents in case of
    *     incompatible schema change.
    */
   SetSchemaResultProto setSchema(SchemaProto schema, boolean ignoreErrorsAndDeleteDocuments);
 
+  /**
+   * Sets the schema for the icing instance.
+   *
+   * @param setSchemaRequest the request proto for setting the schema.
+   */
+  SetSchemaResultProto setSchemaWithRequestProto(SetSchemaRequestProto setSchemaRequest);
+
   /** Gets the schema for the icing instance. */
   GetSchemaResultProto getSchema();
 
diff --git a/java/tests/instrumentation/src/com/google/android/icing/IcingSearchEngineTest.java b/java/tests/instrumentation/src/com/google/android/icing/IcingSearchEngineTest.java
index 1f294c2..975699c 100644
--- a/java/tests/instrumentation/src/com/google/android/icing/IcingSearchEngineTest.java
+++ b/java/tests/instrumentation/src/com/google/android/icing/IcingSearchEngineTest.java
@@ -51,6 +51,7 @@
 import com.google.android.icing.proto.ScoringSpecProto;
 import com.google.android.icing.proto.SearchResultProto;
 import com.google.android.icing.proto.SearchSpecProto;
+import com.google.android.icing.proto.SetSchemaRequestProto;
 import com.google.android.icing.proto.SetSchemaResultProto;
 import com.google.android.icing.proto.SnippetMatchProto;
 import com.google.android.icing.proto.SnippetProto;
@@ -211,6 +212,9 @@
     assertThat(getSchemaTypeResultProto.getSchemaTypeConfig()).isEqualTo(emailTypeConfig);
   }
 
+  // TODO: b/383379132 - Re-enable this test once the JNI API is pre-registered and dropped back
+  // into g3.
+  @Ignore
   @Test
   public void setAndGetSchemaWithDatabase_ok() throws Exception {
     IcingSearchEngineOptions options =
@@ -228,11 +232,23 @@
     SchemaProto db2Schema =
         SchemaProto.newBuilder().addTypes(createEmailTypeConfigWithDatabase(db2)).build();
 
+    SetSchemaRequestProto requestProto1 =
+        SetSchemaRequestProto.newBuilder()
+            .setSchema(db1Schema)
+            .setDatabase(db1)
+            .setIgnoreErrorsAndDeleteDocuments(false)
+            .build();
     SetSchemaResultProto setSchemaResultProto =
-        icingSearchEngine.setSchema(db1Schema, /* ignoreErrorsAndDeleteDocuments= */ false);
+        icingSearchEngine.setSchemaWithRequestProto(requestProto1);
     assertStatusOk(setSchemaResultProto.getStatus());
-    setSchemaResultProto =
-        icingSearchEngine.setSchema(db2Schema, /* ignoreErrorsAndDeleteDocuments= */ false);
+
+    SetSchemaRequestProto requestProto2 =
+        SetSchemaRequestProto.newBuilder()
+            .setSchema(db2Schema)
+            .setDatabase(db2)
+            .setIgnoreErrorsAndDeleteDocuments(false)
+            .build();
+    setSchemaResultProto = icingSearchEngine.setSchemaWithRequestProto(requestProto2);
     assertStatusOk(setSchemaResultProto.getStatus());
 
     // Get schema for individual databases.
@@ -305,6 +321,11 @@
     assertThat(batchPutResultProto.getPutResultProtos(1).getUri()).isEqualTo("uri2");
     assertStatusOk(batchPutResultProto.getPutResultProtos(1).getStatus());
 
+    // PersistToDiskResultProto should not be set if persist_type is not set in the
+    // PutDocumentRequest.
+    assertThat(batchPutResultProto.getPersistToDiskResultProto().getStatus().getCode())
+        .isEqualTo(StatusProto.Code.UNKNOWN);
+
     // Check document 1
     GetResultProto getResultProto =
         icingSearchEngine.get("namespace", "uri1", GetResultSpecProto.getDefaultInstance());
@@ -348,6 +369,11 @@
     assertThat(batchPutResultProto.getPutResultProtos(1).getUri()).isEqualTo("uri");
     assertStatusOk(batchPutResultProto.getPutResultProtos(1).getStatus());
     assertThat(batchPutResultProto.getPutResultProtos(1).getWasReplacement()).isTrue();
+
+    // PersistToDiskResultProto should not be set if persist_type is not set in the
+    // PutDocumentRequest.
+    assertThat(batchPutResultProto.getPersistToDiskResultProto().getStatus().getCode())
+        .isEqualTo(StatusProto.Code.UNKNOWN);
   }
 
   @Test
@@ -368,6 +394,11 @@
 
     BatchPutResultProto expected = BatchPutResultProto.getDefaultInstance();
     assertThat(batchPutResultProto).isEqualTo(expected);
+
+    // PersistToDiskResultProto should not be set if persist_type is not set in the
+    // PutDocumentRequest.
+    assertThat(batchPutResultProto.getPersistToDiskResultProto().getStatus().getCode())
+        .isEqualTo(StatusProto.Code.UNKNOWN);
   }
 
   @Test
@@ -402,6 +433,11 @@
     assertThat(batchPutResultProto.getPutResultProtos(1).getUri()).isEqualTo("uri2");
     assertStatusOk(batchPutResultProto.getPutResultProtos(1).getStatus());
 
+    // PersistToDiskResultProto should not be set if persist_type is not set in the
+    // PutDocumentRequest.
+    assertThat(batchPutResultProto.getPersistToDiskResultProto().getStatus().getCode())
+        .isEqualTo(StatusProto.Code.UNKNOWN);
+
     // Check document 1
     GetResultProto getResultProto =
         icingSearchEngine.get("namespace", "uri1", GetResultSpecProto.getDefaultInstance());
@@ -416,6 +452,38 @@
   }
 
   @Test
+  public void testBatchPutWithPersistToDisk() throws Exception {
+    assertStatusOk(icingSearchEngine.initialize().getStatus());
+
+    SchemaTypeConfigProto emailTypeConfig = createEmailTypeConfig();
+    SchemaProto schema = SchemaProto.newBuilder().addTypes(emailTypeConfig).build();
+    assertThat(
+            icingSearchEngine
+                .setSchema(schema, /* ignoreErrorsAndDeleteDocuments= */ false)
+                .getStatus()
+                .getCode())
+        .isEqualTo(StatusProto.Code.OK);
+
+    DocumentProto emailDocument1 = createEmailDocument("namespace", "uri1");
+    DocumentProto emailDocument2 = createEmailDocument("namespace", "uri2");
+    PutDocumentRequest putDocumentRequest =
+        PutDocumentRequest.newBuilder()
+            .addDocuments(emailDocument1)
+            .addDocuments(emailDocument2)
+            .setPersistType(PersistType.Code.FULL)
+            .build();
+    BatchPutResultProto batchPutResultProto = icingSearchEngine.batchPut(putDocumentRequest);
+
+    assertThat(batchPutResultProto.getPutResultProtos(0).getUri()).isEqualTo("uri1");
+    assertStatusOk(batchPutResultProto.getPutResultProtos(0).getStatus());
+    assertThat(batchPutResultProto.getPutResultProtos(1).getUri()).isEqualTo("uri2");
+    assertStatusOk(batchPutResultProto.getPutResultProtos(1).getStatus());
+
+    // PersistToDisk should be called if persist_type is set in the PutDocumentRequest.
+    assertStatusOk(batchPutResultProto.getPersistToDiskResultProto().getStatus());
+  }
+
+  @Test
   public void testSearch() throws Exception {
     assertStatusOk(icingSearchEngine.initialize().getStatus());
 
diff --git a/proto/icing/proto/document.proto b/proto/icing/proto/document.proto
index 9a92c13..da82d60 100644
--- a/proto/icing/proto/document.proto
+++ b/proto/icing/proto/document.proto
@@ -17,6 +17,7 @@
 package icing.lib;
 
 import "icing/proto/logging.proto";
+import "icing/proto/persist.proto";
 import "icing/proto/status.proto";
 
 option java_package = "com.google.android.icing.proto";
@@ -24,9 +25,13 @@
 option objc_class_prefix = "ICNG";
 
 // Holds a list of DocumentProto.
-// Next tag: 2
+// Next tag: 3
 message PutDocumentRequest {
   repeated DocumentProto documents = 1;
+
+  // The persist type used to call PersistToDisk at the end of the Put request.
+  // If not specified, PersistToDisk will not be called.
+  optional PersistType.Code persist_type = 2;
 }
 
 // Defines a unit of data understood by the IcingSearchEngine.
@@ -125,9 +130,13 @@
 }
 
 // Holds a list of PutResultProto.
-// Next tag: 2
+// Next tag: 3
 message BatchPutResultProto {
   repeated PutResultProto put_result_protos = 1;
+
+  // The result of calling PersistToDisk at the end of the BatchPut request if
+  // PutDocumentRequest.persist_type is set.
+  optional PersistToDiskResultProto persist_to_disk_result_proto = 2;
 }
 
 // Result of a call to IcingSearchEngine.Put
diff --git a/proto/icing/proto/schema.proto b/proto/icing/proto/schema.proto
index 2a461bf..1f7eb4b 100644
--- a/proto/icing/proto/schema.proto
+++ b/proto/icing/proto/schema.proto
@@ -423,6 +423,37 @@
   repeated SchemaTypeConfigProto types = 1;
 }
 
+// Request for a call to IcingSearchEngine.SetSchema
+// Next tag: 4
+message SetSchemaRequestProto {
+  // REQUIRED: The new schema to set. This will replace the existing schema
+  // stored in IcingSearchEngine.
+  //
+  // schema.types is allowed to be empty. In this case, the SetSchema call will
+  // try to delete all types and indexed documents for the provided database,
+  // which is only allowed if ignore_errors_and_delete_documents=true
+  optional SchemaProto schema = 1;
+
+  // OPTIONAL: The database for the set schema request. Only schema types for
+  // this database will be modified.
+  //
+  // For a valid set schema request, this must match the database fields of
+  // schema.types.
+  //
+  // If unset, the default empty database is assumed for the set schema request.
+  optional string database = 2;
+
+  // OPTIONAL: Whether to ignore errors and delete documents when setting the
+  // schema.
+  //
+  // If true, then Icing will try to set the schema even if it is incompatible.
+  // In that case, documents that are invalidated by the new schema would be
+  // deleted from Icing. This cannot be used to force set an invalid schema.
+  //
+  // The default value is false.
+  optional bool ignore_errors_and_delete_documents = 3;
+}
+
 // Result of a call to IcingSearchEngine.SetSchema
 // Next tag: 9
 message SetSchemaResultProto {
diff --git a/synced_AOSP_CL_number.txt b/synced_AOSP_CL_number.txt
index 4172074..ba3548f 100644
--- a/synced_AOSP_CL_number.txt
+++ b/synced_AOSP_CL_number.txt
@@ -1 +1 @@
-set(synced_AOSP_CL_number=732166129)
+set(synced_AOSP_CL_number=733916508)