Moving all DB management logic from LauncherProvider into a separate class

This would make it easier to move the controller to LauncherModel

Bug: 277345535
Test: Presubmit
Flag: N/A
Change-Id: I4d044cf41361f400968ef65e18de5d3976fcdec7
diff --git a/src/com/android/launcher3/AutoInstallsLayout.java b/src/com/android/launcher3/AutoInstallsLayout.java
index 5367d80..197aa5a 100644
--- a/src/com/android/launcher3/AutoInstallsLayout.java
+++ b/src/com/android/launcher3/AutoInstallsLayout.java
@@ -81,7 +81,7 @@
     private static final String FORMATTED_LAYOUT_RES = "default_layout_%dx%d";
     private static final String LAYOUT_RES = "default_layout";
 
-    static AutoInstallsLayout get(Context context, LauncherWidgetHolder appWidgetHolder,
+    public static AutoInstallsLayout get(Context context, LauncherWidgetHolder appWidgetHolder,
             LayoutParserCallback callback) {
         Partner partner = Partner.get(context.getPackageManager(), ACTION_LAUNCHER_CUSTOMIZATION);
         if (partner == null) {
diff --git a/src/com/android/launcher3/InvariantDeviceProfile.java b/src/com/android/launcher3/InvariantDeviceProfile.java
index 4c34648..5e07a3c 100644
--- a/src/com/android/launcher3/InvariantDeviceProfile.java
+++ b/src/com/android/launcher3/InvariantDeviceProfile.java
@@ -183,7 +183,7 @@
 
     public String dbFile;
     public int defaultLayoutId;
-    int demoModeLayoutId;
+    public int demoModeLayoutId;
     public boolean[] inlineQsb = new boolean[COUNT_SIZES];
 
     /**
diff --git a/src/com/android/launcher3/LauncherProvider.java b/src/com/android/launcher3/LauncherProvider.java
index f4892b2..dee3205 100644
--- a/src/com/android/launcher3/LauncherProvider.java
+++ b/src/com/android/launcher3/LauncherProvider.java
@@ -16,10 +16,6 @@
 
 package com.android.launcher3;
 
-import static com.android.launcher3.DefaultLayoutParser.RES_PARTNER_DEFAULT_LAYOUT;
-import static com.android.launcher3.provider.LauncherDbUtils.copyTable;
-import static com.android.launcher3.provider.LauncherDbUtils.tableExists;
-
 import android.annotation.TargetApi;
 import android.appwidget.AppWidgetManager;
 import android.content.ComponentName;
@@ -28,61 +24,33 @@
 import android.content.ContentProviderResult;
 import android.content.ContentUris;
 import android.content.ContentValues;
-import android.content.Context;
 import android.content.OperationApplicationException;
-import android.content.SharedPreferences;
-import android.content.pm.ProviderInfo;
 import android.database.Cursor;
-import android.database.SQLException;
-import android.database.sqlite.SQLiteDatabase;
 import android.database.sqlite.SQLiteQueryBuilder;
 import android.net.Uri;
 import android.os.Binder;
 import android.os.Build;
 import android.os.Bundle;
 import android.os.Process;
-import android.os.UserManager;
-import android.provider.Settings;
 import android.text.TextUtils;
 import android.util.Log;
-import android.util.Xml;
 
 import com.android.launcher3.LauncherSettings.Favorites;
 import com.android.launcher3.config.FeatureFlags;
-import com.android.launcher3.model.DatabaseHelper;
-import com.android.launcher3.provider.LauncherDbUtils;
+import com.android.launcher3.model.ModelDbController;
 import com.android.launcher3.provider.LauncherDbUtils.SQLiteTransaction;
-import com.android.launcher3.provider.RestoreDbTask;
-import com.android.launcher3.util.IOUtils;
-import com.android.launcher3.util.IntArray;
-import com.android.launcher3.util.Partner;
-import com.android.launcher3.util.Thunk;
 import com.android.launcher3.widget.LauncherWidgetHolder;
 
-import org.xmlpull.v1.XmlPullParser;
-
 import java.io.FileDescriptor;
-import java.io.InputStream;
 import java.io.PrintWriter;
-import java.io.StringReader;
 import java.util.ArrayList;
-import java.util.function.Supplier;
 
 public class LauncherProvider extends ContentProvider {
     private static final String TAG = "LauncherProvider";
 
     public static final String AUTHORITY = BuildConfig.APPLICATION_ID + ".settings";
 
-    private static final int TEST_WORKSPACE_LAYOUT_RES_XML = R.xml.default_test_workspace;
-    private static final int TEST2_WORKSPACE_LAYOUT_RES_XML = R.xml.default_test2_workspace;
-    private static final int TAPL_WORKSPACE_LAYOUT_RES_XML = R.xml.default_tapl_test_workspace;
-
-    public static final String EMPTY_DATABASE_CREATED = "EMPTY_DATABASE_CREATED";
-
-    protected DatabaseHelper mOpenHelper;
-    protected String mProviderAuthority;
-
-    private int mDefaultWorkspaceLayoutOverride = 0;
+    protected ModelDbController mModelDbController;
 
     /**
      * $ adb shell dumpsys activity provider com.android.launcher3
@@ -101,6 +69,7 @@
         if (FeatureFlags.IS_STUDIO_BUILD) {
             Log.d(TAG, "Launcher process started");
         }
+        mModelDbController = new ModelDbController(getContext());
 
         // The content provider exists for the entire duration of the launcher main process and
         // is the first component to get created.
@@ -118,49 +87,17 @@
         }
     }
 
-    /**
-     * Overridden in tests
-     */
-    protected synchronized void createDbIfNotExists() {
-        if (mOpenHelper == null) {
-            mOpenHelper = DatabaseHelper.createDatabaseHelper(
-                    getContext(), false /* forMigration */);
-
-            RestoreDbTask.restoreIfNeeded(getContext(), mOpenHelper);
-        }
-    }
-
-    private synchronized boolean prepForMigration(String dbFile, String targetTableName,
-            Supplier<DatabaseHelper> src, Supplier<DatabaseHelper> dst) {
-        if (TextUtils.equals(dbFile, mOpenHelper.getDatabaseName())) {
-            Log.e(TAG, "prepForMigration - target db is same as current: " + dbFile);
-            return false;
-        }
-
-        final DatabaseHelper helper = src.get();
-        mOpenHelper = dst.get();
-        copyTable(helper.getReadableDatabase(), Favorites.TABLE_NAME,
-                mOpenHelper.getWritableDatabase(), targetTableName, getContext());
-        helper.close();
-        return true;
-    }
-
     @Override
     public Cursor query(Uri uri, String[] projection, String selection,
             String[] selectionArgs, String sortOrder) {
-        createDbIfNotExists();
 
         SqlArguments args = new SqlArguments(uri, selection, selectionArgs);
         SQLiteQueryBuilder qb = new SQLiteQueryBuilder();
         qb.setTables(args.table);
 
-        SQLiteDatabase db = mOpenHelper.getWritableDatabase();
-        Cursor result = qb.query(db, projection, args.where, args.args, null, null, sortOrder);
-        final Bundle extra = new Bundle();
-        extra.putString(LauncherSettings.Settings.EXTRA_DB_NAME, mOpenHelper.getDatabaseName());
-        result.setExtras(extra);
+        Cursor result = mModelDbController.query(
+                args.table, projection, args.where, args.args, sortOrder);
         result.setNotificationUri(getContext().getContentResolver(), uri);
-
         return result;
     }
 
@@ -175,9 +112,6 @@
 
     @Override
     public Uri insert(Uri uri, ContentValues initialValues) {
-        createDbIfNotExists();
-        SqlArguments args = new SqlArguments(uri);
-
         // In very limited cases, we support system|signature permission apps to modify the db.
         if (Binder.getCallingPid() != Process.myPid()) {
             if (!initializeExternalAdd(initialValues)) {
@@ -185,11 +119,9 @@
             }
         }
 
-        SQLiteDatabase db = mOpenHelper.getWritableDatabase();
-        addModifiedTime(initialValues);
-        final int rowId = mOpenHelper.dbInsertAndCheck(db, args.table, initialValues);
+        SqlArguments args = new SqlArguments(uri);
+        int rowId = mModelDbController.insert(args.table, initialValues);
         if (rowId < 0) return null;
-        onAddOrDeleteOp(db);
 
         uri = ContentUris.withAppendedId(uri, rowId);
         reloadLauncherIfExternal();
@@ -198,7 +130,7 @@
 
     private boolean initializeExternalAdd(ContentValues values) {
         // 1. Ensure that externally added items have a valid item id
-        int id = mOpenHelper.generateNewItemId();
+        int id = mModelDbController.generateNewItemId();
         values.put(LauncherSettings.Favorites._ID, id);
 
         // 2. In the case of an app widget, and if no app widget id is specified, we
@@ -213,7 +145,7 @@
                     values.getAsString(Favorites.APPWIDGET_PROVIDER));
 
             if (cn != null) {
-                LauncherWidgetHolder widgetHolder = mOpenHelper.newLauncherWidgetHolder();
+                LauncherWidgetHolder widgetHolder = LauncherWidgetHolder.newInstance(getContext());
                 try {
                     int appWidgetId = widgetHolder.allocateAppWidgetId();
                     values.put(LauncherSettings.Favorites.APPWIDGET_ID, appWidgetId);
@@ -238,22 +170,8 @@
 
     @Override
     public int bulkInsert(Uri uri, ContentValues[] values) {
-        createDbIfNotExists();
         SqlArguments args = new SqlArguments(uri);
-
-        SQLiteDatabase db = mOpenHelper.getWritableDatabase();
-        try (SQLiteTransaction t = new SQLiteTransaction(db)) {
-            int numValues = values.length;
-            for (int i = 0; i < numValues; i++) {
-                addModifiedTime(values[i]);
-                if (mOpenHelper.dbInsertAndCheck(db, args.table, values[i]) < 0) {
-                    return 0;
-                }
-            }
-            onAddOrDeleteOp(db);
-            t.commit();
-        }
-
+        mModelDbController.bulkInsert(args.table, values);
         reloadLauncherIfExternal();
         return values.length;
     }
@@ -262,23 +180,13 @@
     @Override
     public ContentProviderResult[] applyBatch(ArrayList<ContentProviderOperation> operations)
             throws OperationApplicationException {
-        createDbIfNotExists();
-        try (SQLiteTransaction t = new SQLiteTransaction(mOpenHelper.getWritableDatabase())) {
-            boolean isAddOrDelete = false;
-
+        try (SQLiteTransaction t = mModelDbController.newTransaction()) {
             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) {
-                onAddOrDeleteOp(t.getDb());
-            }
-
             t.commit();
             reloadLauncherIfExternal();
             return results;
@@ -287,18 +195,9 @@
 
     @Override
     public int delete(Uri uri, String selection, String[] selectionArgs) {
-        createDbIfNotExists();
         SqlArguments args = new SqlArguments(uri, selection, selectionArgs);
-
-        SQLiteDatabase db = mOpenHelper.getWritableDatabase();
-
-        if (Binder.getCallingPid() != Process.myPid()
-                && Favorites.TABLE_NAME.equalsIgnoreCase(args.table)) {
-            mOpenHelper.removeGhostWidgets(mOpenHelper.getWritableDatabase());
-        }
-        int count = db.delete(args.table, args.where, args.args);
+        int count = mModelDbController.delete(args.table, args.where, args.args);
         if (count > 0) {
-            onAddOrDeleteOp(db);
             reloadLauncherIfExternal();
         }
         return count;
@@ -306,12 +205,8 @@
 
     @Override
     public int update(Uri uri, ContentValues values, String selection, String[] selectionArgs) {
-        createDbIfNotExists();
         SqlArguments args = new SqlArguments(uri, selection, selectionArgs);
-
-        addModifiedTime(values);
-        SQLiteDatabase db = mOpenHelper.getWritableDatabase();
-        int count = db.update(args.table, values, args.where, args.args);
+        int count = mModelDbController.update(args.table, values, args.where, args.args);
         reloadLauncherIfExternal();
         return count;
     }
@@ -321,260 +216,76 @@
         if (Binder.getCallingUid() != Process.myUid()) {
             return null;
         }
-        createDbIfNotExists();
 
         switch (method) {
             case LauncherSettings.Settings.METHOD_CLEAR_EMPTY_DB_FLAG: {
-                clearFlagEmptyDbCreated();
+                mModelDbController.clearEmptyDbFlag();
                 return null;
             }
             case LauncherSettings.Settings.METHOD_DELETE_EMPTY_FOLDERS: {
                 Bundle result = new Bundle();
-                result.putIntArray(LauncherSettings.Settings.EXTRA_VALUE, deleteEmptyFolders()
-                        .toArray());
+                result.putIntArray(LauncherSettings.Settings.EXTRA_VALUE,
+                        mModelDbController.deleteEmptyFolders().toArray());
                 return result;
             }
             case LauncherSettings.Settings.METHOD_NEW_ITEM_ID: {
                 Bundle result = new Bundle();
                 result.putInt(LauncherSettings.Settings.EXTRA_VALUE,
-                        mOpenHelper.generateNewItemId());
+                        mModelDbController.generateNewItemId());
                 return result;
             }
             case LauncherSettings.Settings.METHOD_NEW_SCREEN_ID: {
                 Bundle result = new Bundle();
                 result.putInt(LauncherSettings.Settings.EXTRA_VALUE,
-                        mOpenHelper.getNewScreenId());
+                        mModelDbController.getNewScreenId());
                 return result;
             }
             case LauncherSettings.Settings.METHOD_CREATE_EMPTY_DB: {
-                mOpenHelper.createEmptyDB(mOpenHelper.getWritableDatabase());
+                mModelDbController.createEmptyDB();
                 return null;
             }
             case LauncherSettings.Settings.METHOD_SET_USE_TEST_WORKSPACE_LAYOUT_FLAG: {
-                switch (arg) {
-                    case LauncherSettings.Settings.ARG_DEFAULT_WORKSPACE_LAYOUT_TEST:
-                        mDefaultWorkspaceLayoutOverride = TEST_WORKSPACE_LAYOUT_RES_XML;
-                        break;
-                    case LauncherSettings.Settings.ARG_DEFAULT_WORKSPACE_LAYOUT_TEST2:
-                        mDefaultWorkspaceLayoutOverride = TEST2_WORKSPACE_LAYOUT_RES_XML;
-                        break;
-                    case LauncherSettings.Settings.ARG_DEFAULT_WORKSPACE_LAYOUT_TAPL:
-                        mDefaultWorkspaceLayoutOverride = TAPL_WORKSPACE_LAYOUT_RES_XML;
-                        break;
-                    default:
-                        mDefaultWorkspaceLayoutOverride = 0;
-                        break;
-                }
+                mModelDbController.setUseTestWorkspaceLayout(arg);
                 return null;
             }
             case LauncherSettings.Settings.METHOD_CLEAR_USE_TEST_WORKSPACE_LAYOUT_FLAG: {
-                mDefaultWorkspaceLayoutOverride = 0;
+                mModelDbController.setUseTestWorkspaceLayout(null);
                 return null;
             }
             case LauncherSettings.Settings.METHOD_LOAD_DEFAULT_FAVORITES: {
-                loadDefaultFavoritesIfNecessary();
+                mModelDbController.loadDefaultFavoritesIfNecessary();
                 return null;
             }
             case LauncherSettings.Settings.METHOD_REMOVE_GHOST_WIDGETS: {
-                mOpenHelper.removeGhostWidgets(mOpenHelper.getWritableDatabase());
+                mModelDbController.removeGhostWidgets();
                 return null;
             }
             case LauncherSettings.Settings.METHOD_NEW_TRANSACTION: {
                 Bundle result = new Bundle();
                 result.putBinder(LauncherSettings.Settings.EXTRA_VALUE,
-                        new SQLiteTransaction(mOpenHelper.getWritableDatabase()));
+                        mModelDbController.newTransaction());
                 return result;
             }
             case LauncherSettings.Settings.METHOD_REFRESH_HOTSEAT_RESTORE_TABLE: {
-                mOpenHelper.mHotseatRestoreTableExists = tableExists(
-                        mOpenHelper.getReadableDatabase(), Favorites.HYBRID_HOTSEAT_BACKUP_TABLE);
+                mModelDbController.refreshHotseatRestoreTable();
                 return null;
             }
             case LauncherSettings.Settings.METHOD_UPDATE_CURRENT_OPEN_HELPER: {
                 Bundle result = new Bundle();
                 result.putBoolean(LauncherSettings.Settings.EXTRA_VALUE,
-                        prepForMigration(
-                                arg /* dbFile */,
-                                Favorites.TMP_TABLE,
-                                () -> mOpenHelper,
-                                () -> DatabaseHelper.createDatabaseHelper(
-                                        getContext(), true /* forMigration */)));
+                        mModelDbController.updateCurrentOpenHelper(arg /* dbFile */));
                 return result;
             }
             case LauncherSettings.Settings.METHOD_PREP_FOR_PREVIEW: {
                 Bundle result = new Bundle();
                 result.putBoolean(LauncherSettings.Settings.EXTRA_VALUE,
-                        prepForMigration(
-                                arg /* dbFile */,
-                                Favorites.PREVIEW_TABLE_NAME,
-                                () -> DatabaseHelper.createDatabaseHelper(
-                                        getContext(), arg, true /* forMigration */),
-                                () -> mOpenHelper));
+                        mModelDbController.prepareForPreview(arg /* dbFile */));
                 return result;
             }
         }
         return null;
     }
 
-    private void onAddOrDeleteOp(SQLiteDatabase db) {
-        mOpenHelper.onAddOrDeleteOp(db);
-    }
-
-    /**
-     * Deletes any empty folder from the DB.
-     * @return Ids of deleted folders.
-     */
-    private IntArray deleteEmptyFolders() {
-        SQLiteDatabase db = mOpenHelper.getWritableDatabase();
-        try (SQLiteTransaction t = new SQLiteTransaction(db)) {
-            // Select folders whose id do not match any container value.
-            String selection = LauncherSettings.Favorites.ITEM_TYPE + " = "
-                    + LauncherSettings.Favorites.ITEM_TYPE_FOLDER + " AND "
-                    + LauncherSettings.Favorites._ID +  " NOT IN (SELECT " +
-                            LauncherSettings.Favorites.CONTAINER + " FROM "
-                                + Favorites.TABLE_NAME + ")";
-
-            IntArray folderIds = LauncherDbUtils.queryIntArray(false, db, Favorites.TABLE_NAME,
-                    Favorites._ID, selection, null, null);
-            if (!folderIds.isEmpty()) {
-                db.delete(Favorites.TABLE_NAME, Utilities.createDbSelectionQuery(
-                        LauncherSettings.Favorites._ID, folderIds), null);
-            }
-            t.commit();
-            return folderIds;
-        } catch (SQLException ex) {
-            Log.e(TAG, ex.getMessage(), ex);
-            return new IntArray();
-        }
-    }
-
-    @Thunk static void addModifiedTime(ContentValues values) {
-        values.put(LauncherSettings.Favorites.MODIFIED, System.currentTimeMillis());
-    }
-
-    private void clearFlagEmptyDbCreated() {
-        LauncherPrefs.getPrefs(getContext()).edit()
-                .remove(mOpenHelper.getKey(EMPTY_DATABASE_CREATED)).commit();
-    }
-
-    /**
-     * Loads the default workspace based on the following priority scheme:
-     *   1) From the app restrictions
-     *   2) From a package provided by play store
-     *   3) From a partner configuration APK, already in the system image
-     *   4) The default configuration for the particular device
-     */
-    synchronized private void loadDefaultFavoritesIfNecessary() {
-        SharedPreferences sp = LauncherPrefs.getPrefs(getContext());
-
-        if (sp.getBoolean(mOpenHelper.getKey(EMPTY_DATABASE_CREATED), false)) {
-            Log.d(TAG, "loading default workspace");
-
-            LauncherWidgetHolder widgetHolder = mOpenHelper.newLauncherWidgetHolder();
-            try {
-                AutoInstallsLayout loader = createWorkspaceLoaderFromAppRestriction(widgetHolder);
-                if (loader == null) {
-                    loader = AutoInstallsLayout.get(getContext(), widgetHolder, mOpenHelper);
-                }
-                if (loader == null) {
-                    final Partner partner = Partner.get(getContext().getPackageManager());
-                    if (partner != null) {
-                        int workspaceResId = partner.getXmlResId(RES_PARTNER_DEFAULT_LAYOUT);
-                        if (workspaceResId != 0) {
-                            loader = new DefaultLayoutParser(getContext(), widgetHolder,
-                                    mOpenHelper, partner.getResources(), workspaceResId);
-                        }
-                    }
-                }
-
-                final boolean usingExternallyProvidedLayout = loader != null;
-                if (loader == null) {
-                    loader = getDefaultLayoutParser(widgetHolder);
-                }
-
-                // There might be some partially restored DB items, due to buggy restore logic in
-                // previous versions of launcher.
-                mOpenHelper.createEmptyDB(mOpenHelper.getWritableDatabase());
-                // Populate favorites table with initial favorites
-                if ((mOpenHelper.loadFavorites(mOpenHelper.getWritableDatabase(), loader) <= 0)
-                        && usingExternallyProvidedLayout) {
-                    // Unable to load external layout. Cleanup and load the internal layout.
-                    mOpenHelper.createEmptyDB(mOpenHelper.getWritableDatabase());
-                    mOpenHelper.loadFavorites(mOpenHelper.getWritableDatabase(),
-                            getDefaultLayoutParser(widgetHolder));
-                }
-                clearFlagEmptyDbCreated();
-            } finally {
-                widgetHolder.destroy();
-            }
-        }
-    }
-
-    /**
-     * Creates workspace loader from an XML resource listed in the app restrictions.
-     *
-     * @return the loader if the restrictions are set and the resource exists; null otherwise.
-     */
-    private AutoInstallsLayout createWorkspaceLoaderFromAppRestriction(
-            LauncherWidgetHolder widgetHolder) {
-        Context ctx = getContext();
-        final String authority;
-        if (!TextUtils.isEmpty(mProviderAuthority)) {
-            authority = mProviderAuthority;
-        } else {
-            authority = Settings.Secure.getString(ctx.getContentResolver(),
-                    "launcher3.layout.provider");
-        }
-        if (TextUtils.isEmpty(authority)) {
-            return null;
-        }
-
-        ProviderInfo pi = ctx.getPackageManager().resolveContentProvider(authority, 0);
-        if (pi == null) {
-            Log.e(TAG, "No provider found for authority " + authority);
-            return null;
-        }
-        Uri uri = getLayoutUri(authority, ctx);
-        try (InputStream in = ctx.getContentResolver().openInputStream(uri)) {
-            // Read the full xml so that we fail early in case of any IO error.
-            String layout = new String(IOUtils.toByteArray(in));
-            XmlPullParser parser = Xml.newPullParser();
-            parser.setInput(new StringReader(layout));
-
-            Log.d(TAG, "Loading layout from " + authority);
-            return new AutoInstallsLayout(ctx, widgetHolder, mOpenHelper,
-                    ctx.getPackageManager().getResourcesForApplication(pi.applicationInfo),
-                    () -> parser, AutoInstallsLayout.TAG_WORKSPACE);
-        } catch (Exception e) {
-            Log.e(TAG, "Error getting layout stream from: " + authority , e);
-            return null;
-        }
-    }
-
-    public static Uri getLayoutUri(String authority, Context ctx) {
-        InvariantDeviceProfile grid = LauncherAppState.getIDP(ctx);
-        return new Uri.Builder().scheme("content").authority(authority).path("launcher_layout")
-                .appendQueryParameter("version", "1")
-                .appendQueryParameter("gridWidth", Integer.toString(grid.numColumns))
-                .appendQueryParameter("gridHeight", Integer.toString(grid.numRows))
-                .appendQueryParameter("hotseatSize", Integer.toString(grid.numDatabaseHotseatIcons))
-                .build();
-    }
-
-    private DefaultLayoutParser getDefaultLayoutParser(LauncherWidgetHolder widgetHolder) {
-        InvariantDeviceProfile idp = LauncherAppState.getIDP(getContext());
-        int defaultLayout = mDefaultWorkspaceLayoutOverride > 0
-                ? mDefaultWorkspaceLayoutOverride : idp.defaultLayoutId;
-
-        if (getContext().getSystemService(UserManager.class).isDemoUser()
-                && idp.demoModeLayoutId != 0) {
-            defaultLayout = idp.demoModeLayoutId;
-        }
-
-        return new DefaultLayoutParser(getContext(), widgetHolder,
-                mOpenHelper, getContext().getResources(), defaultLayout);
-    }
-
     static class SqlArguments {
         public final String table;
         public final String where;
diff --git a/src/com/android/launcher3/model/DatabaseHelper.java b/src/com/android/launcher3/model/DatabaseHelper.java
index 1840b75..dc5fcf7 100644
--- a/src/com/android/launcher3/model/DatabaseHelper.java
+++ b/src/com/android/launcher3/model/DatabaseHelper.java
@@ -39,7 +39,6 @@
 import com.android.launcher3.InvariantDeviceProfile;
 import com.android.launcher3.LauncherFiles;
 import com.android.launcher3.LauncherPrefs;
-import com.android.launcher3.LauncherProvider;
 import com.android.launcher3.LauncherSettings;
 import com.android.launcher3.LauncherSettings.Favorites;
 import com.android.launcher3.Utilities;
@@ -77,6 +76,7 @@
     private static final boolean LOGD = false;
 
     private static final String DOWNGRADE_SCHEMA_FILE = "downgrade_schema.json";
+    public static final String EMPTY_DATABASE_CREATED = "EMPTY_DATABASE_CREATED";
 
     private final Context mContext;
     private final boolean mForMigration;
@@ -165,8 +165,7 @@
      */
     protected void onEmptyDbCreated() {
         // Set the flag for empty DB
-        LauncherPrefs.getPrefs(mContext).edit().putBoolean(getKey(
-                        LauncherProvider.EMPTY_DATABASE_CREATED), true)
+        LauncherPrefs.getPrefs(mContext).edit().putBoolean(getKey(EMPTY_DATABASE_CREATED), true)
                 .commit();
     }
 
diff --git a/src/com/android/launcher3/model/ModelDbController.java b/src/com/android/launcher3/model/ModelDbController.java
new file mode 100644
index 0000000..7452bcd
--- /dev/null
+++ b/src/com/android/launcher3/model/ModelDbController.java
@@ -0,0 +1,462 @@
+/*
+ * Copyright (C) 2023 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.DefaultLayoutParser.RES_PARTNER_DEFAULT_LAYOUT;
+import static com.android.launcher3.model.DatabaseHelper.EMPTY_DATABASE_CREATED;
+import static com.android.launcher3.provider.LauncherDbUtils.copyTable;
+import static com.android.launcher3.provider.LauncherDbUtils.tableExists;
+
+import android.content.ContentValues;
+import android.content.Context;
+import android.content.SharedPreferences;
+import android.content.pm.ProviderInfo;
+import android.database.Cursor;
+import android.database.SQLException;
+import android.database.sqlite.SQLiteDatabase;
+import android.net.Uri;
+import android.os.Binder;
+import android.os.Bundle;
+import android.os.Process;
+import android.os.UserManager;
+import android.provider.Settings;
+import android.text.TextUtils;
+import android.util.Log;
+import android.util.Xml;
+
+import androidx.annotation.Nullable;
+
+import com.android.launcher3.AutoInstallsLayout;
+import com.android.launcher3.DefaultLayoutParser;
+import com.android.launcher3.InvariantDeviceProfile;
+import com.android.launcher3.LauncherAppState;
+import com.android.launcher3.LauncherPrefs;
+import com.android.launcher3.LauncherSettings;
+import com.android.launcher3.LauncherSettings.Favorites;
+import com.android.launcher3.R;
+import com.android.launcher3.Utilities;
+import com.android.launcher3.provider.LauncherDbUtils;
+import com.android.launcher3.provider.LauncherDbUtils.SQLiteTransaction;
+import com.android.launcher3.provider.RestoreDbTask;
+import com.android.launcher3.util.IOUtils;
+import com.android.launcher3.util.IntArray;
+import com.android.launcher3.util.Partner;
+import com.android.launcher3.widget.LauncherWidgetHolder;
+
+import org.xmlpull.v1.XmlPullParser;
+
+import java.io.InputStream;
+import java.io.StringReader;
+import java.util.function.Supplier;
+
+/**
+ * Utility class which maintains an instance of Launcher database and provides utility methods
+ * around it.
+ */
+public class ModelDbController {
+    private static final String TAG = "LauncherProvider";
+
+    private static final int TEST_WORKSPACE_LAYOUT_RES_XML = R.xml.default_test_workspace;
+    private static final int TEST2_WORKSPACE_LAYOUT_RES_XML = R.xml.default_test2_workspace;
+    private static final int TAPL_WORKSPACE_LAYOUT_RES_XML = R.xml.default_tapl_test_workspace;
+
+    protected DatabaseHelper mOpenHelper;
+    protected String mProviderAuthority;
+
+    private int mDefaultWorkspaceLayoutOverride = 0;
+
+    private final Context mContext;
+
+    public ModelDbController(Context context) {
+        mContext = context;
+    }
+
+    private synchronized void createDbIfNotExists() {
+        if (mOpenHelper == null) {
+            mOpenHelper = DatabaseHelper.createDatabaseHelper(
+                    mContext, false /* forMigration */);
+
+            RestoreDbTask.restoreIfNeeded(mContext, mOpenHelper);
+        }
+    }
+
+    private synchronized boolean prepForMigration(String dbFile, String targetTableName,
+            Supplier<DatabaseHelper> src, Supplier<DatabaseHelper> dst) {
+        if (TextUtils.equals(dbFile, mOpenHelper.getDatabaseName())) {
+            Log.e(TAG, "prepForMigration - target db is same as current: " + dbFile);
+            return false;
+        }
+
+        final DatabaseHelper helper = src.get();
+        mOpenHelper = dst.get();
+        copyTable(helper.getReadableDatabase(), Favorites.TABLE_NAME,
+                mOpenHelper.getWritableDatabase(), targetTableName, mContext);
+        helper.close();
+        return true;
+    }
+
+    /**
+     * Refer {@link SQLiteDatabase#query}
+     */
+    public Cursor query(String table, String[] projection, String selection,
+            String[] selectionArgs, String sortOrder) {
+        createDbIfNotExists();
+        SQLiteDatabase db = mOpenHelper.getWritableDatabase();
+        Cursor result = db.query(
+                table, projection, selection, selectionArgs, null, null, sortOrder);
+
+        final Bundle extra = new Bundle();
+        extra.putString(LauncherSettings.Settings.EXTRA_DB_NAME, mOpenHelper.getDatabaseName());
+        result.setExtras(extra);
+        return result;
+    }
+
+    /**
+     * Refer {@link SQLiteDatabase#insert(String, String, ContentValues)}
+     */
+    public int insert(String table, ContentValues initialValues) {
+        createDbIfNotExists();
+
+        SQLiteDatabase db = mOpenHelper.getWritableDatabase();
+        addModifiedTime(initialValues);
+        int rowId = mOpenHelper.dbInsertAndCheck(db, table, initialValues);
+        if (rowId >= 0) {
+            onAddOrDeleteOp(db);
+        }
+        return rowId;
+    }
+
+    /**
+     * Similar to insert but for adding multiple values in a transaction.
+     */
+    public int bulkInsert(String table, ContentValues[] values) {
+        createDbIfNotExists();
+
+        SQLiteDatabase db = mOpenHelper.getWritableDatabase();
+        try (SQLiteTransaction t = new SQLiteTransaction(db)) {
+            int numValues = values.length;
+            for (int i = 0; i < numValues; i++) {
+                addModifiedTime(values[i]);
+                if (mOpenHelper.dbInsertAndCheck(db, table, values[i]) < 0) {
+                    return 0;
+                }
+            }
+            onAddOrDeleteOp(db);
+            t.commit();
+        }
+        return values.length;
+    }
+
+    /**
+     * Refer {@link SQLiteDatabase#delete(String, String, String[])}
+     */
+    public int delete(String table, String selection, String[] selectionArgs) {
+        createDbIfNotExists();
+        SQLiteDatabase db = mOpenHelper.getWritableDatabase();
+
+        if (Binder.getCallingPid() != Process.myPid()
+                && Favorites.TABLE_NAME.equalsIgnoreCase(table)) {
+            mOpenHelper.removeGhostWidgets(mOpenHelper.getWritableDatabase());
+        }
+        int count = db.delete(table, selection, selectionArgs);
+        if (count > 0) {
+            onAddOrDeleteOp(db);
+        }
+        return count;
+    }
+
+    /**
+     * Refer {@link SQLiteDatabase#update(String, ContentValues, String, String[])}
+     */
+    public int update(String table, ContentValues values,
+            String selection, String[] selectionArgs) {
+        createDbIfNotExists();
+
+        addModifiedTime(values);
+        SQLiteDatabase db = mOpenHelper.getWritableDatabase();
+        int count = db.update(table, values, selection, selectionArgs);
+        return count;
+    }
+
+    /**
+     * Clears a previously set flag corresponding to empty db creation
+     */
+    public void clearEmptyDbFlag() {
+        createDbIfNotExists();
+        clearFlagEmptyDbCreated();
+    }
+
+    /**
+     * Generates an id to be used for new item in the favorites table
+     */
+    public int generateNewItemId() {
+        createDbIfNotExists();
+        return mOpenHelper.generateNewItemId();
+    }
+
+    /**
+     * Generates an id to be used for new workspace screen
+     */
+    public int getNewScreenId() {
+        createDbIfNotExists();
+        return mOpenHelper.getNewScreenId();
+    }
+
+    /**
+     * Creates an empty DB clearing all existing data
+     */
+    public void createEmptyDB() {
+        createDbIfNotExists();
+        mOpenHelper.createEmptyDB(mOpenHelper.getWritableDatabase());
+    }
+
+    /**
+     * Overrides the default xml to be used for setting up workspace
+     */
+    public void setUseTestWorkspaceLayout(@Nullable String layout) {
+        if (LauncherSettings.Settings.ARG_DEFAULT_WORKSPACE_LAYOUT_TEST.equals(layout)) {
+            mDefaultWorkspaceLayoutOverride = TEST_WORKSPACE_LAYOUT_RES_XML;
+        } else if (LauncherSettings.Settings.ARG_DEFAULT_WORKSPACE_LAYOUT_TEST2.equals(layout)) {
+            mDefaultWorkspaceLayoutOverride = TEST2_WORKSPACE_LAYOUT_RES_XML;
+        } else if (LauncherSettings.Settings.ARG_DEFAULT_WORKSPACE_LAYOUT_TAPL.equals(layout)) {
+            mDefaultWorkspaceLayoutOverride = TAPL_WORKSPACE_LAYOUT_RES_XML;
+        } else {
+            mDefaultWorkspaceLayoutOverride = 0;
+        }
+    }
+
+    /**
+     * Removes any widget which are present in the framework, but not in out internal DB
+     */
+    public void removeGhostWidgets() {
+        createDbIfNotExists();
+        mOpenHelper.removeGhostWidgets(mOpenHelper.getWritableDatabase());
+    }
+
+    /**
+     * Returns a new {@link SQLiteTransaction}
+     */
+    public SQLiteTransaction newTransaction() {
+        createDbIfNotExists();
+        return new SQLiteTransaction(mOpenHelper.getWritableDatabase());
+    }
+
+    /**
+     * Refreshes the internal state corresponding to presence of hotseat table
+     */
+    public void refreshHotseatRestoreTable() {
+        createDbIfNotExists();
+        mOpenHelper.mHotseatRestoreTableExists = tableExists(
+                mOpenHelper.getReadableDatabase(), Favorites.HYBRID_HOTSEAT_BACKUP_TABLE);
+    }
+
+    /**
+     * Updates the current DB and copies all the existing data to the temp table
+     * @param dbFile name of the target db file name
+     */
+    public boolean updateCurrentOpenHelper(String dbFile) {
+        createDbIfNotExists();
+        return prepForMigration(
+                dbFile,
+                Favorites.TMP_TABLE,
+                () -> mOpenHelper,
+                () -> DatabaseHelper.createDatabaseHelper(
+                        mContext, true /* forMigration */));
+    }
+
+    /**
+     * Returns the current DatabaseHelper.
+     * Only for tests
+     */
+    public DatabaseHelper getDatabaseHelper() {
+        createDbIfNotExists();
+        return mOpenHelper;
+    }
+
+    /**
+     * Prepares the DB for preview by copying all existing data to preview table
+     */
+    public boolean prepareForPreview(String dbFile) {
+        createDbIfNotExists();
+        return prepForMigration(
+                dbFile,
+                Favorites.PREVIEW_TABLE_NAME,
+                () -> DatabaseHelper.createDatabaseHelper(
+                        mContext, dbFile, true /* forMigration */),
+                () -> mOpenHelper);
+    }
+
+    private void onAddOrDeleteOp(SQLiteDatabase db) {
+        mOpenHelper.onAddOrDeleteOp(db);
+    }
+
+    /**
+     * Deletes any empty folder from the DB.
+     * @return Ids of deleted folders.
+     */
+    public IntArray deleteEmptyFolders() {
+        createDbIfNotExists();
+
+        SQLiteDatabase db = mOpenHelper.getWritableDatabase();
+        try (SQLiteTransaction t = new SQLiteTransaction(db)) {
+            // Select folders whose id do not match any container value.
+            String selection = LauncherSettings.Favorites.ITEM_TYPE + " = "
+                    + LauncherSettings.Favorites.ITEM_TYPE_FOLDER + " AND "
+                    + LauncherSettings.Favorites._ID +  " NOT IN (SELECT "
+                    + LauncherSettings.Favorites.CONTAINER + " FROM "
+                    + Favorites.TABLE_NAME + ")";
+
+            IntArray folderIds = LauncherDbUtils.queryIntArray(false, db, Favorites.TABLE_NAME,
+                    Favorites._ID, selection, null, null);
+            if (!folderIds.isEmpty()) {
+                db.delete(Favorites.TABLE_NAME, Utilities.createDbSelectionQuery(
+                        LauncherSettings.Favorites._ID, folderIds), null);
+            }
+            t.commit();
+            return folderIds;
+        } catch (SQLException ex) {
+            Log.e(TAG, ex.getMessage(), ex);
+            return new IntArray();
+        }
+    }
+
+    private static void addModifiedTime(ContentValues values) {
+        values.put(LauncherSettings.Favorites.MODIFIED, System.currentTimeMillis());
+    }
+
+    private void clearFlagEmptyDbCreated() {
+        LauncherPrefs.getPrefs(mContext).edit()
+                .remove(mOpenHelper.getKey(EMPTY_DATABASE_CREATED)).commit();
+    }
+
+    /**
+     * Loads the default workspace based on the following priority scheme:
+     *   1) From the app restrictions
+     *   2) From a package provided by play store
+     *   3) From a partner configuration APK, already in the system image
+     *   4) The default configuration for the particular device
+     */
+    public synchronized void loadDefaultFavoritesIfNecessary() {
+        createDbIfNotExists();
+        SharedPreferences sp = LauncherPrefs.getPrefs(mContext);
+
+        if (sp.getBoolean(mOpenHelper.getKey(EMPTY_DATABASE_CREATED), false)) {
+            Log.d(TAG, "loading default workspace");
+
+            LauncherWidgetHolder widgetHolder = mOpenHelper.newLauncherWidgetHolder();
+            try {
+                AutoInstallsLayout loader = createWorkspaceLoaderFromAppRestriction(widgetHolder);
+                if (loader == null) {
+                    loader = AutoInstallsLayout.get(mContext, widgetHolder, mOpenHelper);
+                }
+                if (loader == null) {
+                    final Partner partner = Partner.get(mContext.getPackageManager());
+                    if (partner != null) {
+                        int workspaceResId = partner.getXmlResId(RES_PARTNER_DEFAULT_LAYOUT);
+                        if (workspaceResId != 0) {
+                            loader = new DefaultLayoutParser(mContext, widgetHolder,
+                                    mOpenHelper, partner.getResources(), workspaceResId);
+                        }
+                    }
+                }
+
+                final boolean usingExternallyProvidedLayout = loader != null;
+                if (loader == null) {
+                    loader = getDefaultLayoutParser(widgetHolder);
+                }
+
+                // There might be some partially restored DB items, due to buggy restore logic in
+                // previous versions of launcher.
+                mOpenHelper.createEmptyDB(mOpenHelper.getWritableDatabase());
+                // Populate favorites table with initial favorites
+                if ((mOpenHelper.loadFavorites(mOpenHelper.getWritableDatabase(), loader) <= 0)
+                        && usingExternallyProvidedLayout) {
+                    // Unable to load external layout. Cleanup and load the internal layout.
+                    mOpenHelper.createEmptyDB(mOpenHelper.getWritableDatabase());
+                    mOpenHelper.loadFavorites(mOpenHelper.getWritableDatabase(),
+                            getDefaultLayoutParser(widgetHolder));
+                }
+                clearFlagEmptyDbCreated();
+            } finally {
+                widgetHolder.destroy();
+            }
+        }
+    }
+
+    /**
+     * Creates workspace loader from an XML resource listed in the app restrictions.
+     *
+     * @return the loader if the restrictions are set and the resource exists; null otherwise.
+     */
+    private AutoInstallsLayout createWorkspaceLoaderFromAppRestriction(
+            LauncherWidgetHolder widgetHolder) {
+        final String authority;
+        if (!TextUtils.isEmpty(mProviderAuthority)) {
+            authority = mProviderAuthority;
+        } else {
+            authority = Settings.Secure.getString(mContext.getContentResolver(),
+                    "launcher3.layout.provider");
+        }
+        if (TextUtils.isEmpty(authority)) {
+            return null;
+        }
+
+        ProviderInfo pi = mContext.getPackageManager().resolveContentProvider(authority, 0);
+        if (pi == null) {
+            Log.e(TAG, "No provider found for authority " + authority);
+            return null;
+        }
+        Uri uri = getLayoutUri(authority, mContext);
+        try (InputStream in = mContext.getContentResolver().openInputStream(uri)) {
+            // Read the full xml so that we fail early in case of any IO error.
+            String layout = new String(IOUtils.toByteArray(in));
+            XmlPullParser parser = Xml.newPullParser();
+            parser.setInput(new StringReader(layout));
+
+            Log.d(TAG, "Loading layout from " + authority);
+            return new AutoInstallsLayout(mContext, widgetHolder, mOpenHelper,
+                    mContext.getPackageManager().getResourcesForApplication(pi.applicationInfo),
+                    () -> parser, AutoInstallsLayout.TAG_WORKSPACE);
+        } catch (Exception e) {
+            Log.e(TAG, "Error getting layout stream from: " + authority , e);
+            return null;
+        }
+    }
+
+    private static Uri getLayoutUri(String authority, Context ctx) {
+        InvariantDeviceProfile grid = LauncherAppState.getIDP(ctx);
+        return new Uri.Builder().scheme("content").authority(authority).path("launcher_layout")
+                .appendQueryParameter("version", "1")
+                .appendQueryParameter("gridWidth", Integer.toString(grid.numColumns))
+                .appendQueryParameter("gridHeight", Integer.toString(grid.numRows))
+                .appendQueryParameter("hotseatSize", Integer.toString(grid.numDatabaseHotseatIcons))
+                .build();
+    }
+
+    private DefaultLayoutParser getDefaultLayoutParser(LauncherWidgetHolder widgetHolder) {
+        InvariantDeviceProfile idp = LauncherAppState.getIDP(mContext);
+        int defaultLayout = mDefaultWorkspaceLayoutOverride > 0
+                ? mDefaultWorkspaceLayoutOverride : idp.defaultLayoutId;
+
+        if (mContext.getSystemService(UserManager.class).isDemoUser()
+                && idp.demoModeLayoutId != 0) {
+            defaultLayout = idp.demoModeLayoutId;
+        }
+
+        return new DefaultLayoutParser(mContext, widgetHolder,
+                mOpenHelper, mContext.getResources(), defaultLayout);
+    }
+}
diff --git a/tests/src/com/android/launcher3/util/LauncherModelHelper.java b/tests/src/com/android/launcher3/util/LauncherModelHelper.java
index fdfeb7d..5548364 100644
--- a/tests/src/com/android/launcher3/util/LauncherModelHelper.java
+++ b/tests/src/com/android/launcher3/util/LauncherModelHelper.java
@@ -63,6 +63,7 @@
 import com.android.launcher3.model.BgDataModel.Callbacks;
 import com.android.launcher3.model.DatabaseHelper;
 import com.android.launcher3.model.ItemInstallQueue;
+import com.android.launcher3.model.ModelDbController;
 import com.android.launcher3.model.data.AppInfo;
 import com.android.launcher3.model.data.ItemInfo;
 import com.android.launcher3.pm.InstallSessionHelper;
@@ -471,16 +472,16 @@
 
         @Override
         public boolean onCreate() {
+            mModelDbController = new ModelDbController(getContext());
             return true;
         }
 
         public SQLiteDatabase getDb() {
-            createDbIfNotExists();
-            return mOpenHelper.getWritableDatabase();
+            return mModelDbController.getDatabaseHelper().getWritableDatabase();
         }
 
         public DatabaseHelper getHelper() {
-            return mOpenHelper;
+            return mModelDbController.getDatabaseHelper();
         }
     }