Merge pull request #7570 from hayaksi1/fix/7509-run-database-migrations

Run the database migrations again when a database is upgraded

libraries/encrypted-db/build.gradle.kts

@@ -5,6 +5,8 @@
  * SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-Element-Commercial.
  * Please see LICENSE files in the repository root for full details.
  */
+import extension.testCommonDependencies
+
 plugins {
     id("io.element.android-library")
 }
@@ -27,4 +29,6 @@ dependencies {
     implementation(libs.google.tink)
 
     implementation(projects.libraries.androidutils)
+
+    testCommonDependencies(libs)
 }

libraries/encrypted-db/src/main/kotlin/io/element/encrypteddb/SqlCipherDriverFactory.kt

@@ -51,12 +51,28 @@ class SqlCipherDriverFactory(
             is ClientSecret.Passphrase -> secret.formattedAsString().hexToByteArray()
         }
         val factory = SupportOpenHelperFactory(key)
-        return AndroidSqliteDriver(schema = schema, context = context, name = name, factory = factory, callback = object : AndroidSqliteDriver.Callback(
+        return AndroidSqliteDriver(
-                schema
+            schema = schema,
-            ) {
+            context = context,
-            override fun onUpgrade(db: SupportSQLiteDatabase, oldVersion: Int, newVersion: Int) {
+            name = name,
-                onUpgradeCallback?.invoke(db, oldVersion, newVersion)
+            factory = factory,
-            }
+            callback = MigrationCallback(schema, onUpgradeCallback),
-        })
+        )
+    }
+}
+
+/**
+ * Runs [onUpgradeCallback] and then the schema migrations themselves.
+ *
+ * @param schema The SQLite DB schema, whose migrations are applied by [AndroidSqliteDriver.Callback.onUpgrade].
+ * @param onUpgradeCallback Called before the migrations run, so that the database key can be replaced first.
+ */
+internal class MigrationCallback(
+    schema: SqlSchema<QueryResult.Value<Unit>>,
+    private val onUpgradeCallback: ((driver: SupportSQLiteDatabase, oldVersion: Int, newVersion: Int) -> Unit)?,
+) : AndroidSqliteDriver.Callback(schema) {
+    override fun onUpgrade(db: SupportSQLiteDatabase, oldVersion: Int, newVersion: Int) {
+        onUpgradeCallback?.invoke(db, oldVersion, newVersion)
+        super.onUpgrade(db, oldVersion, newVersion)
     }
 }

libraries/encrypted-db/src/test/kotlin/io/element/encrypteddb/MigrationCallbackTest.kt

@@ -0,0 +1,55 @@
+/*
+ * Copyright (c) 2026 Element Creations Ltd.
+ *
+ * SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-Element-Commercial.
+ * Please see LICENSE files in the repository root for full details.
+ */
+
+package io.element.encrypteddb
+
+import app.cash.sqldelight.db.AfterVersion
+import app.cash.sqldelight.db.QueryResult
+import app.cash.sqldelight.db.SqlDriver
+import app.cash.sqldelight.db.SqlSchema
+import com.google.common.truth.Truth.assertThat
+import io.mockk.mockk
+import org.junit.Test
+
+class MigrationCallbackTest {
+    @Test
+    fun `onUpgrade migrates the schema`() {
+        val schema = RecordingSchema()
+        MigrationCallback(schema, onUpgradeCallback = null).onUpgrade(mockk(relaxed = true), 9, 11)
+        assertThat(schema.migrations).containsExactly(9L to 11L)
+    }
+
+    @Test
+    fun `onUpgrade replaces the database key before migrating the schema`() {
+        val order = mutableListOf<String>()
+        val schema = RecordingSchema(onMigrate = { order.add("migrate") })
+        MigrationCallback(schema, onUpgradeCallback = { _, _, _ -> order.add("rekey") })
+            .onUpgrade(mockk(relaxed = true), 9, 11)
+        assertThat(order).containsExactly("rekey", "migrate").inOrder()
+    }
+}
+
+private class RecordingSchema(
+    private val onMigrate: () -> Unit = {},
+) : SqlSchema<QueryResult.Value<Unit>> {
+    val migrations = mutableListOf<Pair<Long, Long>>()
+
+    override val version = 11L
+
+    override fun create(driver: SqlDriver) = QueryResult.Unit
+
+    override fun migrate(
+        driver: SqlDriver,
+        oldVersion: Long,
+        newVersion: Long,
+        vararg callbacks: AfterVersion,
+    ): QueryResult.Value<Unit> {
+        migrations.add(oldVersion to newVersion)
+        onMigrate()
+        return QueryResult.Unit
+    }
+}