Skip to content

Commit f6b201b

Browse files
jakekaplanÅsmund Østvold
authored and
Åsmund Østvold
committed
Add pragma statements automatically if sqlite (PrefectHQ#9169)
1 parent 5b3c871 commit f6b201b

File tree

2 files changed

+11
-3
lines changed

2 files changed

+11
-3
lines changed

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

+2-2
Original file line numberDiff line numberDiff line change
@@ -91,8 +91,8 @@ def dry_run_migrations() -> None:
9191
render_as_batch=dialect.name == "sqlite",
9292
# Each migration is its own transaction
9393
transaction_per_migration=True,
94+
template_args={"dialect": dialect.name},
9495
)
95-
9696
with context.begin_transaction():
9797
context.run_migrations()
9898

@@ -104,7 +104,6 @@ def do_run_migrations(connection: AsyncEngine) -> None:
104104
Args:
105105
connection: a database engine.
106106
"""
107-
108107
context.configure(
109108
connection=connection,
110109
target_metadata=target_metadata,
@@ -123,6 +122,7 @@ def do_run_migrations(connection: AsyncEngine) -> None:
123122
render_as_batch=dialect.name == "sqlite",
124123
# Each migration is its own transaction
125124
transaction_per_migration=True,
125+
template_args={"dialect": dialect.name},
126126
)
127127

128128
with context.begin_transaction():

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

+9-1
Original file line numberDiff line numberDiff line change
@@ -16,10 +16,18 @@ down_revision = ${repr(down_revision)}
1616
branch_labels = ${repr(branch_labels)}
1717
depends_on = ${repr(depends_on)}
1818

19+
<%
20+
sqlite = dialect == "sqlite"
21+
postgresql = dialect == "postgresql"
22+
%>
23+
1924

2025
def upgrade():
26+
${'op.execute("PRAGMA foreign_keys=OFF")' if sqlite else ""}
2127
${upgrades if upgrades else "pass"}
22-
28+
${'op.execute("PRAGMA foreign_keys=ON")' if sqlite else ""}
2329

2430
def downgrade():
31+
${'op.execute("PRAGMA foreign_keys=OFF")' if sqlite else ""}
2532
${downgrades if downgrades else "pass"}
33+
${'op.execute("PRAGMA foreign_keys=ON")' if sqlite else ""}

0 commit comments

Comments
 (0)