Skip to content

Commit 4b0fc3b

Browse files
rakduttakevalmahajan
authored andcommitted
alembic
Signed-off-by: rakdutta <rakhibiswas@yahoo.com>
1 parent f9c9446 commit 4b0fc3b

File tree

1 file changed

+264
-31
lines changed

1 file changed

+264
-31
lines changed

mcpgateway/alembic/versions/356a2d4eed6f_uuid_change_for_prompt_and_resources.py

Lines changed: 264 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -66,24 +66,123 @@ def upgrade() -> None:
6666
)
6767

6868
# 3) Copy data from prompts into prompts_tmp using id_new as id
69-
copy_cols = (
70-
"id, name, description, template, argument_schema, created_at, updated_at, enabled, tags,"
71-
" created_by, created_from_ip, created_via, created_user_agent, modified_by, modified_from_ip,"
72-
" modified_via, modified_user_agent, import_batch_id, federation_source, version, gateway_id, team_id, owner_email, visibility"
69+
# Use SQLAlchemy Core to move rows from `prompts` -> `prompts_tmp` without
70+
# composing SQL text. This avoids dynamic string concat while keeping the
71+
# same column mapping (id_new -> id, is_active -> enabled).
72+
prompts_src = sa.table(
73+
"prompts",
74+
sa.column("id_new"),
75+
sa.column("name"),
76+
sa.column("description"),
77+
sa.column("template"),
78+
sa.column("argument_schema"),
79+
sa.column("created_at"),
80+
sa.column("updated_at"),
81+
sa.column("is_active"),
82+
sa.column("tags"),
83+
sa.column("created_by"),
84+
sa.column("created_from_ip"),
85+
sa.column("created_via"),
86+
sa.column("created_user_agent"),
87+
sa.column("modified_by"),
88+
sa.column("modified_from_ip"),
89+
sa.column("modified_via"),
90+
sa.column("modified_user_agent"),
91+
sa.column("import_batch_id"),
92+
sa.column("federation_source"),
93+
sa.column("version"),
94+
sa.column("gateway_id"),
95+
sa.column("team_id"),
96+
sa.column("owner_email"),
97+
sa.column("visibility"),
7398
)
74-
conn.execute(
75-
text(
76-
(
77-
"INSERT INTO prompts_tmp (" + copy_cols + ") "
78-
"SELECT id_new, name, description, template, argument_schema, created_at, "
79-
"updated_at, is_active, tags, created_by, created_from_ip, created_via, "
80-
"created_user_agent, modified_by, modified_from_ip, modified_via, "
81-
"modified_user_agent, import_batch_id, federation_source, version, gateway_id, "
82-
"team_id, owner_email, visibility FROM prompts"
83-
)
84-
)
99+
100+
prompts_tgt = sa.table(
101+
"prompts_tmp",
102+
sa.column("id"),
103+
sa.column("name"),
104+
sa.column("description"),
105+
sa.column("template"),
106+
sa.column("argument_schema"),
107+
sa.column("created_at"),
108+
sa.column("updated_at"),
109+
sa.column("enabled"),
110+
sa.column("tags"),
111+
sa.column("created_by"),
112+
sa.column("created_from_ip"),
113+
sa.column("created_via"),
114+
sa.column("created_user_agent"),
115+
sa.column("modified_by"),
116+
sa.column("modified_from_ip"),
117+
sa.column("modified_via"),
118+
sa.column("modified_user_agent"),
119+
sa.column("import_batch_id"),
120+
sa.column("federation_source"),
121+
sa.column("version"),
122+
sa.column("gateway_id"),
123+
sa.column("team_id"),
124+
sa.column("owner_email"),
125+
sa.column("visibility"),
85126
)
86127

128+
target_cols = [
129+
"id",
130+
"name",
131+
"description",
132+
"template",
133+
"argument_schema",
134+
"created_at",
135+
"updated_at",
136+
"enabled",
137+
"tags",
138+
"created_by",
139+
"created_from_ip",
140+
"created_via",
141+
"created_user_agent",
142+
"modified_by",
143+
"modified_from_ip",
144+
"modified_via",
145+
"modified_user_agent",
146+
"import_batch_id",
147+
"federation_source",
148+
"version",
149+
"gateway_id",
150+
"team_id",
151+
"owner_email",
152+
"visibility",
153+
]
154+
155+
select_exprs = [
156+
prompts_src.c.id_new,
157+
prompts_src.c.name,
158+
prompts_src.c.description,
159+
prompts_src.c.template,
160+
prompts_src.c.argument_schema,
161+
prompts_src.c.created_at,
162+
prompts_src.c.updated_at,
163+
prompts_src.c.is_active,
164+
prompts_src.c.tags,
165+
prompts_src.c.created_by,
166+
prompts_src.c.created_from_ip,
167+
prompts_src.c.created_via,
168+
prompts_src.c.created_user_agent,
169+
prompts_src.c.modified_by,
170+
prompts_src.c.modified_from_ip,
171+
prompts_src.c.modified_via,
172+
prompts_src.c.modified_user_agent,
173+
prompts_src.c.import_batch_id,
174+
prompts_src.c.federation_source,
175+
prompts_src.c.version,
176+
prompts_src.c.gateway_id,
177+
prompts_src.c.team_id,
178+
prompts_src.c.owner_email,
179+
prompts_src.c.visibility,
180+
]
181+
182+
stmt = sa.select(*select_exprs)
183+
ins = sa.insert(prompts_tgt).from_select(target_cols, stmt)
184+
conn.execute(ins)
185+
87186
# 4) Create new prompt_metrics table with prompt_id varchar(36)
88187
op.create_table(
89188
"prompt_metrics_tmp",
@@ -174,22 +273,137 @@ def upgrade() -> None:
174273
sa.PrimaryKeyConstraint("id", name="pk_resources"),
175274
)
176275

177-
# Copy data into resources_tmp using id_new
178-
res_copy_cols = "id, uri, name, description, mime_type, size, uri_template, created_at, updated_at, enabled, tags, text_content, binary_content, created_by, created_from_ip, created_via, created_user_agent, modified_by, modified_from_ip, modified_via, modified_user_agent, import_batch_id, federation_source, version, gateway_id, team_id, owner_email, visibility"
179-
conn.execute(
180-
text(
181-
(
182-
"INSERT INTO resources_tmp (" + res_copy_cols + ") "
183-
"SELECT id_new, uri, name, description, mime_type, size, uri_template, "
184-
"created_at, updated_at, is_active, tags, text_content, binary_content, "
185-
"created_by, created_from_ip, created_via, created_user_agent, modified_by, "
186-
"modified_from_ip, modified_via, modified_user_agent, import_batch_id, "
187-
"federation_source, version, gateway_id, team_id, owner_email, visibility "
188-
"FROM resources"
189-
)
190-
)
276+
# Copy data into resources_tmp using id_new via SQLAlchemy Core
277+
resources_src = sa.table(
278+
"resources",
279+
sa.column("id_new"),
280+
sa.column("uri"),
281+
sa.column("name"),
282+
sa.column("description"),
283+
sa.column("mime_type"),
284+
sa.column("size"),
285+
sa.column("uri_template"),
286+
sa.column("created_at"),
287+
sa.column("updated_at"),
288+
sa.column("is_active"),
289+
sa.column("tags"),
290+
sa.column("text_content"),
291+
sa.column("binary_content"),
292+
sa.column("created_by"),
293+
sa.column("created_from_ip"),
294+
sa.column("created_via"),
295+
sa.column("created_user_agent"),
296+
sa.column("modified_by"),
297+
sa.column("modified_from_ip"),
298+
sa.column("modified_via"),
299+
sa.column("modified_user_agent"),
300+
sa.column("import_batch_id"),
301+
sa.column("federation_source"),
302+
sa.column("version"),
303+
sa.column("gateway_id"),
304+
sa.column("team_id"),
305+
sa.column("owner_email"),
306+
sa.column("visibility"),
307+
)
308+
309+
resources_tgt = sa.table(
310+
"resources_tmp",
311+
sa.column("id"),
312+
sa.column("uri"),
313+
sa.column("name"),
314+
sa.column("description"),
315+
sa.column("mime_type"),
316+
sa.column("size"),
317+
sa.column("uri_template"),
318+
sa.column("created_at"),
319+
sa.column("updated_at"),
320+
sa.column("enabled"),
321+
sa.column("tags"),
322+
sa.column("text_content"),
323+
sa.column("binary_content"),
324+
sa.column("created_by"),
325+
sa.column("created_from_ip"),
326+
sa.column("created_via"),
327+
sa.column("created_user_agent"),
328+
sa.column("modified_by"),
329+
sa.column("modified_from_ip"),
330+
sa.column("modified_via"),
331+
sa.column("modified_user_agent"),
332+
sa.column("import_batch_id"),
333+
sa.column("federation_source"),
334+
sa.column("version"),
335+
sa.column("gateway_id"),
336+
sa.column("team_id"),
337+
sa.column("owner_email"),
338+
sa.column("visibility"),
191339
)
192340

341+
target_cols_res = [
342+
"id",
343+
"uri",
344+
"name",
345+
"description",
346+
"mime_type",
347+
"size",
348+
"uri_template",
349+
"created_at",
350+
"updated_at",
351+
"enabled",
352+
"tags",
353+
"text_content",
354+
"binary_content",
355+
"created_by",
356+
"created_from_ip",
357+
"created_via",
358+
"created_user_agent",
359+
"modified_by",
360+
"modified_from_ip",
361+
"modified_via",
362+
"modified_user_agent",
363+
"import_batch_id",
364+
"federation_source",
365+
"version",
366+
"gateway_id",
367+
"team_id",
368+
"owner_email",
369+
"visibility",
370+
]
371+
372+
select_exprs_res = [
373+
resources_src.c.id_new,
374+
resources_src.c.uri,
375+
resources_src.c.name,
376+
resources_src.c.description,
377+
resources_src.c.mime_type,
378+
resources_src.c.size,
379+
resources_src.c.uri_template,
380+
resources_src.c.created_at,
381+
resources_src.c.updated_at,
382+
resources_src.c.is_active,
383+
resources_src.c.tags,
384+
resources_src.c.text_content,
385+
resources_src.c.binary_content,
386+
resources_src.c.created_by,
387+
resources_src.c.created_from_ip,
388+
resources_src.c.created_via,
389+
resources_src.c.created_user_agent,
390+
resources_src.c.modified_by,
391+
resources_src.c.modified_from_ip,
392+
resources_src.c.modified_via,
393+
resources_src.c.modified_user_agent,
394+
resources_src.c.import_batch_id,
395+
resources_src.c.federation_source,
396+
resources_src.c.version,
397+
resources_src.c.gateway_id,
398+
resources_src.c.team_id,
399+
resources_src.c.owner_email,
400+
resources_src.c.visibility,
401+
]
402+
403+
stmt_res = sa.select(*select_exprs_res)
404+
ins_res = sa.insert(resources_tgt).from_select(target_cols_res, stmt_res)
405+
conn.execute(ins_res)
406+
193407
# resource_metrics_tmp with resource_id varchar(32)
194408
op.create_table(
195409
"resource_metrics_tmp",
@@ -264,6 +478,7 @@ def upgrade() -> None:
264478
existing_nullable=False,
265479
)
266480

481+
267482
def downgrade() -> None:
268483
"""Downgrade schema."""
269484
conn = op.get_bind()
@@ -304,7 +519,16 @@ def downgrade() -> None:
304519
# We'll preserve uniqueness by using the team_id/owner_email/name triple to later remap.
305520
conn.execute(
306521
text(
307-
"INSERT INTO prompts_old (name, description, template, argument_schema, created_at, updated_at, is_active, tags, created_by, created_from_ip, created_via, created_user_agent, modified_by, modified_from_ip, modified_via, modified_user_agent, import_batch_id, federation_source, version, gateway_id, team_id, owner_email, visibility) SELECT name, description, template, argument_schema, created_at, updated_at, enabled, tags, created_by, created_from_ip, created_via, created_user_agent, modified_by, modified_from_ip, modified_via, modified_user_agent, import_batch_id, federation_source, version, gateway_id, team_id, owner_email, visibility FROM prompts"
522+
(
523+
"INSERT INTO prompts_old (name, description, template, argument_schema, created_at, updated_at, "
524+
"is_active, tags, created_by, created_from_ip, created_via, created_user_agent, modified_by, "
525+
"modified_from_ip, modified_via, modified_user_agent, import_batch_id, federation_source, version, "
526+
"gateway_id, team_id, owner_email, visibility) "
527+
"SELECT name, description, template, argument_schema, created_at, updated_at, enabled, tags, "
528+
"created_by, created_from_ip, created_via, created_user_agent, modified_by, modified_from_ip, "
529+
"modified_via, modified_user_agent, import_batch_id, federation_source, version, gateway_id, "
530+
"team_id, owner_email, visibility FROM prompts"
531+
)
308532
)
309533
)
310534

@@ -426,7 +650,16 @@ def downgrade() -> None:
426650
# 2) Insert rows from current resources into resources_old letting id autoincrement.
427651
conn.execute(
428652
text(
429-
"INSERT INTO resources_old (uri, name, description, mime_type, size, uri_template, created_at, updated_at, is_active, tags, text_content, binary_content, created_by, created_from_ip, created_via, created_user_agent, modified_by, modified_from_ip, modified_via, modified_user_agent, import_batch_id, federation_source, version, gateway_id, team_id, owner_email, visibility) SELECT uri, name, description, mime_type, size, uri_template, created_at, updated_at, enabled, tags, text_content, binary_content, created_by, created_from_ip, created_via, created_user_agent, modified_by, modified_from_ip, modified_via, modified_user_agent, import_batch_id, federation_source, version, gateway_id, team_id, owner_email, visibility FROM resources"
653+
(
654+
"INSERT INTO resources_old (uri, name, description, mime_type, size, uri_template, created_at, "
655+
"updated_at, is_active, tags, text_content, binary_content, created_by, created_from_ip, "
656+
"created_via, created_user_agent, modified_by, modified_from_ip, modified_via, modified_user_agent, "
657+
"import_batch_id, federation_source, version, gateway_id, team_id, owner_email, visibility) "
658+
"SELECT uri, name, description, mime_type, size, uri_template, created_at, updated_at, enabled, tags, "
659+
"text_content, binary_content, created_by, created_from_ip, created_via, created_user_agent, modified_by, "
660+
"modified_from_ip, modified_via, modified_user_agent, import_batch_id, federation_source, version, gateway_id, "
661+
"team_id, owner_email, visibility FROM resources"
662+
)
430663
)
431664
)
432665

0 commit comments

Comments
 (0)