Add pragma statements automatically if sqlite (PrefectHQ#9169)

Author: jakekaplan

src/prefect/server/database/migrations/env.py

@@ -91,8 +91,8 @@ def dry_run_migrations() -> None:
         render_as_batch=dialect.name == "sqlite",
         # Each migration is its own transaction
         transaction_per_migration=True,
+        template_args={"dialect": dialect.name},
     )
-
     with context.begin_transaction():
         context.run_migrations()
 
@@ -104,7 +104,6 @@ def do_run_migrations(connection: AsyncEngine) -> None:
     Args:
         connection: a database engine.
     """
-
     context.configure(
         connection=connection,
         target_metadata=target_metadata,
@@ -123,6 +122,7 @@ def do_run_migrations(connection: AsyncEngine) -> None:
         render_as_batch=dialect.name == "sqlite",
         # Each migration is its own transaction
         transaction_per_migration=True,
+        template_args={"dialect": dialect.name},
     )
 
     with context.begin_transaction():

src/prefect/server/database/migrations/script.py.mako

@@ -16,10 +16,18 @@ down_revision = ${repr(down_revision)}
 branch_labels = ${repr(branch_labels)}
 depends_on = ${repr(depends_on)}
 
+<%
+    sqlite = dialect == "sqlite"
+    postgresql = dialect == "postgresql"
+%>
+
 
 def upgrade():
+    ${'op.execute("PRAGMA foreign_keys=OFF")' if sqlite else ""}
     ${upgrades if upgrades else "pass"}
-
+    ${'op.execute("PRAGMA foreign_keys=ON")' if sqlite else ""}
 
 def downgrade():
+    ${'op.execute("PRAGMA foreign_keys=OFF")' if sqlite else ""}
     ${downgrades if downgrades else "pass"}
+    ${'op.execute("PRAGMA foreign_keys=ON")' if sqlite else ""}