From c929f5ed7aca41ece7453dda22abe99ed6b58fb5 Mon Sep 17 00:00:00 2001 From: Joe Li Date: Tue, 3 Sep 2024 15:58:39 -0700 Subject: [PATCH] fix: pass if table is already removed on upgrade (#30017) --- superset/migrations/shared/constraints.py | 16 ++++++++++++---- superset/migrations/shared/utils.py | 14 ++++++++++++++ ...f4f7811799_remove_sl_dataset_columns_table.py | 12 +++++++----- ...39549add7bfc_remove_sl_table_columns_table.py | 12 +++++++----- ...5-23_38f4144e8558_remove_sl_dataset_tables.py | 12 +++++++----- ...15-27_e53fd48cc078_remove_sl_dataset_users.py | 12 +++++++----- ...08-13_15-29_a6b32d2d07b1_remove_sl_columns.py | 12 +++++++----- ...-08-13_15-31_007a1abffe7e_remove_sl_tables.py | 12 +++++++----- ...8-13_15-33_48cbb571fa3a_remove_sl_datasets.py | 12 +++++++----- 9 files changed, 75 insertions(+), 39 deletions(-) diff --git a/superset/migrations/shared/constraints.py b/superset/migrations/shared/constraints.py index ab80536dbc1f5..5ae63f57ce611 100644 --- a/superset/migrations/shared/constraints.py +++ b/superset/migrations/shared/constraints.py @@ -19,8 +19,10 @@ from dataclasses import dataclass from alembic import op +from sqlalchemy.dialects.sqlite.base import SQLiteDialect # noqa: E402 from sqlalchemy.engine.reflection import Inspector +from superset.migrations.shared.utils import has_table from superset.utils.core import generic_find_fk_constraint_name @@ -75,13 +77,19 @@ def redefine( def drop_fks_for_table(table_name: str) -> None: """ - Drop all foreign key constraints for a table. + Drop all foreign key constraints for a table if it exist and the database + is not sqlite. :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") + if isinstance(connection.dialect, SQLiteDialect): + return # sqlite doesn't like constraints + + if has_table(table_name): + 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/shared/utils.py b/superset/migrations/shared/utils.py index 17e5b4128e46b..dba9ceef6393c 100644 --- a/superset/migrations/shared/utils.py +++ b/superset/migrations/shared/utils.py @@ -168,3 +168,17 @@ def try_load_json(data: Optional[str]) -> dict[str, Any]: except json.JSONDecodeError: print(f"Failed to parse: {data}") return {} + + +def has_table(table_name: str) -> bool: + """ + Check if a table exists in the database. + + :param table_name: The table name + :returns: True if the table exists + """ + + insp = inspect(op.get_context().bind) + table_exists = insp.has_table(table_name) + + return table_exists diff --git a/superset/migrations/versions/2024-05-24_11-31_02f4f7811799_remove_sl_dataset_columns_table.py b/superset/migrations/versions/2024-05-24_11-31_02f4f7811799_remove_sl_dataset_columns_table.py index f3c4f8c2c344b..03158fabd97b1 100644 --- a/superset/migrations/versions/2024-05-24_11-31_02f4f7811799_remove_sl_dataset_columns_table.py +++ b/superset/migrations/versions/2024-05-24_11-31_02f4f7811799_remove_sl_dataset_columns_table.py @@ -26,22 +26,24 @@ from alembic import op from superset.migrations.shared.constraints import drop_fks_for_table +from superset.migrations.shared.utils import has_table # revision identifiers, used by Alembic. revision = "02f4f7811799" down_revision = "f7b6750b67e8" +table_name = "sl_dataset_columns" + def upgrade(): - connection = op.get_bind() - if connection.dialect.name != "sqlite": - drop_fks_for_table("sl_dataset_columns") - op.drop_table("sl_dataset_columns") + if has_table(table_name): + drop_fks_for_table(table_name) + op.drop_table(table_name) def downgrade(): op.create_table( - "sl_dataset_columns", + table_name, sa.Column("dataset_id", sa.Integer(), nullable=False), sa.Column("column_id", sa.Integer(), nullable=False), sa.ForeignKeyConstraint( diff --git a/superset/migrations/versions/2024-08-13_15-17_39549add7bfc_remove_sl_table_columns_table.py b/superset/migrations/versions/2024-08-13_15-17_39549add7bfc_remove_sl_table_columns_table.py index 9fa8e8097b314..1ec41f5e1ca87 100644 --- a/superset/migrations/versions/2024-08-13_15-17_39549add7bfc_remove_sl_table_columns_table.py +++ b/superset/migrations/versions/2024-08-13_15-17_39549add7bfc_remove_sl_table_columns_table.py @@ -26,22 +26,24 @@ from alembic import op from superset.migrations.shared.constraints import drop_fks_for_table +from superset.migrations.shared.utils import has_table # revision identifiers, used by Alembic. revision = "39549add7bfc" down_revision = "02f4f7811799" +table_name = "sl_table_columns" + def upgrade(): - connection = op.get_bind() - if connection.dialect.name != "sqlite": - drop_fks_for_table("sl_table_columns") - op.drop_table("sl_table_columns") + if has_table(table_name): + drop_fks_for_table(table_name) + op.drop_table(table_name) def downgrade(): op.create_table( - "sl_table_columns", + table_name, sa.Column("table_id", sa.Integer(), nullable=False), sa.Column("column_id", sa.Integer(), nullable=False), sa.ForeignKeyConstraint( diff --git a/superset/migrations/versions/2024-08-13_15-23_38f4144e8558_remove_sl_dataset_tables.py b/superset/migrations/versions/2024-08-13_15-23_38f4144e8558_remove_sl_dataset_tables.py index 7f5b75395bf2a..4931456a33830 100644 --- a/superset/migrations/versions/2024-08-13_15-23_38f4144e8558_remove_sl_dataset_tables.py +++ b/superset/migrations/versions/2024-08-13_15-23_38f4144e8558_remove_sl_dataset_tables.py @@ -26,22 +26,24 @@ from alembic import op from superset.migrations.shared.constraints import drop_fks_for_table +from superset.migrations.shared.utils import has_table # revision identifiers, used by Alembic. revision = "38f4144e8558" down_revision = "39549add7bfc" +table_name = "sl_dataset_tables" + def upgrade(): - connection = op.get_bind() - if connection.dialect.name != "sqlite": - drop_fks_for_table("sl_dataset_tables") - op.drop_table("sl_dataset_tables") + if has_table(table_name): + drop_fks_for_table(table_name) + op.drop_table(table_name) def downgrade(): op.create_table( - "sl_dataset_tables", + table_name, sa.Column("dataset_id", sa.Integer(), nullable=False), sa.Column("table_id", sa.Integer(), nullable=False), sa.ForeignKeyConstraint( diff --git a/superset/migrations/versions/2024-08-13_15-27_e53fd48cc078_remove_sl_dataset_users.py b/superset/migrations/versions/2024-08-13_15-27_e53fd48cc078_remove_sl_dataset_users.py index e7e3881d318ad..359b05dcd4948 100644 --- a/superset/migrations/versions/2024-08-13_15-27_e53fd48cc078_remove_sl_dataset_users.py +++ b/superset/migrations/versions/2024-08-13_15-27_e53fd48cc078_remove_sl_dataset_users.py @@ -26,22 +26,24 @@ from alembic import op from superset.migrations.shared.constraints import drop_fks_for_table +from superset.migrations.shared.utils import has_table # revision identifiers, used by Alembic. revision = "e53fd48cc078" down_revision = "38f4144e8558" +table_name = "sl_dataset_users" + def upgrade(): - connection = op.get_bind() - if connection.dialect.name != "sqlite": - drop_fks_for_table("sl_dataset_users") - op.drop_table("sl_dataset_users") + if has_table(table_name): + drop_fks_for_table(table_name) + op.drop_table(table_name) def downgrade(): op.create_table( - "sl_dataset_users", + table_name, sa.Column("dataset_id", sa.Integer(), nullable=False), sa.Column("user_id", sa.Integer(), nullable=False), sa.ForeignKeyConstraint( diff --git a/superset/migrations/versions/2024-08-13_15-29_a6b32d2d07b1_remove_sl_columns.py b/superset/migrations/versions/2024-08-13_15-29_a6b32d2d07b1_remove_sl_columns.py index d5fbf860abaad..1562ed962a0c4 100644 --- a/superset/migrations/versions/2024-08-13_15-29_a6b32d2d07b1_remove_sl_columns.py +++ b/superset/migrations/versions/2024-08-13_15-29_a6b32d2d07b1_remove_sl_columns.py @@ -26,22 +26,24 @@ from alembic import op from superset.migrations.shared.constraints import drop_fks_for_table +from superset.migrations.shared.utils import has_table # revision identifiers, used by Alembic. revision = "a6b32d2d07b1" down_revision = "e53fd48cc078" +table_name = "sl_columns" + def upgrade(): - connection = op.get_bind() - if connection.dialect.name != "sqlite": - drop_fks_for_table("sl_columns") - op.drop_table("sl_columns") + if has_table(table_name): + drop_fks_for_table(table_name) + op.drop_table(table_name) def downgrade(): op.create_table( - "sl_columns", + table_name, sa.Column("uuid", sa.Numeric(precision=16), nullable=True), sa.Column("created_on", sa.DateTime(), nullable=True), sa.Column("changed_on", sa.DateTime(), nullable=True), diff --git a/superset/migrations/versions/2024-08-13_15-31_007a1abffe7e_remove_sl_tables.py b/superset/migrations/versions/2024-08-13_15-31_007a1abffe7e_remove_sl_tables.py index bf5982572fd72..106cd7a1704a0 100644 --- a/superset/migrations/versions/2024-08-13_15-31_007a1abffe7e_remove_sl_tables.py +++ b/superset/migrations/versions/2024-08-13_15-31_007a1abffe7e_remove_sl_tables.py @@ -26,22 +26,24 @@ from alembic import op from superset.migrations.shared.constraints import drop_fks_for_table +from superset.migrations.shared.utils import has_table # revision identifiers, used by Alembic. revision = "007a1abffe7e" down_revision = "a6b32d2d07b1" +table_name = "sl_tables" + def upgrade(): - connection = op.get_bind() - if connection.dialect.name != "sqlite": - drop_fks_for_table("sl_tables") - op.drop_table("sl_tables") + if has_table(table_name): + drop_fks_for_table(table_name) + op.drop_table(table_name) def downgrade(): op.create_table( - "sl_tables", + table_name, sa.Column("uuid", sa.Numeric(precision=16), nullable=True), sa.Column("created_on", sa.DateTime(), nullable=True), sa.Column("changed_on", sa.DateTime(), nullable=True), diff --git a/superset/migrations/versions/2024-08-13_15-33_48cbb571fa3a_remove_sl_datasets.py b/superset/migrations/versions/2024-08-13_15-33_48cbb571fa3a_remove_sl_datasets.py index 56598fc4d2228..c35d3d7572e56 100644 --- a/superset/migrations/versions/2024-08-13_15-33_48cbb571fa3a_remove_sl_datasets.py +++ b/superset/migrations/versions/2024-08-13_15-33_48cbb571fa3a_remove_sl_datasets.py @@ -26,22 +26,24 @@ from alembic import op from superset.migrations.shared.constraints import drop_fks_for_table +from superset.migrations.shared.utils import has_table # revision identifiers, used by Alembic. revision = "48cbb571fa3a" down_revision = "007a1abffe7e" +table_name = "sl_datasets" + def upgrade(): - connection = op.get_bind() - if connection.dialect.name != "sqlite": - drop_fks_for_table("sl_datasets") - op.drop_table("sl_datasets") + if has_table(table_name): + drop_fks_for_table(table_name) + op.drop_table(table_name) def downgrade(): op.create_table( - "sl_datasets", + table_name, sa.Column("uuid", sa.Numeric(precision=16), nullable=True), sa.Column("created_on", sa.DateTime(), nullable=True), sa.Column("changed_on", sa.DateTime(), nullable=True),