diff --git a/src/com/android/launcher3/LauncherProvider.java b/src/com/android/launcher3/LauncherProvider.java
index e004a4b..9891a60 100644
--- a/src/com/android/launcher3/LauncherProvider.java
+++ b/src/com/android/launcher3/LauncherProvider.java
@@ -158,6 +158,7 @@
     private synchronized boolean prepForMigration(String dbFile, String targetTableName,
             Supplier<DatabaseHelper> src, Supplier<DatabaseHelper> dst) {
         if (TextUtils.equals(dbFile, mOpenHelper.getDatabaseName())) {
+            Log.e("b/198965093", "prepForMigration - target db is same as current: " + dbFile);
             return false;
         }
 
@@ -427,7 +428,7 @@
                 Bundle result = new Bundle();
                 result.putBoolean(LauncherSettings.Settings.EXTRA_VALUE,
                         prepForMigration(
-                                InvariantDeviceProfile.INSTANCE.get(getContext()).dbFile,
+                                arg /* dbFile */,
                                 Favorites.TMP_TABLE,
                                 () -> mOpenHelper,
                                 () -> DatabaseHelper.createDatabaseHelper(
diff --git a/src/com/android/launcher3/model/DeviceGridState.java b/src/com/android/launcher3/model/DeviceGridState.java
index 3e49d79..35fcb78 100644
--- a/src/com/android/launcher3/model/DeviceGridState.java
+++ b/src/com/android/launcher3/model/DeviceGridState.java
@@ -43,15 +43,18 @@
     public static final String KEY_WORKSPACE_SIZE = "migration_src_workspace_size";
     public static final String KEY_HOTSEAT_COUNT = "migration_src_hotseat_count";
     public static final String KEY_DEVICE_TYPE = "migration_src_device_type";
+    public static final String KEY_DB_FILE = "migration_src_db_file";
 
     private final String mGridSizeString;
     private final int mNumHotseat;
     private final @DeviceType int mDeviceType;
+    private final String mDbFile;
 
     public DeviceGridState(InvariantDeviceProfile idp) {
         mGridSizeString = String.format(Locale.ENGLISH, "%d,%d", idp.numColumns, idp.numRows);
         mNumHotseat = idp.numDatabaseHotseatIcons;
         mDeviceType = idp.deviceType;
+        mDbFile = idp.dbFile;
     }
 
     public DeviceGridState(Context context) {
@@ -59,6 +62,7 @@
         mGridSizeString = prefs.getString(KEY_WORKSPACE_SIZE, "");
         mNumHotseat = prefs.getInt(KEY_HOTSEAT_COUNT, -1);
         mDeviceType = prefs.getInt(KEY_DEVICE_TYPE, TYPE_PHONE);
+        mDbFile = prefs.getString(KEY_DB_FILE, "");
     }
 
     /**
@@ -69,6 +73,20 @@
     }
 
     /**
+     * Returns the databaseFile for the grid.
+     */
+    public String getDbFile() {
+        return mDbFile;
+    }
+
+    /**
+     * Returns the number of hotseat icons.
+     */
+    public int getNumHotseat() {
+        return mNumHotseat;
+    }
+
+    /**
      * Stores the device state to shared preferences
      */
     public void writeToPrefs(Context context) {
@@ -76,6 +94,7 @@
                 .putString(KEY_WORKSPACE_SIZE, mGridSizeString)
                 .putInt(KEY_HOTSEAT_COUNT, mNumHotseat)
                 .putInt(KEY_DEVICE_TYPE, mDeviceType)
+                .putString(KEY_DB_FILE, mDbFile)
                 .apply();
     }
 
@@ -106,6 +125,7 @@
                 + "mGridSizeString='" + mGridSizeString + '\''
                 + ", mNumHotseat=" + mNumHotseat
                 + ", mDeviceType=" + mDeviceType
+                + ", mDbFile=" + mDbFile
                 + '}';
     }
 
diff --git a/src/com/android/launcher3/model/GridSizeMigrationTaskV2.java b/src/com/android/launcher3/model/GridSizeMigrationTaskV2.java
index 74b0a6f..0369b21 100644
--- a/src/com/android/launcher3/model/GridSizeMigrationTaskV2.java
+++ b/src/com/android/launcher3/model/GridSizeMigrationTaskV2.java
@@ -22,7 +22,6 @@
 import android.content.ContentValues;
 import android.content.Context;
 import android.content.Intent;
-import android.content.SharedPreferences;
 import android.content.pm.PackageInfo;
 import android.content.pm.PackageManager;
 import android.database.Cursor;
@@ -104,13 +103,16 @@
      * Check given a new IDP, if migration is necessary.
      */
     public static boolean needsToMigrate(Context context, InvariantDeviceProfile idp) {
-        DeviceGridState idpGridState = new DeviceGridState(idp);
-        DeviceGridState contextGridState = new DeviceGridState(context);
-        boolean needsToMigrate = !idpGridState.isCompatible(contextGridState);
+        return needsToMigrate(new DeviceGridState(context), new DeviceGridState(idp));
+    }
+
+    private static boolean needsToMigrate(
+            DeviceGridState srcDeviceState, DeviceGridState destDeviceState) {
+        boolean needsToMigrate = !destDeviceState.isCompatible(srcDeviceState);
         // TODO(b/198965093): Revert this change after bug is fixed
         if (needsToMigrate) {
-            Log.d("b/198965093", "Migration is needed. idpGridState: " + idpGridState
-                    + ", contextGridState: " + contextGridState);
+            Log.d("b/198965093", "Migration is needed. destDeviceState: " + destDeviceState
+                    + ", srcDeviceState: " + srcDeviceState);
         }
         return needsToMigrate;
     }
@@ -143,23 +145,26 @@
             idp = LauncherAppState.getIDP(context);
         }
 
-        if (!needsToMigrate(context, idp)) {
+        DeviceGridState srcDeviceState = new DeviceGridState(context);
+        DeviceGridState destDeviceState = new DeviceGridState(idp);
+        if (!needsToMigrate(srcDeviceState, destDeviceState)) {
             return true;
         }
 
-        SharedPreferences prefs = Utilities.getPrefs(context);
         HashSet<String> validPackages = getValidPackages(context);
 
         if (migrateForPreview) {
             if (!LauncherSettings.Settings.call(
                     context.getContentResolver(),
-                    LauncherSettings.Settings.METHOD_PREP_FOR_PREVIEW, idp.dbFile).getBoolean(
+                    LauncherSettings.Settings.METHOD_PREP_FOR_PREVIEW,
+                    destDeviceState.getDbFile()).getBoolean(
                     LauncherSettings.Settings.EXTRA_VALUE)) {
                 return false;
             }
         } else if (!LauncherSettings.Settings.call(
                 context.getContentResolver(),
-                LauncherSettings.Settings.METHOD_UPDATE_CURRENT_OPEN_HELPER).getBoolean(
+                LauncherSettings.Settings.METHOD_UPDATE_CURRENT_OPEN_HELPER,
+                destDeviceState.getDbFile()).getBoolean(
                 LauncherSettings.Settings.EXTRA_VALUE)) {
             return false;
         }
@@ -179,10 +184,10 @@
                             : LauncherSettings.Favorites.TABLE_NAME,
                     context, validPackages);
 
-            Point targetSize = new Point(idp.numColumns, idp.numRows);
+            Point targetSize = new Point(destDeviceState.getColumns(), destDeviceState.getRows());
             GridSizeMigrationTaskV2 task = new GridSizeMigrationTaskV2(context, t.getDb(),
-                    srcReader, destReader, idp.numDatabaseHotseatIcons, targetSize);
-            task.migrate(idp);
+                    srcReader, destReader, destDeviceState.getNumHotseat(), targetSize);
+            task.migrate(srcDeviceState, destDeviceState);
 
             if (!migrateForPreview) {
                 dropTable(t.getDb(), LauncherSettings.Favorites.TMP_TABLE);
@@ -200,13 +205,13 @@
 
             if (!migrateForPreview) {
                 // Save current configuration, so that the migration does not run again.
-                new DeviceGridState(idp).writeToPrefs(context);
+                destDeviceState.writeToPrefs(context);
             }
         }
     }
 
     @VisibleForTesting
-    protected boolean migrate(InvariantDeviceProfile idp) {
+    protected boolean migrate(DeviceGridState srcDeviceState, DeviceGridState destDeviceState) {
         if (mHotseatDiff.isEmpty() && mWorkspaceDiff.isEmpty()) {
             return false;
         }
@@ -228,8 +233,6 @@
 
         boolean preservePages = false;
         if (screens.isEmpty() && FeatureFlags.ENABLE_NEW_MIGRATION_LOGIC.get()) {
-            DeviceGridState srcDeviceState = new DeviceGridState(mContext);
-            DeviceGridState destDeviceState = new DeviceGridState(idp);
             preservePages = destDeviceState.compareTo(srcDeviceState) >= 0
                     && destDeviceState.getColumns() - srcDeviceState.getColumns() <= 2;
         }