diff --git a/superset/migrations/shared/constraints.py b/superset/migrations/shared/constraints.py index d3bc140fd0bab..ab80536dbc1f5 100644 --- a/superset/migrations/shared/constraints.py +++ b/superset/migrations/shared/constraints.py @@ -71,3 +71,17 @@ def redefine( ondelete=on_delete, onupdate=on_update, ) + + +def drop_fks_for_table(table_name: str) -> None: + """ + Drop all foreign key constraints for a table. + + :param table_name: The table name to drop foreign key constraints for + """ + connection = op.get_bind() + inspector = Inspector.from_engine(connection) + foreign_keys = inspector.get_foreign_keys(table_name) + + for fk in foreign_keys: + op.drop_constraint(fk["name"], table_name, type_="foreignkey") diff --git a/superset/migrations/versions/2024-03-20_16-02_678eefb4ab44_add_access_token_table.py b/superset/migrations/versions/2024-03-20_16-02_678eefb4ab44_add_access_token_table.py index ff3f06f077e6b..0e92884e58c99 100644 --- a/superset/migrations/versions/2024-03-20_16-02_678eefb4ab44_add_access_token_table.py +++ b/superset/migrations/versions/2024-03-20_16-02_678eefb4ab44_add_access_token_table.py @@ -30,6 +30,8 @@ from alembic import op # noqa: E402 from sqlalchemy_utils import EncryptedType # noqa: E402 +from superset.migrations.shared.constraints import drop_fks_for_table # noqa: E402 + def upgrade(): op.create_table( @@ -80,5 +82,6 @@ def upgrade(): def downgrade(): + drop_fks_for_table("database_user_oauth2_tokens") op.drop_index("idx_user_id_database_id", table_name="database_user_oauth2_tokens") op.drop_table("database_user_oauth2_tokens")