diff --git a/src/com/android/launcher3/LauncherBackupAgent.java b/src/com/android/launcher3/LauncherBackupAgent.java
index b2f5c57..b3e73f7 100644
--- a/src/com/android/launcher3/LauncherBackupAgent.java
+++ b/src/com/android/launcher3/LauncherBackupAgent.java
@@ -3,24 +3,12 @@
 import android.app.backup.BackupAgent;
 import android.app.backup.BackupDataInput;
 import android.app.backup.BackupDataOutput;
-import android.content.ContentValues;
-import android.database.Cursor;
-import android.database.sqlite.SQLiteDatabase;
 import android.os.ParcelFileDescriptor;
 
-import com.android.launcher3.LauncherProvider.DatabaseHelper;
-import com.android.launcher3.LauncherSettings.Favorites;
-import com.android.launcher3.logging.FileLog;
-
-import java.io.InvalidObjectException;
+import com.android.launcher3.provider.RestoreDbTask;
 
 public class LauncherBackupAgent extends BackupAgent {
 
-    private static final String TAG = "LauncherBackupAgent";
-
-    private static final String INFO_COLUMN_NAME = "name";
-    private static final String INFO_COLUMN_DEFAULT_VALUE = "dflt_value";
-
     @Override
     public void onRestore(
             BackupDataInput data, int appVersionCode, ParcelFileDescriptor newState) {
@@ -35,99 +23,6 @@
 
     @Override
     public void onRestoreFinished() {
-        DatabaseHelper helper = new DatabaseHelper(this, null, LauncherFiles.LAUNCHER_DB);
-
-        if (!sanitizeDBSafely(helper)) {
-            helper.createEmptyDB(helper.getWritableDatabase());
-        }
-
-        try {
-            // Flush all logs before the process is killed.
-            FileLog.flushAll(null);
-        } catch (Exception e) { }
-    }
-
-    private boolean sanitizeDBSafely(DatabaseHelper helper) {
-        SQLiteDatabase db = helper.getWritableDatabase();
-        db.beginTransaction();
-        try {
-            sanitizeDB(helper, db);
-            db.setTransactionSuccessful();
-            return true;
-        } catch (Exception e) {
-            FileLog.e(TAG, "Failed to verify db", e);
-            return false;
-        } finally {
-            db.endTransaction();
-        }
-    }
-
-    /**
-     * Makes the following changes in the provider DB.
-     *   1. Removes all entries belonging to a managed profile as managed profiles
-     *      cannot be restored.
-     *   2. Marks all entries as restored. The flags are updated during first load or as
-     *      the restored apps get installed.
-     *   3. If the user serial for primary profile is different than that of the previous device,
-     *      update the entries to the new profile id.
-     */
-    private void sanitizeDB(DatabaseHelper helper, SQLiteDatabase db) throws Exception {
-        long oldProfileId = getDefaultProfileId(db);
-        // Delete all entries which do not belong to the main user
-        int itemsDeleted = db.delete(
-                Favorites.TABLE_NAME, "profileId != ?", new String[]{Long.toString(oldProfileId)});
-        if (itemsDeleted > 0) {
-            FileLog.d(TAG, itemsDeleted + " items belonging to a managed profile, were deleted");
-        }
-
-        // Mark all items as restored.
-        ContentValues values = new ContentValues();
-        values.put(Favorites.RESTORED, 1);
-        db.update(Favorites.TABLE_NAME, values, null, null);
-
-        // Mark widgets with appropriate restore flag
-        values.put(Favorites.RESTORED,
-                LauncherAppWidgetInfo.FLAG_ID_NOT_VALID |
-                LauncherAppWidgetInfo.FLAG_PROVIDER_NOT_READY |
-                LauncherAppWidgetInfo.FLAG_UI_NOT_READY);
-        db.update(Favorites.TABLE_NAME, values, "itemType = ?",
-                new String[]{Integer.toString(Favorites.ITEM_TYPE_APPWIDGET)});
-
-        long myProfileId = helper.getDefaultUserSerial();
-        if (Utilities.longCompare(oldProfileId, myProfileId) != 0) {
-            FileLog.d(TAG, "Changing primary user id from " + oldProfileId + " to " + myProfileId);
-            migrateProfileId(db, myProfileId);
-        }
-    }
-
-    /**
-     * Updates profile id of all entries and changes the default value for the column.
-     */
-    protected void migrateProfileId(SQLiteDatabase db, long newProfileId) {
-        // Update existing entries.
-        ContentValues values = new ContentValues();
-        values.put(Favorites.PROFILE_ID, newProfileId);
-        db.update(Favorites.TABLE_NAME, values, null, null);
-
-        // Change default value of the column.
-        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;");
-    }
-
-    /**
-     * Returns the profile id for used in the favorites table of the provided db.
-     */
-    protected long getDefaultProfileId(SQLiteDatabase db) throws Exception {
-        try (Cursor c = db.rawQuery("PRAGMA table_info (favorites)", null)){
-            int nameIndex = c.getColumnIndex(INFO_COLUMN_NAME);
-            while (c.moveToNext()) {
-                if (Favorites.PROFILE_ID.equals(c.getString(nameIndex))) {
-                    return c.getLong(c.getColumnIndex(INFO_COLUMN_DEFAULT_VALUE));
-                }
-            }
-            throw new InvalidObjectException("Table does not have a profile id column");
-        }
+        RestoreDbTask.setPending(this, true);
     }
 }
diff --git a/src/com/android/launcher3/LauncherProvider.java b/src/com/android/launcher3/LauncherProvider.java
index 49ce06a..dfb8ba2 100644
--- a/src/com/android/launcher3/LauncherProvider.java
+++ b/src/com/android/launcher3/LauncherProvider.java
@@ -46,7 +46,6 @@
 import android.os.Message;
 import android.os.Process;
 import android.os.UserManager;
-import android.provider.BaseColumns;
 import android.text.TextUtils;
 import android.util.Log;
 import android.util.SparseArray;
@@ -58,6 +57,7 @@
 import com.android.launcher3.compat.UserManagerCompat;
 import com.android.launcher3.config.ProviderConfig;
 import com.android.launcher3.dynamicui.ExtractionUtils;
+import com.android.launcher3.provider.RestoreDbTask;
 import com.android.launcher3.util.ManagedProfileHeuristic;
 import com.android.launcher3.util.NoLocaleSqliteContext;
 import com.android.launcher3.util.Preconditions;
@@ -117,6 +117,15 @@
     protected synchronized void createDbIfNotExists() {
         if (mOpenHelper == null) {
             mOpenHelper = new DatabaseHelper(getContext(), mListenerHandler);
+
+            if (RestoreDbTask.isPending(getContext())) {
+                if (!RestoreDbTask.performRestore(mOpenHelper)) {
+                    mOpenHelper.createEmptyDB(mOpenHelper.getWritableDatabase());
+                }
+                // Set is pending to false irrespective of the result, so that it doesn't get
+                // executed again.
+                RestoreDbTask.setPending(getContext(), false);
+            }
         }
     }
 
@@ -626,7 +635,7 @@
                     mContext);
         }
 
-        protected long getDefaultUserSerial() {
+        public long getDefaultUserSerial() {
             return UserManagerCompat.getInstance(mContext).getSerialNumberForUser(
                     UserHandleCompat.myUserHandle());
         }
diff --git a/src/com/android/launcher3/provider/RestoreDbTask.java b/src/com/android/launcher3/provider/RestoreDbTask.java
new file mode 100644
index 0000000..9d8b6b3
--- /dev/null
+++ b/src/com/android/launcher3/provider/RestoreDbTask.java
@@ -0,0 +1,137 @@
+/*
+ * Copyright (C) 2016 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.provider;
+
+import android.content.ContentValues;
+import android.content.Context;
+import android.database.Cursor;
+import android.database.sqlite.SQLiteDatabase;
+
+import com.android.launcher3.LauncherAppWidgetInfo;
+import com.android.launcher3.LauncherProvider.DatabaseHelper;
+import com.android.launcher3.LauncherSettings.Favorites;
+import com.android.launcher3.Utilities;
+import com.android.launcher3.logging.FileLog;
+
+import java.io.InvalidObjectException;
+
+/**
+ * Utility class to update DB schema after it has been restored.
+ *
+ * This task is executed when Launcher starts for the first time and not immediately after restore.
+ * This helps keep the model consistent if the launcher updates between restore and first startup.
+ */
+public class RestoreDbTask {
+
+    private static final String TAG = "RestoreDbTask";
+    private static final String RESTORE_TASK_PENDING = "restore_task_pending";
+
+    private static final String INFO_COLUMN_NAME = "name";
+    private static final String INFO_COLUMN_DEFAULT_VALUE = "dflt_value";
+
+    public static boolean performRestore(DatabaseHelper helper) {
+        SQLiteDatabase db = helper.getWritableDatabase();
+        db.beginTransaction();
+        try {
+            new RestoreDbTask().sanitizeDB(helper, db);
+            db.setTransactionSuccessful();
+            return true;
+        } catch (Exception e) {
+            FileLog.e(TAG, "Failed to verify db", e);
+            return false;
+        } finally {
+            db.endTransaction();
+        }
+    }
+
+    /**
+     * Makes the following changes in the provider DB.
+     *   1. Removes all entries belonging to a managed profile as managed profiles
+     *      cannot be restored.
+     *   2. Marks all entries as restored. The flags are updated during first load or as
+     *      the restored apps get installed.
+     *   3. If the user serial for primary profile is different than that of the previous device,
+     *      update the entries to the new profile id.
+     */
+    private void sanitizeDB(DatabaseHelper helper, SQLiteDatabase db) throws Exception {
+        long oldProfileId = getDefaultProfileId(db);
+        // Delete all entries which do not belong to the main user
+        int itemsDeleted = db.delete(
+                Favorites.TABLE_NAME, "profileId != ?", new String[]{Long.toString(oldProfileId)});
+        if (itemsDeleted > 0) {
+            FileLog.d(TAG, itemsDeleted + " items belonging to a managed profile, were deleted");
+        }
+
+        // Mark all items as restored.
+        ContentValues values = new ContentValues();
+        values.put(Favorites.RESTORED, 1);
+        db.update(Favorites.TABLE_NAME, values, null, null);
+
+        // Mark widgets with appropriate restore flag
+        values.put(Favorites.RESTORED,
+                LauncherAppWidgetInfo.FLAG_ID_NOT_VALID |
+                        LauncherAppWidgetInfo.FLAG_PROVIDER_NOT_READY |
+                        LauncherAppWidgetInfo.FLAG_UI_NOT_READY);
+        db.update(Favorites.TABLE_NAME, values, "itemType = ?",
+                new String[]{Integer.toString(Favorites.ITEM_TYPE_APPWIDGET)});
+
+        long myProfileId = helper.getDefaultUserSerial();
+        if (Utilities.longCompare(oldProfileId, myProfileId) != 0) {
+            FileLog.d(TAG, "Changing primary user id from " + oldProfileId + " to " + myProfileId);
+            migrateProfileId(db, myProfileId);
+        }
+    }
+
+    /**
+     * Updates profile id of all entries and changes the default value for the column.
+     */
+    protected void migrateProfileId(SQLiteDatabase db, long newProfileId) {
+        // Update existing entries.
+        ContentValues values = new ContentValues();
+        values.put(Favorites.PROFILE_ID, newProfileId);
+        db.update(Favorites.TABLE_NAME, values, null, null);
+
+        // Change default value of the column.
+        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;");
+    }
+
+    /**
+     * Returns the profile id used in the favorites table of the provided db.
+     */
+    protected long getDefaultProfileId(SQLiteDatabase db) throws Exception {
+        try (Cursor c = db.rawQuery("PRAGMA table_info (favorites)", null)){
+            int nameIndex = c.getColumnIndex(INFO_COLUMN_NAME);
+            while (c.moveToNext()) {
+                if (Favorites.PROFILE_ID.equals(c.getString(nameIndex))) {
+                    return c.getLong(c.getColumnIndex(INFO_COLUMN_DEFAULT_VALUE));
+                }
+            }
+            throw new InvalidObjectException("Table does not have a profile id column");
+        }
+    }
+
+    public static boolean isPending(Context context) {
+        return Utilities.getPrefs(context).getBoolean(RESTORE_TASK_PENDING, false);
+    }
+
+    public static void setPending(Context context, boolean isPending) {
+        Utilities.getPrefs(context).edit().putBoolean(RESTORE_TASK_PENDING, isPending).commit();
+    }
+}
diff --git a/tests/src/com/android/launcher3/LauncherBackupAgentTest.java b/tests/src/com/android/launcher3/provider/RestoreDbTaskTest.java
similarity index 88%
rename from tests/src/com/android/launcher3/LauncherBackupAgentTest.java
rename to tests/src/com/android/launcher3/provider/RestoreDbTaskTest.java
index 020a557..29f738b 100644
--- a/tests/src/com/android/launcher3/LauncherBackupAgentTest.java
+++ b/tests/src/com/android/launcher3/provider/RestoreDbTaskTest.java
@@ -1,4 +1,4 @@
-package com.android.launcher3;
+package com.android.launcher3.provider;
 
 import android.content.ContentValues;
 import android.database.Cursor;
@@ -10,14 +10,14 @@
 import com.android.launcher3.LauncherSettings.Favorites;
 
 /**
- * Tests for {@link LauncherBackupAgent}
+ * Tests for {@link RestoreDbTask}
  */
 @MediumTest
-public class LauncherBackupAgentTest extends AndroidTestCase {
+public class RestoreDbTaskTest extends AndroidTestCase {
 
     public void testGetProfileId() throws Exception {
         SQLiteDatabase db = new MyDatabaseHelper(23).getWritableDatabase();
-        assertEquals(23, new LauncherBackupAgent().getDefaultProfileId(db));
+        assertEquals(23, new RestoreDbTask().getDefaultProfileId(db));
     }
 
     public void testMigrateProfileId() throws Exception {
@@ -32,7 +32,7 @@
         // Verify item add
         assertEquals(5, getCount(db, "select * from favorites where profileId = 42"));
 
-        new LauncherBackupAgent().migrateProfileId(db, 33);
+        new RestoreDbTask().migrateProfileId(db, 33);
 
         // verify data migrated
         assertEquals(0, getCount(db, "select * from favorites where profileId = 42"));
diff --git a/tests/src/com/android/launcher3/util/TestLauncherProvider.java b/tests/src/com/android/launcher3/util/TestLauncherProvider.java
index bd3e86c..6ca2121 100644
--- a/tests/src/com/android/launcher3/util/TestLauncherProvider.java
+++ b/tests/src/com/android/launcher3/util/TestLauncherProvider.java
@@ -37,7 +37,7 @@
         }
 
         @Override
-        protected long getDefaultUserSerial() {
+        public long getDefaultUserSerial() {
             return 0;
         }