Add RoomDatabase.Callback

Users can now add Callbacks to RoomDatabase.

Bug: 62699324
Test: DatabaseCallbackTest
Change-Id: I0fa38ba97614e1ad721594d238960b8183e96769
diff --git a/room/compiler/src/main/kotlin/android/arch/persistence/room/writer/SQLiteOpenHelperWriter.kt b/room/compiler/src/main/kotlin/android/arch/persistence/room/writer/SQLiteOpenHelperWriter.kt
index 50dcb32..afb12e9 100644
--- a/room/compiler/src/main/kotlin/android/arch/persistence/room/writer/SQLiteOpenHelperWriter.kt
+++ b/room/compiler/src/main/kotlin/android/arch/persistence/room/writer/SQLiteOpenHelperWriter.kt
@@ -16,7 +16,6 @@
 
 package android.arch.persistence.room.writer
 
-import android.support.annotation.VisibleForTesting
 import android.arch.persistence.room.ext.L
 import android.arch.persistence.room.ext.N
 import android.arch.persistence.room.ext.RoomTypeNames
@@ -26,6 +25,7 @@
 import android.arch.persistence.room.solver.CodeGenScope
 import android.arch.persistence.room.vo.Database
 import android.arch.persistence.room.vo.Entity
+import android.support.annotation.VisibleForTesting
 import com.squareup.javapoet.MethodSpec
 import com.squareup.javapoet.ParameterSpec
 import com.squareup.javapoet.TypeSpec
@@ -67,7 +67,8 @@
             superclass(RoomTypeNames.OPEN_HELPER_DELEGATE)
             addMethod(createCreateAllTables())
             addMethod(createDropAllTables())
-            addMethod(createOnOpen())
+            addMethod(createOnCreate(scope.fork()))
+            addMethod(createOnOpen(scope.fork()))
             addMethod(createValidateMigration(scope.fork()))
         }.build()
     }
@@ -85,7 +86,15 @@
         }.build()
     }
 
-    private fun createOnOpen(): MethodSpec {
+    private fun createOnCreate(scope: CodeGenScope): MethodSpec {
+        return MethodSpec.methodBuilder("onCreate").apply {
+            addModifiers(PROTECTED)
+            addParameter(SupportDbTypeNames.DB, "_db")
+            invokeCallbacks(scope, "onCreate")
+        }.build()
+    }
+
+    private fun createOnOpen(scope: CodeGenScope): MethodSpec {
         return MethodSpec.methodBuilder("onOpen").apply {
             addModifiers(PUBLIC)
             addParameter(SupportDbTypeNames.DB, "_db")
@@ -94,6 +103,7 @@
                 addStatement("_db.execSQL($S)", "PRAGMA foreign_keys = ON")
             }
             addStatement("internalInitInvalidationTracker(_db)")
+            invokeCallbacks(scope, "onOpen")
         }.build()
     }
 
@@ -117,6 +127,19 @@
         }.build()
     }
 
+    private fun MethodSpec.Builder.invokeCallbacks(scope: CodeGenScope, methodName: String) {
+        val iVar = scope.getTmpVar("_i")
+        val sizeVar = scope.getTmpVar("_size")
+        beginControlFlow("if (mCallbacks != null)").apply {
+            beginControlFlow("for (int $N = 0, $N = mCallbacks.size(); $N < $N; $N++)",
+                    iVar, sizeVar, iVar, sizeVar, iVar).apply {
+                addStatement("mCallbacks.get($N).$N(_db)", iVar, methodName)
+            }
+            endControlFlow()
+        }
+        endControlFlow()
+    }
+
     @VisibleForTesting
     fun createQuery(entity : Entity) : String {
         return entity.createTableQuery
diff --git a/room/compiler/src/test/data/databasewriter/output/ComplexDatabase.java b/room/compiler/src/test/data/databasewriter/output/ComplexDatabase.java
index b7c4657..8ed676b 100644
--- a/room/compiler/src/test/data/databasewriter/output/ComplexDatabase.java
+++ b/room/compiler/src/test/data/databasewriter/output/ComplexDatabase.java
@@ -32,9 +32,22 @@
                 _db.execSQL("DROP TABLE IF EXISTS `User`");
             }
 
+            protected void onCreate(SupportSQLiteDatabase _db) {
+                if (mCallbacks != null) {
+                    for (int _i = 0, _size = mCallbacks.size(); _i < _size; _i++) {
+                        mCallbacks.get(_i).onCreate(_db);
+                    }
+                }
+            }
+
             public void onOpen(SupportSQLiteDatabase _db) {
                 mDatabase = _db;
                 internalInitInvalidationTracker(_db);
+                if (mCallbacks != null) {
+                    for (int _i = 0, _size = mCallbacks.size(); _i < _size; _i++) {
+                        mCallbacks.get(_i).onOpen(_db);
+                    }
+                }
             }
 
             protected void validateMigration(SupportSQLiteDatabase _db) {
diff --git a/room/integration-tests/testapp/src/androidTest/java/android/arch/persistence/room/integration/testapp/test/DatabaseCallbackTest.java b/room/integration-tests/testapp/src/androidTest/java/android/arch/persistence/room/integration/testapp/test/DatabaseCallbackTest.java
new file mode 100644
index 0000000..579b3e4
--- /dev/null
+++ b/room/integration-tests/testapp/src/androidTest/java/android/arch/persistence/room/integration/testapp/test/DatabaseCallbackTest.java
@@ -0,0 +1,119 @@
+/*
+ * Copyright (C) 2017 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 android.arch.persistence.room.integration.testapp.test;
+
+import static org.hamcrest.CoreMatchers.is;
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.hamcrest.Matchers.empty;
+import static org.hamcrest.core.IsCollectionContaining.hasItem;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+
+import android.arch.persistence.db.SupportSQLiteDatabase;
+import android.arch.persistence.room.Room;
+import android.arch.persistence.room.RoomDatabase;
+import android.arch.persistence.room.integration.testapp.TestDatabase;
+import android.arch.persistence.room.integration.testapp.vo.User;
+import android.content.Context;
+import android.database.Cursor;
+import android.support.annotation.NonNull;
+import android.support.test.InstrumentationRegistry;
+import android.support.test.filters.MediumTest;
+import android.support.test.filters.SmallTest;
+import android.support.test.runner.AndroidJUnit4;
+
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+import java.util.ArrayList;
+import java.util.List;
+
+@RunWith(AndroidJUnit4.class)
+public class DatabaseCallbackTest {
+
+    @Test
+    @MediumTest
+    public void createAndOpen() {
+        Context context = InstrumentationRegistry.getTargetContext();
+        TestDatabaseCallback callback1 = new TestDatabaseCallback();
+        TestDatabase db1 = Room.databaseBuilder(context, TestDatabase.class, "test")
+                .addCallback(callback1)
+                .build();
+        assertFalse(callback1.mCreated);
+        assertFalse(callback1.mOpened);
+        User user1 = TestUtil.createUser(3);
+        user1.setName("george");
+        db1.getUserDao().insert(user1);
+        assertTrue(callback1.mCreated);
+        assertTrue(callback1.mOpened);
+        TestDatabaseCallback callback2 = new TestDatabaseCallback();
+        TestDatabase db2 = Room.databaseBuilder(context, TestDatabase.class, "test")
+                .addCallback(callback2)
+                .build();
+        assertFalse(callback2.mCreated);
+        assertFalse(callback2.mOpened);
+        User user2 = db2.getUserDao().load(3);
+        assertThat(user2.getName(), is("george"));
+        assertFalse(callback2.mCreated); // Not called; already created by db1
+        assertTrue(callback2.mOpened);
+    }
+
+    @Test
+    @SmallTest
+    public void writeOnCreate() {
+        Context context = InstrumentationRegistry.getTargetContext();
+        TestDatabase db = Room.inMemoryDatabaseBuilder(context, TestDatabase.class)
+                .addCallback(new RoomDatabase.Callback() {
+                    @Override
+                    public void onCreate(@NonNull SupportSQLiteDatabase db) {
+                        Cursor cursor = null;
+                        try {
+                            cursor = db.query(
+                                    "SELECT name FROM sqlite_master WHERE type = 'table'");
+                            ArrayList<String> names = new ArrayList<>();
+                            while (cursor.moveToNext()) {
+                                names.add(cursor.getString(0));
+                            }
+                            assertThat(names, hasItem("User"));
+                        } finally {
+                            if (cursor != null) {
+                                cursor.close();
+                            }
+                        }
+                    }
+                })
+                .build();
+        List<Integer> ids = db.getUserDao().loadIds();
+        assertThat(ids, is(empty()));
+    }
+
+    public static class TestDatabaseCallback extends RoomDatabase.Callback {
+
+        boolean mCreated;
+        boolean mOpened;
+
+        @Override
+        public void onCreate(@NonNull SupportSQLiteDatabase db) {
+            mCreated = true;
+        }
+
+        @Override
+        public void onOpen(@NonNull SupportSQLiteDatabase db) {
+            mOpened = true;
+        }
+    }
+}
diff --git a/room/runtime/src/main/java/android/arch/persistence/room/DatabaseConfiguration.java b/room/runtime/src/main/java/android/arch/persistence/room/DatabaseConfiguration.java
index e5ac1c0..e3d0c63 100644
--- a/room/runtime/src/main/java/android/arch/persistence/room/DatabaseConfiguration.java
+++ b/room/runtime/src/main/java/android/arch/persistence/room/DatabaseConfiguration.java
@@ -22,6 +22,8 @@
 import android.support.annotation.Nullable;
 import android.support.annotation.RestrictTo;
 
+import java.util.List;
+
 /**
  * Configuration class for a {@link RoomDatabase}.
  */
@@ -49,6 +51,9 @@
     @NonNull
     public final RoomDatabase.MigrationContainer migrationContainer;
 
+    @Nullable
+    public final List<RoomDatabase.Callback> callbacks;
+
     /**
      * Whether Room should throw an exception for queries run on the main thread.
      */
@@ -61,6 +66,7 @@
      * @param name Name of the database, can be null if it is in memory.
      * @param sqliteOpenHelperFactory The open helper factory to use.
      * @param migrationContainer The migration container for migrations.
+     * @param callbacks The list of callbacks for database events.
      * @param allowMainThreadQueries Whether to allow main thread reads/writes or not.
      *
      * @hide
@@ -69,11 +75,13 @@
     public DatabaseConfiguration(@NonNull Context context, @Nullable String name,
             @NonNull SupportSQLiteOpenHelper.Factory sqliteOpenHelperFactory,
             @NonNull RoomDatabase.MigrationContainer migrationContainer,
+            @Nullable List<RoomDatabase.Callback> callbacks,
             boolean allowMainThreadQueries) {
         this.sqliteOpenHelperFactory = sqliteOpenHelperFactory;
         this.context = context;
         this.name = name;
         this.migrationContainer = migrationContainer;
+        this.callbacks = callbacks;
         this.allowMainThreadQueries = allowMainThreadQueries;
     }
 }
diff --git a/room/runtime/src/main/java/android/arch/persistence/room/RoomDatabase.java b/room/runtime/src/main/java/android/arch/persistence/room/RoomDatabase.java
index 7846999..cc265c0 100644
--- a/room/runtime/src/main/java/android/arch/persistence/room/RoomDatabase.java
+++ b/room/runtime/src/main/java/android/arch/persistence/room/RoomDatabase.java
@@ -56,6 +56,9 @@
     private final InvalidationTracker mInvalidationTracker;
     private boolean mAllowMainThreadQueries;
 
+    @Nullable
+    protected List<Callback> mCallbacks;
+
     /**
      * Creates a RoomDatabase.
      * <p>
@@ -75,6 +78,7 @@
     @CallSuper
     public void init(DatabaseConfiguration configuration) {
         mOpenHelper = createOpenHelper(configuration);
+        mCallbacks = configuration.callbacks;
         mAllowMainThreadQueries = configuration.allowMainThreadQueries;
     }
 
@@ -285,6 +289,7 @@
         private final Class<T> mDatabaseClass;
         private final String mName;
         private final Context mContext;
+        private ArrayList<Callback> mCallbacks;
 
         private SupportSQLiteOpenHelper.Factory mFactory;
         private boolean mInMemory;
@@ -355,6 +360,20 @@
         }
 
         /**
+         * Adds a {@link Callback} to this database.
+         *
+         * @param callback The callback.
+         * @return this
+         */
+        public Builder<T> addCallback(@NonNull Callback callback) {
+            if (mCallbacks == null) {
+                mCallbacks = new ArrayList<>();
+            }
+            mCallbacks.add(callback);
+            return this;
+        }
+
+        /**
          * Creates the databases and initializes it.
          * <p>
          * By default, all RoomDatabases use in memory storage for TEMP tables and enables recursive
@@ -377,7 +396,7 @@
             }
             DatabaseConfiguration configuration =
                     new DatabaseConfiguration(mContext, mName, mFactory, mMigrationContainer,
-                            mAllowMainThreadQueries);
+                            mCallbacks, mAllowMainThreadQueries);
             T db = Room.getGeneratedImplementation(mDatabaseClass, DB_IMPL_SUFFIX);
             db.init(configuration);
             return db;
@@ -475,4 +494,27 @@
             return result;
         }
     }
+
+    /**
+     * Callback for {@link RoomDatabase}.
+     */
+    public abstract static class Callback {
+
+        /**
+         * Called when the database is created for the first time. This is called after all the
+         * tables are created.
+         *
+         * @param db The database.
+         */
+        public void onCreate(@NonNull SupportSQLiteDatabase db) {
+        }
+
+        /**
+         * Called when the database has been opened.
+         *
+         * @param db The database.
+         */
+        public void onOpen(@NonNull SupportSQLiteDatabase db) {
+        }
+    }
 }
diff --git a/room/runtime/src/main/java/android/arch/persistence/room/RoomOpenHelper.java b/room/runtime/src/main/java/android/arch/persistence/room/RoomOpenHelper.java
index 7ac73df..a5b8150 100644
--- a/room/runtime/src/main/java/android/arch/persistence/room/RoomOpenHelper.java
+++ b/room/runtime/src/main/java/android/arch/persistence/room/RoomOpenHelper.java
@@ -58,6 +58,7 @@
     public void onCreate(SupportSQLiteDatabase db) {
         updateIdentity(db);
         mDelegate.createAllTables(db);
+        mDelegate.onCreate(db);
     }
 
     @Override
@@ -134,6 +135,8 @@
 
         protected abstract void onOpen(SupportSQLiteDatabase database);
 
+        protected abstract void onCreate(SupportSQLiteDatabase database);
+
         /**
          * Called after a migration run to validate database integrity.
          *
diff --git a/room/testing/src/main/java/android/arch/persistence/room/testing/MigrationTestHelper.java b/room/testing/src/main/java/android/arch/persistence/room/testing/MigrationTestHelper.java
index 904706e..2f49f39 100644
--- a/room/testing/src/main/java/android/arch/persistence/room/testing/MigrationTestHelper.java
+++ b/room/testing/src/main/java/android/arch/persistence/room/testing/MigrationTestHelper.java
@@ -128,7 +128,7 @@
         SchemaBundle schemaBundle = loadSchema(version);
         RoomDatabase.MigrationContainer container = new RoomDatabase.MigrationContainer();
         DatabaseConfiguration configuration = new DatabaseConfiguration(
-                mInstrumentation.getTargetContext(), name, mOpenFactory, container, true);
+                mInstrumentation.getTargetContext(), name, mOpenFactory, container, null, true);
         RoomOpenHelper roomOpenHelper = new RoomOpenHelper(configuration,
                 new CreatingDelegate(schemaBundle.getDatabase()),
                 schemaBundle.getDatabase().getIdentityHash());
@@ -170,7 +170,7 @@
         RoomDatabase.MigrationContainer container = new RoomDatabase.MigrationContainer();
         container.addMigrations(migrations);
         DatabaseConfiguration configuration = new DatabaseConfiguration(
-                mInstrumentation.getTargetContext(), name, mOpenFactory, container, true);
+                mInstrumentation.getTargetContext(), name, mOpenFactory, container, null, true);
         RoomOpenHelper roomOpenHelper = new RoomOpenHelper(configuration,
                 new MigratingDelegate(schemaBundle.getDatabase(), validateDroppedTables),
                 schemaBundle.getDatabase().getIdentityHash());
@@ -394,8 +394,11 @@
         }
 
         @Override
-        protected void onOpen(SupportSQLiteDatabase database) {
+        protected void onCreate(SupportSQLiteDatabase database) {
+        }
 
+        @Override
+        protected void onOpen(SupportSQLiteDatabase database) {
         }
     }
 }