« prev ^ index » next coverage.py v7.10.6, created at 2025-12-05 13:38 +0000
1"""Work queue data migration
3Revision ID: 1678f2fb8b33
4Revises: b9bda9f142f1
5Create Date: 2023-01-31 10:54:42.747849
7"""
9import sqlalchemy as sa 1a
10from alembic import op 1a
12# revision identifiers, used by Alembic.
13revision = "1678f2fb8b33" 1a
14down_revision = "b9bda9f142f1" 1a
15branch_labels = None 1a
16depends_on = None 1a
19def upgrade(): 1a
20 # Create temporary indexes for migration
21 op.execute( 1a
22 "CREATE INDEX IF NOT EXISTS ix_flow_run__work_queue_id_work_queue_name ON"
23 " flow_run (work_queue_id, work_queue_name)"
24 )
25 op.execute( 1a
26 "CREATE INDEX IF NOT EXISTS ix_deployment__work_queue_id_work_queue_name ON"
27 " deployment (work_queue_id, work_queue_name)"
28 )
30 # Create default agent work pool and associate all existing queues with it
31 connection = op.get_bind() 1a
33 n_work_queues_without_a_work_pool = connection.execute( 1a
34 sa.text("SELECT COUNT(*) FROM work_queue WHERE work_pool_id IS NULL")
35 ).fetchone()[0]
37 if n_work_queues_without_a_work_pool > 0: 37 ↛ 38line 37 didn't jump to line 38 because the condition on line 37 was never true1a
38 connection.execute(
39 sa.text(
40 "INSERT INTO work_pool (name, type) VALUES ('default-agent-pool',"
41 " 'prefect-agent')"
42 )
43 )
45 default_pool_id = connection.execute(
46 sa.text("SELECT id FROM work_pool WHERE name = 'default-agent-pool'")
47 ).fetchone()[0]
49 default_queue = connection.execute(
50 sa.text("SELECT id FROM work_queue WHERE name = 'default'")
51 ).fetchone()
53 if not default_queue:
54 connection.execute(
55 sa.text(
56 "INSERT INTO work_queue (name, work_pool_id) VALUES ('default',"
57 " :default_pool_id)"
58 ).params({"default_pool_id": default_pool_id}),
59 )
61 connection.execute(
62 sa.text(
63 "UPDATE work_queue SET work_pool_id = :default_pool_id WHERE work_pool_id"
64 " IS NULL"
65 ).params({"default_pool_id": default_pool_id}),
66 )
68 default_queue_id = connection.execute(
69 sa.text(
70 "SELECT id FROM work_queue WHERE name = 'default' and work_pool_id ="
71 " :default_pool_id"
72 ).params({"default_pool_id": default_pool_id}),
73 ).fetchone()[0]
75 connection.execute(
76 sa.text(
77 "UPDATE work_pool SET default_queue_id = :default_queue_id WHERE id ="
78 " :default_pool_id"
79 ).params(
80 {
81 "default_pool_id": default_pool_id,
82 "default_queue_id": default_queue_id,
83 }
84 ),
85 )
87 # Set priority on all queues and update flow runs and deployments
88 queue_rows = connection.execute(
89 sa.text(
90 "SELECT id, name FROM work_queue WHERE work_pool_id = :default_pool_id"
91 ).params({"default_pool_id": default_pool_id}),
92 ).fetchall()
94 with op.get_context().autocommit_block():
95 for enumeration, row in enumerate(queue_rows):
96 connection.execute(
97 sa.text(
98 "UPDATE work_queue SET priority = :priority WHERE id = :id"
99 ).params({"priority": enumeration + 1, "id": row[0]}),
100 )
102 batch_size = 250
104 while True:
105 result = connection.execute(
106 sa.text(
107 """
108 UPDATE flow_run
109 SET work_queue_id=:id
110 WHERE flow_run.id in (
111 SELECT id
112 FROM flow_run
113 WHERE flow_run.work_queue_id IS NULL and flow_run.work_queue_name=:name
114 LIMIT :batch_size
115 )
116 """
117 ).params(
118 {"id": row[0], "name": row[1], "batch_size": batch_size}
119 ),
120 )
121 if result.rowcount <= batch_size:
122 break
124 while True:
125 result = connection.execute(
126 sa.text(
127 """
128 UPDATE deployment
129 SET work_queue_id=:id
130 WHERE deployment.id in (
131 SELECT id
132 FROM deployment
133 WHERE deployment.work_queue_id IS NULL and deployment.work_queue_name=:name
134 LIMIT :batch_size
135 )
136 """
137 ).params(
138 {"id": row[0], "name": row[1], "batch_size": batch_size}
139 ),
140 )
141 if result.rowcount <= batch_size:
142 break
144 with op.batch_alter_table("work_queue", schema=None) as batch_op: 1a
145 batch_op.drop_constraint("uq_work_queue__name") 1a
146 batch_op.create_unique_constraint( 1a
147 op.f("uq_work_queue__work_pool_id_name"), ["work_pool_id", "name"]
148 )
149 batch_op.alter_column("work_pool_id", nullable=False) 1a
151 op.execute("DROP INDEX IF EXISTS ix_flow_run__work_queue_id_work_queue_name") 1a
152 op.execute("DROP INDEX IF EXISTS ix_deployment__work_queue_id_work_queue_name") 1a
155def downgrade(): 1a
156 connection = op.get_bind()
157 # Create temporary indexes for migration
158 op.execute(
159 "CREATE INDEX IF NOT EXISTS ix_flow_run__work_queue_id_work_queue_name ON"
160 " flow_run (work_queue_id, work_queue_name)"
161 )
162 op.execute(
163 "CREATE INDEX IF NOT EXISTS ix_deployment__work_queue_id_work_queue_name ON"
164 " deployment (work_queue_id, work_queue_name)"
165 )
167 with op.batch_alter_table("work_queue", schema=None) as batch_op:
168 batch_op.alter_column("work_pool_id", nullable=True)
170 # Delete all non-default queues and pools
171 default_pool_id_result = connection.execute(
172 sa.text("SELECT id FROM work_pool WHERE name = 'default-agent-pool'")
173 ).fetchone()
174 if default_pool_id_result:
175 default_pool_id = default_pool_id_result[0]
176 connection.execute(
177 sa.text(
178 "DELETE FROM work_queue WHERE work_pool_id != :default_pool_id"
179 ).params({"default_pool_id": default_pool_id})
180 )
181 queue_rows = connection.execute(
182 sa.text("SELECT id, name FROM work_queue"),
183 ).fetchall()
185 with op.get_context().autocommit_block():
186 for row in queue_rows:
187 batch_size = 250
189 while True:
190 result = connection.execute(
191 sa.text(
192 """
193 UPDATE flow_run
194 SET work_queue_id=NULL
195 WHERE flow_run.id in (
196 SELECT id
197 FROM flow_run
198 WHERE flow_run.work_queue_id IS NOT NULL and flow_run.work_queue_id=:id
199 LIMIT :batch_size
200 )
201 """
202 ).params({"id": row[0], "batch_size": batch_size}),
203 )
204 if result.rowcount <= batch_size:
205 break
207 while True:
208 result = connection.execute(
209 sa.text(
210 """
211 UPDATE deployment
212 SET work_queue_id=NULL
213 WHERE deployment.id in (
214 SELECT id
215 FROM deployment
216 WHERE deployment.work_queue_id IS NOT NULL and deployment.work_queue_id=:id
217 LIMIT :batch_size
218 )
219 """
220 ).params({"id": row[0], "batch_size": batch_size}),
221 )
222 if result.rowcount <= batch_size:
223 break
225 connection.execute(sa.text("UPDATE work_queue SET work_pool_id = NULL"))
227 connection.execute(sa.text("DELETE FROM work_pool"))
229 with op.batch_alter_table("work_queue", schema=None) as batch_op:
230 batch_op.drop_constraint("uq_work_queue__work_pool_id_name")
231 batch_op.create_unique_constraint("uq_work_queue__name", ["name"])
233 op.execute("DROP INDEX IF EXISTS ix_flow_run__work_queue_id_work_queue_name")
234 op.execute("DROP INDEX IF EXISTS ix_deployment__work_queue_id_work_queue_name")