|
| 1 | +# -*- coding: utf-8 -*- |
| 2 | +"""Location: ./mcpgateway/alembic/versions/f3a3a3d901b8_remove_gateway_url_unique_constraint.py |
| 3 | +Copyright 2025 |
| 4 | +SPDX-License-Identifier: Apache-2.0 |
| 5 | +Authors: Keval Mahajan |
| 6 | +
|
| 7 | +Alembic migration to remove unique constraint on gateway URL. |
| 8 | +An improved alternative duplication check has been implemented for gateway duplication prevention. |
| 9 | +
|
| 10 | +Revision ID: f3a3a3d901b8 |
| 11 | +Revises: aac21d6f9522 |
| 12 | +Create Date: 2025-11-11 22:30:05.474282 |
| 13 | +
|
| 14 | +""" |
| 15 | + |
| 16 | +# Standard |
| 17 | +from typing import Sequence, Union |
| 18 | + |
| 19 | +# Third-Party |
| 20 | +from alembic import op |
| 21 | +from sqlalchemy.engine import Inspector |
| 22 | + |
| 23 | +# revision identifiers, used by Alembic. |
| 24 | +revision: str = "f3a3a3d901b8" |
| 25 | +down_revision: Union[str, Sequence[str], None] = "aac21d6f9522" |
| 26 | +branch_labels: Union[str, Sequence[str], None] = None |
| 27 | +depends_on: Union[str, Sequence[str], None] = None |
| 28 | + |
| 29 | + |
| 30 | +def constraint_exists(inspector, table_name, constraint_name): |
| 31 | + """ |
| 32 | + Check if a specific unique constraint exists on a given table. |
| 33 | +
|
| 34 | + This function queries the database using the provided SQLAlchemy |
| 35 | + inspector to determine if a constraint with the given name exists. |
| 36 | + If the check fails due to an exception (e.g., database connectivity issues), |
| 37 | + it conservatively assumes that the constraint exists. |
| 38 | +
|
| 39 | + Args: |
| 40 | + inspector (sqlalchemy.engine.reflection.Inspector): SQLAlchemy inspector |
| 41 | + instance for database introspection. |
| 42 | + table_name (str): Name of the table to inspect. |
| 43 | + constraint_name (str): Name of the unique constraint to check. |
| 44 | +
|
| 45 | + Returns: |
| 46 | + bool: True if the constraint exists or if the check could not be performed, |
| 47 | + False if the constraint does not exist. |
| 48 | + """ |
| 49 | + try: |
| 50 | + unique_constraints = inspector.get_unique_constraints(table_name) |
| 51 | + return any(uc["name"] == constraint_name for uc in unique_constraints) |
| 52 | + except Exception: |
| 53 | + # Fallback: assume constraint exists if we can't check |
| 54 | + return True |
| 55 | + |
| 56 | + |
| 57 | +def upgrade(): |
| 58 | + """Remove the unique constraint on (team_id, owner_email, url) from gateway table.""" |
| 59 | + |
| 60 | + conn = op.get_bind() |
| 61 | + inspector = Inspector.from_engine(conn) |
| 62 | + |
| 63 | + # Check if constraint exists before attempting to drop |
| 64 | + if not constraint_exists(inspector, "gateways", "uq_team_owner_url_gateway"): |
| 65 | + print("Constraint 'uq_team_owner_url_gateway' does not exist, skipping drop.") |
| 66 | + return |
| 67 | + |
| 68 | + if conn.dialect.name == "sqlite": |
| 69 | + # SQLite: Use batch mode to recreate table without the constraint |
| 70 | + with op.batch_alter_table("gateways", schema=None) as batch_op: |
| 71 | + batch_op.drop_constraint("uq_team_owner_url_gateway", type_="unique") |
| 72 | + else: |
| 73 | + # PostgreSQL, MySQL, etc.: Direct constraint drop |
| 74 | + op.drop_constraint("uq_team_owner_url_gateway", "gateways", type_="unique") |
| 75 | + |
| 76 | + print("Successfully removed constraint 'uq_team_owner_url_gateway' from gateway table.") |
| 77 | + |
| 78 | + |
| 79 | +def downgrade(): |
| 80 | + """Re-add the unique constraint on (team_id, owner_email, url) to gateway table.""" |
| 81 | + |
| 82 | + conn = op.get_bind() |
| 83 | + inspector = Inspector.from_engine(conn) |
| 84 | + |
| 85 | + # Check if constraint already exists before attempting to create |
| 86 | + if constraint_exists(inspector, "gateways", "uq_team_owner_url_gateway"): |
| 87 | + print("Constraint 'uq_team_owner_url_gateway' already exists, skipping creation.") |
| 88 | + return |
| 89 | + |
| 90 | + if conn.dialect.name == "sqlite": |
| 91 | + # SQLite: Use batch mode to recreate table with the constraint |
| 92 | + with op.batch_alter_table("gateways", schema=None) as batch_op: |
| 93 | + batch_op.create_unique_constraint("uq_team_owner_url_gateway", ["team_id", "owner_email", "url"]) |
| 94 | + else: |
| 95 | + # PostgreSQL, MySQL, etc.: Direct constraint creation |
| 96 | + op.create_unique_constraint("uq_team_owner_url_constraint", "gateways", ["team_id", "owner_email", "url"]) |
| 97 | + |
| 98 | + print("Successfully re-added constraint 'uq_team_owner_url_gateway' to gateways table.") |
0 commit comments