1 files changed, 56 insertions, 1 deletions

diff --git a/src/com/android/messaging/datamodel/DatabaseUpgradeHelper.java b/src/com/android/messaging/datamodel/DatabaseUpgradeHelper.java
index d112533..96aba86 100644
--- a/src/com/android/messaging/datamodel/DatabaseUpgradeHelper.java
+++ b/src/com/android/messaging/datamodel/DatabaseUpgradeHelper.java
@@ -15,8 +15,10 @@
  */
 package com.android.messaging.datamodel;
 
+import android.content.Context;
 import android.database.sqlite.SQLiteDatabase;
 
+import com.android.messaging.Factory;
 import com.android.messaging.util.Assert;
 import com.android.messaging.util.LogUtil;
 
@@ -30,8 +32,61 @@ public class DatabaseUpgradeHelper {
         }
 
         LogUtil.i(TAG, "Database upgrade started from version " + oldVersion + " to " + newVersion);
+        try {
+            doUpgradeWithExceptions(db, oldVersion, newVersion);
+            LogUtil.i(TAG, "Finished database upgrade");
+        } catch (final Exception ex) {
+            LogUtil.e(TAG, "Failed to perform db upgrade from version " +
+                    oldVersion + " to version " + newVersion, ex);
+            DatabaseHelper.rebuildTables(db);
+        }
+    }
+
+    public void doUpgradeWithExceptions(final SQLiteDatabase db, final int oldVersion,
+            final int newVersion) throws Exception {
+        int currentVersion = oldVersion;
+        if (currentVersion < 2) {
+            currentVersion = upgradeToVersion2(db);
+        }
+        // Rebuild all the views
+        final Context context = Factory.get().getApplicationContext();
+        DatabaseHelper.dropAllViews(db);
+        DatabaseHelper.rebuildAllViews(new DatabaseWrapper(context, db));
+        // Finally, check if we have arrived at the final version.
+        checkAndUpdateVersionAtReleaseEnd(currentVersion, Integer.MAX_VALUE, newVersion);
+    }
 
-        // Add future upgrade code here
+    private int upgradeToVersion2(final SQLiteDatabase db) {
+        db.execSQL("ALTER TABLE " + DatabaseHelper.CONVERSATIONS_TABLE + " ADD COLUMN " +
+                DatabaseHelper.ConversationColumns.IS_ENTERPRISE + " INT DEFAULT(0)");
+        LogUtil.i(TAG, "Ugraded database to version 2");
+        return 2;
+    }
+
+    /**
+     * Checks db version correctness at the end of each milestone release. If target database
+     * version lies beyond the version range that the current release may handle, we snap the
+     * current version to the end of the release, so that we may go on to the next release' upgrade
+     * path. Otherwise, if target version is within reach of the current release, but we are not
+     * at the target version, then throw an exception to force a table rebuild.
+     */
+    private int checkAndUpdateVersionAtReleaseEnd(final int currentVersion,
+            final int maxVersionForRelease, final int targetVersion) throws Exception {
+        if (maxVersionForRelease < targetVersion) {
+            // Target version is beyond the current release. Snap to max version for the
+            // current release so we can go on to the upgrade path for the next release.
+            return maxVersionForRelease;
+        }
+
+        // If we are here, this means the current release' upgrade handler should upgrade to
+        // target version...
+        if (currentVersion != targetVersion) {
+            // No more upgrade handlers. So we can't possibly upgrade to the final version.
+            throw new Exception("Missing upgrade handler from version " +
+                    currentVersion + " to version " + targetVersion);
+        }
+        // Upgrade succeeded.
+        return targetVersion;
     }
 
     public void onDowngrade(final SQLiteDatabase db, final int oldVersion, final int newVersion) {