Implement schema migration to another type in framework.

Changes included:
*05b7d39:Change schema version from per AppSearchSchema to overall.
*effe024:Support schema migration to another type.
*1be54dc:Minor fix for where we set version in schema.

Bug: 182620003
Test: AppSearchSchemaMigrationCtsTest
Change-Id: Ifabfaa65c32fc3684d7fc7a09dd021671ff3252d

5 files changed, 115 insertions, 139 deletions

diff --git a/apex/appsearch/framework/java/android/app/appsearch/AppSearchMigrationHelper.java b/apex/appsearch/framework/java/android/app/appsearch/AppSearchMigrationHelper.java
index e585d9147895..4357905f0371 100644
--- a/apex/appsearch/framework/java/android/app/appsearch/AppSearchMigrationHelper.java
+++ b/apex/appsearch/framework/java/android/app/appsearch/AppSearchMigrationHelper.java
@@ -16,6 +16,7 @@
 
 package android.app.appsearch;
 
+import static android.app.appsearch.AppSearchResult.RESULT_INVALID_SCHEMA;
 import static android.os.ParcelFileDescriptor.MODE_READ_ONLY;
 import static android.os.ParcelFileDescriptor.MODE_WRITE_ONLY;
 
@@ -27,6 +28,7 @@ import android.os.Bundle;
 import android.os.Parcel;
 import android.os.ParcelFileDescriptor;
 import android.os.RemoteException;
+import android.util.ArraySet;
 
 import com.android.internal.infra.AndroidFuture;
 
@@ -39,8 +41,8 @@ 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.Set;
 import java.util.concurrent.ExecutionException;
 
 /**
@@ -55,23 +57,23 @@ public class AppSearchMigrationHelper implements Closeable {
     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 final Set<String> mDestinationTypes;
     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 {
+            @NonNull String databaseName,
+            @NonNull Set<AppSearchSchema> newSchemas) 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);
+        mDestinationTypes = new ArraySet<>(newSchemas.size());
+        for (AppSearchSchema newSchema : newSchemas) {
+            mDestinationTypes.add(newSchema.getSchemaType());
+        }
     }
 
     /**
@@ -87,7 +89,8 @@ public class AppSearchMigrationHelper implements Closeable {
      *     GenericDocument} to new version.
      */
     @WorkerThread
-    public void queryAndTransform(@NonNull String schemaType, @NonNull Migrator migrator)
+    public void queryAndTransform(@NonNull String schemaType, @NonNull Migrator migrator,
+            int currentVersion, int finalVersion)
             throws IOException, AppSearchException, InterruptedException, ExecutionException {
         File queryFile = File.createTempFile(/*prefix=*/"appsearch", /*suffix=*/null);
         try (ParcelFileDescriptor fileDescriptor =
@@ -111,7 +114,7 @@ public class AppSearchMigrationHelper implements Closeable {
             if (!result.isSuccess()) {
                 throw new AppSearchException(result.getResultCode(), result.getErrorMessage());
             }
-            readAndTransform(queryFile, migrator);
+            readAndTransform(queryFile, migrator, currentVersion, finalVersion);
         } catch (RemoteException e) {
             throw e.rethrowFromSystemServer();
         } finally {
@@ -173,8 +176,9 @@ public class AppSearchMigrationHelper implements Closeable {
      *
      * <p>Save migrated {@link GenericDocument}s to the {@link #mMigratedFile}.
      */
-    private void readAndTransform(@NonNull File file, @NonNull Migrator migrator)
-            throws IOException {
+    private void readAndTransform(@NonNull File file, @NonNull Migrator migrator,
+            int currentVersion, int finalVersion)
+            throws IOException, AppSearchException {
         try (DataInputStream inputStream = new DataInputStream(new FileInputStream(file));
              DataOutputStream outputStream = new DataOutputStream(new FileOutputStream(
                      mMigratedFile, /*append=*/ true))) {
@@ -187,9 +191,6 @@ public class AppSearchMigrationHelper implements Closeable {
                     // 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);
@@ -197,6 +198,18 @@ public class AppSearchMigrationHelper implements Closeable {
                     // currentVersion == finalVersion case won't trigger migration and get here.
                     newDocument = migrator.onDowngrade(currentVersion, finalVersion, document);
                 }
+
+                if (!mDestinationTypes.contains(newDocument.getSchemaType())) {
+                    // we exit before the new schema has been set to AppSearch. So no
+                    // observable changes will be applied to stored schemas and documents.
+                    // And the temp file will be deleted at close(), which will be triggered at
+                    // the end of try-with-resources block of SearchSessionImpl.
+                    throw new AppSearchException(
+                            RESULT_INVALID_SCHEMA,
+                            "Receive a migrated document with schema type: "
+                                    + newDocument.getSchemaType()
+                                    + ". But the schema types doesn't exist in the request");
+                }
                 writeBundleToOutputStream(outputStream, newDocument.getBundle());
             }
             mAreDocumentsMigrated = true;
diff --git a/apex/appsearch/framework/java/android/app/appsearch/AppSearchSession.java b/apex/appsearch/framework/java/android/app/appsearch/AppSearchSession.java
index 4dd1b79b536c..367748951902 100644
--- a/apex/appsearch/framework/java/android/app/appsearch/AppSearchSession.java
+++ b/apex/appsearch/framework/java/android/app/appsearch/AppSearchSession.java
@@ -19,7 +19,6 @@ 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;
@@ -647,8 +646,8 @@ public final class AppSearchSession implements Closeable {
                     new ArrayList<>(request.getSchemasNotDisplayedBySystem()),
                     schemasPackageAccessibleBundles,
                     request.isForceOverride(),
-                    mUserId,
                     request.getVersion(),
+                    mUserId,
                     new IAppSearchResultCallback.Stub() {
                         public void onResult(AppSearchResult result) {
                             executor.execute(() -> {
@@ -661,7 +660,7 @@ public final class AppSearchSession implements Closeable {
                                             // Throw exception if there is any deleted types or
                                             // incompatible types. That's the only case we swallowed
                                             // in the AppSearchImpl#setSchema().
-                                            checkDeletedAndIncompatible(
+                                            SchemaMigrationUtil.checkDeletedAndIncompatible(
                                                     setSchemaResponse.getDeletedTypes(),
                                                     setSchemaResponse.getIncompatibleTypes());
                                         }
@@ -698,7 +697,7 @@ public final class AppSearchSession implements Closeable {
         workExecutor.execute(() -> {
             try {
                 // Migration process
-                // 1. Generate the current and the final version map.
+                // 1. Validate and retrieve all active migrators.
                 AndroidFuture<AppSearchResult<GetSchemaResponse>> getSchemaFuture =
                         new AndroidFuture<>();
                 getSchema(callbackExecutor, getSchemaFuture::complete);
@@ -709,11 +708,18 @@ public final class AppSearchSession implements Closeable {
                     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());
+                int currentVersion = getSchemaResponse.getVersion();
+                int finalVersion = request.getVersion();
+                Map<String, Migrator> activeMigrators = SchemaMigrationUtil.getActiveMigrators(
+                        getSchemaResponse.getSchemas(), request.getMigrators(), currentVersion,
+                        finalVersion);
+
+                // No need to trigger migration if no migrator is active.
+                if (activeMigrators.isEmpty()) {
+                    setSchemaNoMigrations(request, schemaBundles, schemasPackageAccessibleBundles,
+                            callbackExecutor, callback);
+                    return;
+                }
 
                 // 2. SetSchema with forceOverride=false, to retrieve the list of
                 // incompatible/deleted types.
@@ -725,8 +731,8 @@ public final class AppSearchSession implements Closeable {
                         new ArrayList<>(request.getSchemasNotDisplayedBySystem()),
                         schemasPackageAccessibleBundles,
                         /*forceOverride=*/ false,
-                        mUserId,
                         request.getVersion(),
+                        mUserId,
                         new IAppSearchResultCallback.Stub() {
                             public void onResult(AppSearchResult result) {
                                 setSchemaFuture.complete(result);
@@ -741,46 +747,27 @@ public final class AppSearchSession implements Closeable {
                 SetSchemaResponse setSchemaResponse =
                         new SetSchemaResponse(setSchemaResult.getResultValue());
 
-                // 1. If forceOverride is false, check that all incompatible types will be migrated.
+                // 3. 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);
+                    SchemaMigrationUtil.checkDeletedAndIncompatibleAfterMigration(setSchemaResponse,
+                            activeMigrators.keySet());
                 }
 
-                try (AppSearchMigrationHelper migrationHelper =
-                             new AppSearchMigrationHelper(
-                                     mService, mUserId, currentVersionMap, finalVersionMap,
-                                     mPackageName, mDatabaseName)) {
-                    Map<String, Migrator> migratorMap = request.getMigrators();
+                try (AppSearchMigrationHelper migrationHelper = new AppSearchMigrationHelper(
+                        mService, mUserId, mPackageName, mDatabaseName, request.getSchemas())) {
 
-                    // 2. Trigger migration for all migrators.
+                    // 4. 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);
-                        }
+                    for (Map.Entry<String, Migrator> entry : activeMigrators.entrySet()) {
+                        migrationHelper.queryAndTransform(/*schemaType=*/ entry.getKey(),
+                                /*migrator=*/ entry.getValue(), currentVersion,
+                                finalVersion);
                     }
 
-                    // 3. SetSchema a second time with forceOverride=true if the first attempted
+                    // 5. SetSchema a second time with forceOverride=true if the first attempted
                     // failed.
                     if (!setSchemaResponse.getIncompatibleTypes().isEmpty()
                             || !setSchemaResponse.getDeletedTypes().isEmpty()) {
@@ -809,13 +796,16 @@ public final class AppSearchSession implements Closeable {
                             // error in the first setSchema call, all other errors will be thrown at
                             // the first time.
                             callbackExecutor.execute(() -> callback.accept(
-                                    AppSearchResult.newFailedResult(setSchemaResult)));
+                                    AppSearchResult.newFailedResult(setSchema2Result)));
                             return;
                         }
                     }
 
                     SetSchemaResponse.Builder responseBuilder = setSchemaResponse.toBuilder()
-                            .addMigratedTypes(migratedTypes);
+                            .addMigratedTypes(activeMigrators.keySet());
+
+                    // 6. Put all the migrated documents into the index, now that the new schema is
+                    // set.
                     AppSearchResult<SetSchemaResponse> putResult =
                             migrationHelper.putMigratedDocuments(responseBuilder);
                     callbackExecutor.execute(() -> callback.accept(putResult));
@@ -826,17 +816,4 @@ public final class AppSearchSession implements Closeable {
             }
         });
     }
-
-    /**  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 a8ac27c11089..4d05ad71a121 100644
--- a/apex/appsearch/framework/java/android/app/appsearch/IAppSearchManager.aidl
+++ b/apex/appsearch/framework/java/android/app/appsearch/IAppSearchManager.aidl
@@ -40,6 +40,7 @@ interface IAppSearchManager {
      *     packages. The value List contains PackageIdentifier Bundles.
      * @param forceOverride Whether to apply the new schema even if it is incompatible. All
      *     incompatible documents will be deleted.
+     * @param schemaVersion  The overall schema version number of the request.
      * @param userId Id of the calling user
      * @param callback {@link IAppSearchResultCallback#onResult} will be called with an
      *     {@link AppSearchResult}&lt;{@link Bundle}&gt;, where the value are
@@ -52,8 +53,8 @@ interface IAppSearchManager {
         in List<String> schemasNotDisplayedBySystem,
         in Map<String, List<Bundle>> schemasPackageAccessibleBundles,
         boolean forceOverride,
-        in int userId,
         in int schemaVersion,
+        in int userId,
         in IAppSearchResultCallback callback);
 
     /**
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 c9473bdeff32..32d7e043e954 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,12 +20,11 @@ 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;
@@ -36,84 +35,70 @@ import java.util.Set;
  * @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.
-     */
+    /** Returns all active {@link Migrator}s that need to be triggered in this migration. */
     @NonNull
-    public static Set<String> getUnmigratedIncompatibleTypes(
-            @NonNull Set<String> incompatibleSchemaTypes,
+    public static Map<String, Migrator> getActiveMigrators(
+            @NonNull Set<AppSearchSchema> existingSchemas,
             @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.shouldMigrate(currentVersion, finalVersion)) {
-                unmigratedSchemaTypes.add(unmigratedSchemaType);
+            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);
             }
         }
-        return Collections.unmodifiableSet(unmigratedSchemaTypes);
+        return activeMigrators;
     }
 
     /**
-     * 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.
+     * Checks the setSchema() call won't delete any types or has incompatible types after all {@link
+     * Migrator} has been triggered..
      */
-    public static boolean shouldTriggerMigration(
-            @NonNull String schemaType,
-            @NonNull Migrator migrator,
-            @NonNull Map<String, Integer> currentVersionMap,
-            @NonNull Map<String, Integer> finalVersionMap)
+    public static void checkDeletedAndIncompatibleAfterMigration(
+            @NonNull SetSchemaResponse setSchemaResponse, @NonNull Set<String> activeMigrators)
             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.shouldMigrate(currentVersion, finalVersion);
+        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);
     }
 
-    /** 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);
+    /** 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);
         }
-        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 5f3a3ae75731..704ee6446e54 100644
--- a/apex/appsearch/service/java/com/android/server/appsearch/AppSearchManagerService.java
+++ b/apex/appsearch/service/java/com/android/server/appsearch/AppSearchManagerService.java
@@ -157,8 +157,8 @@ public class AppSearchManagerService extends SystemService {
                 @NonNull List<String> schemasNotDisplayedBySystem,
                 @NonNull Map<String, List<Bundle>> schemasPackageAccessibleBundles,
                 boolean forceOverride,
-                @UserIdInt int userId,
                 int schemaVersion,
+                @UserIdInt int userId,
                 @NonNull IAppSearchResultCallback callback) {
             Preconditions.checkNotNull(packageName);
             Preconditions.checkNotNull(databaseName);