Adding support for backing up favorites table

Favorites table is copied as a separate table name during the first grid migration.
On subsequent migrations this backup table is used if it exists, otherwise new
backup is created. The backup table is also removed if there is any insert or
delete operation on the db (outside of the migration operation itself).

Bug: 111850268
Bug: 121048571
Change-Id: I6f02f4a355c369ee99d89430971be258f7516f6e
diff --git a/src/com/android/launcher3/LauncherProvider.java b/src/com/android/launcher3/LauncherProvider.java
index 8ed3314..24dc17a 100644
--- a/src/com/android/launcher3/LauncherProvider.java
+++ b/src/com/android/launcher3/LauncherProvider.java
@@ -16,6 +16,9 @@
 
 package com.android.launcher3;
 
+import static com.android.launcher3.provider.LauncherDbUtils.dropTable;
+import static com.android.launcher3.provider.LauncherDbUtils.tableExists;
+
 import android.annotation.TargetApi;
 import android.appwidget.AppWidgetHost;
 import android.appwidget.AppWidgetManager;
@@ -210,6 +213,7 @@
         addModifiedTime(initialValues);
         final int rowId = dbInsertAndCheck(mOpenHelper, db, args.table, null, initialValues);
         if (rowId < 0) return null;
+        mOpenHelper.onAddOrDeleteOp(db);
 
         uri = ContentUris.withAppendedId(uri, rowId);
         notifyListeners();
@@ -282,6 +286,7 @@
                     return 0;
                 }
             }
+            mOpenHelper.onAddOrDeleteOp(db);
             t.commit();
         }
 
@@ -290,15 +295,30 @@
         return values.length;
     }
 
+    @TargetApi(Build.VERSION_CODES.M)
     @Override
     public ContentProviderResult[] applyBatch(ArrayList<ContentProviderOperation> operations)
             throws OperationApplicationException {
         createDbIfNotExists();
         try (SQLiteTransaction t = new SQLiteTransaction(mOpenHelper.getWritableDatabase())) {
-            ContentProviderResult[] result =  super.applyBatch(operations);
+            boolean isAddOrDelete = !Utilities.ATLEAST_MARSHMALLOW;
+
+            final int numOperations = operations.size();
+            final ContentProviderResult[] results = new ContentProviderResult[numOperations];
+            for (int i = 0; i < numOperations; i++) {
+                ContentProviderOperation op = operations.get(i);
+                results[i] = op.apply(this, results, i);
+
+                isAddOrDelete |= (op.isInsert() || op.isDelete()) &&
+                        results[i].count != null && results[i].count > 0;
+            }
+            if (isAddOrDelete) {
+                mOpenHelper.onAddOrDeleteOp(t.getDb());
+            }
+
             t.commit();
             reloadLauncherIfExternal();
-            return result;
+            return results;
         }
     }
 
@@ -315,6 +335,7 @@
         }
         int count = db.delete(args.table, args.where, args.args);
         if (count > 0) {
+            mOpenHelper.onAddOrDeleteOp(db);
             notifyListeners();
             reloadLauncherIfExternal();
         }
@@ -381,6 +402,17 @@
                 mOpenHelper.removeGhostWidgets(mOpenHelper.getWritableDatabase());
                 return null;
             }
+            case LauncherSettings.Settings.METHOD_NEW_TRANSACTION: {
+                Bundle result = new Bundle();
+                result.putBinder(LauncherSettings.Settings.EXTRA_VALUE,
+                        new SQLiteTransaction(mOpenHelper.getWritableDatabase()));
+                return result;
+            }
+            case LauncherSettings.Settings.METHOD_REFRESH_BACKUP_TABLE: {
+                mOpenHelper.mBackupTableExists =
+                        tableExists(mOpenHelper.getReadableDatabase(), Favorites.BACKUP_TABLE_NAME);
+                return null;
+            }
         }
         return null;
     }
@@ -528,17 +560,19 @@
         private final Context mContext;
         private int mMaxItemId = -1;
         private int mMaxScreenId = -1;
+        private boolean mBackupTableExists;
 
         DatabaseHelper(Context context, Handler widgetHostResetHandler) {
             this(context, widgetHostResetHandler, 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.
-            if (!tableExists(Favorites.TABLE_NAME)) {
+            if (!tableExists(getReadableDatabase(), Favorites.TABLE_NAME)) {
                 Log.e(TAG, "Tables are missing after onCreate has been called. Trying to recreate");
                 // This operation is a no-op if the table already exists.
                 addFavoritesTable(getWritableDatabase(), true);
             }
+            mBackupTableExists = tableExists(getReadableDatabase(), Favorites.BACKUP_TABLE_NAME);
 
             initIds();
         }
@@ -564,18 +598,6 @@
             }
         }
 
-        private boolean tableExists(String tableName) {
-            Cursor c = getReadableDatabase().query(
-                    true, "sqlite_master", new String[] {"tbl_name"},
-                    "tbl_name = ?", new String[] {tableName},
-                    null, null, null, null, null);
-            try {
-                return c.getCount() > 0;
-            } finally {
-                c.close();
-            }
-        }
-
         @Override
         public void onCreate(SQLiteDatabase db) {
             if (LOGD) Log.d(TAG, "creating new launcher database");
@@ -590,6 +612,13 @@
             onEmptyDbCreated();
         }
 
+        protected void onAddOrDeleteOp(SQLiteDatabase db) {
+            if (mBackupTableExists) {
+                dropTable(db, Favorites.BACKUP_TABLE_NAME);
+                mBackupTableExists = false;
+            }
+        }
+
         /**
          * Overriden in tests.
          */
@@ -733,7 +762,7 @@
                                 Favorites.CONTAINER, Favorites.CONTAINER_DESKTOP);
                         db.execSQL(query);
                     }
-                    db.execSQL("DROP TABLE IF EXISTS workspaceScreens");
+                    dropTable(db, "workspaceScreens");
                 }
                 case 28:
                     // DB Upgraded successfully
@@ -762,8 +791,8 @@
          */
         public void createEmptyDB(SQLiteDatabase db) {
             try (SQLiteTransaction t = new SQLiteTransaction(db)) {
-                db.execSQL("DROP TABLE IF EXISTS " + Favorites.TABLE_NAME);
-                db.execSQL("DROP TABLE IF EXISTS workspaceScreens");
+                dropTable(db, Favorites.TABLE_NAME);
+                dropTable(db, "workspaceScreens");
                 onCreate(db);
                 t.commit();
             }
diff --git a/src/com/android/launcher3/LauncherSettings.java b/src/com/android/launcher3/LauncherSettings.java
index 79c8208..e248ba0 100644
--- a/src/com/android/launcher3/LauncherSettings.java
+++ b/src/com/android/launcher3/LauncherSettings.java
@@ -89,6 +89,11 @@
         public static final String TABLE_NAME = "favorites";
 
         /**
+         * Backup table created when when the favorites table is modified during grid migration
+         */
+        public static final String BACKUP_TABLE_NAME = "favorites_bakup";
+
+        /**
          * The content:// style URL for this table
          */
         public static final Uri CONTENT_URI = Uri.parse("content://" +
@@ -231,8 +236,13 @@
         public static final String OPTIONS = "options";
 
         public static void addTableToDb(SQLiteDatabase db, long myProfileId, boolean optional) {
+            addTableToDb(db, myProfileId, optional, TABLE_NAME);
+        }
+
+        public static void addTableToDb(SQLiteDatabase db, long myProfileId, boolean optional,
+                String tableName) {
             String ifNotExists = optional ? " IF NOT EXISTS " : "";
-            db.execSQL("CREATE TABLE " + ifNotExists + TABLE_NAME + " (" +
+            db.execSQL("CREATE TABLE " + ifNotExists + tableName + " (" +
                     "_id INTEGER PRIMARY KEY," +
                     "title TEXT," +
                     "intent TEXT," +
@@ -279,6 +289,10 @@
 
         public static final String METHOD_REMOVE_GHOST_WIDGETS = "remove_ghost_widgets";
 
+        public static final String METHOD_NEW_TRANSACTION = "new_db_transaction";
+
+        public static final String METHOD_REFRESH_BACKUP_TABLE = "refresh_backup_table";
+
         public static final String EXTRA_VALUE = "value";
 
         public static Bundle call(ContentResolver cr, String method) {
diff --git a/src/com/android/launcher3/model/GridBackupTable.java b/src/com/android/launcher3/model/GridBackupTable.java
new file mode 100644
index 0000000..804a040
--- /dev/null
+++ b/src/com/android/launcher3/model/GridBackupTable.java
@@ -0,0 +1,135 @@
+/*
+ * Copyright (C) 2018 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 static com.android.launcher3.LauncherSettings.Favorites.BACKUP_TABLE_NAME;
+import static com.android.launcher3.provider.LauncherDbUtils.dropTable;
+import static com.android.launcher3.provider.LauncherDbUtils.tableExists;
+
+import android.content.ContentValues;
+import android.content.Context;
+import android.database.Cursor;
+import android.database.sqlite.SQLiteDatabase;
+import android.graphics.Point;
+import android.os.Process;
+import android.util.Log;
+
+import com.android.launcher3.LauncherSettings.Favorites;
+import com.android.launcher3.LauncherSettings.Settings;
+import com.android.launcher3.compat.UserManagerCompat;
+
+/**
+ * Helper class to backup and restore Favorites table into a separate table
+ * within the same data base.
+ */
+public class GridBackupTable {
+    private static final String TAG = "GridBackupTable";
+
+    private static final int ID_PROPERTY = -1;
+
+    private static final String KEY_HOTSEAT_SIZE = Favorites.SCREEN;
+    private static final String KEY_GRID_X_SIZE = Favorites.SPANX;
+    private static final String KEY_GRID_Y_SIZE = Favorites.SPANY;
+    private static final String KEY_DB_VERSION = Favorites.RANK;
+
+    private final Context mContext;
+    private final SQLiteDatabase mDb;
+
+    private final int mOldHotseatSize;
+    private final int mOldGridX;
+    private final int mOldGridY;
+
+    private int mRestoredHotseatSize;
+    private int mRestoredGridX;
+    private int mRestoredGridY;
+
+    public GridBackupTable(Context context, SQLiteDatabase db,
+            int hotseatSize, int gridX, int gridY) {
+        mContext = context;
+        mDb = db;
+
+        mOldHotseatSize = hotseatSize;
+        mOldGridX = gridX;
+        mOldGridY = gridY;
+    }
+
+    public boolean backupOrRestoreAsNeeded() {
+        // Check if backup table exists
+        if (!tableExists(mDb, 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.
+                return false;
+            }
+
+            copyTable(Favorites.TABLE_NAME, BACKUP_TABLE_NAME);
+            encodeDBProperties();
+            return false;
+        }
+
+        if (!loadDbProperties()) {
+            return false;
+        }
+        copyTable(BACKUP_TABLE_NAME, Favorites.TABLE_NAME);
+        Log.d(TAG, "Backup table found");
+        return true;
+    }
+
+    public int getRestoreHotseatAndGridSize(Point outGridSize) {
+        outGridSize.set(mRestoredGridX, mRestoredGridY);
+        return mRestoredHotseatSize;
+    }
+
+    private void copyTable(String from, String to) {
+        long userSerial = UserManagerCompat.getInstance(mContext).getSerialNumberForUser(
+                Process.myUserHandle());
+        dropTable(mDb, to);
+        Favorites.addTableToDb(mDb, userSerial, false, to);
+        mDb.execSQL("INSERT INTO " + to + " SELECT * FROM " + from + " where _id > " + ID_PROPERTY);
+    }
+
+    private void encodeDBProperties() {
+        ContentValues values = new ContentValues();
+        values.put(Favorites._ID, ID_PROPERTY);
+        values.put(KEY_DB_VERSION, mDb.getVersion());
+        values.put(KEY_GRID_X_SIZE, mOldGridX);
+        values.put(KEY_GRID_Y_SIZE, mOldGridY);
+        values.put(KEY_HOTSEAT_SIZE, mOldHotseatSize);
+        mDb.insert(BACKUP_TABLE_NAME, null, values);
+    }
+
+    private boolean loadDbProperties() {
+        try (Cursor c = mDb.query(BACKUP_TABLE_NAME, new String[] {
+                        KEY_DB_VERSION,     // 0
+                        KEY_GRID_X_SIZE,    // 1
+                        KEY_GRID_Y_SIZE,    // 2
+                        KEY_HOTSEAT_SIZE},  // 3
+                "_id=" + ID_PROPERTY, null, null, null, null)) {
+            if (!c.moveToNext()) {
+                Log.e(TAG, "Meta data not found in backup table");
+                return false;
+            }
+            if (mDb.getVersion() != c.getInt(0)) {
+                return false;
+            }
+
+            mRestoredGridX = c.getInt(1);
+            mRestoredGridY = c.getInt(2);
+            mRestoredHotseatSize = c.getInt(3);
+            return true;
+        }
+    }
+}
diff --git a/src/com/android/launcher3/model/GridSizeMigrationTask.java b/src/com/android/launcher3/model/GridSizeMigrationTask.java
index a9ddccb..1c7bffa 100644
--- a/src/com/android/launcher3/model/GridSizeMigrationTask.java
+++ b/src/com/android/launcher3/model/GridSizeMigrationTask.java
@@ -1,10 +1,10 @@
 package com.android.launcher3.model;
 
+import static com.android.launcher3.LauncherSettings.Settings.EXTRA_VALUE;
 import static com.android.launcher3.Utilities.getPointString;
 import static com.android.launcher3.Utilities.parsePoint;
 
 import android.content.ComponentName;
-import android.content.ContentProviderOperation;
 import android.content.ContentValues;
 import android.content.Context;
 import android.content.Intent;
@@ -12,24 +12,27 @@
 import android.content.pm.PackageInfo;
 import android.content.pm.PackageManager;
 import android.database.Cursor;
+import android.database.sqlite.SQLiteDatabase;
 import android.graphics.Point;
 import android.util.Log;
+import android.util.SparseArray;
 
 import com.android.launcher3.InvariantDeviceProfile;
 import com.android.launcher3.ItemInfo;
 import com.android.launcher3.LauncherAppState;
 import com.android.launcher3.LauncherAppWidgetProviderInfo;
-import com.android.launcher3.LauncherProvider;
 import com.android.launcher3.LauncherSettings;
 import com.android.launcher3.LauncherSettings.Favorites;
+import com.android.launcher3.LauncherSettings.Settings;
 import com.android.launcher3.Utilities;
 import com.android.launcher3.Workspace;
 import com.android.launcher3.compat.AppWidgetManagerCompat;
 import com.android.launcher3.compat.PackageInstallerCompat;
 import com.android.launcher3.config.FeatureFlags;
+import com.android.launcher3.provider.LauncherDbUtils;
+import com.android.launcher3.provider.LauncherDbUtils.SQLiteTransaction;
 import com.android.launcher3.util.GridOccupancy;
 import com.android.launcher3.util.IntArray;
-import com.android.launcher3.util.IntSet;
 import com.android.launcher3.util.IntSparseArrayMap;
 
 import java.util.ArrayList;
@@ -60,13 +63,13 @@
     private static final float WT_WIDGET_FACTOR = 0.6f;
     private static final float WT_FOLDER_FACTOR = 0.5f;
 
-    private final Context mContext;
-    private final InvariantDeviceProfile mIdp;
+    protected final SQLiteDatabase mDb;
+    protected final Context mContext;
 
-    private final ContentValues mTempValues = new ContentValues();
     protected final IntArray mEntryToRemove = new IntArray();
-    private final ArrayList<ContentProviderOperation> mUpdateOperations = new ArrayList<>();
     protected final ArrayList<DbEntry> mCarryOver = new ArrayList<>();
+
+    private final SparseArray<ContentValues> mUpdateOperations = new SparseArray<>();
     private final HashSet<String> mValidPackages;
 
     private final int mSrcX, mSrcY;
@@ -76,11 +79,11 @@
     private final int mSrcHotseatSize;
     private final int mDestHotseatSize;
 
-    protected GridSizeMigrationTask(Context context, InvariantDeviceProfile idp,
+    protected GridSizeMigrationTask(Context context, SQLiteDatabase db,
             HashSet<String> validPackages, Point sourceSize, Point targetSize) {
         mContext = context;
+        mDb = db;
         mValidPackages = validPackages;
-        mIdp = idp;
 
         mSrcX = sourceSize.x;
         mSrcY = sourceSize.y;
@@ -95,11 +98,10 @@
         mSrcHotseatSize = mDestHotseatSize = -1;
     }
 
-    protected GridSizeMigrationTask(Context context,
-            InvariantDeviceProfile idp, HashSet<String> validPackages,
-            int srcHotseatSize, int destHotseatSize) {
+    protected GridSizeMigrationTask(Context context, SQLiteDatabase db,
+            HashSet<String> validPackages, int srcHotseatSize, int destHotseatSize) {
         mContext = context;
-        mIdp = idp;
+        mDb = db;
         mValidPackages = validPackages;
 
         mSrcHotseatSize = srcHotseatSize;
@@ -118,20 +120,21 @@
      */
     private boolean applyOperations() throws Exception {
         // Update items
-        if (!mUpdateOperations.isEmpty()) {
-            mContext.getContentResolver().applyBatch(LauncherProvider.AUTHORITY, mUpdateOperations);
+        int updateCount = mUpdateOperations.size();
+        for (int i = 0; i < updateCount; i++) {
+            mDb.update(Favorites.TABLE_NAME, mUpdateOperations.valueAt(i),
+                    "_id=" + mUpdateOperations.keyAt(i), null);
         }
 
         if (!mEntryToRemove.isEmpty()) {
             if (DEBUG) {
                 Log.d(TAG, "Removing items: " + mEntryToRemove.toConcatString());
             }
-            mContext.getContentResolver().delete(LauncherSettings.Favorites.CONTENT_URI,
-                    Utilities.createDbSelectionQuery(
-                            LauncherSettings.Favorites._ID, mEntryToRemove), null);
+            mDb.delete(Favorites.TABLE_NAME, Utilities.createDbSelectionQuery(
+                    Favorites._ID, mEntryToRemove), null);
         }
 
-        return !mUpdateOperations.isEmpty() || !mEntryToRemove.isEmpty();
+        return updateCount > 0 || !mEntryToRemove.isEmpty();
     }
 
     /**
@@ -181,24 +184,17 @@
     }
 
     @VisibleForTesting
-    static IntArray getWorkspaceScreenIds(Context context) {
-        IntSet set = new IntSet();
-        try (Cursor c = context.getContentResolver().query(Favorites.CONTENT_URI,
-                new String[] {Favorites.SCREEN},
+    static IntArray getWorkspaceScreenIds(SQLiteDatabase db) {
+        return LauncherDbUtils.queryIntArray(db, Favorites.TABLE_NAME, Favorites.SCREEN,
                 Favorites.CONTAINER + " = " + Favorites.CONTAINER_DESKTOP,
-                null, Favorites.SCREEN)) {
-            while (c.moveToNext()) {
-                set.add(c.getInt(0));
-            }
-        }
-        return set.getArray();
+                Favorites.SCREEN, Favorites.SCREEN);
     }
 
     /**
      * @return true if any DB change was made
      */
     protected boolean migrateWorkspace() throws Exception {
-        IntArray allScreens = getWorkspaceScreenIds(mContext);
+        IntArray allScreens = getWorkspaceScreenIds(mDb);
         if (allScreens.isEmpty()) {
             throw new Exception("Unable to get workspace screens");
         }
@@ -230,8 +226,7 @@
                     int newScreenId = LauncherSettings.Settings.call(
                             mContext.getContentResolver(),
                             LauncherSettings.Settings.METHOD_NEW_SCREEN_ID)
-                                .getInt(LauncherSettings.Settings.EXTRA_VALUE);
-
+                            .getInt(EXTRA_VALUE);
                     for (DbEntry item : placement.finalPlacedItems) {
                         if (!mCarryOver.remove(itemMap.get(item.id))) {
                             throw new Exception("Unable to find matching items");
@@ -362,11 +357,9 @@
      * Updates an item in the DB.
      */
     protected void update(DbEntry item) {
-        mTempValues.clear();
-        item.addToContentValues(mTempValues);
-        mUpdateOperations.add(ContentProviderOperation
-                .newUpdate(LauncherSettings.Favorites.getContentUri(item.id))
-                .withValues(mTempValues).build());
+        ContentValues values = new ContentValues();
+        item.addToContentValues(values);
+        mUpdateOperations.put(item.id, values);
     }
 
     /**
@@ -612,13 +605,13 @@
     }
 
     private ArrayList<DbEntry> loadHotseatEntries() {
-        Cursor c =  mContext.getContentResolver().query(LauncherSettings.Favorites.CONTENT_URI,
+        Cursor c =  queryWorkspace(
                 new String[]{
                         Favorites._ID,                  // 0
                         Favorites.ITEM_TYPE,            // 1
                         Favorites.INTENT,               // 2
                         Favorites.SCREEN},              // 3
-                Favorites.CONTAINER + " = " + Favorites.CONTAINER_HOTSEAT, null, null, null);
+                Favorites.CONTAINER + " = " + Favorites.CONTAINER_HOTSEAT);
 
         final int indexId = c.getColumnIndexOrThrow(Favorites._ID);
         final int indexItemType = c.getColumnIndexOrThrow(Favorites.ITEM_TYPE);
@@ -796,8 +789,7 @@
     }
 
     protected Cursor queryWorkspace(String[] columns, String where) {
-        return mContext.getContentResolver().query(LauncherSettings.Favorites.CONTENT_URI,
-                columns, where, null, null, null);
+        return mDb.query(Favorites.TABLE_NAME, columns, where, null, null, null, null);
     }
 
     /**
@@ -864,11 +856,11 @@
         }
 
         public void addToContentValues(ContentValues values) {
-            values.put(LauncherSettings.Favorites.SCREEN, screenId);
-            values.put(LauncherSettings.Favorites.CELLX, cellX);
-            values.put(LauncherSettings.Favorites.CELLY, cellY);
-            values.put(LauncherSettings.Favorites.SPANX, spanX);
-            values.put(LauncherSettings.Favorites.SPANY, spanY);
+            values.put(Favorites.SCREEN, screenId);
+            values.put(Favorites.CELLX, cellX);
+            values.put(Favorites.CELLY, cellY);
+            values.put(Favorites.SPANX, spanX);
+            values.put(Favorites.SPANY, spanY);
         }
     }
 
@@ -907,34 +899,43 @@
         }
 
         long migrationStartTime = System.currentTimeMillis();
-        try {
+        try (SQLiteTransaction transaction = (SQLiteTransaction) Settings.call(
+                context.getContentResolver(), Settings.METHOD_NEW_TRANSACTION)
+                .getBinder(Settings.EXTRA_VALUE)) {
+
+            int srcHotseatCount = prefs.getInt(KEY_MIGRATION_SRC_HOTSEAT_COUNT,
+                    idp.numHotseatIcons);
+            Point sourceSize = parsePoint(prefs.getString(
+                    KEY_MIGRATION_SRC_WORKSPACE_SIZE, gridSizeString));
+
             boolean dbChanged = false;
 
+            GridBackupTable backupTable = new GridBackupTable(context, transaction.getDb(),
+                    srcHotseatCount, sourceSize.x, sourceSize.y);
+            if (backupTable.backupOrRestoreAsNeeded()) {
+                dbChanged = true;
+                srcHotseatCount = backupTable.getRestoreHotseatAndGridSize(sourceSize);
+            }
+
             HashSet<String> validPackages = getValidPackages(context);
             // Hotseat
-            int srcHotseatCount = prefs.getInt(KEY_MIGRATION_SRC_HOTSEAT_COUNT,
-                    idp.numHotseatIcons);
             if (srcHotseatCount != idp.numHotseatIcons) {
                 // Migrate hotseat.
-
-                dbChanged = new GridSizeMigrationTask(context, LauncherAppState.getIDP(context),
+                dbChanged = new GridSizeMigrationTask(context, transaction.getDb(),
                         validPackages, srcHotseatCount, idp.numHotseatIcons).migrateHotseat();
             }
 
             // Grid size
             Point targetSize = new Point(idp.numColumns, idp.numRows);
-            Point sourceSize = parsePoint(prefs.getString(
-                    KEY_MIGRATION_SRC_WORKSPACE_SIZE, gridSizeString));
-
-            if (new MultiStepMigrationTask(validPackages, context).migrate(sourceSize,
-                    targetSize)) {
+            if (new MultiStepMigrationTask(validPackages, context, transaction.getDb())
+                    .migrate(sourceSize, targetSize)) {
                 dbChanged = true;
             }
 
             if (dbChanged) {
                 // Make sure we haven't removed everything.
                 final Cursor c = context.getContentResolver().query(
-                        LauncherSettings.Favorites.CONTENT_URI, null, null, null, null);
+                        Favorites.CONTENT_URI, null, null, null, null);
                 boolean hasData = c.moveToNext();
                 c.close();
                 if (!hasData) {
@@ -942,6 +943,8 @@
                 }
             }
 
+            transaction.commit();
+            Settings.call(context.getContentResolver(), Settings.METHOD_REFRESH_BACKUP_TABLE);
             return true;
         } catch (Exception e) {
             Log.e(TAG, "Error during grid migration", e);
@@ -982,19 +985,24 @@
      */
     public static IntSparseArrayMap<Object> removeBrokenHotseatItems(Context context)
             throws Exception {
-        GridSizeMigrationTask task = new GridSizeMigrationTask(
-                context, LauncherAppState.getIDP(context), getValidPackages(context),
-                Integer.MAX_VALUE, Integer.MAX_VALUE);
+        try (SQLiteTransaction transaction = (SQLiteTransaction) Settings.call(
+                context.getContentResolver(), Settings.METHOD_NEW_TRANSACTION)
+                .getBinder(Settings.EXTRA_VALUE)) {
+            GridSizeMigrationTask task = new GridSizeMigrationTask(
+                    context, transaction.getDb(), getValidPackages(context),
+                    Integer.MAX_VALUE, Integer.MAX_VALUE);
 
-        // Load all the valid entries
-        ArrayList<DbEntry> items = task.loadHotseatEntries();
-        // Delete any entry marked for deletion by above load.
-        task.applyOperations();
-        IntSparseArrayMap<Object> positions = new IntSparseArrayMap<>();
-        for (DbEntry item : items) {
-            positions.put(item.screenId, item);
+            // Load all the valid entries
+            ArrayList<DbEntry> items = task.loadHotseatEntries();
+            // Delete any entry marked for deletion by above load.
+            task.applyOperations();
+            IntSparseArrayMap<Object> positions = new IntSparseArrayMap<>();
+            for (DbEntry item : items) {
+                positions.put(item.screenId, item);
+            }
+            transaction.commit();
+            return positions;
         }
-        return positions;
     }
 
     /**
@@ -1003,10 +1011,13 @@
     protected static class MultiStepMigrationTask {
         private final HashSet<String> mValidPackages;
         private final Context mContext;
+        private final SQLiteDatabase mDb;
 
-        public MultiStepMigrationTask(HashSet<String> validPackages, Context context) {
+        public MultiStepMigrationTask(HashSet<String> validPackages, Context context,
+                SQLiteDatabase db) {
             mValidPackages = validPackages;
             mContext = context;
+            mDb = db;
         }
 
         public boolean migrate(Point sourceSize, Point targetSize) throws Exception {
@@ -1042,7 +1053,7 @@
         }
 
         protected boolean runStepTask(Point sourceSize, Point nextSize) throws Exception {
-            return new GridSizeMigrationTask(mContext, LauncherAppState.getIDP(mContext),
+            return new GridSizeMigrationTask(mContext, mDb,
                     mValidPackages, sourceSize, nextSize).migrateWorkspace();
         }
     }
diff --git a/src/com/android/launcher3/provider/LauncherDbUtils.java b/src/com/android/launcher3/provider/LauncherDbUtils.java
index b79478a..2c843f9 100644
--- a/src/com/android/launcher3/provider/LauncherDbUtils.java
+++ b/src/com/android/launcher3/provider/LauncherDbUtils.java
@@ -21,6 +21,7 @@
 import android.database.Cursor;
 import android.database.DatabaseUtils;
 import android.database.sqlite.SQLiteDatabase;
+import android.os.Binder;
 import android.util.Log;
 
 import com.android.launcher3.LauncherAppState;
@@ -103,10 +104,22 @@
         return out;
     }
 
+    public static boolean tableExists(SQLiteDatabase db, String tableName) {
+        try (Cursor c = db.query(true, "sqlite_master", new String[] {"tbl_name"},
+                "tbl_name = ?", new String[] {tableName},
+                null, null, null, null, null)) {
+            return c.getCount() > 0;
+        }
+    }
+
+    public static void dropTable(SQLiteDatabase db, String tableName) {
+        db.execSQL("DROP TABLE IF EXISTS " + tableName);
+    }
+
     /**
      * Utility class to simplify managing sqlite transactions
      */
-    public static class SQLiteTransaction implements AutoCloseable {
+    public static class SQLiteTransaction extends Binder implements AutoCloseable {
         private final SQLiteDatabase mDb;
 
         public SQLiteTransaction(SQLiteDatabase db) {
@@ -122,5 +135,9 @@
         public void close() {
             mDb.endTransaction();
         }
+
+        public SQLiteDatabase getDb() {
+            return mDb;
+        }
     }
 }
diff --git a/src/com/android/launcher3/provider/LossyScreenMigrationTask.java b/src/com/android/launcher3/provider/LossyScreenMigrationTask.java
index 9166b83..6d839f3 100644
--- a/src/com/android/launcher3/provider/LossyScreenMigrationTask.java
+++ b/src/com/android/launcher3/provider/LossyScreenMigrationTask.java
@@ -37,29 +37,21 @@
  */
 public class LossyScreenMigrationTask extends GridSizeMigrationTask {
 
-    private final SQLiteDatabase mDb;
-
     private final IntSparseArrayMap<DbEntry> mOriginalItems;
     private final IntSparseArrayMap<DbEntry> mUpdates;
 
     protected LossyScreenMigrationTask(
             Context context, InvariantDeviceProfile idp, SQLiteDatabase db) {
         // Decrease the rows count by 1
-        super(context, idp, getValidPackages(context),
+        super(context, db, getValidPackages(context),
                 new Point(idp.numColumns, idp.numRows + 1),
                 new Point(idp.numColumns, idp.numRows));
 
-        mDb = db;
         mOriginalItems = new IntSparseArrayMap<>();
         mUpdates = new IntSparseArrayMap<>();
     }
 
     @Override
-    protected Cursor queryWorkspace(String[] columns, String where) {
-        return mDb.query(Favorites.TABLE_NAME, columns, where, null, null, null, null);
-    }
-
-    @Override
     protected void update(DbEntry item) {
         mUpdates.put(item.id, item.copy());
     }
diff --git a/src/com/android/launcher3/provider/RestoreDbTask.java b/src/com/android/launcher3/provider/RestoreDbTask.java
index 17c66b4..bcca4d8 100644
--- a/src/com/android/launcher3/provider/RestoreDbTask.java
+++ b/src/com/android/launcher3/provider/RestoreDbTask.java
@@ -16,6 +16,8 @@
 
 package com.android.launcher3.provider;
 
+import static com.android.launcher3.provider.LauncherDbUtils.dropTable;
+
 import android.content.ContentValues;
 import android.content.Context;
 import android.database.Cursor;
@@ -111,7 +113,7 @@
         db.execSQL("ALTER TABLE favorites RENAME TO favorites_old;");
         Favorites.addTableToDb(db, newProfileId, false);
         db.execSQL("INSERT INTO favorites SELECT * FROM favorites_old;");
-        db.execSQL("DROP TABLE favorites_old;");
+        dropTable(db, "favorites_old");
     }
 
     /**