Re-enable foreign key support if needed when a migration fails

library/src/commonMain/kotlin/com/eygraber/sqldelight/androidx/driver/AndroidxSqliteDriver.kt

@@ -483,40 +483,62 @@ private inline fun SQLiteConnection.withDeferredForeignKeyChecks(
     prepare("PRAGMA foreign_keys = OFF;").use(SQLiteStatement::step)
   }
 
-  block()
+  try {
+    block()
 
-  if(configuration.isForeignKeyConstraintsEnabled) {
-    prepare("PRAGMA foreign_keys = ON;").use(SQLiteStatement::step)
-
-    if(configuration.isForeignKeyConstraintsCheckedAfterCreateOrUpdate) {
-      prepare("PRAGMA foreign_key_check;").use { check ->
-        val violations = mutableListOf<AndroidxSqliteDriver.ForeignKeyConstraintViolation>()
-        var count = 0
-        while(check.step() && count++ < configuration.maxMigrationForeignKeyConstraintViolationsToReport) {
-          violations.add(
-            AndroidxSqliteDriver.ForeignKeyConstraintViolation(
-              referencingTable = check.getText(0),
-              referencingRowId = check.getInt(1),
-              referencedTable = check.getText(2),
-              referencingConstraintIndex = check.getInt(3),
-            ),
-          )
-        }
+    if(configuration.isForeignKeyConstraintsEnabled) {
+      prepare("PRAGMA foreign_keys = ON;").use(SQLiteStatement::step)
+
+      if(configuration.isForeignKeyConstraintsCheckedAfterCreateOrUpdate) {
+        reportForeignKeyViolations(
+          configuration.maxMigrationForeignKeyConstraintViolationsToReport,
+        )
+      }
+    }
+  } catch(e: Throwable) {
+    // An exception happened during creation / migration.
+    // We will try to re-enable foreign keys, and if that also fails,
+    // we will add it as a suppressed exception to the original one.
+    try {
+      if(configuration.isForeignKeyConstraintsEnabled) {
+        prepare("PRAGMA foreign_keys = ON;").use(SQLiteStatement::step)
+      }
+    } catch(fkException: Throwable) {
+      e.addSuppressed(fkException)
+    }
+    throw e
+  }
+}
+
+private fun SQLiteConnection.reportForeignKeyViolations(
+  maxMigrationForeignKeyConstraintViolationsToReport: Int,
+) {
+  prepare("PRAGMA foreign_key_check;").use { check ->
+    val violations = mutableListOf<AndroidxSqliteDriver.ForeignKeyConstraintViolation>()
+    var count = 0
+    while(check.step() && count++ < maxMigrationForeignKeyConstraintViolationsToReport) {
+      violations.add(
+        AndroidxSqliteDriver.ForeignKeyConstraintViolation(
+          referencingTable = check.getText(0),
+          referencingRowId = check.getInt(1),
+          referencedTable = check.getText(2),
+          referencingConstraintIndex = check.getInt(3),
+        ),
+      )
+    }
 
-        if(violations.isNotEmpty()) {
-          val unprintedViolationsCount = violations.size - 5
-          val unprintedDisclaimer = if(unprintedViolationsCount > 0) " ($unprintedViolationsCount not shown)" else ""
+    if(violations.isNotEmpty()) {
+      val unprintedViolationsCount = violations.size - 5
+      val unprintedDisclaimer = if(unprintedViolationsCount > 0) " ($unprintedViolationsCount not shown)" else ""
 
-          throw AndroidxSqliteDriver.ForeignKeyConstraintCheckException(
-            violations = violations,
-            message = """
+      throw AndroidxSqliteDriver.ForeignKeyConstraintCheckException(
+        violations = violations,
+        message = """
             |The following foreign key constraints are violated$unprintedDisclaimer:
             |
             |${violations.take(5).joinToString(separator = "\n\n")}
-            """.trimMargin(),
-          )
-        }
-      }
+        """.trimMargin(),
+      )
     }
   }
 }

library/src/commonTest/kotlin/com/eygraber/sqldelight/androidx/driver/AndroidxSqliteCreationTest.kt

@@ -10,6 +10,7 @@ import kotlin.test.Test
 import kotlin.test.assertContentEquals
 import kotlin.test.assertEquals
 import kotlin.test.assertFailsWith
+import kotlin.test.assertTrue
 
 abstract class AndroidxSqliteCreationTest {
   private fun getSchema(
@@ -429,4 +430,72 @@ abstract class AndroidxSqliteCreationTest {
       execute(null, "PRAGMA user_version;", 0, null)
     }
   }
+
+  @Test
+  fun `exceptions thrown during creation are propagated to the caller`() {
+    val schema = getSchema {
+      throw RuntimeException("Test")
+    }
+    val dbName = Random.nextULong().toHexString()
+
+    withDatabase(
+      schema = schema,
+      dbName = dbName,
+      onCreate = {},
+      onUpdate = { _, _ -> },
+      onOpen = {},
+      onConfigure = {},
+      configuration = AndroidxSqliteConfiguration(
+        isForeignKeyConstraintsEnabled = true,
+        isForeignKeyConstraintsCheckedAfterCreateOrUpdate = false,
+      ),
+    ) {
+      val message = assertFailsWith<RuntimeException> {
+        execute(null, "PRAGMA user_version;", 0, null)
+      }.message
+
+      assertEquals("Test", message)
+    }
+  }
+
+  @Test
+  fun `foreign keys are re-enabled after an exception is thrown during creation`() {
+    val schema = getSchema {
+      throw RuntimeException("Test")
+    }
+    val dbName = Random.nextULong().toHexString()
+
+    withDatabase(
+      schema = schema,
+      dbName = dbName,
+      onCreate = {},
+      onUpdate = { _, _ -> },
+      onOpen = {},
+      onConfigure = {},
+      configuration = AndroidxSqliteConfiguration(
+        isForeignKeyConstraintsEnabled = true,
+        isForeignKeyConstraintsCheckedAfterCreateOrUpdate = false,
+      ),
+    ) {
+      assertFailsWith<RuntimeException> {
+        execute(null, "PRAGMA user_version;", 0, null)
+      }
+
+      assertTrue {
+        executeQuery(
+          identifier = null,
+          sql = "PRAGMA foreign_keys;",
+          mapper = { cursor ->
+            QueryResult.Value(
+              when {
+                cursor.next().value -> cursor.getLong(0)
+                else -> 0L
+              },
+            )
+          },
+          parameters = 0,
+        ).value == 1L
+      }
+    }
+  }
 }

library/src/commonTest/kotlin/com/eygraber/sqldelight/androidx/driver/AndroidxSqliteMigrationTest.kt

@@ -10,6 +10,7 @@ import kotlin.test.Test
 import kotlin.test.assertContentEquals
 import kotlin.test.assertEquals
 import kotlin.test.assertFailsWith
+import kotlin.test.assertTrue
 
 abstract class AndroidxSqliteMigrationTest {
   private fun getSchema(
@@ -596,4 +597,111 @@ abstract class AndroidxSqliteMigrationTest {
     assertEquals(0, create)
     assertEquals(1, update)
   }
+
+  @Test
+  fun `exceptions thrown during migration are propagated to the caller`() {
+    val schema = getSchema {
+      throw RuntimeException("Test")
+    }
+    val dbName = Random.nextULong().toHexString()
+
+    // trigger creation
+    withDatabase(
+      schema = schema,
+      dbName = dbName,
+      onCreate = {},
+      onUpdate = { _, _ -> },
+      onOpen = {},
+      onConfigure = {},
+      deleteDbAfterRun = false,
+      configuration = AndroidxSqliteConfiguration(
+        isForeignKeyConstraintsEnabled = true,
+        isForeignKeyConstraintsCheckedAfterCreateOrUpdate = false,
+      ),
+    ) {
+      execute(null, "PRAGMA user_version;", 0, null)
+    }
+
+    schema.version++
+
+    withDatabase(
+      schema = schema,
+      dbName = dbName,
+      onCreate = {},
+      onUpdate = { _, _ -> },
+      onOpen = {},
+      onConfigure = {},
+      deleteDbBeforeRun = false,
+      configuration = AndroidxSqliteConfiguration(
+        isForeignKeyConstraintsEnabled = true,
+        isForeignKeyConstraintsCheckedAfterCreateOrUpdate = false,
+      ),
+    ) {
+      val message = assertFailsWith<RuntimeException> {
+        execute(null, "PRAGMA user_version;", 0, null)
+      }.message
+      assertEquals("Test", message)
+    }
+  }
+
+  @Test
+  fun `foreign keys are re-enabled after an exception is thrown during migration`() {
+    val schema = getSchema {
+      throw RuntimeException("Test")
+    }
+    val dbName = Random.nextULong().toHexString()
+
+    // trigger creation
+    withDatabase(
+      schema = schema,
+      dbName = dbName,
+      onCreate = {},
+      onUpdate = { _, _ -> },
+      onOpen = {},
+      onConfigure = {},
+      deleteDbAfterRun = false,
+      configuration = AndroidxSqliteConfiguration(
+        isForeignKeyConstraintsEnabled = true,
+        isForeignKeyConstraintsCheckedAfterCreateOrUpdate = false,
+      ),
+    ) {
+      execute(null, "PRAGMA user_version;", 0, null)
+    }
+
+    schema.version++
+
+    withDatabase(
+      schema = schema,
+      dbName = dbName,
+      onCreate = {},
+      onUpdate = { _, _ -> },
+      onOpen = {},
+      onConfigure = {},
+      deleteDbBeforeRun = false,
+      configuration = AndroidxSqliteConfiguration(
+        isForeignKeyConstraintsEnabled = true,
+        isForeignKeyConstraintsCheckedAfterCreateOrUpdate = false,
+      ),
+    ) {
+      assertFailsWith<RuntimeException> {
+        execute(null, "PRAGMA user_version;", 0, null)
+      }
+
+      assertTrue {
+        executeQuery(
+          identifier = null,
+          sql = "PRAGMA foreign_keys;",
+          mapper = { cursor ->
+            QueryResult.Value(
+              when {
+                cursor.next().value -> cursor.getLong(0)
+                else -> 0L
+              },
+            )
+          },
+          parameters = 0,
+        ).value == 1L
+      }
+    }
+  }
 }