Updated ROOM

app/src/main/java/androidx/room/SQLiteCopyOpenHelper.java

@@ -28,20 +28,23 @@ import androidx.room.util.DBUtil;
 import androidx.room.util.FileUtil;
 import androidx.sqlite.db.SupportSQLiteDatabase;
 import androidx.sqlite.db.SupportSQLiteOpenHelper;
+import androidx.sqlite.db.framework.FrameworkSQLiteOpenHelperFactory;
 
 import java.io.File;
 import java.io.FileInputStream;
 import java.io.FileOutputStream;
 import java.io.IOException;
+import java.io.InputStream;
 import java.nio.channels.Channels;
 import java.nio.channels.FileChannel;
 import java.nio.channels.ReadableByteChannel;
+import java.util.concurrent.Callable;
 
 /**
  * An open helper that will copy & open a pre-populated database if it doesn't exists in internal
  * storage.
  */
-class SQLiteCopyOpenHelper implements SupportSQLiteOpenHelper {
+class SQLiteCopyOpenHelper implements SupportSQLiteOpenHelper, DelegatingOpenHelper {
 
     @NonNull
     private final Context mContext;
@@ -49,6 +52,8 @@ class SQLiteCopyOpenHelper implements SupportSQLiteOpenHelper {
     private final String mCopyFromAssetPath;
     @Nullable
     private final File mCopyFromFile;
+    @Nullable
+    private final Callable<InputStream> mCopyFromInputStream;
     private final int mDatabaseVersion;
     @NonNull
     private final SupportSQLiteOpenHelper mDelegate;
@@ -61,11 +66,13 @@ class SQLiteCopyOpenHelper implements SupportSQLiteOpenHelper {
             @NonNull Context context,
             @Nullable String copyFromAssetPath,
             @Nullable File copyFromFile,
+            @Nullable Callable<InputStream> copyFromInputStream,
             int databaseVersion,
             @NonNull SupportSQLiteOpenHelper supportSQLiteOpenHelper) {
         mContext = context;
         mCopyFromAssetPath = copyFromAssetPath;
         mCopyFromFile = copyFromFile;
+        mCopyFromInputStream = copyFromInputStream;
         mDatabaseVersion = databaseVersion;
         mDelegate = supportSQLiteOpenHelper;
     }
@@ -84,7 +91,7 @@ class SQLiteCopyOpenHelper implements SupportSQLiteOpenHelper {
     @Override
     public synchronized SupportSQLiteDatabase getWritableDatabase() {
         if (!mVerified) {
-            verifyDatabaseFile();
+            verifyDatabaseFile(true);
             mVerified = true;
         }
         return mDelegate.getWritableDatabase();
@@ -93,7 +100,7 @@ class SQLiteCopyOpenHelper implements SupportSQLiteOpenHelper {
     @Override
     public synchronized SupportSQLiteDatabase getReadableDatabase() {
         if (!mVerified) {
-            verifyDatabaseFile();
+            verifyDatabaseFile(false);
             mVerified = true;
         }
         return mDelegate.getReadableDatabase();
@@ -105,13 +112,19 @@ class SQLiteCopyOpenHelper implements SupportSQLiteOpenHelper {
         mVerified = false;
     }
 
+    @Override
+    @NonNull
+    public SupportSQLiteOpenHelper getDelegate() {
+        return mDelegate;
+    }
+
     // Can't be constructor param because the factory is needed by the database builder which in
     // turn is the one that actually builds the configuration.
     void setDatabaseConfiguration(@Nullable DatabaseConfiguration databaseConfiguration) {
         mDatabaseConfiguration = databaseConfiguration;
     }
 
-    private void verifyDatabaseFile() {
+    private void verifyDatabaseFile(boolean writable) {
         String databaseName = getDatabaseName();
         File databaseFile = mContext.getDatabasePath(databaseName);
         boolean processLevelLock = mDatabaseConfiguration == null
@@ -125,7 +138,7 @@ class SQLiteCopyOpenHelper implements SupportSQLiteOpenHelper {
             if (!databaseFile.exists()) {
                 try {
                     // No database file found, copy and be done.
-                    copyDatabaseFile(databaseFile);
+                    copyDatabaseFile(databaseFile, writable);
                     return;
                 } catch (IOException e) {
                     throw new RuntimeException("Unable to copy database file.", e);
@@ -157,7 +170,7 @@ class SQLiteCopyOpenHelper implements SupportSQLiteOpenHelper {
 
             if (mContext.deleteDatabase(databaseName)) {
                 try {
-                    copyDatabaseFile(databaseFile);
+                    copyDatabaseFile(databaseFile, writable);
                 } catch (IOException e) {
                     // We are more forgiving copying a database on a destructive migration since
                     // there is already a database file that can be opened.
@@ -172,14 +185,23 @@ class SQLiteCopyOpenHelper implements SupportSQLiteOpenHelper {
         }
     }
 
-    private void copyDatabaseFile(File destinationFile) throws IOException {
+    private void copyDatabaseFile(File destinationFile, boolean writable) throws IOException {
         ReadableByteChannel input;
         if (mCopyFromAssetPath != null) {
             input = Channels.newChannel(mContext.getAssets().open(mCopyFromAssetPath));
         } else if (mCopyFromFile != null) {
             input = new FileInputStream(mCopyFromFile).getChannel();
+        } else if (mCopyFromInputStream != null) {
+            final InputStream inputStream;
+            try {
+                inputStream = mCopyFromInputStream.call();
+            } catch (Exception e) {
+                throw new IOException("inputStreamCallable exception on call", e);
+            }
+            input = Channels.newChannel(inputStream);
         } else {
-            throw new IllegalStateException("copyFromAssetPath and copyFromFile == null!");
+            throw new IllegalStateException("copyFromAssetPath, copyFromFile and "
+                    + "copyFromInputStream are all null!");
         }
 
         // An intermediate file is used so that we never end up with a half-copied database file
@@ -196,10 +218,58 @@ class SQLiteCopyOpenHelper implements SupportSQLiteOpenHelper {
                     + destinationFile.getAbsolutePath());
         }
 
+        // Temporarily open intermediate database file using FrameworkSQLiteOpenHelper and dispatch
+        // the open pre-packaged callback. If it fails then intermediate file won't be copied making
+        // invoking pre-packaged callback a transactional operation.
+        dispatchOnOpenPrepackagedDatabase(intermediateFile, writable);
+
         if (!intermediateFile.renameTo(destinationFile)) {
             throw new IOException("Failed to move intermediate file ("
                     + intermediateFile.getAbsolutePath() + ") to destination ("
                     + destinationFile.getAbsolutePath() + ").");
         }
     }
+
+    private void dispatchOnOpenPrepackagedDatabase(File databaseFile, boolean writable) {
+        if (mDatabaseConfiguration == null
+                || mDatabaseConfiguration.prepackagedDatabaseCallback == null) {
+            return;
+        }
+
+        SupportSQLiteOpenHelper helper = createFrameworkOpenHelper(databaseFile);
+        try {
+            SupportSQLiteDatabase db = writable ? helper.getWritableDatabase() :
+                    helper.getReadableDatabase();
+            mDatabaseConfiguration.prepackagedDatabaseCallback.onOpenPrepackagedDatabase(db);
+        } finally {
+            // Close the db and let Room re-open it through a normal path
+            helper.close();
+        }
+    }
+
+    private SupportSQLiteOpenHelper createFrameworkOpenHelper(File databaseFile) {
+        String databaseName = databaseFile.getName();
+        int version;
+        try {
+            version = DBUtil.readVersion(databaseFile);
+        } catch (IOException e) {
+            throw new RuntimeException("Malformed database file, unable to read version.", e);
+        }
+
+        FrameworkSQLiteOpenHelperFactory factory = new FrameworkSQLiteOpenHelperFactory();
+        Configuration configuration = Configuration.builder(mContext)
+                .name(databaseName)
+                .callback(new Callback(version) {
+                    @Override
+                    public void onCreate(@NonNull SupportSQLiteDatabase db) {
+                    }
+
+                    @Override
+                    public void onUpgrade(@NonNull SupportSQLiteDatabase db, int oldVersion,
+                            int newVersion) {
+                    }
+                })
+                .build();
+        return factory.create(configuration);
+    }
 }