Merge "Update Framework from Jetpack." into sc-dev

21 files changed, 610 insertions, 369 deletions

diff --git a/apex/appsearch/framework/api/current.txt b/apex/appsearch/framework/api/current.txt
index f6bfaa356799..e08d22ce194c 100644
--- a/apex/appsearch/framework/api/current.txt
+++ b/apex/appsearch/framework/api/current.txt
@@ -24,14 +24,6 @@ package android.app.appsearch {
     method @Deprecated @NonNull public android.app.appsearch.AppSearchManager.SearchContext.Builder setDatabaseName(@NonNull String);
   }
 
-  public interface AppSearchMigrationHelper {
-    method public void queryAndTransform(@NonNull String, @NonNull android.app.appsearch.AppSearchMigrationHelper.Transformer) throws java.lang.Exception;
-  }
-
-  public static interface AppSearchMigrationHelper.Transformer {
-    method @NonNull public android.app.appsearch.GenericDocument transform(int, int, @NonNull android.app.appsearch.GenericDocument) throws java.lang.Exception;
-  }
-
   public final class AppSearchResult<ValueType> {
     method @Nullable public String getErrorMessage();
     method public int getResultCode();
@@ -109,11 +101,6 @@ package android.app.appsearch {
     method @NonNull public android.app.appsearch.AppSearchSchema.Int64PropertyConfig.Builder setCardinality(int);
   }
 
-  public static interface AppSearchSchema.Migrator {
-    method public default void onDowngrade(int, int, @NonNull android.app.appsearch.AppSearchMigrationHelper) throws java.lang.Exception;
-    method public default void onUpgrade(int, int, @NonNull android.app.appsearch.AppSearchMigrationHelper) throws java.lang.Exception;
-  }
-
   public abstract static class AppSearchSchema.PropertyConfig {
     method public int getCardinality();
     method @NonNull public String getName();
@@ -221,6 +208,13 @@ package android.app.appsearch {
     method @NonNull public android.app.appsearch.SearchResults search(@NonNull String, @NonNull android.app.appsearch.SearchSpec);
   }
 
+  public abstract class Migrator {
+    ctor public Migrator();
+    ctor public Migrator(int);
+    method @NonNull @WorkerThread public abstract android.app.appsearch.GenericDocument onDowngrade(int, int, @NonNull android.app.appsearch.GenericDocument);
+    method @NonNull @WorkerThread public abstract android.app.appsearch.GenericDocument onUpgrade(int, int, @NonNull android.app.appsearch.GenericDocument);
+  }
+
   public class PackageIdentifier {
     ctor public PackageIdentifier(@NonNull String, @NonNull byte[]);
     method @NonNull public String getPackageName();
@@ -357,7 +351,7 @@ package android.app.appsearch {
   }
 
   public final class SetSchemaRequest {
-    method @NonNull public java.util.Map<java.lang.String,android.app.appsearch.AppSearchSchema.Migrator> getMigrators();
+    method @NonNull public java.util.Map<java.lang.String,android.app.appsearch.Migrator> getMigrators();
     method @NonNull public java.util.Set<android.app.appsearch.AppSearchSchema> getSchemas();
     method @NonNull public java.util.Set<java.lang.String> getSchemasNotDisplayedBySystem();
     method @Deprecated @NonNull public java.util.Set<java.lang.String> getSchemasNotVisibleToSystemUi();
@@ -371,7 +365,7 @@ package android.app.appsearch {
     method @NonNull public android.app.appsearch.SetSchemaRequest.Builder addSchemas(@NonNull java.util.Collection<android.app.appsearch.AppSearchSchema>);
     method @NonNull public android.app.appsearch.SetSchemaRequest build();
     method @NonNull public android.app.appsearch.SetSchemaRequest.Builder setForceOverride(boolean);
-    method @NonNull public android.app.appsearch.SetSchemaRequest.Builder setMigrator(@NonNull String, @NonNull android.app.appsearch.AppSearchSchema.Migrator);
+    method @NonNull public android.app.appsearch.SetSchemaRequest.Builder setMigrator(@NonNull String, @NonNull android.app.appsearch.Migrator);
     method @NonNull public android.app.appsearch.SetSchemaRequest.Builder setSchemaTypeDisplayedBySystem(@NonNull String, boolean);
     method @NonNull public android.app.appsearch.SetSchemaRequest.Builder setSchemaTypeVisibilityForPackage(@NonNull String, boolean, @NonNull android.app.appsearch.PackageIdentifier);
     method @Deprecated @NonNull public android.app.appsearch.SetSchemaRequest.Builder setSchemaTypeVisibilityForSystemUi(@NonNull String, boolean);
diff --git a/apex/appsearch/framework/java/android/app/appsearch/AppSearchSession.java b/apex/appsearch/framework/java/android/app/appsearch/AppSearchSession.java
index 486acb4a3d7e..9ea73a9773bc 100644
--- a/apex/appsearch/framework/java/android/app/appsearch/AppSearchSession.java
+++ b/apex/appsearch/framework/java/android/app/appsearch/AppSearchSession.java
@@ -114,8 +114,6 @@ public final class AppSearchSession implements Closeable {
      * @param executor Executor on which to invoke the callback.
      * @param callback Callback to receive errors resulting from setting the schema. If the
      *                 operation succeeds, the callback will be invoked with {@code null}.
-     * @see android.app.appsearch.AppSearchSchema.Migrator
-     * @see android.app.appsearch.AppSearchMigrationHelper.Transformer
      */
     // TODO(b/169883602): Change @code references to @link when setPlatformSurfaceable APIs are
     //  exposed.
diff --git a/apex/appsearch/framework/java/external/android/app/appsearch/AppSearchMigrationHelper.java b/apex/appsearch/framework/java/external/android/app/appsearch/AppSearchMigrationHelper.java
deleted file mode 100644
index 37943fc24ded..000000000000
--- a/apex/appsearch/framework/java/external/android/app/appsearch/AppSearchMigrationHelper.java
+++ /dev/null
@@ -1,65 +0,0 @@
-/*
- * Copyright 2021 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package android.app.appsearch;
-
-import android.annotation.NonNull;
-import android.annotation.SuppressLint;
-
-/**
- * The helper class for {@link AppSearchSchema} migration.
- *
- * <p>It will query and migrate {@link GenericDocument} in given type to a new version.
- */
-public interface AppSearchMigrationHelper {
-
-    /**
-     * Queries all documents that need to be migrated to the different version, and transform
-     * documents to that version by passing them to the provided {@link Transformer}.
-     *
-     * @param schemaType The schema that need be updated and migrated {@link GenericDocument} under
-     *     this type.
-     * @param transformer The {@link Transformer} that will upgrade or downgrade a {@link
-     *     GenericDocument} to new version.
-     * @see Transformer#transform
-     */
-    // Rethrow the Generic Exception thrown from the Transformer.
-    @SuppressLint("GenericException")
-    void queryAndTransform(@NonNull String schemaType, @NonNull Transformer transformer)
-            throws Exception;
-
-    /** The class to migrate {@link GenericDocument} between different version. */
-    interface Transformer {
-
-        /**
-         * Translates a {@link GenericDocument} from a version to a different version.
-         *
-         * <p>If the uri, schema type or namespace is changed via the transform, it will apply to
-         * the new {@link GenericDocument}.
-         *
-         * @param currentVersion The current version of the document's schema.
-         * @param finalVersion The final version that documents need to be migrated to.
-         * @param document The {@link GenericDocument} need to be translated to new version.
-         * @return A {@link GenericDocument} in new version.
-         */
-        @NonNull
-        // This method will be overridden by users, allow them to throw any customer Exceptions.
-        @SuppressLint("GenericException")
-        GenericDocument transform(
-                int currentVersion, int finalVersion, @NonNull GenericDocument document)
-                throws Exception;
-    }
-}
diff --git a/apex/appsearch/framework/java/external/android/app/appsearch/AppSearchSchema.java b/apex/appsearch/framework/java/external/android/app/appsearch/AppSearchSchema.java
index 2cf52716cffc..55f0c80b16d0 100644
--- a/apex/appsearch/framework/java/external/android/app/appsearch/AppSearchSchema.java
+++ b/apex/appsearch/framework/java/external/android/app/appsearch/AppSearchSchema.java
@@ -20,7 +20,6 @@ import android.annotation.IntDef;
 import android.annotation.IntRange;
 import android.annotation.NonNull;
 import android.annotation.Nullable;
-import android.annotation.SuppressLint;
 import android.app.appsearch.exceptions.IllegalSchemaException;
 import android.app.appsearch.util.BundleUtil;
 import android.os.Bundle;
@@ -179,7 +178,7 @@ public final class AppSearchSchema {
          * @throws IllegalStateException if the version is negative or the builder has already been
          *     used.
          * @see AppSearchSession#setSchema
-         * @see AppSearchSchema.Migrator
+         * @see Migrator
          * @see SetSchemaRequest.Builder#setMigrator
          */
         @NonNull
@@ -861,43 +860,4 @@ public final class AppSearchSchema {
             }
         }
     }
-
-    /**
-     * A migrator class to translate {@link GenericDocument} from different version of {@link
-     * AppSearchSchema}
-     */
-    public interface Migrator {
-
-        /**
-         * Migrates {@link GenericDocument} to a newer version of {@link AppSearchSchema}.
-         *
-         * <p>This methods will be invoked only if the {@link SetSchemaRequest} is setting a higher
-         * version number than the current {@link AppSearchSchema} saved in AppSearch.
-         *
-         * @param currentVersion The current version of the document's schema.
-         * @param targetVersion The final version that documents need to be migrated to.
-         * @param helper The helper class could help to query all documents need to be migrated.
-         */
-        // This method will be overridden by users, allow them to throw any customer Exceptions.
-        @SuppressLint("GenericException")
-        default void onUpgrade(
-                int currentVersion, int targetVersion, @NonNull AppSearchMigrationHelper helper)
-                throws Exception {}
-
-        /**
-         * Migrates {@link GenericDocument} to an older version of {@link AppSearchSchema}.
-         *
-         * <p>The methods will be invoked only if the {@link SetSchemaRequest} is setting a higher
-         * version number than the current {@link AppSearchSchema} saved in AppSearch.
-         *
-         * @param currentVersion The current version of the document's schema.
-         * @param targetVersion The final version that documents need to be migrated to.
-         * @param helper The helper class could help to query all documents need to be migrated.
-         */
-        // This method will be overridden by users, allow them to throw any customer Exceptions.
-        @SuppressLint("GenericException")
-        default void onDowngrade(
-                int currentVersion, int targetVersion, @NonNull AppSearchMigrationHelper helper)
-                throws Exception {}
-    }
 }
diff --git a/apex/appsearch/framework/java/external/android/app/appsearch/Migrator.java b/apex/appsearch/framework/java/external/android/app/appsearch/Migrator.java
new file mode 100644
index 000000000000..5ae9a4186a52
--- /dev/null
+++ b/apex/appsearch/framework/java/external/android/app/appsearch/Migrator.java
@@ -0,0 +1,107 @@
+/*
+ * Copyright 2021 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package android.app.appsearch;
+
+import android.annotation.NonNull;
+import android.annotation.WorkerThread;
+
+import com.android.internal.util.Preconditions;
+
+/**
+ * A migrator class to translate {@link GenericDocument} from different version of {@link
+ * AppSearchSchema}
+ *
+ * <p>Make non-backwards-compatible changes will delete all stored documents in old schema. You can
+ * save your documents by setting {@link Migrator} via the {@link
+ * SetSchemaRequest.Builder#setMigrator} for each type and target version you want to save.
+ *
+ * <p>{@link #onDowngrade} or {@link #onUpgrade} will be triggered if the version number of the
+ * schema stored in AppSearch is different with the version in the request.
+ *
+ * <p>If any error or Exception occurred in the {@link #onDowngrade} or {@link #onUpgrade}, all the
+ * setSchema request will be rejected unless the schema changes are backwards-compatible, and stored
+ * documents won't have any observable changes.
+ */
+public abstract class Migrator {
+    private final int mStartVersion;
+
+    /**
+     * Creates a {@link Migrator} will trigger migration for any version less than the final version
+     * in the new schema.
+     */
+    public Migrator() {
+        this(/*startVersion=*/ 0);
+    }
+
+    /**
+     * Creates a {@link Migrator} with a non-negative start version.
+     *
+     * <p>Providing 0 will trigger migration for any version less than the final version in the new
+     * schema.
+     *
+     * @param startVersion The migration will be only triggered for those versions greater or equal
+     *     to the given startVersion.
+     */
+    public Migrator(int startVersion) {
+        Preconditions.checkArgumentNonnegative(startVersion);
+        mStartVersion = startVersion;
+    }
+
+    /**
+     * @return {@code True} if the current version need to be migrated.
+     * @hide
+     */
+    public boolean shouldMigrateToFinalVersion(int currentVersion, int finalVersion) {
+        return currentVersion >= mStartVersion && currentVersion != finalVersion;
+    }
+
+    /**
+     * Migrates {@link GenericDocument} to a newer version of {@link AppSearchSchema}.
+     *
+     * <p>This method will be invoked only if the {@link SetSchemaRequest} is setting a higher
+     * version number than the current {@link AppSearchSchema} saved in AppSearch.
+     *
+     * <p>This method will be invoked on the background worker thread.
+     *
+     * @param currentVersion The current version of the document's schema.
+     * @param targetVersion The final version that documents need to be migrated to.
+     * @param document The {@link GenericDocument} need to be translated to new version.
+     * @return A {@link GenericDocument} in new version.
+     */
+    @WorkerThread
+    @NonNull
+    public abstract GenericDocument onUpgrade(
+            int currentVersion, int targetVersion, @NonNull GenericDocument document);
+
+    /**
+     * Migrates {@link GenericDocument} to an older version of {@link AppSearchSchema}.
+     *
+     * <p>This method will be invoked only if the {@link SetSchemaRequest} is setting a lower
+     * version number than the current {@link AppSearchSchema} saved in AppSearch.
+     *
+     * <p>This method will be invoked on the background worker thread.
+     *
+     * @param currentVersion The current version of the document's schema.
+     * @param targetVersion The final version that documents need to be migrated to.
+     * @param document The {@link GenericDocument} need to be translated to new version.
+     * @return A {@link GenericDocument} in new version.
+     */
+    @WorkerThread
+    @NonNull
+    public abstract GenericDocument onDowngrade(
+            int currentVersion, int targetVersion, @NonNull GenericDocument document);
+}
diff --git a/apex/appsearch/framework/java/external/android/app/appsearch/SetSchemaRequest.java b/apex/appsearch/framework/java/external/android/app/appsearch/SetSchemaRequest.java
index c05406358995..c1eedcd63dbf 100644
--- a/apex/appsearch/framework/java/external/android/app/appsearch/SetSchemaRequest.java
+++ b/apex/appsearch/framework/java/external/android/app/appsearch/SetSchemaRequest.java
@@ -63,24 +63,33 @@ import java.util.Set;
  * android.app.appsearch.exceptions.AppSearchException}, with a message describing the
  * incompatibility. As a result, the previously set schema will remain unchanged.
  *
- * <p>Backward incompatible changes can be made by setting {@link
- * SetSchemaRequest.Builder#setForceOverride} method to {@code true}. This deletes all documents
- * that are incompatible with the new schema. The new schema is then saved and persisted to disk.
+ * <p>Backward incompatible changes can be made by :
+ *
+ * <ul>
+ *   <li>setting {@link SetSchemaRequest.Builder#setForceOverride} method to {@code true}. This
+ *       deletes all documents that are incompatible with the new schema. The new schema is then
+ *       saved and persisted to disk.
+ *   <li>Add a {@link Migrator} for each incompatible type and make no deletion. The migrator will
+ *       migrate documents from it's old schema version to the new version. Migrated types will be
+ *       set into both {@link SetSchemaResponse#getIncompatibleTypes()} and {@link
+ *       SetSchemaResponse#getMigratedTypes()}. See the migration section below.
+ * </ul>
  *
  * @see AppSearchSession#setSchema
+ * @see Migrator
  */
 public final class SetSchemaRequest {
     private final Set<AppSearchSchema> mSchemas;
     private final Set<String> mSchemasNotDisplayedBySystem;
     private final Map<String, Set<PackageIdentifier>> mSchemasVisibleToPackages;
-    private final Map<String, AppSearchSchema.Migrator> mMigrators;
+    private final Map<String, Migrator> mMigrators;
     private final boolean mForceOverride;
 
     SetSchemaRequest(
             @NonNull Set<AppSearchSchema> schemas,
             @NonNull Set<String> schemasNotDisplayedBySystem,
             @NonNull Map<String, Set<PackageIdentifier>> schemasVisibleToPackages,
-            @NonNull Map<String, AppSearchSchema.Migrator> migrators,
+            @NonNull Map<String, Migrator> migrators,
             boolean forceOverride) {
         mSchemas = Preconditions.checkNotNull(schemas);
         mSchemasNotDisplayedBySystem = Preconditions.checkNotNull(schemasNotDisplayedBySystem);
@@ -129,9 +138,12 @@ public final class SetSchemaRequest {
         return copy;
     }
 
-    /** Returns the map of {@link android.app.appsearch.AppSearchSchema.Migrator}. */
+    /**
+     * Returns the map of {@link Migrator}, the key will be the schema type of the {@link Migrator}
+     * associated with.
+     */
     @NonNull
-    public Map<String, AppSearchSchema.Migrator> getMigrators() {
+    public Map<String, Migrator> getMigrators() {
         return Collections.unmodifiableMap(mMigrators);
     }
 
@@ -164,7 +176,7 @@ public final class SetSchemaRequest {
         private final Set<String> mSchemasNotDisplayedBySystem = new ArraySet<>();
         private final Map<String, Set<PackageIdentifier>> mSchemasVisibleToPackages =
                 new ArrayMap<>();
-        private final Map<String, AppSearchSchema.Migrator> mMigrators = new ArrayMap<>();
+        private final Map<String, Migrator> mMigrators = new ArrayMap<>();
         private boolean mForceOverride = false;
         private boolean mBuilt = false;
 
@@ -295,7 +307,7 @@ public final class SetSchemaRequest {
         }
 
         /**
-         * Sets the {@link android.app.appsearch.AppSearchSchema.Migrator}.
+         * Sets the {@link Migrator}.
          *
          * @param schemaType The schema type to set migrator on.
          * @param migrator The migrator translate a document from it's old version to a new
@@ -303,8 +315,7 @@ public final class SetSchemaRequest {
          */
         @NonNull
         @SuppressLint("MissingGetterMatchingBuilder") // Getter return plural objects.
-        public Builder setMigrator(
-                @NonNull String schemaType, @NonNull AppSearchSchema.Migrator migrator) {
+        public Builder setMigrator(@NonNull String schemaType, @NonNull Migrator migrator) {
             Preconditions.checkNotNull(schemaType);
             Preconditions.checkNotNull(migrator);
             mMigrators.put(schemaType, migrator);
diff --git a/apex/appsearch/framework/java/external/android/app/appsearch/SetSchemaResponse.java b/apex/appsearch/framework/java/external/android/app/appsearch/SetSchemaResponse.java
index 98cd49b79737..d63e4372f61f 100644
--- a/apex/appsearch/framework/java/external/android/app/appsearch/SetSchemaResponse.java
+++ b/apex/appsearch/framework/java/external/android/app/appsearch/SetSchemaResponse.java
@@ -128,8 +128,8 @@ public class SetSchemaResponse {
      * Returns a {@link Set} of schema type whose new definitions set in the {@link
      * AppSearchSession#setSchema} call were incompatible with the pre-existing schema.
      *
-     * <p>If a {@link android.app.appsearch.AppSearchSchema.Migrator} is provided for this type and
-     * the migration is success triggered. The type will also appear in {@link #getMigratedTypes()}.
+     * <p>If a {@link Migrator} is provided for this type and the migration is success triggered.
+     * The type will also appear in {@link #getMigratedTypes()}.
      *
      * @see AppSearchSession#setSchema
      * @see SetSchemaRequest.Builder#setForceOverride
diff --git a/apex/appsearch/framework/java/external/android/app/appsearch/util/SchemaMigrationUtil.java b/apex/appsearch/framework/java/external/android/app/appsearch/util/SchemaMigrationUtil.java
new file mode 100644
index 000000000000..fae8ad48e61c
--- /dev/null
+++ b/apex/appsearch/framework/java/external/android/app/appsearch/util/SchemaMigrationUtil.java
@@ -0,0 +1,118 @@
+/*
+ * Copyright 2021 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package android.app.appsearch.util;
+
+import android.annotation.NonNull;
+import android.app.appsearch.AppSearchResult;
+import android.app.appsearch.AppSearchSchema;
+import android.app.appsearch.Migrator;
+import android.app.appsearch.exceptions.AppSearchException;
+import android.util.ArrayMap;
+import android.util.ArraySet;
+import android.util.Log;
+
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Map;
+import java.util.Set;
+
+/**
+ * Utilities for schema migration.
+ *
+ * @hide
+ */
+public final class SchemaMigrationUtil {
+    private static final String TAG = "AppSearchMigrateUtil";
+
+    private SchemaMigrationUtil() {}
+
+    /**
+     * Finds out which incompatible schema type won't be migrated by comparing its current and final
+     * version number.
+     */
+    @NonNull
+    public static Set<String> getUnmigratedIncompatibleTypes(
+            @NonNull Set<String> incompatibleSchemaTypes,
+            @NonNull Map<String, Migrator> migrators,
+            @NonNull Map<String, Integer> currentVersionMap,
+            @NonNull Map<String, Integer> finalVersionMap)
+            throws AppSearchException {
+        Set<String> unmigratedSchemaTypes = new ArraySet<>();
+        for (String unmigratedSchemaType : incompatibleSchemaTypes) {
+            Integer currentVersion = currentVersionMap.get(unmigratedSchemaType);
+            Integer finalVersion = finalVersionMap.get(unmigratedSchemaType);
+            if (currentVersion == null) {
+                // impossible, we have done something wrong.
+                throw new AppSearchException(
+                        AppSearchResult.RESULT_UNKNOWN_ERROR,
+                        "Cannot find the current version number for schema type: "
+                                + unmigratedSchemaType);
+            }
+            if (finalVersion == null) {
+                // The schema doesn't exist in the SetSchemaRequest.
+                unmigratedSchemaTypes.add(unmigratedSchemaType);
+                continue;
+            }
+            // we don't have migrator or won't trigger migration for this schema type.
+            Migrator migrator = migrators.get(unmigratedSchemaType);
+            if (migrator == null
+                    || !migrator.shouldMigrateToFinalVersion(currentVersion, finalVersion)) {
+                unmigratedSchemaTypes.add(unmigratedSchemaType);
+            }
+        }
+        return Collections.unmodifiableSet(unmigratedSchemaTypes);
+    }
+
+    /**
+     * Triggers upgrade or downgrade migration for the given schema type if its version stored in
+     * AppSearch is different with the version in the request.
+     *
+     * @return {@code True} if we trigger the migration for the given type.
+     */
+    public static boolean shouldTriggerMigration(
+            @NonNull String schemaType,
+            @NonNull Migrator migrator,
+            @NonNull Map<String, Integer> currentVersionMap,
+            @NonNull Map<String, Integer> finalVersionMap)
+            throws AppSearchException {
+        Integer currentVersion = currentVersionMap.get(schemaType);
+        Integer finalVersion = finalVersionMap.get(schemaType);
+        if (currentVersion == null) {
+            Log.d(TAG, "The SchemaType: " + schemaType + " not present in AppSearch.");
+            return false;
+        }
+        if (finalVersion == null) {
+            throw new AppSearchException(
+                    AppSearchResult.RESULT_INVALID_ARGUMENT,
+                    "Receive a migrator for schema type : "
+                            + schemaType
+                            + ", but the schema doesn't exist in the request.");
+        }
+        return migrator.shouldMigrateToFinalVersion(currentVersion, finalVersion);
+    }
+
+    /** Builds a Map of SchemaType and its version of given set of {@link AppSearchSchema}. */
+    @NonNull
+    public static Map<String, Integer> buildVersionMap(
+            @NonNull Collection<AppSearchSchema> schemas) {
+        Map<String, Integer> currentVersionMap = new ArrayMap<>(schemas.size());
+        for (AppSearchSchema currentSchema : schemas) {
+            currentVersionMap.put(currentSchema.getSchemaType(), currentSchema.getVersion());
+        }
+        return currentVersionMap;
+    }
+}
diff --git a/apex/appsearch/service/java/com/android/server/appsearch/AppSearchManagerService.java b/apex/appsearch/service/java/com/android/server/appsearch/AppSearchManagerService.java
index 27c9ccb2cca6..77aa14c37d79 100644
--- a/apex/appsearch/service/java/com/android/server/appsearch/AppSearchManagerService.java
+++ b/apex/appsearch/service/java/com/android/server/appsearch/AppSearchManagerService.java
@@ -193,7 +193,7 @@ public class AppSearchManagerService extends SystemService {
                     try {
                         // TODO(b/173451571): reduce burden of binder thread by enqueue request onto
                         // a separate thread.
-                        impl.putDocument(packageName, databaseName, document);
+                        impl.putDocument(packageName, databaseName, document, /*logger=*/ null);
                         resultBuilder.setSuccess(document.getUri(), /*result=*/ null);
                     } catch (Throwable t) {
                         resultBuilder.setResult(document.getUri(), throwableToFailedResult(t));
diff --git a/apex/appsearch/service/java/com/android/server/appsearch/VisibilityStore.java b/apex/appsearch/service/java/com/android/server/appsearch/VisibilityStore.java
index 7c92456bea49..ad94a0a57d27 100644
--- a/apex/appsearch/service/java/com/android/server/appsearch/VisibilityStore.java
+++ b/apex/appsearch/service/java/com/android/server/appsearch/VisibilityStore.java
@@ -368,7 +368,8 @@ public class VisibilityStore {
                     packageAccessibleDocuments.toArray(new GenericDocument[0]));
         }
 
-        mAppSearchImpl.putDocument(PACKAGE_NAME, DATABASE_NAME, visibilityDocument.build());
+        mAppSearchImpl.putDocument(
+                PACKAGE_NAME, DATABASE_NAME, visibilityDocument.build(), /*logger=*/ null);
 
         // Update derived data structures.
         mNotPlatformSurfaceableMap.put(prefix, schemasNotPlatformSurfaceable);
diff --git a/apex/appsearch/service/java/com/android/server/appsearch/external/localstorage/AppSearchImpl.java b/apex/appsearch/service/java/com/android/server/appsearch/external/localstorage/AppSearchImpl.java
index e2c211b7d303..5e8760ec35c2 100644
--- a/apex/appsearch/service/java/com/android/server/appsearch/external/localstorage/AppSearchImpl.java
+++ b/apex/appsearch/service/java/com/android/server/appsearch/external/localstorage/AppSearchImpl.java
@@ -17,6 +17,7 @@
 package com.android.server.appsearch.external.localstorage;
 
 import android.annotation.NonNull;
+import android.annotation.Nullable;
 import android.annotation.WorkerThread;
 import android.app.appsearch.AppSearchResult;
 import android.app.appsearch.AppSearchSchema;
@@ -29,6 +30,7 @@ import android.app.appsearch.SetSchemaResponse;
 import android.app.appsearch.exceptions.AppSearchException;
 import android.content.Context;
 import android.os.Bundle;
+import android.os.SystemClock;
 import android.util.ArrayMap;
 import android.util.ArraySet;
 import android.util.Log;
@@ -43,6 +45,7 @@ import com.android.server.appsearch.external.localstorage.converter.SearchResult
 import com.android.server.appsearch.external.localstorage.converter.SearchSpecToProtoConverter;
 import com.android.server.appsearch.external.localstorage.converter.SetSchemaResponseToProtoConverter;
 import com.android.server.appsearch.external.localstorage.converter.TypePropertyPathToProtoConverter;
+import com.android.server.appsearch.external.localstorage.stats.PutDocumentStats;
 
 import com.google.android.icing.IcingSearchEngine;
 import com.google.android.icing.proto.DeleteByQueryResultProto;
@@ -449,23 +452,65 @@ public final class AppSearchImpl implements Closeable {
     public void putDocument(
             @NonNull String packageName,
             @NonNull String databaseName,
-            @NonNull GenericDocument document)
+            @NonNull GenericDocument document,
+            @Nullable AppSearchLogger logger)
             throws AppSearchException {
+        PutDocumentStats.Builder pStatsBuilder = null;
+        if (logger != null) {
+            pStatsBuilder = new PutDocumentStats.Builder(packageName, databaseName);
+        }
+        long totalStartTimeMillis = SystemClock.elapsedRealtime();
+
         mReadWriteLock.writeLock().lock();
         try {
             throwIfClosedLocked();
 
+            // Generate Document Proto
+            long generateDocumentProtoStartTimeMillis = SystemClock.elapsedRealtime();
             DocumentProto.Builder documentBuilder =
                     GenericDocumentToProtoConverter.toDocumentProto(document).toBuilder();
+            long generateDocumentProtoEndTimeMillis = SystemClock.elapsedRealtime();
+
+            // Rewrite Document Type
+            long rewriteDocumentTypeStartTimeMillis = SystemClock.elapsedRealtime();
             String prefix = createPrefix(packageName, databaseName);
             addPrefixToDocument(documentBuilder, prefix);
+            long rewriteDocumentTypeEndTimeMillis = SystemClock.elapsedRealtime();
 
             PutResultProto putResultProto = mIcingSearchEngineLocked.put(documentBuilder.build());
             addToMap(mNamespaceMapLocked, prefix, documentBuilder.getNamespace());
 
+            // Logging stats
+            if (logger != null) {
+                pStatsBuilder
+                        .getGeneralStatsBuilder()
+                        .setStatusCode(
+                                statusProtoToAppSearchException(putResultProto.getStatus())
+                                        .getResultCode());
+                pStatsBuilder
+                        .setGenerateDocumentProtoLatencyMillis(
+                                (int)
+                                        (generateDocumentProtoEndTimeMillis
+                                                - generateDocumentProtoStartTimeMillis))
+                        .setRewriteDocumentTypesLatencyMillis(
+                                (int)
+                                        (rewriteDocumentTypeEndTimeMillis
+                                                - rewriteDocumentTypeStartTimeMillis));
+                AppSearchLoggerHelper.copyNativeStats(
+                        putResultProto.getPutDocumentStats(), pStatsBuilder);
+            }
+
             checkSuccess(putResultProto.getStatus());
         } finally {
             mReadWriteLock.writeLock().unlock();
+
+            if (logger != null) {
+                long totalEndTimeMillis = SystemClock.elapsedRealtime();
+                pStatsBuilder
+                        .getGeneralStatsBuilder()
+                        .setTotalLatencyMillis((int) (totalEndTimeMillis - totalStartTimeMillis));
+                logger.logStats(pStatsBuilder.build());
+            }
         }
     }
 
diff --git a/apex/appsearch/service/java/com/android/server/appsearch/external/localstorage/AppSearchLoggerHelper.java b/apex/appsearch/service/java/com/android/server/appsearch/external/localstorage/AppSearchLoggerHelper.java
new file mode 100644
index 000000000000..5680670629ae
--- /dev/null
+++ b/apex/appsearch/service/java/com/android/server/appsearch/external/localstorage/AppSearchLoggerHelper.java
@@ -0,0 +1,58 @@
+/*
+ * Copyright 2021 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.server.appsearch.external.localstorage;
+
+import android.annotation.NonNull;
+
+import com.android.internal.util.Preconditions;
+import com.android.server.appsearch.external.localstorage.stats.PutDocumentStats;
+
+import com.google.android.icing.proto.PutDocumentStatsProto;
+
+/**
+ * Class contains helper functions for logging.
+ *
+ * <p>E.g. we need to have helper functions to copy numbers from IcingLib to stats classes.
+ *
+ * @hide
+ */
+public final class AppSearchLoggerHelper {
+    private AppSearchLoggerHelper() {}
+
+    /**
+     * Copies native stats to builder.
+     *
+     * @param fromNativeStats stats copied from
+     * @param toStatsBuilder stats copied to
+     */
+    static void copyNativeStats(
+            @NonNull PutDocumentStatsProto fromNativeStats,
+            @NonNull PutDocumentStats.Builder toStatsBuilder) {
+        Preconditions.checkNotNull(fromNativeStats);
+        Preconditions.checkNotNull(toStatsBuilder);
+        toStatsBuilder
+                .setNativeLatencyMillis(fromNativeStats.getLatencyMs())
+                .setNativeDocumentStoreLatencyMillis(fromNativeStats.getDocumentStoreLatencyMs())
+                .setNativeIndexLatencyMillis(fromNativeStats.getIndexLatencyMs())
+                .setNativeIndexMergeLatencyMillis(fromNativeStats.getIndexMergeLatencyMs())
+                .setNativeDocumentSizeBytes(fromNativeStats.getDocumentSize())
+                .setNativeNumTokensIndexed(
+                        fromNativeStats.getTokenizationStats().getNumTokensIndexed())
+                .setNativeExceededMaxNumTokens(
+                        fromNativeStats.getTokenizationStats().getExceededMaxTokenNum());
+    }
+}
diff --git a/apex/appsearch/service/java/com/android/server/appsearch/external/localstorage/AppSearchMigrationHelperImpl.java b/apex/appsearch/service/java/com/android/server/appsearch/external/localstorage/AppSearchMigrationHelperImpl.java
deleted file mode 100644
index 4b8ce6d2c1d1..000000000000
--- a/apex/appsearch/service/java/com/android/server/appsearch/external/localstorage/AppSearchMigrationHelperImpl.java
+++ /dev/null
@@ -1,168 +0,0 @@
-/*
- * Copyright 2021 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package com.android.server.appsearch.external.localstorage;
-
-import static android.app.appsearch.AppSearchResult.throwableToFailedResult;
-
-import android.annotation.NonNull;
-import android.app.appsearch.AppSearchBatchResult;
-import android.app.appsearch.AppSearchMigrationHelper;
-import android.app.appsearch.GenericDocument;
-import android.app.appsearch.SearchResultPage;
-import android.app.appsearch.SearchSpec;
-import android.app.appsearch.SetSchemaResponse;
-import android.app.appsearch.exceptions.AppSearchException;
-import android.os.Bundle;
-import android.os.Parcel;
-
-import com.android.internal.util.Preconditions;
-
-import com.google.protobuf.CodedInputStream;
-import com.google.protobuf.CodedOutputStream;
-
-import java.io.File;
-import java.io.FileInputStream;
-import java.io.FileOutputStream;
-import java.io.IOException;
-import java.io.InputStream;
-import java.util.Map;
-
-/**
- * An implementation of {@link AppSearchMigrationHelper} which query document and save post-migrated
- * documents to locally in the app's storage space.
- */
-class AppSearchMigrationHelperImpl implements AppSearchMigrationHelper {
-    private final AppSearchImpl mAppSearchImpl;
-    private final String mPackageName;
-    private final String mDatabaseName;
-    private final File mFile;
-    private final Map<String, Integer> mCurrentVersionMap;
-    private final Map<String, Integer> mFinalVersionMap;
-
-    AppSearchMigrationHelperImpl(
-            @NonNull AppSearchImpl appSearchImpl,
-            @NonNull Map<String, Integer> currentVersionMap,
-            @NonNull Map<String, Integer> finalVersionMap,
-            @NonNull String packageName,
-            @NonNull String databaseName)
-            throws IOException {
-        mAppSearchImpl = Preconditions.checkNotNull(appSearchImpl);
-        mCurrentVersionMap = Preconditions.checkNotNull(currentVersionMap);
-        mFinalVersionMap = Preconditions.checkNotNull(finalVersionMap);
-        mPackageName = Preconditions.checkNotNull(packageName);
-        mDatabaseName = Preconditions.checkNotNull(databaseName);
-        mFile = File.createTempFile(/*prefix=*/ "appsearch", /*suffix=*/ null);
-    }
-
-    @Override
-    public void queryAndTransform(
-            @NonNull String schemaType, @NonNull AppSearchMigrationHelper.Transformer migrator)
-            throws Exception {
-        Preconditions.checkState(mFile.exists(), "Internal temp file does not exist.");
-        int currentVersion = mCurrentVersionMap.get(schemaType);
-        int finalVersion = mFinalVersionMap.get(schemaType);
-        try (FileOutputStream outputStream = new FileOutputStream(mFile)) {
-            // TODO(b/151178558) change the output stream so that we can use it in platform
-            CodedOutputStream codedOutputStream = CodedOutputStream.newInstance(outputStream);
-            SearchResultPage searchResultPage =
-                    mAppSearchImpl.query(
-                            mPackageName,
-                            mDatabaseName,
-                            /*queryExpression=*/ "",
-                            new SearchSpec.Builder()
-                                    .addFilterSchemas(schemaType)
-                                    .setTermMatch(SearchSpec.TERM_MATCH_PREFIX)
-                                    .build());
-            while (!searchResultPage.getResults().isEmpty()) {
-                for (int i = 0; i < searchResultPage.getResults().size(); i++) {
-                    GenericDocument newDocument =
-                            migrator.transform(
-                                    currentVersion,
-                                    finalVersion,
-                                    searchResultPage.getResults().get(i).getGenericDocument());
-                    Bundle bundle = newDocument.getBundle();
-                    Parcel parcel = Parcel.obtain();
-                    parcel.writeBundle(bundle);
-                    byte[] serializedMessage = parcel.marshall();
-                    parcel.recycle();
-                    codedOutputStream.writeByteArrayNoTag(serializedMessage);
-                }
-                codedOutputStream.flush();
-                searchResultPage = mAppSearchImpl.getNextPage(searchResultPage.getNextPageToken());
-                outputStream.flush();
-            }
-        }
-    }
-
-    /**
-     * Reads {@link GenericDocument} from the temperate file and saves them to AppSearch.
-     *
-     * <p>This method should be only called once.
-     *
-     * @return the {@link AppSearchBatchResult} for migration documents.
-     */
-    @NonNull
-    public SetSchemaResponse readAndPutDocuments(SetSchemaResponse.Builder responseBuilder)
-            throws IOException, AppSearchException {
-        Preconditions.checkState(mFile.exists(), "Internal temp file does not exist.");
-        try (InputStream inputStream = new FileInputStream(mFile)) {
-            CodedInputStream codedInputStream = CodedInputStream.newInstance(inputStream);
-            while (!codedInputStream.isAtEnd()) {
-                GenericDocument document = readDocumentFromInputStream(codedInputStream);
-                try {
-                    mAppSearchImpl.putDocument(mPackageName, mDatabaseName, document);
-                } catch (Throwable t) {
-                    responseBuilder.addMigrationFailure(
-                            new SetSchemaResponse.MigrationFailure.Builder()
-                                    .setNamespace(document.getNamespace())
-                                    .setSchemaType(document.getSchemaType())
-                                    .setUri(document.getUri())
-                                    .setAppSearchResult(throwableToFailedResult(t))
-                                    .build());
-                }
-            }
-            mAppSearchImpl.persistToDisk();
-            return responseBuilder.build();
-        } finally {
-            mFile.delete();
-        }
-    }
-
-    void deleteTempFile() {
-        mFile.delete();
-    }
-
-    /**
-     * Reads {@link GenericDocument} from given {@link CodedInputStream}.
-     *
-     * @param codedInputStream The codedInputStream to read from
-     * @throws IOException on File operation error.
-     */
-    @NonNull
-    private static GenericDocument readDocumentFromInputStream(
-            @NonNull CodedInputStream codedInputStream) throws IOException {
-        byte[] serializedMessage = codedInputStream.readByteArray();
-
-        Parcel parcel = Parcel.obtain();
-        parcel.unmarshall(serializedMessage, 0, serializedMessage.length);
-        parcel.setDataPosition(0);
-        Bundle bundle = parcel.readBundle();
-        parcel.recycle();
-
-        return new GenericDocument(bundle);
-    }
-}
diff --git a/apex/appsearch/service/java/com/android/server/appsearch/external/localstorage/stats/CallStats.java b/apex/appsearch/service/java/com/android/server/appsearch/external/localstorage/stats/CallStats.java
index 81a5067c9fa1..a724f95ad8fa 100644
--- a/apex/appsearch/service/java/com/android/server/appsearch/external/localstorage/stats/CallStats.java
+++ b/apex/appsearch/service/java/com/android/server/appsearch/external/localstorage/stats/CallStats.java
@@ -76,7 +76,7 @@ public class CallStats {
 
     CallStats(@NonNull Builder builder) {
         Preconditions.checkNotNull(builder);
-        mGeneralStats = Preconditions.checkNotNull(builder.mGeneralStats);
+        mGeneralStats = Preconditions.checkNotNull(builder.mGeneralStatsBuilder).build();
         mCallType = builder.mCallType;
         mEstimatedBinderLatencyMillis = builder.mEstimatedBinderLatencyMillis;
         mNumOperationsSucceeded = builder.mNumOperationsSucceeded;
@@ -132,15 +132,23 @@ public class CallStats {
 
     /** Builder for {@link CallStats}. */
     public static class Builder {
-        @NonNull final GeneralStats mGeneralStats;
+        @NonNull final GeneralStats.Builder mGeneralStatsBuilder;
         @CallType int mCallType;
         int mEstimatedBinderLatencyMillis;
         int mNumOperationsSucceeded;
         int mNumOperationsFailed;
 
-        /** Builder takes {@link GeneralStats} to hold general stats. */
-        public Builder(@NonNull GeneralStats generalStats) {
-            mGeneralStats = Preconditions.checkNotNull(generalStats);
+        /** Builder takes {@link GeneralStats.Builder}. */
+        public Builder(@NonNull String packageName, @NonNull String database) {
+            Preconditions.checkNotNull(packageName);
+            Preconditions.checkNotNull(database);
+            mGeneralStatsBuilder = new GeneralStats.Builder(packageName, database);
+        }
+
+        /** Returns {@link GeneralStats.Builder}. */
+        @NonNull
+        public GeneralStats.Builder getGeneralStatsBuilder() {
+            return mGeneralStatsBuilder;
         }
 
         /** Sets type of the call. */
diff --git a/apex/appsearch/service/java/com/android/server/appsearch/external/localstorage/stats/GeneralStats.java b/apex/appsearch/service/java/com/android/server/appsearch/external/localstorage/stats/GeneralStats.java
index d2a45d5304f9..8ce8eda3d7a7 100644
--- a/apex/appsearch/service/java/com/android/server/appsearch/external/localstorage/stats/GeneralStats.java
+++ b/apex/appsearch/service/java/com/android/server/appsearch/external/localstorage/stats/GeneralStats.java
@@ -82,18 +82,18 @@ public final class GeneralStats {
     public static class Builder {
         @NonNull final String mPackageName;
         @NonNull final String mDatabase;
-        @AppSearchResult.ResultCode int mStatusCode;
+        @AppSearchResult.ResultCode int mStatusCode = AppSearchResult.RESULT_UNKNOWN_ERROR;
         int mTotalLatencyMillis;
 
         /**
          * Constructor
          *
          * @param packageName name of the package logging stats
-         * @param dataBase name of the database logging stats
+         * @param database name of the database logging stats
          */
-        public Builder(@NonNull String packageName, @NonNull String dataBase) {
+        public Builder(@NonNull String packageName, @NonNull String database) {
             mPackageName = Preconditions.checkNotNull(packageName);
-            mDatabase = Preconditions.checkNotNull(dataBase);
+            mDatabase = Preconditions.checkNotNull(database);
         }
 
         /** Sets status code returned from {@link AppSearchResult#getResultCode()} */
diff --git a/apex/appsearch/service/java/com/android/server/appsearch/external/localstorage/stats/PutDocumentStats.java b/apex/appsearch/service/java/com/android/server/appsearch/external/localstorage/stats/PutDocumentStats.java
index b1b643b66859..c1f6fb118797 100644
--- a/apex/appsearch/service/java/com/android/server/appsearch/external/localstorage/stats/PutDocumentStats.java
+++ b/apex/appsearch/service/java/com/android/server/appsearch/external/localstorage/stats/PutDocumentStats.java
@@ -54,12 +54,14 @@ public final class PutDocumentStats {
     /** Number of tokens added to the index. */
     private final int mNativeNumTokensIndexed;
 
-    /** Number of tokens clipped for exceeding the max number. */
-    private final int mNativeNumTokensClipped;
+    /**
+     * Whether the number of tokens to be indexed exceeded the max number of tokens per document.
+     */
+    private final boolean mNativeExceededMaxNumTokens;
 
     PutDocumentStats(@NonNull Builder builder) {
         Preconditions.checkNotNull(builder);
-        mGeneralStats = Preconditions.checkNotNull(builder.mGeneralStats);
+        mGeneralStats = Preconditions.checkNotNull(builder.mGeneralStatsBuilder).build();
         mGenerateDocumentProtoLatencyMillis = builder.mGenerateDocumentProtoLatencyMillis;
         mRewriteDocumentTypesLatencyMillis = builder.mRewriteDocumentTypesLatencyMillis;
         mNativeLatencyMillis = builder.mNativeLatencyMillis;
@@ -68,7 +70,7 @@ public final class PutDocumentStats {
         mNativeIndexMergeLatencyMillis = builder.mNativeIndexMergeLatencyMillis;
         mNativeDocumentSizeBytes = builder.mNativeDocumentSizeBytes;
         mNativeNumTokensIndexed = builder.mNativeNumTokensIndexed;
-        mNativeNumTokensClipped = builder.mNativeNumTokensClipped;
+        mNativeExceededMaxNumTokens = builder.mNativeExceededMaxNumTokens;
     }
 
     /** Returns the {@link GeneralStats} object attached to this instance. */
@@ -117,14 +119,17 @@ public final class PutDocumentStats {
         return mNativeNumTokensIndexed;
     }
 
-    /** Returns number of tokens clipped for exceeding the max number. */
-    public int getNativeNumTokensClipped() {
-        return mNativeNumTokensClipped;
+    /**
+     * Returns whether the number of tokens to be indexed exceeded the max number of tokens per
+     * document.
+     */
+    public boolean getNativeExceededMaxNumTokens() {
+        return mNativeExceededMaxNumTokens;
     }
 
     /** Builder for {@link PutDocumentStats}. */
     public static class Builder {
-        @NonNull final GeneralStats mGeneralStats;
+        @NonNull final GeneralStats.Builder mGeneralStatsBuilder;
         int mGenerateDocumentProtoLatencyMillis;
         int mRewriteDocumentTypesLatencyMillis;
         int mNativeLatencyMillis;
@@ -133,11 +138,19 @@ public final class PutDocumentStats {
         int mNativeIndexMergeLatencyMillis;
         int mNativeDocumentSizeBytes;
         int mNativeNumTokensIndexed;
-        int mNativeNumTokensClipped;
+        boolean mNativeExceededMaxNumTokens;
 
-        /** Builder takes {@link GeneralStats} to hold general stats. */
-        public Builder(@NonNull GeneralStats generalStats) {
-            mGeneralStats = Preconditions.checkNotNull(generalStats);
+        /** Builder takes {@link GeneralStats.Builder}. */
+        public Builder(@NonNull String packageName, @NonNull String database) {
+            Preconditions.checkNotNull(packageName);
+            Preconditions.checkNotNull(database);
+            mGeneralStatsBuilder = new GeneralStats.Builder(packageName, database);
+        }
+
+        /** Returns {@link GeneralStats.Builder}. */
+        @NonNull
+        public GeneralStats.Builder getGeneralStatsBuilder() {
+            return mGeneralStatsBuilder;
         }
 
         /** Sets how much time we spend for generating document proto, in milliseconds. */
@@ -200,10 +213,13 @@ public final class PutDocumentStats {
             return this;
         }
 
-        /** Sets number of tokens clipped for exceeding the max number. */
+        /**
+         * Sets whether the number of tokens to be indexed exceeded the max number of tokens per
+         * document.
+         */
         @NonNull
-        public Builder setNativeNumTokensClipped(int nativeNumTokensClipped) {
-            mNativeNumTokensClipped = nativeNumTokensClipped;
+        public Builder setNativeExceededMaxNumTokens(boolean nativeExceededMaxNumTokens) {
+            mNativeExceededMaxNumTokens = nativeExceededMaxNumTokens;
             return this;
         }
 
diff --git a/apex/appsearch/synced_jetpack_changeid.txt b/apex/appsearch/synced_jetpack_changeid.txt
index 68531b6ee1e5..09522152cefb 100644
--- a/apex/appsearch/synced_jetpack_changeid.txt
+++ b/apex/appsearch/synced_jetpack_changeid.txt
@@ -1 +1 @@
-I1926fb1d13628607f7a513c8149b65dd86c98dd6
+I723a9d7b5e64329ab25b6d7627f3b2d222c31ac7
diff --git a/services/tests/servicestests/src/com/android/server/appsearch/AppSearchImplPlatformTest.java b/services/tests/servicestests/src/com/android/server/appsearch/AppSearchImplPlatformTest.java
index b5f49124483e..ad22cbaef575 100644
--- a/services/tests/servicestests/src/com/android/server/appsearch/AppSearchImplPlatformTest.java
+++ b/services/tests/servicestests/src/com/android/server/appsearch/AppSearchImplPlatformTest.java
@@ -103,12 +103,12 @@ public class AppSearchImplPlatformTest {
         // Insert package1 document
         GenericDocument document1 =
                 new GenericDocument.Builder<>("uri", "schema1").setNamespace("namespace").build();
-        mAppSearchImpl.putDocument("package1", "database1", document1);
+        mAppSearchImpl.putDocument("package1", "database1", document1, /*logger=*/ null);
 
         // Insert package2 document
         GenericDocument document2 =
                 new GenericDocument.Builder<>("uri", "schema2").setNamespace("namespace").build();
-        mAppSearchImpl.putDocument("package2", "database2", document2);
+        mAppSearchImpl.putDocument("package2", "database2", document2, /*logger=*/ null);
 
         // No query filters specified, global query can retrieve all documents.
         SearchSpec searchSpec =
@@ -155,12 +155,12 @@ public class AppSearchImplPlatformTest {
         // Insert package1 document
         GenericDocument document1 =
                 new GenericDocument.Builder<>("uri", "schema1").setNamespace("namespace").build();
-        mAppSearchImpl.putDocument("package1", "database1", document1);
+        mAppSearchImpl.putDocument("package1", "database1", document1, /*logger=*/ null);
 
         // Insert package2 document
         GenericDocument document2 =
                 new GenericDocument.Builder<>("uri", "schema2").setNamespace("namespace").build();
-        mAppSearchImpl.putDocument("package2", "database2", document2);
+        mAppSearchImpl.putDocument("package2", "database2", document2, /*logger=*/ null);
 
         // "package1" filter specified
         SearchSpec searchSpec =
diff --git a/services/tests/servicestests/src/com/android/server/appsearch/external/localstorage/AppSearchImplTest.java b/services/tests/servicestests/src/com/android/server/appsearch/external/localstorage/AppSearchImplTest.java
index c84c1cf8c8f1..e0cdeddc200e 100644
--- a/services/tests/servicestests/src/com/android/server/appsearch/external/localstorage/AppSearchImplTest.java
+++ b/services/tests/servicestests/src/com/android/server/appsearch/external/localstorage/AppSearchImplTest.java
@@ -416,7 +416,7 @@ public class AppSearchImplTest {
                 i++) {
             GenericDocument document =
                     new GenericDocument.Builder<>("namespace", "uri" + i, "type").build();
-            mAppSearchImpl.putDocument("package", "database", document);
+            mAppSearchImpl.putDocument("package", "database", document, /*logger=*/ null);
         }
 
         // Check optimize() will release 0 docs since there is no deletion.
@@ -476,7 +476,7 @@ public class AppSearchImplTest {
         // Insert document
         GenericDocument document =
                 new GenericDocument.Builder<>("namespace", "uri", "type").build();
-        mAppSearchImpl.putDocument("package", "database", document);
+        mAppSearchImpl.putDocument("package", "database", document, /*logger=*/ null);
 
         // Rewrite SearchSpec
         mAppSearchImpl.rewriteSearchSpecForPrefixesLocked(
@@ -516,11 +516,11 @@ public class AppSearchImplTest {
         // Insert documents
         GenericDocument document1 =
                 new GenericDocument.Builder<>("namespace", "uri", "typeA").build();
-        mAppSearchImpl.putDocument("package", "database1", document1);
+        mAppSearchImpl.putDocument("package", "database1", document1, /*logger=*/ null);
 
         GenericDocument document2 =
                 new GenericDocument.Builder<>("namespace", "uri", "typeB").build();
-        mAppSearchImpl.putDocument("package", "database2", document2);
+        mAppSearchImpl.putDocument("package", "database2", document2, /*logger=*/ null);
 
         // Rewrite SearchSpec
         mAppSearchImpl.rewriteSearchSpecForPrefixesLocked(
@@ -560,7 +560,7 @@ public class AppSearchImplTest {
         // Insert document
         GenericDocument document =
                 new GenericDocument.Builder<>("namespace", "uri", "type").build();
-        mAppSearchImpl.putDocument("package", "database", document);
+        mAppSearchImpl.putDocument("package", "database", document, /*logger=*/ null);
 
         // If 'allowedPrefixedSchemas' is empty, this returns false since there's nothing to
         // search over. Despite the searchSpecProto having schema type filters.
@@ -613,7 +613,7 @@ public class AppSearchImplTest {
         // Insert package1 document
         GenericDocument document =
                 new GenericDocument.Builder<>("namespace", "uri", "schema1").build();
-        mAppSearchImpl.putDocument("package1", "database1", document);
+        mAppSearchImpl.putDocument("package1", "database1", document, /*logger=*/ null);
 
         // No query filters specified, package2 shouldn't be able to query for package1's documents.
         SearchSpec searchSpec =
@@ -624,7 +624,7 @@ public class AppSearchImplTest {
 
         // Insert package2 document
         document = new GenericDocument.Builder<>("namespace", "uri", "schema2").build();
-        mAppSearchImpl.putDocument("package2", "database2", document);
+        mAppSearchImpl.putDocument("package2", "database2", document, /*logger=*/ null);
 
         // No query filters specified. package2 should only get its own documents back.
         searchResultPage = mAppSearchImpl.query("package2", "database2", "", searchSpec);
@@ -663,7 +663,7 @@ public class AppSearchImplTest {
         // Insert package1 document
         GenericDocument document =
                 new GenericDocument.Builder<>("namespace", "uri", "schema1").build();
-        mAppSearchImpl.putDocument("package1", "database1", document);
+        mAppSearchImpl.putDocument("package1", "database1", document, /*logger=*/ null);
 
         // "package1" filter specified, but package2 shouldn't be able to query for package1's
         // documents.
@@ -678,7 +678,7 @@ public class AppSearchImplTest {
 
         // Insert package2 document
         document = new GenericDocument.Builder<>("namespace", "uri", "schema2").build();
-        mAppSearchImpl.putDocument("package2", "database2", document);
+        mAppSearchImpl.putDocument("package2", "database2", document, /*logger=*/ null);
 
         // "package2" filter specified, package2 should only get its own documents back.
         searchSpec =
@@ -1124,7 +1124,8 @@ public class AppSearchImplTest {
                     appSearchImpl.putDocument(
                             "package",
                             "database",
-                            new GenericDocument.Builder<>("namespace", "uri", "type").build());
+                            new GenericDocument.Builder<>("namespace", "uri", "type").build(),
+                            /*logger=*/ null);
                 });
 
         expectThrows(
diff --git a/services/tests/servicestests/src/com/android/server/appsearch/external/localstorage/AppSearchLoggerTest.java b/services/tests/servicestests/src/com/android/server/appsearch/external/localstorage/AppSearchLoggerTest.java
new file mode 100644
index 000000000000..467ede4efc1c
--- /dev/null
+++ b/services/tests/servicestests/src/com/android/server/appsearch/external/localstorage/AppSearchLoggerTest.java
@@ -0,0 +1,144 @@
+/*
+ * Copyright 2021 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.server.appsearch.external.localstorage;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import android.annotation.NonNull;
+import android.annotation.Nullable;
+import android.app.appsearch.AppSearchResult;
+import android.app.appsearch.AppSearchSchema;
+import android.app.appsearch.GenericDocument;
+import android.content.Context;
+
+import androidx.test.core.app.ApplicationProvider;
+
+import com.android.server.appsearch.external.localstorage.stats.CallStats;
+import com.android.server.appsearch.external.localstorage.stats.PutDocumentStats;
+import com.android.server.appsearch.proto.PutDocumentStatsProto;
+
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.TemporaryFolder;
+
+import java.util.Collections;
+import java.util.List;
+
+public class AppSearchLoggerTest {
+    @Rule public TemporaryFolder mTemporaryFolder = new TemporaryFolder();
+    private AppSearchImpl mAppSearchImpl;
+    private TestLogger mLogger;
+
+    @Before
+    public void setUp() throws Exception {
+        Context context = ApplicationProvider.getApplicationContext();
+
+        // Give ourselves global query permissions
+        mAppSearchImpl =
+                AppSearchImpl.create(
+                        mTemporaryFolder.newFolder(),
+                        context,
+                        VisibilityStore.NO_OP_USER_ID,
+                        /*globalQuerierPackage=*/ context.getPackageName());
+        mLogger = new TestLogger();
+    }
+
+    // Test only not thread safe.
+    public class TestLogger implements AppSearchLogger {
+        @Nullable PutDocumentStats mPutDocumentStats;
+
+        @Override
+        public void logStats(@NonNull CallStats stats) {
+            throw new UnsupportedOperationException();
+        }
+
+        @Override
+        public void logStats(@NonNull PutDocumentStats stats) {
+            mPutDocumentStats = stats;
+        }
+    }
+
+    @Test
+    public void testAppSearchLoggerHelper_testCopyNativeStats_putDocument() {
+        final int nativeLatencyMillis = 3;
+        final int nativeDocumentStoreLatencyMillis = 4;
+        final int nativeIndexLatencyMillis = 5;
+        final int nativeIndexMergeLatencyMillis = 6;
+        final int nativeDocumentSize = 7;
+        final int nativeNumTokensIndexed = 8;
+        final boolean nativeExceededMaxNumTokens = true;
+        PutDocumentStatsProto nativePutDocumentStats =
+                PutDocumentStatsProto.newBuilder()
+                        .setLatencyMs(nativeLatencyMillis)
+                        .setDocumentStoreLatencyMs(nativeDocumentStoreLatencyMillis)
+                        .setIndexLatencyMs(nativeIndexLatencyMillis)
+                        .setIndexMergeLatencyMs(nativeIndexMergeLatencyMillis)
+                        .setDocumentSize(nativeDocumentSize)
+                        .setTokenizationStats(
+                                PutDocumentStatsProto.TokenizationStats.newBuilder()
+                                        .setNumTokensIndexed(nativeNumTokensIndexed)
+                                        .setExceededMaxTokenNum(nativeExceededMaxNumTokens)
+                                        .build())
+                        .build();
+        PutDocumentStats.Builder pBuilder = new PutDocumentStats.Builder("packageName", "database");
+
+        AppSearchLoggerHelper.copyNativeStats(nativePutDocumentStats, pBuilder);
+
+        PutDocumentStats pStats = pBuilder.build();
+        assertThat(pStats.getNativeLatencyMillis()).isEqualTo(nativeLatencyMillis);
+        assertThat(pStats.getNativeDocumentStoreLatencyMillis())
+                .isEqualTo(nativeDocumentStoreLatencyMillis);
+        assertThat(pStats.getNativeIndexLatencyMillis()).isEqualTo(nativeIndexLatencyMillis);
+        assertThat(pStats.getNativeIndexMergeLatencyMillis())
+                .isEqualTo(nativeIndexMergeLatencyMillis);
+        assertThat(pStats.getNativeDocumentSizeBytes()).isEqualTo(nativeDocumentSize);
+        assertThat(pStats.getNativeNumTokensIndexed()).isEqualTo(nativeNumTokensIndexed);
+        assertThat(pStats.getNativeExceededMaxNumTokens()).isEqualTo(nativeExceededMaxNumTokens);
+    }
+
+    //
+    // Testing actual logging
+    //
+    @Test
+    public void testLoggingStats_putDocument() throws Exception {
+        // Insert schema
+        final String testPackageName = "testPackage";
+        final String testDatabase = "testDatabase";
+        List<AppSearchSchema> schemas =
+                Collections.singletonList(new AppSearchSchema.Builder("type").build());
+        mAppSearchImpl.setSchema(
+                testPackageName,
+                testDatabase,
+                schemas,
+                /*schemasNotPlatformSurfaceable=*/ Collections.emptyList(),
+                /*schemasPackageAccessible=*/ Collections.emptyMap(),
+                /*forceOverride=*/ false);
+        GenericDocument document =
+                new GenericDocument.Builder<>("namespace", "uri", "type").build();
+
+        mAppSearchImpl.putDocument(testPackageName, testDatabase, document, mLogger);
+
+        PutDocumentStats pStats = mLogger.mPutDocumentStats;
+        assertThat(pStats).isNotNull();
+        assertThat(pStats.getGeneralStats().getPackageName()).isEqualTo(testPackageName);
+        assertThat(pStats.getGeneralStats().getDatabase()).isEqualTo(testDatabase);
+        assertThat(pStats.getGeneralStats().getStatusCode()).isEqualTo(AppSearchResult.RESULT_OK);
+        // The rest of native stats have been tested in testCopyNativeStats
+        assertThat(pStats.getNativeDocumentSizeBytes()).isGreaterThan(0);
+    }
+}
diff --git a/services/tests/servicestests/src/com/android/server/appsearch/external/localstorage/stats/AppSearchStatsTest.java b/services/tests/servicestests/src/com/android/server/appsearch/external/localstorage/stats/AppSearchStatsTest.java
index 4308885faaad..8dbf249c4e96 100644
--- a/services/tests/servicestests/src/com/android/server/appsearch/external/localstorage/stats/AppSearchStatsTest.java
+++ b/services/tests/servicestests/src/com/android/server/appsearch/external/localstorage/stats/AppSearchStatsTest.java
@@ -18,12 +18,14 @@ package com.android.server.appsearch.external.localstorage.stats;
 
 import static com.google.common.truth.Truth.assertThat;
 
+import android.app.appsearch.AppSearchResult;
+
 import org.junit.Test;
 
 public class AppSearchStatsTest {
     static final String TEST_PACKAGE_NAME = "com.google.test";
     static final String TEST_DATA_BASE = "testDataBase";
-    static final int TEST_STATUS_CODE = 2;
+    static final int TEST_STATUS_CODE = AppSearchResult.RESULT_INTERNAL_ERROR;
     static final int TEST_TOTAL_LATENCY_MILLIS = 20;
 
     @Test
@@ -40,25 +42,38 @@ public class AppSearchStatsTest {
         assertThat(gStats.getTotalLatencyMillis()).isEqualTo(TEST_TOTAL_LATENCY_MILLIS);
     }
 
+    /** Make sure status code is UNKNOWN if not set in {@link GeneralStats} */
     @Test
-    public void testAppSearchStats_CallStats() {
-        final int estimatedBinderLatencyMillis = 1;
-        final int numOperationsSucceeded = 2;
-        final int numOperationsFailed = 3;
-
+    public void testAppSearchStats_GeneralStats_defaultStatsCode_Unknown() {
         final GeneralStats gStats =
                 new GeneralStats.Builder(TEST_PACKAGE_NAME, TEST_DATA_BASE)
-                        .setStatusCode(TEST_STATUS_CODE)
                         .setTotalLatencyMillis(TEST_TOTAL_LATENCY_MILLIS)
                         .build();
+
+        assertThat(gStats.getPackageName()).isEqualTo(TEST_PACKAGE_NAME);
+        assertThat(gStats.getDatabase()).isEqualTo(TEST_DATA_BASE);
+        assertThat(gStats.getStatusCode()).isEqualTo(AppSearchResult.RESULT_UNKNOWN_ERROR);
+        assertThat(gStats.getTotalLatencyMillis()).isEqualTo(TEST_TOTAL_LATENCY_MILLIS);
+    }
+
+    @Test
+    public void testAppSearchStats_CallStats() {
+        final int estimatedBinderLatencyMillis = 1;
+        final int numOperationsSucceeded = 2;
+        final int numOperationsFailed = 3;
         final @CallStats.CallType int callType = CallStats.CALL_TYPE_PUT_DOCUMENTS;
-        final CallStats cStats =
-                new CallStats.Builder(gStats)
+
+        final CallStats.Builder cStatsBuilder =
+                new CallStats.Builder(TEST_PACKAGE_NAME, TEST_DATA_BASE)
                         .setCallType(callType)
                         .setEstimatedBinderLatencyMillis(estimatedBinderLatencyMillis)
                         .setNumOperationsSucceeded(numOperationsSucceeded)
-                        .setNumOperationsFailed(numOperationsFailed)
-                        .build();
+                        .setNumOperationsFailed(numOperationsFailed);
+        cStatsBuilder
+                .getGeneralStatsBuilder()
+                .setStatusCode(TEST_STATUS_CODE)
+                .setTotalLatencyMillis(TEST_TOTAL_LATENCY_MILLIS);
+        final CallStats cStats = cStatsBuilder.build();
 
         assertThat(cStats.getGeneralStats().getPackageName()).isEqualTo(TEST_PACKAGE_NAME);
         assertThat(cStats.getGeneralStats().getDatabase()).isEqualTo(TEST_DATA_BASE);
@@ -82,15 +97,9 @@ public class AppSearchStatsTest {
         final int nativeIndexMergeLatencyMillis = 6;
         final int nativeDocumentSize = 7;
         final int nativeNumTokensIndexed = 8;
-        final int nativeNumTokensClipped = 9;
-
-        final GeneralStats gStats =
-                new GeneralStats.Builder(TEST_PACKAGE_NAME, TEST_DATA_BASE)
-                        .setStatusCode(TEST_STATUS_CODE)
-                        .setTotalLatencyMillis(TEST_TOTAL_LATENCY_MILLIS)
-                        .build();
-        final PutDocumentStats pStats =
-                new PutDocumentStats.Builder(gStats)
+        final boolean nativeExceededMaxNumTokens = true;
+        final PutDocumentStats.Builder pStatsBuilder =
+                new PutDocumentStats.Builder(TEST_PACKAGE_NAME, TEST_DATA_BASE)
                         .setGenerateDocumentProtoLatencyMillis(generateDocumentProtoLatencyMillis)
                         .setRewriteDocumentTypesLatencyMillis(rewriteDocumentTypesLatencyMillis)
                         .setNativeLatencyMillis(nativeLatencyMillis)
@@ -99,8 +108,12 @@ public class AppSearchStatsTest {
                         .setNativeIndexMergeLatencyMillis(nativeIndexMergeLatencyMillis)
                         .setNativeDocumentSizeBytes(nativeDocumentSize)
                         .setNativeNumTokensIndexed(nativeNumTokensIndexed)
-                        .setNativeNumTokensClipped(nativeNumTokensClipped)
-                        .build();
+                        .setNativeExceededMaxNumTokens(nativeExceededMaxNumTokens);
+        pStatsBuilder
+                .getGeneralStatsBuilder()
+                .setStatusCode(TEST_STATUS_CODE)
+                .setTotalLatencyMillis(TEST_TOTAL_LATENCY_MILLIS);
+        final PutDocumentStats pStats = pStatsBuilder.build();
 
         assertThat(pStats.getGeneralStats().getPackageName()).isEqualTo(TEST_PACKAGE_NAME);
         assertThat(pStats.getGeneralStats().getDatabase()).isEqualTo(TEST_DATA_BASE);
@@ -119,6 +132,6 @@ public class AppSearchStatsTest {
                 .isEqualTo(nativeIndexMergeLatencyMillis);
         assertThat(pStats.getNativeDocumentSizeBytes()).isEqualTo(nativeDocumentSize);
         assertThat(pStats.getNativeNumTokensIndexed()).isEqualTo(nativeNumTokensIndexed);
-        assertThat(pStats.getNativeNumTokensClipped()).isEqualTo(nativeNumTokensClipped);
+        assertThat(pStats.getNativeExceededMaxNumTokens()).isEqualTo(nativeExceededMaxNumTokens);
     }
 }