Coverage for mcpgateway/alembic/versions/b77ca9d2de7e_uuid_pk_and_slug_refactor.py: 7%

175 statements  

« prev     ^ index     » next       coverage.py v7.9.2, created at 2025-07-09 11:03 +0100

1# -*- coding: utf-8 -*- 

2"""uuid-pk_and_slug_refactor 

3 

4Revision ID: b77ca9d2de7e 

5Revises: 

6Create Date: 2025-06-26 21:29:59.117140 

7 

8""" 

9 

10# Standard 

11from typing import Sequence, Union 

12import uuid 

13 

14# Third-Party 

15import sqlalchemy as sa 

16from sqlalchemy.orm import Session 

17 

18# First-Party 

19from alembic import op 

20from mcpgateway.config import settings 

21from mcpgateway.utils.create_slug import slugify 

22 

23# revision identifiers, used by Alembic. 

24revision: str = "b77ca9d2de7e" 

25down_revision: Union[str, Sequence[str], None] = None 

26branch_labels: Union[str, Sequence[str], None] = None 

27depends_on: Union[str, Sequence[str], None] = None 

28 

29 

30# ────────────────────────────────────────────────────────────────────────────── 

31# Helpers 

32# ────────────────────────────────────────────────────────────────────────────── 

33def _use_batch() -> bool: 

34 return op.get_bind().dialect.name == "sqlite" 

35 

36 

37# ────────────────────────────────────────────────────────────────────────────── 

38# Upgrade 

39# ────────────────────────────────────────────────────────────────────────────── 

40def upgrade() -> None: 

41 bind = op.get_bind() 

42 sess = Session(bind=bind) 

43 inspector = sa.inspect(bind) 

44 

45 if not inspector.has_table("gateways"): 

46 print("Fresh database detected. Skipping migration.") 

47 return 

48 

49 print("Existing installation detected. Starting data and schema migration...") 

50 

51 # ── STAGE 1: ADD NEW NULLABLE COLUMNS AS PLACEHOLDERS ───────────────── 

52 op.add_column("gateways", sa.Column("slug", sa.String(), nullable=True)) 

53 op.add_column("gateways", sa.Column("id_new", sa.String(36), nullable=True)) 

54 

55 op.add_column("tools", sa.Column("id_new", sa.String(36), nullable=True)) 

56 op.add_column("tools", sa.Column("original_name", sa.String(), nullable=True)) 

57 op.add_column("tools", sa.Column("original_name_slug", sa.String(), nullable=True)) 

58 op.add_column("tools", sa.Column("name_new", sa.String(), nullable=True)) 

59 op.add_column("tools", sa.Column("gateway_id_new", sa.String(36), nullable=True)) 

60 

61 op.add_column("resources", sa.Column("gateway_id_new", sa.String(36), nullable=True)) 

62 op.add_column("prompts", sa.Column("gateway_id_new", sa.String(36), nullable=True)) 

63 

64 op.add_column("servers", sa.Column("id_new", sa.String(36), nullable=True)) 

65 

66 op.add_column("server_tool_association", sa.Column("server_id_new", sa.String(36), nullable=True)) 

67 op.add_column("server_tool_association", sa.Column("tool_id_new", sa.String(36), nullable=True)) 

68 

69 op.add_column("tool_metrics", sa.Column("tool_id_new", sa.String(36), nullable=True)) 

70 op.add_column("server_metrics", sa.Column("server_id_new", sa.String(36), nullable=True)) 

71 op.add_column("server_resource_association", sa.Column("server_id_new", sa.String(36), nullable=True)) 

72 op.add_column("server_prompt_association", sa.Column("server_id_new", sa.String(36), nullable=True)) 

73 

74 # ── STAGE 2: POPULATE THE NEW COLUMNS (DATA MIGRATION) ─────────────── 

75 gateways = sess.execute(sa.select(sa.text("id, name")).select_from(sa.text("gateways"))).all() 

76 for gid, gname in gateways: 

77 g_uuid = uuid.uuid4().hex 

78 sess.execute( 

79 sa.text("UPDATE gateways SET id_new=:u, slug=:s WHERE id=:i"), 

80 {"u": g_uuid, "s": slugify(gname), "i": gid}, 

81 ) 

82 

83 tools = sess.execute(sa.select(sa.text("id, name, gateway_id")).select_from(sa.text("tools"))).all() 

84 for tid, tname, g_old in tools: 

85 t_uuid = uuid.uuid4().hex 

86 tool_slug = slugify(tname) 

87 sess.execute( 

88 sa.text( 

89 """ 

90 UPDATE tools 

91 SET id_new=:u, 

92 original_name=:on, 

93 original_name_slug=:ons, 

94 name_new = CASE 

95 WHEN :g IS NOT NULL THEN (SELECT slug FROM gateways WHERE id = :g) || :sep || :ons 

96 ELSE :ons 

97 END, 

98 gateway_id_new=(SELECT id_new FROM gateways WHERE id=:g) 

99 WHERE id=:i 

100 """ 

101 ), 

102 { 

103 "u": t_uuid, 

104 "on": tname, 

105 "ons": tool_slug, 

106 "sep": settings.gateway_tool_name_separator, 

107 "g": g_old, 

108 "i": tid, 

109 }, 

110 ) 

111 

112 servers = sess.execute(sa.select(sa.text("id")).select_from(sa.text("servers"))).all() 

113 for (sid,) in servers: 

114 sess.execute( 

115 sa.text("UPDATE servers SET id_new=:u WHERE id=:i"), 

116 {"u": uuid.uuid4().hex, "i": sid}, 

117 ) 

118 

119 # Populate all dependent tables 

120 resources = sess.execute(sa.select(sa.text("id, gateway_id")).select_from(sa.text("resources"))).all() 

121 for rid, g_old in resources: 

122 sess.execute(sa.text("UPDATE resources SET gateway_id_new=(SELECT id_new FROM gateways WHERE id=:g) WHERE id=:i"), {"g": g_old, "i": rid}) 

123 prompts = sess.execute(sa.select(sa.text("id, gateway_id")).select_from(sa.text("prompts"))).all() 

124 for pid, g_old in prompts: 

125 sess.execute(sa.text("UPDATE prompts SET gateway_id_new=(SELECT id_new FROM gateways WHERE id=:g) WHERE id=:i"), {"g": g_old, "i": pid}) 

126 sta = sess.execute(sa.select(sa.text("server_id, tool_id")).select_from(sa.text("server_tool_association"))).all() 

127 for s_old, t_old in sta: 

128 sess.execute( 

129 sa.text("UPDATE server_tool_association SET server_id_new=(SELECT id_new FROM servers WHERE id=:s), tool_id_new=(SELECT id_new FROM tools WHERE id=:t) WHERE server_id=:s AND tool_id=:t"), 

130 {"s": s_old, "t": t_old}, 

131 ) 

132 tool_metrics = sess.execute(sa.select(sa.text("id, tool_id")).select_from(sa.text("tool_metrics"))).all() 

133 for tmid, t_old in tool_metrics: 

134 sess.execute(sa.text("UPDATE tool_metrics SET tool_id_new=(SELECT id_new FROM tools WHERE id=:t) WHERE id=:i"), {"t": t_old, "i": tmid}) 

135 server_metrics = sess.execute(sa.select(sa.text("id, server_id")).select_from(sa.text("server_metrics"))).all() 

136 for smid, s_old in server_metrics: 

137 sess.execute(sa.text("UPDATE server_metrics SET server_id_new=(SELECT id_new FROM servers WHERE id=:s) WHERE id=:i"), {"s": s_old, "i": smid}) 

138 server_resource_assoc = sess.execute(sa.select(sa.text("server_id, resource_id")).select_from(sa.text("server_resource_association"))).all() 

139 for s_old, r_id in server_resource_assoc: 

140 sess.execute(sa.text("UPDATE server_resource_association SET server_id_new=(SELECT id_new FROM servers WHERE id=:s) WHERE server_id=:s AND resource_id=:r"), {"s": s_old, "r": r_id}) 

141 server_prompt_assoc = sess.execute(sa.select(sa.text("server_id, prompt_id")).select_from(sa.text("server_prompt_association"))).all() 

142 for s_old, p_id in server_prompt_assoc: 

143 sess.execute(sa.text("UPDATE server_prompt_association SET server_id_new=(SELECT id_new FROM servers WHERE id=:s) WHERE server_id=:s AND prompt_id=:p"), {"s": s_old, "p": p_id}) 

144 

145 sess.commit() 

146 

147 # ── STAGE 3: FINALIZE SCHEMA (CORRECTED ORDER) ─────────────────────── 

148 # First, rebuild all tables that depend on `servers` and `gateways`. 

149 # This implicitly drops their old foreign key constraints. 

150 with op.batch_alter_table("server_tool_association") as batch_op: 

151 batch_op.drop_column("server_id") 

152 batch_op.drop_column("tool_id") 

153 batch_op.alter_column("server_id_new", new_column_name="server_id", nullable=False) 

154 batch_op.alter_column("tool_id_new", new_column_name="tool_id", nullable=False) 

155 batch_op.create_primary_key("pk_server_tool_association", ["server_id", "tool_id"]) 

156 

157 with op.batch_alter_table("server_resource_association") as batch_op: 

158 batch_op.drop_column("server_id") 

159 batch_op.alter_column("server_id_new", new_column_name="server_id", nullable=False) 

160 

161 with op.batch_alter_table("server_prompt_association") as batch_op: 

162 batch_op.drop_column("server_id") 

163 batch_op.alter_column("server_id_new", new_column_name="server_id", nullable=False) 

164 

165 with op.batch_alter_table("server_metrics") as batch_op: 

166 batch_op.drop_column("server_id") 

167 batch_op.alter_column("server_id_new", new_column_name="server_id", nullable=False) 

168 

169 with op.batch_alter_table("tool_metrics") as batch_op: 

170 batch_op.drop_column("tool_id") 

171 batch_op.alter_column("tool_id_new", new_column_name="tool_id", nullable=False) 

172 

173 with op.batch_alter_table("tools") as batch_op: 

174 batch_op.drop_column("id") 

175 batch_op.alter_column("id_new", new_column_name="id", nullable=False) 

176 batch_op.create_primary_key("pk_tools", ["id"]) 

177 batch_op.drop_column("gateway_id") 

178 batch_op.alter_column("gateway_id_new", new_column_name="gateway_id", nullable=True) 

179 batch_op.drop_column("name") 

180 batch_op.alter_column("name_new", new_column_name="name", nullable=True) 

181 batch_op.alter_column("original_name", nullable=False) 

182 batch_op.alter_column("original_name_slug", nullable=False) 

183 batch_op.create_unique_constraint("uq_tools_name", ["name"]) 

184 batch_op.create_unique_constraint("uq_gateway_id__original_name", ["gateway_id", "original_name"]) 

185 

186 with op.batch_alter_table("resources") as batch_op: 

187 batch_op.drop_column("gateway_id") 

188 batch_op.alter_column("gateway_id_new", new_column_name="gateway_id", nullable=True) 

189 

190 with op.batch_alter_table("prompts") as batch_op: 

191 batch_op.drop_column("gateway_id") 

192 batch_op.alter_column("gateway_id_new", new_column_name="gateway_id", nullable=True) 

193 

194 # Second, now that no tables point to their old IDs, rebuild `gateways` and `servers`. 

195 with op.batch_alter_table("gateways") as batch_op: 

196 batch_op.drop_column("id") 

197 batch_op.alter_column("id_new", new_column_name="id", nullable=False) 

198 batch_op.create_primary_key("pk_gateways", ["id"]) 

199 batch_op.alter_column("slug", nullable=False) 

200 batch_op.create_unique_constraint("uq_gateways_slug", ["slug"]) 

201 batch_op.create_unique_constraint("uq_gateways_url", ["url"]) 

202 

203 with op.batch_alter_table("servers") as batch_op: 

204 batch_op.drop_column("id") 

205 batch_op.alter_column("id_new", new_column_name="id", nullable=False) 

206 batch_op.create_primary_key("pk_servers", ["id"]) 

207 

208 # Finally, recreate all the foreign key constraints in batch mode for SQLite compatibility. 

209 # The redundant `source_table` argument has been removed from each call. 

210 with op.batch_alter_table("tools") as batch_op: 

211 batch_op.create_foreign_key("fk_tools_gateway_id", "gateways", ["gateway_id"], ["id"]) 

212 with op.batch_alter_table("resources") as batch_op: 

213 batch_op.create_foreign_key("fk_resources_gateway_id", "gateways", ["gateway_id"], ["id"]) 

214 with op.batch_alter_table("prompts") as batch_op: 

215 batch_op.create_foreign_key("fk_prompts_gateway_id", "gateways", ["gateway_id"], ["id"]) 

216 with op.batch_alter_table("server_tool_association") as batch_op: 

217 batch_op.create_foreign_key("fk_server_tool_association_servers", "servers", ["server_id"], ["id"]) 

218 batch_op.create_foreign_key("fk_server_tool_association_tools", "tools", ["tool_id"], ["id"]) 

219 with op.batch_alter_table("tool_metrics") as batch_op: 

220 batch_op.create_foreign_key("fk_tool_metrics_tool_id", "tools", ["tool_id"], ["id"]) 

221 with op.batch_alter_table("server_metrics") as batch_op: 

222 batch_op.create_foreign_key("fk_server_metrics_server_id", "servers", ["server_id"], ["id"]) 

223 with op.batch_alter_table("server_resource_association") as batch_op: 

224 batch_op.create_foreign_key("fk_server_resource_association_server_id", "servers", ["server_id"], ["id"]) 

225 with op.batch_alter_table("server_prompt_association") as batch_op: 

226 batch_op.create_foreign_key("fk_server_prompt_association_server_id", "servers", ["server_id"], ["id"]) 

227 

228 

229# def upgrade() -> None: 

230# bind = op.get_bind() 

231# sess = Session(bind=bind) 

232# inspector = sa.inspect(bind) 

233 

234# if not inspector.has_table("gateways"): 

235# print("Fresh database detected. Skipping migration.") 

236# return 

237 

238# print("Existing installation detected. Starting data and schema migration...") 

239 

240# # ── STAGE 1: ADD NEW NULLABLE COLUMNS AS PLACEHOLDERS ───────────────── 

241# op.add_column("gateways", sa.Column("slug", sa.String(), nullable=True)) 

242# op.add_column("gateways", sa.Column("id_new", sa.String(36), nullable=True)) 

243 

244# op.add_column("tools", sa.Column("id_new", sa.String(36), nullable=True)) 

245# op.add_column("tools", sa.Column("original_name", sa.String(), nullable=True)) 

246# op.add_column("tools", sa.Column("original_name_slug", sa.String(), nullable=True)) 

247# op.add_column("tools", sa.Column("name_new", sa.String(), nullable=True)) 

248# op.add_column("tools", sa.Column("gateway_id_new", sa.String(36), nullable=True)) 

249 

250# op.add_column("resources", sa.Column("gateway_id_new", sa.String(36), nullable=True)) 

251# op.add_column("prompts", sa.Column("gateway_id_new", sa.String(36), nullable=True)) 

252 

253# op.add_column("servers", sa.Column("id_new", sa.String(36), nullable=True)) 

254 

255# op.add_column("server_tool_association", sa.Column("server_id_new", sa.String(36), nullable=True)) 

256# op.add_column("server_tool_association", sa.Column("tool_id_new", sa.String(36), nullable=True)) 

257 

258# op.add_column("tool_metrics", sa.Column("tool_id_new", sa.String(36), nullable=True)) 

259 

260# # Add columns for the new server dependencies 

261# op.add_column("server_metrics", sa.Column("server_id_new", sa.String(36), nullable=True)) 

262# op.add_column("server_resource_association", sa.Column("server_id_new", sa.String(36), nullable=True)) 

263# op.add_column("server_prompt_association", sa.Column("server_id_new", sa.String(36), nullable=True)) 

264 

265 

266# # ── STAGE 2: POPULATE THE NEW COLUMNS (DATA MIGRATION) ─────────────── 

267# gateways = sess.execute(sa.select(sa.text("id, name")).select_from(sa.text("gateways"))).all() 

268# for gid, gname in gateways: 

269# g_uuid = uuid.uuid4().hex 

270# sess.execute( 

271# sa.text("UPDATE gateways SET id_new=:u, slug=:s WHERE id=:i"), 

272# {"u": g_uuid, "s": slugify(gname), "i": gid}, 

273# ) 

274 

275# tools = sess.execute( 

276# sa.select(sa.text("id, name, gateway_id")).select_from(sa.text("tools")) 

277# ).all() 

278# for tid, tname, g_old in tools: 

279# t_uuid = uuid.uuid4().hex 

280# tool_slug = slugify(tname) 

281# sess.execute( 

282# sa.text( 

283# """ 

284# UPDATE tools 

285# SET id_new=:u, 

286# original_name=:on, 

287# original_name_slug=:ons, 

288# name_new = CASE 

289# WHEN :g IS NOT NULL THEN (SELECT slug FROM gateways WHERE id = :g) || :sep || :ons 

290# ELSE :ons 

291# END, 

292# gateway_id_new=(SELECT id_new FROM gateways WHERE id=:g) 

293# WHERE id=:i 

294# """ 

295# ), 

296# { 

297# "u": t_uuid, "on": tname, "ons": tool_slug, 

298# "sep": settings.gateway_tool_name_separator, "g": g_old, "i": tid, 

299# }, 

300# ) 

301 

302# servers = sess.execute(sa.select(sa.text("id")).select_from(sa.text("servers"))).all() 

303# for (sid,) in servers: 

304# sess.execute( 

305# sa.text("UPDATE servers SET id_new=:u WHERE id=:i"), 

306# {"u": uuid.uuid4().hex, "i": sid}, 

307# ) 

308 

309# # Populate all dependent tables 

310# resources = sess.execute(sa.select(sa.text("id, gateway_id")).select_from(sa.text("resources"))).all() 

311# for rid, g_old in resources: 

312# sess.execute(sa.text("UPDATE resources SET gateway_id_new=(SELECT id_new FROM gateways WHERE id=:g) WHERE id=:i"), {"g": g_old, "i": rid}) 

313# prompts = sess.execute(sa.select(sa.text("id, gateway_id")).select_from(sa.text("prompts"))).all() 

314# for pid, g_old in prompts: 

315# sess.execute(sa.text("UPDATE prompts SET gateway_id_new=(SELECT id_new FROM gateways WHERE id=:g) WHERE id=:i"), {"g": g_old, "i": pid}) 

316# sta = sess.execute(sa.select(sa.text("server_id, tool_id")).select_from(sa.text("server_tool_association"))).all() 

317# for s_old, t_old in sta: 

318# sess.execute(sa.text("UPDATE server_tool_association SET server_id_new=(SELECT id_new FROM servers WHERE id=:s), tool_id_new=(SELECT id_new FROM tools WHERE id=:t) WHERE server_id=:s AND tool_id=:t"), {"s": s_old, "t": t_old}) 

319# tool_metrics = sess.execute(sa.select(sa.text("id, tool_id")).select_from(sa.text("tool_metrics"))).all() 

320# for tmid, t_old in tool_metrics: 

321# sess.execute(sa.text("UPDATE tool_metrics SET tool_id_new=(SELECT id_new FROM tools WHERE id=:t) WHERE id=:i"), {"t": t_old, "i": tmid}) 

322# server_metrics = sess.execute(sa.select(sa.text("id, server_id")).select_from(sa.text("server_metrics"))).all() 

323# for smid, s_old in server_metrics: 

324# sess.execute(sa.text("UPDATE server_metrics SET server_id_new=(SELECT id_new FROM servers WHERE id=:s) WHERE id=:i"), {"s": s_old, "i": smid}) 

325# server_resource_assoc = sess.execute(sa.select(sa.text("server_id, resource_id")).select_from(sa.text("server_resource_association"))).all() 

326# for s_old, r_id in server_resource_assoc: 

327# sess.execute(sa.text("UPDATE server_resource_association SET server_id_new=(SELECT id_new FROM servers WHERE id=:s) WHERE server_id=:s AND resource_id=:r"), {"s": s_old, "r": r_id}) 

328# server_prompt_assoc = sess.execute(sa.select(sa.text("server_id, prompt_id")).select_from(sa.text("server_prompt_association"))).all() 

329# for s_old, p_id in server_prompt_assoc: 

330# sess.execute(sa.text("UPDATE server_prompt_association SET server_id_new=(SELECT id_new FROM servers WHERE id=:s) WHERE server_id=:s AND prompt_id=:p"), {"s": s_old, "p": p_id}) 

331 

332# sess.commit() 

333 

334# # ── STAGE 3: FINALIZE SCHEMA (CORRECTED ORDER) ─────────────────────── 

335# with op.batch_alter_table("server_tool_association") as batch_op: 

336# batch_op.drop_column("server_id") 

337# batch_op.drop_column("tool_id") 

338# batch_op.alter_column("server_id_new", new_column_name="server_id", nullable=False) 

339# batch_op.alter_column("tool_id_new", new_column_name="tool_id", nullable=False) 

340# batch_op.create_primary_key("pk_server_tool_association", ["server_id", "tool_id"]) 

341 

342# with op.batch_alter_table("server_resource_association") as batch_op: 

343# batch_op.drop_column("server_id") 

344# batch_op.alter_column("server_id_new", new_column_name="server_id", nullable=False) 

345 

346# with op.batch_alter_table("server_prompt_association") as batch_op: 

347# batch_op.drop_column("server_id") 

348# batch_op.alter_column("server_id_new", new_column_name="server_id", nullable=False) 

349 

350# with op.batch_alter_table("server_metrics") as batch_op: 

351# batch_op.drop_column("server_id") 

352# batch_op.alter_column("server_id_new", new_column_name="server_id", nullable=False) 

353 

354# with op.batch_alter_table("tool_metrics") as batch_op: 

355# batch_op.drop_column("tool_id") 

356# batch_op.alter_column("tool_id_new", new_column_name="tool_id", nullable=False) 

357 

358# with op.batch_alter_table("tools") as batch_op: 

359# batch_op.drop_column("id") 

360# batch_op.alter_column("id_new", new_column_name="id", nullable=False) 

361# batch_op.create_primary_key("pk_tools", ["id"]) 

362# batch_op.drop_column("gateway_id") 

363# batch_op.alter_column("gateway_id_new", new_column_name="gateway_id", nullable=True) 

364# batch_op.drop_column("name") 

365# batch_op.alter_column("name_new", new_column_name="name", nullable=False) 

366# batch_op.alter_column("original_name", nullable=False) 

367# batch_op.alter_column("original_name_slug", nullable=False) 

368# batch_op.create_unique_constraint("uq_tools_name", ["name"]) 

369# batch_op.create_unique_constraint("uq_gateway_id__original_name", ["gateway_id", "original_name"]) 

370 

371# with op.batch_alter_table("resources") as batch_op: 

372# batch_op.drop_column("gateway_id") 

373# batch_op.alter_column("gateway_id_new", new_column_name="gateway_id", nullable=True) 

374 

375# with op.batch_alter_table("prompts") as batch_op: 

376# batch_op.drop_column("gateway_id") 

377# batch_op.alter_column("gateway_id_new", new_column_name="gateway_id", nullable=True) 

378 

379# with op.batch_alter_table("gateways") as batch_op: 

380# batch_op.drop_column("id") 

381# batch_op.alter_column("id_new", new_column_name="id", nullable=False) 

382# batch_op.create_primary_key("pk_gateways", ["id"]) 

383# batch_op.alter_column("slug", nullable=False) 

384# batch_op.create_unique_constraint("uq_gateways_slug", ["slug"]) 

385# batch_op.create_unique_constraint("uq_gateways_url", ["url"]) 

386 

387# with op.batch_alter_table("servers") as batch_op: 

388# batch_op.drop_column("id") 

389# batch_op.alter_column("id_new", new_column_name="id", nullable=False) 

390# batch_op.create_primary_key("pk_servers", ["id"]) 

391 

392# # Finally, recreate all the foreign key constraints 

393# op.create_foreign_key("fk_tools_gateway_id", "tools", "gateways", ["gateway_id"], ["id"]) 

394# op.create_foreign_key("fk_resources_gateway_id", "resources", "gateways", ["gateway_id"], ["id"]) 

395# op.create_foreign_key("fk_prompts_gateway_id", "prompts", "gateways", ["gateway_id"], ["id"]) 

396# op.create_foreign_key("fk_server_tool_association_servers", "server_tool_association", "servers", ["server_id"], ["id"]) 

397# op.create_foreign_key("fk_server_tool_association_tools", "server_tool_association", "tools", ["tool_id"], ["id"]) 

398# op.create_foreign_key("fk_tool_metrics_tool_id", "tool_metrics", "tools", ["tool_id"], ["id"]) 

399# op.create_foreign_key("fk_server_metrics_server_id", "server_metrics", "servers", ["server_id"], ["id"]) 

400# op.create_foreign_key("fk_server_resource_association_server_id", "server_resource_association", "servers", ["server_id"], ["id"]) 

401# op.create_foreign_key("fk_server_prompt_association_server_id", "server_prompt_association", "servers", ["server_id"], ["id"]) 

402 

403 

404def downgrade() -> None: 

405 # ── STAGE 1 (REVERSE): Revert Schema to original state ───────────────── 

406 # This reverses the operations from STAGE 3 of the upgrade. 

407 # Data from the new columns will be lost, which is expected. 

408 

409 with op.batch_alter_table("server_tool_association") as batch_op: 

410 # Drop new constraints 

411 batch_op.drop_constraint("fk_server_tool_association_tools", type_="foreignkey") 

412 batch_op.drop_constraint("fk_server_tool_association_servers", type_="foreignkey") 

413 batch_op.drop_constraint("pk_server_tool_association", type_="primarykey") 

414 # Rename final columns back to temporary names 

415 batch_op.alter_column("server_id", new_column_name="server_id_new") 

416 batch_op.alter_column("tool_id", new_column_name="tool_id_new") 

417 # Add back old integer columns (data is not restored) 

418 batch_op.add_column(sa.Column("server_id", sa.Integer(), nullable=True)) 

419 batch_op.add_column(sa.Column("tool_id", sa.Integer(), nullable=True)) 

420 

421 with op.batch_alter_table("tools") as batch_op: 

422 # Drop new constraints 

423 batch_op.drop_constraint("fk_tools_gateway_id", type_="foreignkey") 

424 batch_op.drop_constraint("uq_gateway_id__original_name", type_="unique") 

425 batch_op.drop_constraint("uq_tools_name", type_="unique") 

426 batch_op.drop_constraint("pk_tools", type_="primarykey") 

427 # Rename final columns back to temporary names 

428 batch_op.alter_column("id", new_column_name="id_new") 

429 batch_op.alter_column("gateway_id", new_column_name="gateway_id_new") 

430 batch_op.alter_column("name", new_column_name="name_new") 

431 # Add back old columns 

432 batch_op.add_column(sa.Column("id", sa.Integer(), nullable=True)) 

433 batch_op.add_column(sa.Column("gateway_id", sa.Integer(), nullable=True)) 

434 batch_op.add_column(sa.Column("name", sa.String(), nullable=True)) 

435 

436 with op.batch_alter_table("servers") as batch_op: 

437 batch_op.drop_constraint("pk_servers", type_="primarykey") 

438 batch_op.alter_column("id", new_column_name="id_new") 

439 batch_op.add_column(sa.Column("id", sa.Integer(), nullable=True)) 

440 

441 with op.batch_alter_table("gateways") as batch_op: 

442 batch_op.drop_constraint("uq_gateways_url", type_="unique") 

443 batch_op.drop_constraint("uq_gateways_slug", type_="unique") 

444 batch_op.drop_constraint("pk_gateways", type_="primarykey") 

445 batch_op.alter_column("id", new_column_name="id_new") 

446 batch_op.add_column(sa.Column("id", sa.Integer(), nullable=True)) 

447 

448 # ── STAGE 2 (REVERSE): Reverse Data Migration (No-Op for Schema) ────── 

449 # Reversing the data population (e.g., creating integer PKs from UUIDs) 

450 # is a complex, stateful operation and is omitted here. At this point, 

451 # the original columns exist but are empty (NULL). 

452 

453 # ── STAGE 3 (REVERSE): Drop the temporary/new columns ──────────────── 

454 # This reverses the operations from STAGE 1 of the upgrade. 

455 op.drop_column("server_tool_association", "tool_id_new") 

456 op.drop_column("server_tool_association", "server_id_new") 

457 op.drop_column("servers", "id_new") 

458 op.drop_column("tools", "gateway_id_new") 

459 op.drop_column("tools", "name_new") 

460 op.drop_column("tools", "original_name_slug") 

461 op.drop_column("tools", "original_name") 

462 op.drop_column("tools", "id_new") 

463 op.drop_column("gateways", "id_new") 

464 op.drop_column("gateways", "slug")