Setup infrastructure (multi-db support) for the new grid migration algorithm

We'll have a db for each grid option and a db for back up / restore.

TODO(pinyaoting): support back up / restore using the new infrastructure, particularly calls to GridBackupTable should use different DBs when the feature flag (NEW_GRID_MIGRATION_ALGORITHM) is on.

Bug: 144052802
Test: N/A

Change-Id: I644a3e70148bd78204a747a337446a3c038f616f
diff --git a/res/values/attrs.xml b/res/values/attrs.xml
index de17eb7..707424c 100644
--- a/res/values/attrs.xml
+++ b/res/values/attrs.xml
@@ -115,6 +115,7 @@
         <attr name="numFolderColumns" format="integer" />
         <!-- numHotseatIcons defaults to numColumns, if not specified -->
         <attr name="numHotseatIcons" format="integer" />
+        <attr name="dbFile" format="string" />
         <attr name="defaultLayoutId" format="reference" />
         <attr name="demoModeLayoutId" format="reference" />
     </declare-styleable>
diff --git a/res/xml/device_profiles.xml b/res/xml/device_profiles.xml
index 82547d5..1c99dfc 100644
--- a/res/xml/device_profiles.xml
+++ b/res/xml/device_profiles.xml
@@ -24,6 +24,7 @@
         launcher:numFolderRows="2"
         launcher:numFolderColumns="3"
         launcher:numHotseatIcons="3"
+        launcher:dbFile="launcher_3_by_3.db"
         launcher:defaultLayoutId="@xml/default_workspace_3x3" >
 
         <display-option
@@ -51,6 +52,7 @@
         launcher:numFolderRows="3"
         launcher:numFolderColumns="4"
         launcher:numHotseatIcons="4"
+        launcher:dbFile="launcher_4_by_4.db"
         launcher:defaultLayoutId="@xml/default_workspace_4x4" >
 
         <display-option
@@ -102,6 +104,7 @@
         launcher:numFolderRows="4"
         launcher:numFolderColumns="4"
         launcher:numHotseatIcons="5"
+        launcher:dbFile="launcher.db"
         launcher:defaultLayoutId="@xml/default_workspace_5x5" >
 
         <display-option
diff --git a/robolectric_tests/src/com/android/launcher3/model/BackupRestoreTest.java b/robolectric_tests/src/com/android/launcher3/model/BackupRestoreTest.java
index 7072adf..50f9494 100644
--- a/robolectric_tests/src/com/android/launcher3/model/BackupRestoreTest.java
+++ b/robolectric_tests/src/com/android/launcher3/model/BackupRestoreTest.java
@@ -134,7 +134,7 @@
                 { APP_ICON, SHORTCUT, SHORTCUT, NO__ICON},
             }}, 2, OLD_WORK_PROFILE_ID);
         // simulates the creation of backup upon restore
-        new GridBackupTable(RuntimeEnvironment.application, mDb, mIdp.numHotseatIcons,
+        new GridBackupTable(RuntimeEnvironment.application, mDb, mDb, mIdp.numHotseatIcons,
                 mIdp.numColumns, mIdp.numRows).doBackup(
                         MY_OLD_PROFILE_ID, GridBackupTable.OPTION_REQUIRES_SANITIZATION);
         // reset favorites table
diff --git a/robolectric_tests/src/com/android/launcher3/model/GridBackupTableTest.java b/robolectric_tests/src/com/android/launcher3/model/GridBackupTableTest.java
index f46b849..36fd96b 100644
--- a/robolectric_tests/src/com/android/launcher3/model/GridBackupTableTest.java
+++ b/robolectric_tests/src/com/android/launcher3/model/GridBackupTableTest.java
@@ -63,7 +63,7 @@
 
     @Test
     public void backupTableCreated() {
-        GridBackupTable backupTable = new GridBackupTable(mContext, mDb, 4, 4, 4);
+        GridBackupTable backupTable = new GridBackupTable(mContext, mDb, mDb, 4, 4, 4);
         assertFalse(backupTable.backupOrRestoreAsNeeded());
         Settings.call(mContext.getContentResolver(), Settings.METHOD_REFRESH_BACKUP_TABLE);
 
@@ -75,14 +75,14 @@
 
     @Test
     public void backupTableRestored() {
-        assertFalse(new GridBackupTable(mContext, mDb, 4, 4, 4).backupOrRestoreAsNeeded());
+        assertFalse(new GridBackupTable(mContext, mDb, mDb, 4, 4, 4).backupOrRestoreAsNeeded());
         Settings.call(mContext.getContentResolver(), Settings.METHOD_REFRESH_BACKUP_TABLE);
 
         // Delete entries
         mDb.delete(TABLE_NAME, null, null);
         assertEquals(0, queryNumEntries(mDb, TABLE_NAME));
 
-        GridBackupTable backupTable = new GridBackupTable(mContext, mDb, 3, 3, 3);
+        GridBackupTable backupTable = new GridBackupTable(mContext, mDb, mDb, 3, 3, 3);
         assertTrue(backupTable.backupOrRestoreAsNeeded());
 
         // Items have been restored
@@ -96,7 +96,7 @@
 
     @Test
     public void backupTableRemovedOnAdd() {
-        assertFalse(new GridBackupTable(mContext, mDb, 4, 4, 4).backupOrRestoreAsNeeded());
+        assertFalse(new GridBackupTable(mContext, mDb, mDb, 4, 4, 4).backupOrRestoreAsNeeded());
         Settings.call(mContext.getContentResolver(), Settings.METHOD_REFRESH_BACKUP_TABLE);
 
         assertTrue(tableExists(mDb, BACKUP_TABLE_NAME));
@@ -107,7 +107,7 @@
 
     @Test
     public void backupTableRemovedOnDelete() {
-        assertFalse(new GridBackupTable(mContext, mDb, 4, 4, 4).backupOrRestoreAsNeeded());
+        assertFalse(new GridBackupTable(mContext, mDb, mDb, 4, 4, 4).backupOrRestoreAsNeeded());
         Settings.call(mContext.getContentResolver(), Settings.METHOD_REFRESH_BACKUP_TABLE);
 
         assertTrue(tableExists(mDb, BACKUP_TABLE_NAME));
@@ -118,7 +118,7 @@
 
     @Test
     public void backupTableRetainedOnUpdate() {
-        assertFalse(new GridBackupTable(mContext, mDb, 4, 4, 4).backupOrRestoreAsNeeded());
+        assertFalse(new GridBackupTable(mContext, mDb, mDb, 4, 4, 4).backupOrRestoreAsNeeded());
         Settings.call(mContext.getContentResolver(), Settings.METHOD_REFRESH_BACKUP_TABLE);
 
         assertTrue(tableExists(mDb, BACKUP_TABLE_NAME));
diff --git a/src/com/android/launcher3/InvariantDeviceProfile.java b/src/com/android/launcher3/InvariantDeviceProfile.java
index a807e4f..857db8e 100644
--- a/src/com/android/launcher3/InvariantDeviceProfile.java
+++ b/src/com/android/launcher3/InvariantDeviceProfile.java
@@ -119,6 +119,7 @@
      */
     public int numAllAppsColumns;
 
+    public String dbFile;
     public int defaultLayoutId;
     int demoModeLayoutId;
 
@@ -146,6 +147,7 @@
         iconTextSize = p.iconTextSize;
         numHotseatIcons = p.numHotseatIcons;
         numAllAppsColumns = p.numAllAppsColumns;
+        dbFile = p.dbFile;
         defaultLayoutId = p.defaultLayoutId;
         demoModeLayoutId = p.demoModeLayoutId;
         mExtraAttrs = p.mExtraAttrs;
@@ -292,6 +294,7 @@
         numRows = closestProfile.numRows;
         numColumns = closestProfile.numColumns;
         numHotseatIcons = closestProfile.numHotseatIcons;
+        dbFile = closestProfile.dbFile;
         defaultLayoutId = closestProfile.defaultLayoutId;
         demoModeLayoutId = closestProfile.demoModeLayoutId;
         numFolderRows = closestProfile.numFolderRows;
@@ -559,6 +562,7 @@
 
         private final int numHotseatIcons;
 
+        private final String dbFile;
         private final int defaultLayoutId;
         private final int demoModeLayoutId;
 
@@ -571,6 +575,7 @@
             numRows = a.getInt(R.styleable.GridDisplayOption_numRows, 0);
             numColumns = a.getInt(R.styleable.GridDisplayOption_numColumns, 0);
 
+            dbFile = a.getString(R.styleable.GridDisplayOption_dbFile);
             defaultLayoutId = a.getResourceId(
                     R.styleable.GridDisplayOption_defaultLayoutId, 0);
             demoModeLayoutId = a.getResourceId(
diff --git a/src/com/android/launcher3/LauncherFiles.java b/src/com/android/launcher3/LauncherFiles.java
index 9c4646b..25afb55 100644
--- a/src/com/android/launcher3/LauncherFiles.java
+++ b/src/com/android/launcher3/LauncherFiles.java
@@ -15,8 +15,13 @@
     private static final String XML = ".xml";
 
     public static final String LAUNCHER_DB = "launcher.db";
+    public static final String LAUNCHER_4_BY_4_DB = "launcher_4_by_4.db";
+    public static final String LAUNCHER_3_BY_3_DB = "launcher_3_by_3.db";
+    public static final String LAUNCHER_2_BY_2_DB = "launcher_2_by_2.db";
+    public static final String BACKUP_DB = "backup.db";
     public static final String SHARED_PREFERENCES_KEY = "com.android.launcher3.prefs";
-    public static final String MANAGED_USER_PREFERENCES_KEY = "com.android.launcher3.managedusers.prefs";
+    public static final String MANAGED_USER_PREFERENCES_KEY =
+            "com.android.launcher3.managedusers.prefs";
     // This preference file is not backed up to cloud.
     public static final String DEVICE_PREFERENCES_KEY = "com.android.launcher3.device.prefs";
 
@@ -25,6 +30,10 @@
 
     public static final List<String> ALL_FILES = Collections.unmodifiableList(Arrays.asList(
             LAUNCHER_DB,
+            LAUNCHER_4_BY_4_DB,
+            LAUNCHER_3_BY_3_DB,
+            LAUNCHER_2_BY_2_DB,
+            BACKUP_DB,
             SHARED_PREFERENCES_KEY + XML,
             WIDGET_PREVIEWS_DB,
             MANAGED_USER_PREFERENCES_KEY + XML,
diff --git a/src/com/android/launcher3/LauncherProvider.java b/src/com/android/launcher3/LauncherProvider.java
index b0ab35c..5544240 100644
--- a/src/com/android/launcher3/LauncherProvider.java
+++ b/src/com/android/launcher3/LauncherProvider.java
@@ -16,6 +16,7 @@
 
 package com.android.launcher3;
 
+import static com.android.launcher3.config.FeatureFlags.MULTI_DB_GRID_MIRATION_ALGO;
 import static com.android.launcher3.provider.LauncherDbUtils.dropTable;
 import static com.android.launcher3.provider.LauncherDbUtils.tableExists;
 import static com.android.launcher3.util.Executors.MODEL_EXECUTOR;
@@ -157,6 +158,17 @@
         }
     }
 
+    private synchronized boolean updateCurrentOpenHelper() {
+        final InvariantDeviceProfile idp = InvariantDeviceProfile.INSTANCE.get(getContext());
+        if (TextUtils.equals(idp.dbFile, mOpenHelper.getDatabaseName())) {
+            return false;
+        }
+
+        mOpenHelper.close();
+        mOpenHelper = new DatabaseHelper(getContext());
+        return true;
+    }
+
     @Override
     public Cursor query(Uri uri, String[] projection, String selection,
             String[] selectionArgs, String sortOrder) {
@@ -210,7 +222,7 @@
         addModifiedTime(initialValues);
         final int rowId = dbInsertAndCheck(mOpenHelper, db, args.table, null, initialValues);
         if (rowId < 0) return null;
-        mOpenHelper.onAddOrDeleteOp(db);
+        onAddOrDeleteOp(db);
 
         uri = ContentUris.withAppendedId(uri, rowId);
         reloadLauncherIfExternal();
@@ -268,7 +280,7 @@
                     return 0;
                 }
             }
-            mOpenHelper.onAddOrDeleteOp(db);
+            onAddOrDeleteOp(db);
             t.commit();
         }
 
@@ -294,7 +306,7 @@
                         results[i].count != null && results[i].count > 0;
             }
             if (isAddOrDelete) {
-                mOpenHelper.onAddOrDeleteOp(t.getDb());
+                onAddOrDeleteOp(t.getDb());
             }
 
             t.commit();
@@ -316,7 +328,7 @@
         }
         int count = db.delete(args.table, args.where, args.args);
         if (count > 0) {
-            mOpenHelper.onAddOrDeleteOp(db);
+            onAddOrDeleteOp(db);
             reloadLauncherIfExternal();
         }
         return count;
@@ -360,12 +372,14 @@
             }
             case LauncherSettings.Settings.METHOD_NEW_ITEM_ID: {
                 Bundle result = new Bundle();
-                result.putInt(LauncherSettings.Settings.EXTRA_VALUE, mOpenHelper.generateNewItemId());
+                result.putInt(LauncherSettings.Settings.EXTRA_VALUE,
+                        mOpenHelper.generateNewItemId());
                 return result;
             }
             case LauncherSettings.Settings.METHOD_NEW_SCREEN_ID: {
                 Bundle result = new Bundle();
-                result.putInt(LauncherSettings.Settings.EXTRA_VALUE, mOpenHelper.generateNewScreenId());
+                result.putInt(LauncherSettings.Settings.EXTRA_VALUE,
+                        mOpenHelper.generateNewScreenId());
                 return result;
             }
             case LauncherSettings.Settings.METHOD_CREATE_EMPTY_DB: {
@@ -387,8 +401,12 @@
                 return result;
             }
             case LauncherSettings.Settings.METHOD_REFRESH_BACKUP_TABLE: {
-                mOpenHelper.mBackupTableExists =
-                        tableExists(mOpenHelper.getReadableDatabase(), Favorites.BACKUP_TABLE_NAME);
+                // TODO(pinyaoting): Update the behavior here.
+                if (!MULTI_DB_GRID_MIRATION_ALGO.get()) {
+                    mOpenHelper.mBackupTableExists =
+                            tableExists(mOpenHelper.getReadableDatabase(),
+                                    Favorites.BACKUP_TABLE_NAME);
+                }
                 return null;
             }
             case LauncherSettings.Settings.METHOD_RESTORE_BACKUP_TABLE: {
@@ -399,10 +417,26 @@
                         TOKEN_RESTORE_BACKUP_TABLE, RESTORE_BACKUP_TABLE_DELAY);
                 return null;
             }
+            case LauncherSettings.Settings.METHOD_UPDATE_CURRENT_OPEN_HELPER: {
+                if (MULTI_DB_GRID_MIRATION_ALGO.get()) {
+                    Bundle result = new Bundle();
+                    result.putBoolean(LauncherSettings.Settings.EXTRA_VALUE,
+                            updateCurrentOpenHelper());
+                    return result;
+                }
+            }
         }
         return null;
     }
 
+    private void onAddOrDeleteOp(SQLiteDatabase db) {
+        if (MULTI_DB_GRID_MIRATION_ALGO.get()) {
+            // TODO(pingyaoting): Implement the behavior here.
+        } else {
+            mOpenHelper.onAddOrDeleteOp(db);
+        }
+    }
+
     /**
      * Deletes any empty folder from the DB.
      * @return Ids of deleted folders.
@@ -551,14 +585,16 @@
     /**
      * The class is subclassed in tests to create an in-memory db.
      */
-    public static class DatabaseHelper extends NoLocaleSQLiteHelper implements LayoutParserCallback {
+    public static class DatabaseHelper extends NoLocaleSQLiteHelper implements
+            LayoutParserCallback {
         private final Context mContext;
         private int mMaxItemId = -1;
         private int mMaxScreenId = -1;
         private boolean mBackupTableExists;
 
         DatabaseHelper(Context context) {
-            this(context, LauncherFiles.LAUNCHER_DB);
+            this(context, MULTI_DB_GRID_MIRATION_ALGO.get() ? InvariantDeviceProfile.INSTANCE.get(
+                    context).dbFile : LauncherFiles.LAUNCHER_DB);
             // Table creation sometimes fails silently, which leads to a crash loop.
             // This way, we will try to create a table every time after crash, so the device
             // would eventually be able to recover.
@@ -567,7 +603,10 @@
                 // This operation is a no-op if the table already exists.
                 addFavoritesTable(getWritableDatabase(), true);
             }
-            mBackupTableExists = tableExists(getReadableDatabase(), Favorites.BACKUP_TABLE_NAME);
+            if (!MULTI_DB_GRID_MIRATION_ALGO.get()) {
+                mBackupTableExists = tableExists(getReadableDatabase(),
+                        Favorites.BACKUP_TABLE_NAME);
+            }
 
             initIds();
         }
@@ -575,8 +614,8 @@
         /**
          * Constructor used in tests and for restore.
          */
-        public DatabaseHelper(Context context, String tableName) {
-            super(context, tableName, SCHEMA_VERSION);
+        public DatabaseHelper(Context context, String dbName) {
+            super(context, dbName, SCHEMA_VERSION);
             mContext = context;
         }
 
@@ -606,7 +645,7 @@
         }
 
         protected void onAddOrDeleteOp(SQLiteDatabase db) {
-            if (mBackupTableExists) {
+            if (!MULTI_DB_GRID_MIRATION_ALGO.get() && mBackupTableExists) {
                 dropTable(db, Favorites.BACKUP_TABLE_NAME);
                 mBackupTableExists = false;
             }
diff --git a/src/com/android/launcher3/LauncherSettings.java b/src/com/android/launcher3/LauncherSettings.java
index 4c5c61c..49831f6 100644
--- a/src/com/android/launcher3/LauncherSettings.java
+++ b/src/com/android/launcher3/LauncherSettings.java
@@ -302,6 +302,8 @@
 
         public static final String METHOD_RESTORE_BACKUP_TABLE = "restore_backup_table";
 
+        public static final String METHOD_UPDATE_CURRENT_OPEN_HELPER = "update_current_open_helper";
+
         public static final String EXTRA_VALUE = "value";
 
         public static Bundle call(ContentResolver cr, String method) {
diff --git a/src/com/android/launcher3/config/FeatureFlags.java b/src/com/android/launcher3/config/FeatureFlags.java
index b1a2c33..cc33965 100644
--- a/src/com/android/launcher3/config/FeatureFlags.java
+++ b/src/com/android/launcher3/config/FeatureFlags.java
@@ -112,6 +112,9 @@
     public static final BooleanFlag ENABLE_DEEP_SHORTCUT_ICON_CACHE = getDebugFlag(
             "ENABLE_DEEP_SHORTCUT_ICON_CACHE", true, "R/W deep shortcut in IconCache");
 
+    public static final BooleanFlag MULTI_DB_GRID_MIRATION_ALGO = getDebugFlag(
+            "MULTI_DB_GRID_MIRATION_ALGO", false, "Use the multi-db grid migration algorithm");
+
     public static final BooleanFlag ENABLE_LAUNCHER_PREVIEW_IN_GRID_PICKER = getDebugFlag(
             "ENABLE_LAUNCHER_PREVIEW_IN_GRID_PICKER", false,
             "Show launcher preview in grid picker");
diff --git a/src/com/android/launcher3/model/GridBackupTable.java b/src/com/android/launcher3/model/GridBackupTable.java
index fc9948e..eb99af5 100644
--- a/src/com/android/launcher3/model/GridBackupTable.java
+++ b/src/com/android/launcher3/model/GridBackupTable.java
@@ -33,6 +33,8 @@
 import com.android.launcher3.LauncherSettings.Settings;
 import com.android.launcher3.pm.UserCache;
 
+import java.util.Objects;
+
 /**
  * Helper class to backup and restore Favorites table into a separate table
  * within the same data base.
@@ -61,7 +63,8 @@
     private static final int STATE_SANITIZED = 2;
 
     private final Context mContext;
-    private final SQLiteDatabase mDb;
+    private final SQLiteDatabase mFavoritesDb;
+    private final SQLiteDatabase mBackupDb;
 
     private final int mOldHotseatSize;
     private final int mOldGridX;
@@ -74,10 +77,11 @@
     @IntDef({STATE_NOT_FOUND, STATE_RAW, STATE_SANITIZED})
     private @interface BackupState { }
 
-    public GridBackupTable(Context context, SQLiteDatabase db,
+    public GridBackupTable(Context context, SQLiteDatabase favoritesDb, SQLiteDatabase backupDb,
             int hotseatSize, int gridX, int gridY) {
         mContext = context;
-        mDb = db;
+        mFavoritesDb = favoritesDb;
+        mBackupDb = backupDb;
 
         mOldHotseatSize = hotseatSize;
         mOldGridX = gridX;
@@ -90,7 +94,7 @@
      */
     public boolean backupOrRestoreAsNeeded() {
         // Check if backup table exists
-        if (!tableExists(mDb, BACKUP_TABLE_NAME)) {
+        if (!tableExists(mBackupDb, BACKUP_TABLE_NAME)) {
             if (Settings.call(mContext.getContentResolver(), Settings.METHOD_WAS_EMPTY_DB_CREATED)
                     .getBoolean(Settings.EXTRA_VALUE, false)) {
                 // No need to copy if empty DB was created.
@@ -105,7 +109,7 @@
         }
         long userSerial = UserCache.INSTANCE.get(mContext).getSerialNumberForUser(
                 Process.myUserHandle());
-        copyTable(BACKUP_TABLE_NAME, Favorites.TABLE_NAME, userSerial);
+        copyTable(mBackupDb, BACKUP_TABLE_NAME, mFavoritesDb, Favorites.TABLE_NAME, userSerial);
         Log.d(TAG, "Backup table found");
         return true;
     }
@@ -118,28 +122,37 @@
     /**
      * Copy valid grid entries from one table to another.
      */
-    private void copyTable(String from, String to, long userSerial) {
-        dropTable(mDb, to);
-        Favorites.addTableToDb(mDb, userSerial, false, to);
-        mDb.execSQL("INSERT INTO " + to + " SELECT * FROM " + from + " where _id > " + ID_PROPERTY);
+    private static void copyTable(SQLiteDatabase fromDb, String fromTable, SQLiteDatabase toDb,
+            String toTable, long userSerial) {
+        dropTable(toDb, toTable);
+        Favorites.addTableToDb(toDb, userSerial, false, toTable);
+        if (fromDb != toDb) {
+            toDb.execSQL("ATTACH DATABASE '" + fromDb.getPath() + "' AS from_db");
+            toDb.execSQL(
+                    "INSERT INTO " + toTable + " SELECT * FROM from_db." + fromTable
+                            + " where _id > " + ID_PROPERTY);
+        } else {
+            toDb.execSQL("INSERT INTO " + toTable + " SELECT * FROM " + fromTable + " where _id > "
+                    + ID_PROPERTY);
+        }
     }
 
     private void encodeDBProperties(int options) {
         ContentValues values = new ContentValues();
         values.put(Favorites._ID, ID_PROPERTY);
-        values.put(KEY_DB_VERSION, mDb.getVersion());
+        values.put(KEY_DB_VERSION, mFavoritesDb.getVersion());
         values.put(KEY_GRID_X_SIZE, mOldGridX);
         values.put(KEY_GRID_Y_SIZE, mOldGridY);
         values.put(KEY_HOTSEAT_SIZE, mOldHotseatSize);
         values.put(Favorites.OPTIONS, options);
-        mDb.insert(BACKUP_TABLE_NAME, null, values);
+        mBackupDb.insert(BACKUP_TABLE_NAME, null, values);
     }
 
     /**
      * Load DB properties from grid backup table.
      */
     public @BackupState int loadDBProperties() {
-        try (Cursor c = mDb.query(BACKUP_TABLE_NAME, new String[] {
+        try (Cursor c = mBackupDb.query(BACKUP_TABLE_NAME, new String[] {
                 KEY_DB_VERSION,     // 0
                 KEY_GRID_X_SIZE,    // 1
                 KEY_GRID_Y_SIZE,    // 2
@@ -150,7 +163,7 @@
                 Log.e(TAG, "Meta data not found in backup table");
                 return STATE_NOT_FOUND;
             }
-            if (!validateDBVersion(mDb.getVersion(), c.getInt(0))) {
+            if (!validateDBVersion(mBackupDb.getVersion(), c.getInt(0))) {
                 return STATE_NOT_FOUND;
             }
 
@@ -166,7 +179,7 @@
      * Restore workspace from raw backup if available.
      */
     public boolean restoreFromRawBackupIfAvailable(long oldProfileId) {
-        if (!tableExists(mDb, Favorites.BACKUP_TABLE_NAME)
+        if (!tableExists(mBackupDb, Favorites.BACKUP_TABLE_NAME)
                 || loadDBProperties() != STATE_RAW
                 || mOldHotseatSize != mRestoredHotseatSize
                 || mOldGridX != mRestoredGridX
@@ -174,7 +187,8 @@
             // skip restore if dimensions in backup table differs from current setup.
             return false;
         }
-        copyTable(Favorites.BACKUP_TABLE_NAME, Favorites.TABLE_NAME, oldProfileId);
+        copyTable(mBackupDb, Favorites.BACKUP_TABLE_NAME, mFavoritesDb, Favorites.TABLE_NAME,
+                oldProfileId);
         Log.d(TAG, "Backup restored");
         return true;
     }
@@ -183,7 +197,8 @@
      * Performs a backup on the workspace layout.
      */
     public void doBackup(long profileId, int options) {
-        copyTable(Favorites.TABLE_NAME, Favorites.BACKUP_TABLE_NAME, profileId);
+        copyTable(mFavoritesDb, Favorites.TABLE_NAME, mBackupDb, Favorites.BACKUP_TABLE_NAME,
+                profileId);
         encodeDBProperties(options);
     }
 
diff --git a/src/com/android/launcher3/model/GridSizeMigrationTask.java b/src/com/android/launcher3/model/GridSizeMigrationTask.java
index c35c4b9..4f92066 100644
--- a/src/com/android/launcher3/model/GridSizeMigrationTask.java
+++ b/src/com/android/launcher3/model/GridSizeMigrationTask.java
@@ -909,7 +909,7 @@
             boolean dbChanged = false;
 
             GridBackupTable backupTable = new GridBackupTable(context, transaction.getDb(),
-                    srcHotseatCount, sourceSize.x, sourceSize.y);
+                    transaction.getDb(), srcHotseatCount, sourceSize.x, sourceSize.y);
             if (backupTable.backupOrRestoreAsNeeded()) {
                 dbChanged = true;
                 srcHotseatCount = backupTable.getRestoreHotseatAndGridSize(sourceSize);
diff --git a/src/com/android/launcher3/model/GridSizeMigrationTaskV2.java b/src/com/android/launcher3/model/GridSizeMigrationTaskV2.java
new file mode 100644
index 0000000..63b7191
--- /dev/null
+++ b/src/com/android/launcher3/model/GridSizeMigrationTaskV2.java
@@ -0,0 +1,40 @@
+/*
+ * Copyright (C) 2020 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.launcher3.model;
+
+import android.content.Context;
+
+/**
+ * This class takes care of shrinking the workspace (by maximum of one row and one column), as a
+ * result of restoring from a larger device or device density change.
+ */
+public class GridSizeMigrationTaskV2 {
+
+    private GridSizeMigrationTaskV2(Context context) {
+
+    }
+
+    /**
+     * Migrates the workspace and hotseat in case their sizes changed.
+     *
+     * @return false if the migration failed.
+     */
+    public static boolean migrateGridIfNeeded(Context context) {
+        // To be implemented.
+        return true;
+    }
+}
diff --git a/src/com/android/launcher3/model/LoaderTask.java b/src/com/android/launcher3/model/LoaderTask.java
index 23ec459..ec23f98 100644
--- a/src/com/android/launcher3/model/LoaderTask.java
+++ b/src/com/android/launcher3/model/LoaderTask.java
@@ -19,6 +19,7 @@
 import static com.android.launcher3.ItemInfoWithIcon.FLAG_DISABLED_LOCKED_USER;
 import static com.android.launcher3.ItemInfoWithIcon.FLAG_DISABLED_SAFEMODE;
 import static com.android.launcher3.ItemInfoWithIcon.FLAG_DISABLED_SUSPENDED;
+import static com.android.launcher3.config.FeatureFlags.MULTI_DB_GRID_MIRATION_ALGO;
 import static com.android.launcher3.model.ModelUtils.filterCurrentWorkspaceItems;
 import static com.android.launcher3.util.Executors.MODEL_EXECUTOR;
 import static com.android.launcher3.util.PackageManagerHelper.hasShortcutsPermission;
@@ -320,7 +321,9 @@
             clearDb = true;
         }
 
-        if (!clearDb && !GridSizeMigrationTask.migrateGridIfNeeded(context)) {
+        if (!clearDb && (MULTI_DB_GRID_MIRATION_ALGO.get()
+                ? !GridSizeMigrationTaskV2.migrateGridIfNeeded(context)
+                : !GridSizeMigrationTask.migrateGridIfNeeded(context))) {
             // Migration failed. Clear workspace.
             clearDb = true;
         }
diff --git a/src/com/android/launcher3/provider/RestoreDbTask.java b/src/com/android/launcher3/provider/RestoreDbTask.java
index 407ff31..842be40 100644
--- a/src/com/android/launcher3/provider/RestoreDbTask.java
+++ b/src/com/android/launcher3/provider/RestoreDbTask.java
@@ -107,15 +107,17 @@
      */
     private void backupWorkspace(Context context, SQLiteDatabase db) throws Exception {
         InvariantDeviceProfile idp = LauncherAppState.getIDP(context);
-        new GridBackupTable(context, db, idp.numHotseatIcons, idp.numColumns, idp.numRows)
+        // TODO(pinyaoting): Support backing up workspace with multiple grid options.
+        new GridBackupTable(context, db, db, idp.numHotseatIcons, idp.numColumns, idp.numRows)
                 .doBackup(getDefaultProfileId(db), GridBackupTable.OPTION_REQUIRES_SANITIZATION);
     }
 
     private void restoreWorkspace(@NonNull Context context, @NonNull SQLiteDatabase db,
             @NonNull DatabaseHelper helper, @NonNull BackupManager backupManager)
             throws Exception {
+        // TODO(pinyaoting): Support restoring workspace with multiple grid options.
         final InvariantDeviceProfile idp = LauncherAppState.getIDP(context);
-        GridBackupTable backupTable = new GridBackupTable(context, db,
+        GridBackupTable backupTable = new GridBackupTable(context, db, db,
                 idp.numHotseatIcons, idp.numColumns, idp.numRows);
         if (backupTable.restoreFromRawBackupIfAvailable(getDefaultProfileId(db))) {
             sanitizeDB(helper, db, backupManager);