1 files changed, 58 insertions, 9 deletions

diff --git a/adservices/service-core/java/com/android/adservices/data/DbHelper.java b/adservices/service-core/java/com/android/adservices/data/DbHelper.java
index 85f29184f4..8925df8963 100644
--- a/adservices/service-core/java/com/android/adservices/data/DbHelper.java
+++ b/adservices/service-core/java/com/android/adservices/data/DbHelper.java
@@ -28,7 +28,11 @@ import com.android.adservices.data.enrollment.EnrollmentTables;
 import com.android.adservices.data.measurement.MeasurementTables;
 import com.android.adservices.data.measurement.migration.IMeasurementDbMigrator;
 import com.android.adservices.data.measurement.migration.MeasurementDbMigratorV2;
+import com.android.adservices.data.measurement.migration.MeasurementDbMigratorV3;
 import com.android.adservices.data.topics.TopicsTables;
+import com.android.adservices.data.topics.migration.ITopicsDbMigrator;
+import com.android.adservices.data.topics.migration.TopicDbMigratorV3;
+import com.android.adservices.service.FlagsFactory;
 import com.android.internal.annotations.VisibleForTesting;
 
 import com.google.common.collect.ImmutableList;
@@ -41,12 +45,16 @@ import java.util.List;
  * get the same reference.
  */
 public class DbHelper extends SQLiteOpenHelper {
+    // Version 3: Add TopicContributors Table for Topics API, guarded by feature flag.
+    public static final int DATABASE_VERSION_V3 = 3;
 
-    static final int LATEST_DATABASE_VERSION = 2;
+    static final int CURRENT_DATABASE_VERSION = 2;
     private static final String DATABASE_NAME = "adservices.db";
 
     private static DbHelper sSingleton = null;
     private final File mDbFile;
+    // The version when the database is actually created
+    private final int mDbVersion;
 
     /**
      * It's only public to unit test.
@@ -59,6 +67,7 @@ public class DbHelper extends SQLiteOpenHelper {
     public DbHelper(@NonNull Context context, @NonNull String dbName, int dbVersion) {
         super(context, dbName, null, dbVersion);
         mDbFile = context.getDatabasePath(dbName);
+        this.mDbVersion = dbVersion;
     }
 
     /** Returns an instance of the DbHelper given a context. */
@@ -66,7 +75,7 @@ public class DbHelper extends SQLiteOpenHelper {
     public static DbHelper getInstance(@NonNull Context ctx) {
         synchronized (DbHelper.class) {
             if (sSingleton == null) {
-                sSingleton = new DbHelper(ctx, DATABASE_NAME, LATEST_DATABASE_VERSION);
+                sSingleton = new DbHelper(ctx, DATABASE_NAME, getDatabaseVersionToCreate());
             }
             return sSingleton;
         }
@@ -74,7 +83,7 @@ public class DbHelper extends SQLiteOpenHelper {
 
     @Override
     public void onCreate(@NonNull SQLiteDatabase db) {
-        LogUtil.d("DbHelper.onCreate.");
+        LogUtil.d("DbHelper.onCreate with version %d. Name: %s", mDbVersion, mDbFile.getName());
         for (String sql : TopicsTables.CREATE_STATEMENTS) {
             db.execSQL(sql);
         }
@@ -89,6 +98,12 @@ public class DbHelper extends SQLiteOpenHelper {
         }
     }
 
+    @Override
+    public void onOpen(SQLiteDatabase db) {
+        super.onOpen(db);
+        db.execSQL("PRAGMA foreign_keys=ON");
+    }
+
     /**
      * Wraps getReadableDatabase to catch SQLiteException and log error.
      */
@@ -102,9 +117,7 @@ public class DbHelper extends SQLiteOpenHelper {
         }
     }
 
-    /**
-     * Wraps getWritableDatabase to catch SQLiteException and log error.
-     */
+    /** Wraps getWritableDatabase to catch SQLiteException and log error. */
     @Nullable
     public SQLiteDatabase safeGetWritableDatabase() {
         try {
@@ -115,18 +128,54 @@ public class DbHelper extends SQLiteOpenHelper {
         }
     }
 
+    // TODO(b/255964885): Consolidate DB Migrator Class across Rubidium
     @Override
     public void onUpgrade(@NonNull SQLiteDatabase db, int oldVersion, int newVersion) {
-        LogUtil.d("DbHelper.onUpgrade.");
+        LogUtil.d(
+                "DbHelper.onUpgrade. Attempting to upgrade version from %d to %d.",
+                oldVersion, newVersion);
         getOrderedDbMigrators()
                 .forEach(dbMigrator -> dbMigrator.performMigration(db, oldVersion, newVersion));
+        try {
+            topicsGetOrderedDbMigrators()
+                    .forEach(dbMigrator -> dbMigrator.performMigration(db, oldVersion, newVersion));
+        } catch (IllegalArgumentException e) {
+            LogUtil.e(
+                    "Topics DB Upgrade is not performed! oldVersion: %d, newVersion: %d.",
+                    oldVersion, newVersion);
+        }
     }
 
     public long getDbFileSize() {
         return mDbFile != null && mDbFile.exists() ? mDbFile.length() : -1;
     }
 
-    private static List<IMeasurementDbMigrator> getOrderedDbMigrators() {
-        return ImmutableList.of(new MeasurementDbMigratorV2());
+    /**
+     * Check whether TopContributors Table is supported in current database. TopContributors is
+     * introduced in Version 3.
+     */
+    public boolean supportsTopicContributorsTable() {
+        return mDbVersion >= DATABASE_VERSION_V3;
+    }
+
+    /** Get Migrators in order for Measurement. */
+    @VisibleForTesting
+    public List<IMeasurementDbMigrator> getOrderedDbMigrators() {
+        return ImmutableList.of(new MeasurementDbMigratorV2(), new MeasurementDbMigratorV3());
+    }
+
+    /** Get Migrators in order for Topics. */
+    @VisibleForTesting
+    public List<ITopicsDbMigrator> topicsGetOrderedDbMigrators() {
+        return ImmutableList.of(new TopicDbMigratorV3());
+    }
+
+    // Get the database version to create. It may be different as LATEST_DATABASE_VERSION, depending
+    // on Flags status.
+    @VisibleForTesting
+    static int getDatabaseVersionToCreate() {
+        return FlagsFactory.getFlags().getEnableDatabaseSchemaVersion3()
+                ? DATABASE_VERSION_V3
+                : CURRENT_DATABASE_VERSION;
     }
 }