Implement schema migration in framework.

Changes included:
* 9d1cf52:Supports schema migration in AppSearch

Bug: 177266929
Test: AppSearchSchemaMigrationCtsTest
Change-Id: Ia5b964baeb413ec7d176bb12c9e4c2af26cf2bf3

6 files changed, 694 insertions, 64 deletions

diff --git a/apex/appsearch/framework/java/android/app/appsearch/AppSearchMigrationHelper.java b/apex/appsearch/framework/java/android/app/appsearch/AppSearchMigrationHelper.java
new file mode 100644
index 000000000000..e585d9147895
--- /dev/null
+++ b/apex/appsearch/framework/java/android/app/appsearch/AppSearchMigrationHelper.java
@@ -0,0 +1,256 @@
+/*
+ * Copyright (C) 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 static android.os.ParcelFileDescriptor.MODE_READ_ONLY;
+import static android.os.ParcelFileDescriptor.MODE_WRITE_ONLY;
+
+import android.annotation.NonNull;
+import android.annotation.UserIdInt;
+import android.annotation.WorkerThread;
+import android.app.appsearch.exceptions.AppSearchException;
+import android.os.Bundle;
+import android.os.Parcel;
+import android.os.ParcelFileDescriptor;
+import android.os.RemoteException;
+
+import com.android.internal.infra.AndroidFuture;
+
+import java.io.Closeable;
+import java.io.DataInputStream;
+import java.io.DataOutputStream;
+import java.io.EOFException;
+import java.io.File;
+import java.io.FileInputStream;
+import java.io.FileOutputStream;
+import java.io.IOException;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.concurrent.ExecutionException;
+
+/**
+ * The helper class for {@link AppSearchSchema} migration.
+ *
+ * <p>It will query and migrate {@link GenericDocument} in given type to a new version.
+ * @hide
+ */
+public class AppSearchMigrationHelper implements Closeable {
+    private final IAppSearchManager mService;
+    private final String mPackageName;
+    private final String mDatabaseName;
+    private final int mUserId;
+    private final File mMigratedFile;
+    private final Map<String, Integer> mCurrentVersionMap;
+    private final Map<String, Integer> mFinalVersionMap;
+    private boolean mAreDocumentsMigrated = false;
+
+    AppSearchMigrationHelper(@NonNull IAppSearchManager service,
+            @UserIdInt int userId,
+            @NonNull Map<String, Integer> currentVersionMap,
+            @NonNull Map<String, Integer> finalVersionMap,
+            @NonNull String packageName,
+            @NonNull String databaseName) throws IOException {
+        mService = Objects.requireNonNull(service);
+        mCurrentVersionMap = Objects.requireNonNull(currentVersionMap);
+        mFinalVersionMap = Objects.requireNonNull(finalVersionMap);
+        mPackageName = Objects.requireNonNull(packageName);
+        mDatabaseName = Objects.requireNonNull(databaseName);
+        mUserId = userId;
+        mMigratedFile = File.createTempFile(/*prefix=*/"appsearch", /*suffix=*/null);
+    }
+
+    /**
+     * Queries all documents that need to be migrated to a different version and transform
+     * documents to that version by passing them to the provided {@link Migrator}.
+     *
+     * <p>The method will be executed on the executor provided to
+     * {@link AppSearchSession#setSchema}.
+     *
+     * @param schemaType The schema type that needs to be updated and whose {@link GenericDocument}
+     *                   need to be migrated.
+     * @param migrator The {@link Migrator} that will upgrade or downgrade a {@link
+     *     GenericDocument} to new version.
+     */
+    @WorkerThread
+    public void queryAndTransform(@NonNull String schemaType, @NonNull Migrator migrator)
+            throws IOException, AppSearchException, InterruptedException, ExecutionException {
+        File queryFile = File.createTempFile(/*prefix=*/"appsearch", /*suffix=*/null);
+        try (ParcelFileDescriptor fileDescriptor =
+                     ParcelFileDescriptor.open(queryFile, MODE_WRITE_ONLY)) {
+            AndroidFuture<AppSearchResult<Void>> androidFuture = new AndroidFuture<>();
+            mService.writeQueryResultsToFile(mPackageName, mDatabaseName,
+                    fileDescriptor,
+                    /*queryExpression=*/ "",
+                    new SearchSpec.Builder()
+                            .addFilterSchemas(schemaType)
+                            .setTermMatch(SearchSpec.TERM_MATCH_EXACT_ONLY)
+                            .build().getBundle(),
+                    mUserId,
+                    new IAppSearchResultCallback.Stub() {
+                        @Override
+                        public void onResult(AppSearchResult result) throws RemoteException {
+                            androidFuture.complete(result);
+                        }
+                    });
+            AppSearchResult<Void> result = androidFuture.get();
+            if (!result.isSuccess()) {
+                throw new AppSearchException(result.getResultCode(), result.getErrorMessage());
+            }
+            readAndTransform(queryFile, migrator);
+        } catch (RemoteException e) {
+            throw e.rethrowFromSystemServer();
+        } finally {
+            queryFile.delete();
+        }
+    }
+
+    /**
+     * Puts all {@link GenericDocument} migrated from the previous call to
+     * {@link #queryAndTransform} into AppSearch.
+     *
+     * <p> This method should be only called once.
+     *
+     * @param responseBuilder a SetSchemaResponse builder whose result will be returned by this
+     *                        function with any
+     *                        {@link android.app.appsearch.SetSchemaResponse.MigrationFailure}
+     *                        added in.
+     * @return the {@link SetSchemaResponse} for {@link AppSearchSession#setSchema} call.
+     */
+    @NonNull
+    AppSearchResult<SetSchemaResponse> putMigratedDocuments(
+            @NonNull SetSchemaResponse.Builder responseBuilder) {
+        if (!mAreDocumentsMigrated) {
+            return AppSearchResult.newSuccessfulResult(responseBuilder.build());
+        }
+        try (ParcelFileDescriptor fileDescriptor =
+                     ParcelFileDescriptor.open(mMigratedFile, MODE_READ_ONLY)) {
+            AndroidFuture<AppSearchResult<List<Bundle>>> androidFuture = new AndroidFuture<>();
+            mService.putDocumentsFromFile(mPackageName, mDatabaseName, fileDescriptor, mUserId,
+                    new IAppSearchResultCallback.Stub() {
+                        @Override
+                        public void onResult(AppSearchResult result) throws RemoteException {
+                            androidFuture.complete(result);
+                        }
+                    });
+            AppSearchResult<List<Bundle>> result = androidFuture.get();
+            if (!result.isSuccess()) {
+                return AppSearchResult.newFailedResult(result);
+            }
+            List<Bundle> migratedFailureBundles = result.getResultValue();
+            for (int i = 0; i < migratedFailureBundles.size(); i++) {
+                responseBuilder.addMigrationFailure(
+                        new SetSchemaResponse.MigrationFailure(migratedFailureBundles.get(i)));
+            }
+        } catch (RemoteException e) {
+            throw e.rethrowFromSystemServer();
+        } catch (Throwable t) {
+            return AppSearchResult.throwableToFailedResult(t);
+        } finally {
+            mMigratedFile.delete();
+        }
+        return AppSearchResult.newSuccessfulResult(responseBuilder.build());
+    }
+
+    /**
+     * Reads all saved {@link GenericDocument}s from the given {@link File}.
+     *
+     * <p>Transforms those {@link GenericDocument}s to the final version.
+     *
+     * <p>Save migrated {@link GenericDocument}s to the {@link #mMigratedFile}.
+     */
+    private void readAndTransform(@NonNull File file, @NonNull Migrator migrator)
+            throws IOException {
+        try (DataInputStream inputStream = new DataInputStream(new FileInputStream(file));
+             DataOutputStream outputStream = new DataOutputStream(new FileOutputStream(
+                     mMigratedFile, /*append=*/ true))) {
+            GenericDocument document;
+            while (true) {
+                try {
+                    document = readDocumentFromInputStream(inputStream);
+                } catch (EOFException e) {
+                    break;
+                    // Nothing wrong. We just finished reading.
+                }
+
+                int currentVersion = mCurrentVersionMap.get(document.getSchemaType());
+                int finalVersion = mFinalVersionMap.get(document.getSchemaType());
+
+                GenericDocument newDocument;
+                if (currentVersion < finalVersion) {
+                    newDocument = migrator.onUpgrade(currentVersion, finalVersion, document);
+                } else {
+                    // currentVersion == finalVersion case won't trigger migration and get here.
+                    newDocument = migrator.onDowngrade(currentVersion, finalVersion, document);
+                }
+                writeBundleToOutputStream(outputStream, newDocument.getBundle());
+            }
+            mAreDocumentsMigrated = true;
+        }
+    }
+
+    /**
+     * Reads the {@link Bundle} of a {@link GenericDocument} from given {@link DataInputStream}.
+     *
+     * @param inputStream The inputStream to read from
+     *
+     * @throws IOException        on read failure.
+     * @throws EOFException       if {@link java.io.InputStream} reaches the end.
+     */
+    @NonNull
+    public static GenericDocument readDocumentFromInputStream(
+            @NonNull DataInputStream inputStream) throws IOException {
+        int length = inputStream.readInt();
+        if (length == 0) {
+            throw new EOFException();
+        }
+        byte[] serializedMessage = new byte[length];
+        inputStream.read(serializedMessage);
+
+        Parcel parcel = Parcel.obtain();
+        try {
+            parcel.unmarshall(serializedMessage, 0, serializedMessage.length);
+            parcel.setDataPosition(0);
+            Bundle bundle = parcel.readBundle();
+            return new GenericDocument(bundle);
+        } finally {
+            parcel.recycle();
+        }
+    }
+
+    /**
+     * Serializes a {@link Bundle} and writes into the given {@link DataOutputStream}.
+     */
+    public static void writeBundleToOutputStream(
+            @NonNull DataOutputStream outputStream, @NonNull Bundle bundle)
+            throws IOException {
+        Parcel parcel = Parcel.obtain();
+        try {
+            parcel.writeBundle(bundle);
+            byte[] serializedMessage = parcel.marshall();
+            outputStream.writeInt(serializedMessage.length);
+            outputStream.write(serializedMessage);
+        } finally {
+            parcel.recycle();
+        }
+    }
+
+    @Override
+    public void close() throws IOException {
+        mMigratedFile.delete();
+    }
+}
diff --git a/apex/appsearch/framework/java/android/app/appsearch/AppSearchSession.java b/apex/appsearch/framework/java/android/app/appsearch/AppSearchSession.java
index bf733ed77442..a0718dc956f7 100644
--- a/apex/appsearch/framework/java/android/app/appsearch/AppSearchSession.java
+++ b/apex/appsearch/framework/java/android/app/appsearch/AppSearchSession.java
@@ -19,6 +19,8 @@ package android.app.appsearch;
 import android.annotation.CallbackExecutor;
 import android.annotation.NonNull;
 import android.annotation.UserIdInt;
+import android.app.appsearch.exceptions.AppSearchException;
+import android.app.appsearch.util.SchemaMigrationUtil;
 import android.os.Bundle;
 import android.os.ParcelableException;
 import android.os.RemoteException;
@@ -26,14 +28,17 @@ import android.util.ArrayMap;
 import android.util.ArraySet;
 import android.util.Log;
 
+import com.android.internal.infra.AndroidFuture;
 import com.android.internal.util.Preconditions;
 
 import java.io.Closeable;
+import java.io.IOException;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
 import java.util.Set;
+import java.util.concurrent.ExecutionException;
 import java.util.concurrent.Executor;
 import java.util.function.Consumer;
 
@@ -55,7 +60,6 @@ public final class AppSearchSession implements Closeable {
     private boolean mIsMutated = false;
     private boolean mIsClosed = false;
 
-
     /**
      * Creates a search session for the client, defined by the {@code userId} and
      * {@code packageName}.
@@ -157,26 +161,65 @@ public final class AppSearchSession implements Closeable {
             }
             schemasPackageAccessibleBundles.put(entry.getKey(), packageIdentifierBundles);
         }
+
+        // No need to trigger migration if user never set migrator
+        if (request.getMigrators().isEmpty()) {
+            setSchemaNoMigrations(
+                    request,
+                    schemaBundles,
+                    schemasPackageAccessibleBundles,
+                    callbackExecutor,
+                    callback);
+            return;
+        }
+
         try {
+            // Migration process
+            // 1. Generate the current and the final version map.
+            // TODO(b/182855402) Release binder thread and move the heavy work into worker thread.
+            AndroidFuture<AppSearchResult<GetSchemaResponse>> future = new AndroidFuture<>();
+            getSchema(callbackExecutor, future::complete);
+            AppSearchResult<GetSchemaResponse> getSchemaResult = future.get();
+            if (!getSchemaResult.isSuccess()) {
+                callback.accept(AppSearchResult.newFailedResult(getSchemaResult));
+                return;
+            }
+            GetSchemaResponse getSchemaResponse = getSchemaResult.getResultValue();
+            Set<AppSearchSchema> currentSchemas = getSchemaResponse.getSchemas();
+            Map<String, Integer> currentVersionMap =
+                    SchemaMigrationUtil.buildVersionMap(currentSchemas,
+                            getSchemaResponse.getVersion());
+            Map<String, Integer> finalVersionMap =
+                    SchemaMigrationUtil.buildVersionMap(request.getSchemas(), request.getVersion());
+
+            // 2. SetSchema with forceOverride=false, to retrieve the list of incompatible/deleted
+            // types.
             mService.setSchema(
                     mPackageName,
                     mDatabaseName,
                     schemaBundles,
                     new ArrayList<>(request.getSchemasNotDisplayedBySystem()),
                     schemasPackageAccessibleBundles,
-                    request.isForceOverride(),
+                    /*forceOverride=*/ false,
                     mUserId,
                     request.getVersion(),
                     new IAppSearchResultCallback.Stub() {
                         public void onResult(AppSearchResult result) {
                             callbackExecutor.execute(() -> {
                                 if (result.isSuccess()) {
-                                    callback.accept(
-                                            // TODO(b/177266929) implement Migration in platform.
-                                            // TODO(b/183177268): once migration is implemented, run
-                                            //  it on workExecutor.
-                                            AppSearchResult.newSuccessfulResult(
-                                                    new SetSchemaResponse.Builder().build()));
+                                    // TODO(b/183177268): once migration is implemented, run
+                                    //  it on workExecutor.
+                                    try {
+                                        Bundle bundle = (Bundle) result.getResultValue();
+                                        SetSchemaResponse setSchemaResponse =
+                                                new SetSchemaResponse(bundle);
+                                        setSchemaMigration(
+                                                request, setSchemaResponse, schemaBundles,
+                                                schemasPackageAccessibleBundles, currentVersionMap,
+                                                finalVersionMap, callback);
+                                    } catch (Throwable t) {
+                                        callback.accept(AppSearchResult.throwableToFailedResult(t));
+                                    }
                                 } else {
                                     callback.accept(result);
                                 }
@@ -186,6 +229,8 @@ public final class AppSearchSession implements Closeable {
             mIsMutated = true;
         } catch (RemoteException e) {
             throw e.rethrowFromSystemServer();
+        }  catch (Throwable t) {
+            callback.accept(AppSearchResult.throwableToFailedResult(t));
         }
     }
 
@@ -627,4 +672,159 @@ public final class AppSearchSession implements Closeable {
             }
         }
     }
+
+    /**
+     * Set schema to Icing for no-migration scenario.
+     *
+     * <p>We only need one time {@link #setSchema} call for no-migration scenario by using the
+     * forceoverride in the request.
+     */
+    private void setSchemaNoMigrations(@NonNull SetSchemaRequest request,
+            @NonNull List<Bundle> schemaBundles,
+            @NonNull Map<String, List<Bundle>> schemasPackageAccessibleBundles,
+            @NonNull @CallbackExecutor Executor executor,
+            @NonNull Consumer<AppSearchResult<SetSchemaResponse>> callback) {
+        try {
+            mService.setSchema(
+                    mPackageName,
+                    mDatabaseName,
+                    schemaBundles,
+                    new ArrayList<>(request.getSchemasNotDisplayedBySystem()),
+                    schemasPackageAccessibleBundles,
+                    request.isForceOverride(),
+                    mUserId,
+                    request.getVersion(),
+                    new IAppSearchResultCallback.Stub() {
+                        public void onResult(AppSearchResult result) {
+                            executor.execute(() -> {
+                                if (result.isSuccess()) {
+                                    try {
+                                        SetSchemaResponse setSchemaResponse =
+                                                new SetSchemaResponse(
+                                                        (Bundle) result.getResultValue());
+                                        if (!request.isForceOverride()) {
+                                            // Throw exception if there is any deleted types or
+                                            // incompatible types. That's the only case we swallowed
+                                            // in the AppSearchImpl#setSchema().
+                                            checkDeletedAndIncompatible(
+                                                    setSchemaResponse.getDeletedTypes(),
+                                                    setSchemaResponse.getIncompatibleTypes());
+                                        }
+                                        callback.accept(AppSearchResult
+                                                .newSuccessfulResult(setSchemaResponse));
+                                    } catch (Throwable t) {
+                                        callback.accept(AppSearchResult.throwableToFailedResult(t));
+                                    }
+                                } else {
+                                    callback.accept(result);
+                                }
+                            });
+                        }
+                    });
+            mIsMutated = true;
+        } catch (RemoteException e) {
+            throw e.rethrowFromSystemServer();
+        }
+    }
+
+    /**
+     * Set schema to Icing for migration scenario.
+     *
+     * <p>First time {@link #setSchema} call with forceOverride is false gives us all incompatible
+     * changes. After trigger migrations, the second time call {@link #setSchema} will actually
+     * apply the changes.
+     *
+     * @param setSchemaResponse the result of the first setSchema call with forceOverride=false.
+     */
+    private void setSchemaMigration(@NonNull SetSchemaRequest request,
+            @NonNull SetSchemaResponse setSchemaResponse,
+            @NonNull List<Bundle> schemaBundles,
+            @NonNull Map<String, List<Bundle>> schemasPackageAccessibleBundles,
+            @NonNull Map<String, Integer> currentVersionMap, Map<String, Integer> finalVersionMap,
+            @NonNull Consumer<AppSearchResult<SetSchemaResponse>> callback)
+            throws AppSearchException, IOException, RemoteException, ExecutionException,
+            InterruptedException {
+        // 1. If forceOverride is false, check that all incompatible types will be migrated.
+        // If some aren't we must throw an error, rather than proceeding and deleting those
+        // types.
+        if (!request.isForceOverride()) {
+            Set<String> unmigratedTypes = SchemaMigrationUtil.getUnmigratedIncompatibleTypes(
+                    setSchemaResponse.getIncompatibleTypes(),
+                    request.getMigrators(),
+                    currentVersionMap,
+                    finalVersionMap);
+            // check if there are any unmigrated types or deleted types. If there are, we will throw
+            // an exception.
+            // Since the force override is false, the schema will not have been set if there are any
+            // incompatible or deleted types.
+            checkDeletedAndIncompatible(setSchemaResponse.getDeletedTypes(),
+                    unmigratedTypes);
+        }
+
+        try (AppSearchMigrationHelper migrationHelper =
+                     new AppSearchMigrationHelper(mService, mUserId, currentVersionMap,
+                             finalVersionMap, mPackageName, mDatabaseName)) {
+            Map<String, Migrator> migratorMap = request.getMigrators();
+
+            // 2. Trigger migration for all migrators.
+            // TODO(b/177266929) trigger migration for all types together rather than separately.
+            Set<String> migratedTypes = new ArraySet<>();
+            for (Map.Entry<String, Migrator> entry : migratorMap.entrySet()) {
+                String schemaType = entry.getKey();
+                Migrator migrator = entry.getValue();
+                if (SchemaMigrationUtil.shouldTriggerMigration(
+                        schemaType, migrator, currentVersionMap, finalVersionMap)) {
+                    migrationHelper.queryAndTransform(schemaType, migrator);
+                    migratedTypes.add(schemaType);
+                }
+            }
+
+            // 3. SetSchema a second time with forceOverride=true if the first attempted failed.
+            if (!setSchemaResponse.getIncompatibleTypes().isEmpty()
+                    || !setSchemaResponse.getDeletedTypes().isEmpty()) {
+                AndroidFuture<AppSearchResult<SetSchemaResponse>> future = new AndroidFuture<>();
+                // only trigger second setSchema() call if the first one is fail.
+                mService.setSchema(
+                        mPackageName,
+                        mDatabaseName,
+                        schemaBundles,
+                        new ArrayList<>(request.getSchemasNotDisplayedBySystem()),
+                        schemasPackageAccessibleBundles,
+                        /*forceOverride=*/ true,
+                        mUserId,
+                        request.getVersion(),
+                        new IAppSearchResultCallback.Stub() {
+                            @Override
+                            public void onResult(AppSearchResult result) throws RemoteException {
+                                future.complete(result);
+                            }
+                        });
+                AppSearchResult<SetSchemaResponse> secondSetSchemaResult = future.get();
+                if (!secondSetSchemaResult.isSuccess()) {
+                    // we failed to set the schema in second time with force override = true, which
+                    // is an impossible case. Since we only swallow the incompatible error in the
+                    // first setSchema call, all other errors will be thrown at the first time.
+                    callback.accept(secondSetSchemaResult);
+                    return;
+                }
+            }
+
+            SetSchemaResponse.Builder responseBuilder = setSchemaResponse.toBuilder()
+                    .addMigratedTypes(migratedTypes);
+            callback.accept(migrationHelper.putMigratedDocuments(responseBuilder));
+        }
+    }
+
+    /**  Checks the setSchema() call won't delete any types or has incompatible types. */
+    //TODO(b/177266929) move this method to util
+    private void checkDeletedAndIncompatible(Set<String> deletedTypes,
+            Set<String> incompatibleTypes)
+            throws AppSearchException {
+        if (!deletedTypes.isEmpty() || !incompatibleTypes.isEmpty()) {
+            String newMessage = "Schema is incompatible."
+                    + "\n  Deleted types: " + deletedTypes
+                    + "\n  Incompatible types: " + incompatibleTypes;
+            throw new AppSearchException(AppSearchResult.RESULT_INVALID_SCHEMA, newMessage);
+        }
+    }
 }
diff --git a/apex/appsearch/framework/java/android/app/appsearch/IAppSearchManager.aidl b/apex/appsearch/framework/java/android/app/appsearch/IAppSearchManager.aidl
index d436488f34f2..48c397f324ec 100644
--- a/apex/appsearch/framework/java/android/app/appsearch/IAppSearchManager.aidl
+++ b/apex/appsearch/framework/java/android/app/appsearch/IAppSearchManager.aidl
@@ -21,6 +21,7 @@ import android.app.appsearch.AppSearchBatchResult;
 import android.app.appsearch.AppSearchResult;
 import android.app.appsearch.IAppSearchBatchResultCallback;
 import android.app.appsearch.IAppSearchResultCallback;
+import android.os.ParcelFileDescriptor;
 import com.android.internal.infra.AndroidFuture;
 
 parcelable SearchResults;
@@ -41,7 +42,8 @@ interface IAppSearchManager {
      *     incompatible documents will be deleted.
      * @param userId Id of the calling user
      * @param callback {@link IAppSearchResultCallback#onResult} will be called with an
-     *     {@link AppSearchResult}&lt;{@link Void}&gt;.
+     *     {@link AppSearchResult}&lt;{@link Bundle}&gt;, where the value are
+     *     {@link SetSchemaResponse} bundle.
      */
     void setSchema(
         in String packageName,
@@ -189,6 +191,47 @@ interface IAppSearchManager {
     void invalidateNextPageToken(in long nextPageToken, in int userId);
 
     /**
+    * Searches a document based on a given specifications.
+    *
+    * <p>Documents will be save to the given ParcelFileDescriptor
+    *
+    * @param packageName The name of the package to query over.
+    * @param databaseName The databaseName this query for.
+    * @param fileDescriptor The ParcelFileDescriptor where documents should be written to.
+    * @param queryExpression String to search for.
+    * @param searchSpecBundle SearchSpec bundle.
+    * @param userId Id of the calling user.
+    * @param callback {@link IAppSearchResultCallback#onResult} will be called with an
+    *        {@link AppSearchResult}&lt;{@code null}&gt;.
+    */
+    void writeQueryResultsToFile(
+        in String packageName,
+        in String databaseName,
+        in ParcelFileDescriptor fileDescriptor,
+        in String queryExpression,
+        in Bundle searchSpecBundle,
+        in int userId,
+        in IAppSearchResultCallback callback);
+
+    /**
+    * Inserts documents from the given file into the index.
+    *
+    * @param packageName The name of the package that owns this document.
+    * @param databaseName  The name of the database where this document lives.
+    * @param fileDescriptor The ParcelFileDescriptor where documents should be read from.
+    * @param userId Id of the calling user.
+    * @param callback {@link IAppSearchResultCallback#onResult} will be called with an
+    *     {@link AppSearchResult}&lt;{@link List}&lt;{@link Bundle}&gt;&gt;, where the value are
+    *     MigrationFailure bundles.
+    */
+    void putDocumentsFromFile(
+        in String packageName,
+        in String databaseName,
+        in ParcelFileDescriptor fileDescriptor,
+        in int userId,
+        in IAppSearchResultCallback callback);
+
+    /**
      * Reports usage of a particular document by URI and namespace.
      *
      * <p>A usage report represents an event in which a user interacted with or viewed a document.
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
index 32d7e043e954..c9473bdeff32 100644
--- a/apex/appsearch/framework/java/external/android/app/appsearch/util/SchemaMigrationUtil.java
+++ b/apex/appsearch/framework/java/external/android/app/appsearch/util/SchemaMigrationUtil.java
@@ -20,11 +20,12 @@ import android.annotation.NonNull;
 import android.app.appsearch.AppSearchResult;
 import android.app.appsearch.AppSearchSchema;
 import android.app.appsearch.Migrator;
-import android.app.appsearch.SetSchemaResponse;
 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;
@@ -35,70 +36,84 @@ import java.util.Set;
  * @hide
  */
 public final class SchemaMigrationUtil {
+    private static final String TAG = "AppSearchMigrateUtil";
+
     private SchemaMigrationUtil() {}
 
-    /** Returns all active {@link Migrator}s that need to be triggered in this migration. */
+    /**
+     * Finds out which incompatible schema type won't be migrated by comparing its current and final
+     * version number.
+     */
     @NonNull
-    public static Map<String, Migrator> getActiveMigrators(
-            @NonNull Set<AppSearchSchema> existingSchemas,
+    public static Set<String> getUnmigratedIncompatibleTypes(
+            @NonNull Set<String> incompatibleSchemaTypes,
             @NonNull Map<String, Migrator> migrators,
-            int currentVersion,
-            int finalVersion) {
-        if (currentVersion == finalVersion) {
-            return Collections.emptyMap();
-        }
-        Set<String> existingTypes = new ArraySet<>(existingSchemas.size());
-        for (AppSearchSchema schema : existingSchemas) {
-            existingTypes.add(schema.getSchemaType());
-        }
-
-        Map<String, Migrator> activeMigrators = new ArrayMap<>();
-        for (Map.Entry<String, Migrator> entry : migrators.entrySet()) {
-            // The device contains the source type, and we should trigger migration for the type.
-            String schemaType = entry.getKey();
-            Migrator migrator = entry.getValue();
-            if (existingTypes.contains(schemaType)
-                    && migrator.shouldMigrate(currentVersion, finalVersion)) {
-                activeMigrators.put(schemaType, migrator);
+            @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.shouldMigrate(currentVersion, finalVersion)) {
+                unmigratedSchemaTypes.add(unmigratedSchemaType);
             }
         }
-        return activeMigrators;
+        return Collections.unmodifiableSet(unmigratedSchemaTypes);
     }
 
     /**
-     * Checks the setSchema() call won't delete any types or has incompatible types after all {@link
-     * Migrator} has been triggered..
+     * 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 void checkDeletedAndIncompatibleAfterMigration(
-            @NonNull SetSchemaResponse setSchemaResponse, @NonNull Set<String> activeMigrators)
+    public static boolean shouldTriggerMigration(
+            @NonNull String schemaType,
+            @NonNull Migrator migrator,
+            @NonNull Map<String, Integer> currentVersionMap,
+            @NonNull Map<String, Integer> finalVersionMap)
             throws AppSearchException {
-        Set<String> unmigratedIncompatibleTypes =
-                new ArraySet<>(setSchemaResponse.getIncompatibleTypes());
-        unmigratedIncompatibleTypes.removeAll(activeMigrators);
-
-        Set<String> unmigratedDeletedTypes = new ArraySet<>(setSchemaResponse.getDeletedTypes());
-        unmigratedDeletedTypes.removeAll(activeMigrators);
-
-        // check if there are any unmigrated incompatible types or deleted types. If there
-        // are, we will getActiveMigratorsthrow an exception. That's the only case we
-        // swallowed in the AppSearchImpl#setSchema().
-        // Since the force override is false, the schema will not have been set if there are
-        // any incompatible or deleted types.
-        checkDeletedAndIncompatible(unmigratedDeletedTypes, unmigratedIncompatibleTypes);
+        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.shouldMigrate(currentVersion, finalVersion);
     }
 
-    /** Checks the setSchema() call won't delete any types or has incompatible types. */
-    public static void checkDeletedAndIncompatible(
-            @NonNull Set<String> deletedTypes, @NonNull Set<String> incompatibleTypes)
-            throws AppSearchException {
-        if (deletedTypes.size() > 0 || incompatibleTypes.size() > 0) {
-            String newMessage =
-                    "Schema is incompatible."
-                            + "\n  Deleted types: "
-                            + deletedTypes
-                            + "\n  Incompatible types: "
-                            + incompatibleTypes;
-            throw new AppSearchException(AppSearchResult.RESULT_INVALID_SCHEMA, newMessage);
+    /** Builds a Map of SchemaType and its version of given set of {@link AppSearchSchema}. */
+    //TODO(b/182620003) remove this method once support migrate to another type
+    @NonNull
+    public static Map<String, Integer> buildVersionMap(
+            @NonNull Collection<AppSearchSchema> schemas, int version) {
+        Map<String, Integer> currentVersionMap = new ArrayMap<>(schemas.size());
+        for (AppSearchSchema currentSchema : schemas) {
+            currentVersionMap.put(currentSchema.getSchemaType(), version);
         }
+        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 6e3fb82ba213..91ed6cd4a638 100644
--- a/apex/appsearch/service/java/com/android/server/appsearch/AppSearchManagerService.java
+++ b/apex/appsearch/service/java/com/android/server/appsearch/AppSearchManagerService.java
@@ -21,6 +21,7 @@ import android.annotation.NonNull;
 import android.annotation.UserIdInt;
 import android.app.ActivityManager;
 import android.app.appsearch.AppSearchBatchResult;
+import android.app.appsearch.AppSearchMigrationHelper;
 import android.app.appsearch.AppSearchResult;
 import android.app.appsearch.AppSearchSchema;
 import android.app.appsearch.GenericDocument;
@@ -31,10 +32,12 @@ import android.app.appsearch.IAppSearchResultCallback;
 import android.app.appsearch.PackageIdentifier;
 import android.app.appsearch.SearchResultPage;
 import android.app.appsearch.SearchSpec;
+import android.app.appsearch.SetSchemaResponse;
 import android.content.Context;
 import android.content.pm.PackageManagerInternal;
 import android.os.Binder;
 import android.os.Bundle;
+import android.os.ParcelFileDescriptor;
 import android.os.ParcelableException;
 import android.os.RemoteException;
 import android.os.UserHandle;
@@ -49,6 +52,11 @@ import com.android.server.LocalServices;
 import com.android.server.SystemService;
 import com.android.server.appsearch.external.localstorage.AppSearchImpl;
 
+import java.io.DataInputStream;
+import java.io.DataOutputStream;
+import java.io.EOFException;
+import java.io.FileInputStream;
+import java.io.FileOutputStream;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
@@ -125,7 +133,7 @@ public class AppSearchManagerService extends SystemService {
                     schemasPackageAccessible.put(entry.getKey(), packageIdentifiers);
                 }
                 AppSearchImpl impl = mImplInstanceManager.getAppSearchImpl(callingUserId);
-                impl.setSchema(
+                SetSchemaResponse setSchemaResponse = impl.setSchema(
                         packageName,
                         databaseName,
                         schemas,
@@ -133,8 +141,8 @@ public class AppSearchManagerService extends SystemService {
                         schemasPackageAccessible,
                         forceOverride,
                         schemaVersion);
-                invokeCallbackOnResult(
-                        callback, AppSearchResult.newSuccessfulResult(/*result=*/ null));
+                invokeCallbackOnResult(callback,
+                        AppSearchResult.newSuccessfulResult(setSchemaResponse.getBundle()));
             } catch (Throwable t) {
                 invokeCallbackOnError(callback, t);
             } finally {
@@ -399,6 +407,98 @@ public class AppSearchManagerService extends SystemService {
         }
 
         @Override
+        public void writeQueryResultsToFile(
+                @NonNull String packageName,
+                @NonNull String databaseName,
+                @NonNull ParcelFileDescriptor fileDescriptor,
+                @NonNull String queryExpression,
+                @NonNull Bundle searchSpecBundle,
+                @UserIdInt int userId,
+                @NonNull IAppSearchResultCallback callback) {
+            int callingUid = Binder.getCallingUid();
+            int callingUserId = handleIncomingUser(userId, callingUid);
+            final long callingIdentity = Binder.clearCallingIdentity();
+            try {
+                verifyCallingPackage(callingUid, packageName);
+                AppSearchImpl impl =
+                        mImplInstanceManager.getAppSearchImpl(callingUserId);
+                // we don't need to append the file. The file is always brand new.
+                try (DataOutputStream outputStream = new DataOutputStream(
+                        new FileOutputStream(fileDescriptor.getFileDescriptor()))) {
+                    SearchResultPage searchResultPage = impl.query(
+                            packageName,
+                            databaseName,
+                            queryExpression,
+                            new SearchSpec(searchSpecBundle));
+                    while (!searchResultPage.getResults().isEmpty()) {
+                        for (int i = 0; i < searchResultPage.getResults().size(); i++) {
+                            AppSearchMigrationHelper.writeBundleToOutputStream(
+                                    outputStream, searchResultPage.getResults().get(i)
+                                            .getGenericDocument().getBundle());
+                        }
+                        searchResultPage = impl.getNextPage(searchResultPage.getNextPageToken());
+                    }
+                }
+                invokeCallbackOnResult(callback, AppSearchResult.newSuccessfulResult(null));
+            } catch (Throwable t) {
+                invokeCallbackOnError(callback, t);
+            } finally {
+                Binder.restoreCallingIdentity(callingIdentity);
+            }
+        }
+
+        @Override
+        public void putDocumentsFromFile(
+                @NonNull String packageName,
+                @NonNull String databaseName,
+                @NonNull ParcelFileDescriptor fileDescriptor,
+                @UserIdInt int userId,
+                @NonNull IAppSearchResultCallback callback) {
+            int callingUid = Binder.getCallingUid();
+            int callingUserId = handleIncomingUser(userId, callingUid);
+            final long callingIdentity = Binder.clearCallingIdentity();
+            try {
+                verifyCallingPackage(callingUid, packageName);
+                AppSearchImpl impl =
+                        mImplInstanceManager.getAppSearchImpl(callingUserId);
+
+                GenericDocument document;
+                ArrayList<Bundle> migrationFailureBundles = new ArrayList<>();
+                try (DataInputStream inputStream = new DataInputStream(
+                        new FileInputStream(fileDescriptor.getFileDescriptor()))) {
+                    while (true) {
+                        try {
+                            document = AppSearchMigrationHelper
+                                    .readDocumentFromInputStream(inputStream);
+                        } catch (EOFException e) {
+                            // nothing wrong, we just finish the reading.
+                            break;
+                        }
+                        try {
+                            impl.putDocument(packageName, databaseName, document, /*logger=*/ null);
+                        } catch (Throwable t) {
+                            migrationFailureBundles.add(
+                                    new SetSchemaResponse.MigrationFailure.Builder()
+                                            .setNamespace(document.getNamespace())
+                                            .setSchemaType(document.getSchemaType())
+                                            .setUri(document.getUri())
+                                            .setAppSearchResult(
+                                                    AppSearchResult.throwableToFailedResult(t))
+                                            .build().getBundle());
+                        }
+                    }
+                }
+                impl.persistToDisk();
+                invokeCallbackOnResult(callback,
+                        AppSearchResult.newSuccessfulResult(migrationFailureBundles));
+            } catch (Throwable t) {
+                invokeCallbackOnError(callback, t);
+            } finally {
+                Binder.restoreCallingIdentity(callingIdentity);
+            }
+        }
+
+        @Override
         public void reportUsage(
                 @NonNull String packageName,
                 @NonNull String databaseName,
diff --git a/services/tests/servicestests/src/com/android/server/pm/BaseShortcutManagerTest.java b/services/tests/servicestests/src/com/android/server/pm/BaseShortcutManagerTest.java
index 212a9c614632..8e2b20700e2c 100644
--- a/services/tests/servicestests/src/com/android/server/pm/BaseShortcutManagerTest.java
+++ b/services/tests/servicestests/src/com/android/server/pm/BaseShortcutManagerTest.java
@@ -89,6 +89,7 @@ import android.os.FileUtils;
 import android.os.Handler;
 import android.os.IBinder;
 import android.os.Looper;
+import android.os.ParcelFileDescriptor;
 import android.os.PersistableBundle;
 import android.os.Process;
 import android.os.RemoteException;
@@ -769,6 +770,21 @@ public abstract class BaseShortcutManagerTest extends InstrumentationTestCase {
         }
 
         @Override
+        public void writeQueryResultsToFile(String packageName, String databaseName,
+                ParcelFileDescriptor fileDescriptor, String queryExpression,
+                Bundle searchSpecBundle, int userId, IAppSearchResultCallback callback)
+                throws RemoteException {
+            ignore(callback);
+        }
+
+        @Override
+        public void putDocumentsFromFile(String packageName, String databaseName,
+                ParcelFileDescriptor fileDescriptor, int userId, IAppSearchResultCallback callback)
+                throws RemoteException {
+            ignore(callback);
+        }
+
+        @Override
         public void reportUsage(String packageName, String databaseName, String namespace,
                 String uri, long usageTimeMillis, boolean systemUsage, int userId,
                 IAppSearchResultCallback callback)