fix: Downgrade of revision 678eefb4ab44 throws error (apache#29799)

preset-io · Aug 1, 2024 · 249f5ec · 249f5ec
1 parent 2cbd945
commit 249f5ec
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 0 deletions.
diff --git 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
@@ -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")