This is an automated email from the ASF dual-hosted git repository. elizabeth pushed a commit to branch elizabeth/index-migration-function in repository https://gitbox.apache.org/repos/asf/superset.git
commit f146e99c74487e9f4ba3d6ea5a784eb474242702 Author: Elizabeth Thompson <[email protected]> AuthorDate: Wed Apr 16 16:49:07 2025 -0700 migrate to more db migration utils --- superset/migrations/shared/utils.py | 12 +++++++-- ...cf3d64daf4_add_user_id_dttm_idx_to_log_model.py | 13 +++++----- ..._90139bf715e4_add_currency_column_to_metrics.py | 17 ++++++------ ...20_65a167d4c62e_add_indexes_to_report_models.py | 30 ++++++++++++---------- ...20_16-02_678eefb4ab44_add_access_token_table.py | 5 ++-- ...0_3dfd0e78650e_add_query_sql_editor_id_index.py | 24 ++++++++--------- 6 files changed, 56 insertions(+), 45 deletions(-) diff --git a/superset/migrations/shared/utils.py b/superset/migrations/shared/utils.py index db97383a56..57aabdfd2c 100644 --- a/superset/migrations/shared/utils.py +++ b/superset/migrations/shared/utils.py @@ -303,7 +303,7 @@ def add_columns(table_name: str, *columns: Column) -> None: """ Adds new columns to an existing database table. - If a column already exist, it logs an informational message and skips the adding process. + If a column already exist, or the table doesn't exist, it logs an informational message and skips the adding process. Otherwise, it proceeds to add the new column to the table. The operation is performed using Alembic's batch_alter_table. @@ -312,6 +312,10 @@ def add_columns(table_name: str, *columns: Column) -> None: :param columns: A list of SQLAlchemy Column objects that define the name, type, and other attributes of the columns to be added. """ # noqa: E501 + if not has_table(table_name=table_name): + logger.info(f"Table {GREEN}{table_name}{RESET} doesn't exist. Skipping...") + return + cols_to_add = [] for col in columns: if table_has_column(table_name=table_name, column_name=col.name): @@ -333,7 +337,7 @@ def drop_columns(table_name: str, *columns: str) -> None: """ Drops specified columns from an existing database table. - If a column does not exist, it logs an informational message and skips the dropping process. + If a column or table does not exist, it logs an informational message and skips the dropping process. Otherwise, it proceeds to remove the column from the table. The operation is performed using Alembic's batch_alter_table. @@ -342,6 +346,10 @@ def drop_columns(table_name: str, *columns: str) -> None: :param columns: A list of column names to be dropped. """ # noqa: E501 + if not has_table(table_name=table_name): + logger.info(f"Table {GREEN}{table_name}{RESET} doesn't exist. Skipping...") + return + cols_to_drop = [] for col in columns: if not table_has_column(table_name=table_name, column_name=col): diff --git a/superset/migrations/versions/2022-07-07_00-00_cdcf3d64daf4_add_user_id_dttm_idx_to_log_model.py b/superset/migrations/versions/2022-07-07_00-00_cdcf3d64daf4_add_user_id_dttm_idx_to_log_model.py index 9a2bfad53b..f0e97b9bd6 100644 --- a/superset/migrations/versions/2022-07-07_00-00_cdcf3d64daf4_add_user_id_dttm_idx_to_log_model.py +++ b/superset/migrations/versions/2022-07-07_00-00_cdcf3d64daf4_add_user_id_dttm_idx_to_log_model.py @@ -22,19 +22,20 @@ Create Date: 2022-04-05 13:27:06.028908 """ +from alembic import op + +from superset.migrations.shared.utils import create_index, drop_index + # revision identifiers, used by Alembic. revision = "cdcf3d64daf4" down_revision = "7fb8bca906d2" -from alembic import op # noqa: E402 - - def upgrade(): - op.create_index( - op.f("ix_logs_user_id_dttm"), "logs", ["user_id", "dttm"], unique=False + create_index( + "logs", op.f("ix_logs_user_id_dttm"), ["user_id", "dttm"], unique=False ) def downgrade(): - op.drop_index(op.f("ix_logs_user_id_dttm"), table_name="logs") + drop_index(index_name=op.f("ix_logs_user_id_dttm"), table_name="logs") diff --git a/superset/migrations/versions/2023-06-21_14-02_90139bf715e4_add_currency_column_to_metrics.py b/superset/migrations/versions/2023-06-21_14-02_90139bf715e4_add_currency_column_to_metrics.py index 005f4ce40d..a59dc329b2 100644 --- a/superset/migrations/versions/2023-06-21_14-02_90139bf715e4_add_currency_column_to_metrics.py +++ b/superset/migrations/versions/2023-06-21_14-02_90139bf715e4_add_currency_column_to_metrics.py @@ -22,21 +22,20 @@ Create Date: 2023-06-21 14:02:08.200955 """ +import sqlalchemy as sa + +from superset.migrations.shared.utils import add_columns, drop_columns + # revision identifiers, used by Alembic. revision = "90139bf715e4" down_revision = "83e1abbe777f" -import sqlalchemy as sa # noqa: E402 -from alembic import op # noqa: E402 - def upgrade(): - op.add_column("metrics", sa.Column("currency", sa.String(128), nullable=True)) - op.add_column("sql_metrics", sa.Column("currency", sa.String(128), nullable=True)) + add_columns("metrics", sa.Column("currency", sa.String(128), nullable=True)) + add_columns("sql_metrics", sa.Column("currency", sa.String(128), nullable=True)) def downgrade(): - with op.batch_alter_table("sql_metrics") as batch_op_sql_metrics: - batch_op_sql_metrics.drop_column("currency") - with op.batch_alter_table("metrics") as batch_op_metrics: - batch_op_metrics.drop_column("currency") + drop_columns("sql_metrics", "currency") + drop_columns("metrics", "currency") diff --git a/superset/migrations/versions/2024-01-05_16-20_65a167d4c62e_add_indexes_to_report_models.py b/superset/migrations/versions/2024-01-05_16-20_65a167d4c62e_add_indexes_to_report_models.py index 19fd18c5e5..ec233f1958 100644 --- a/superset/migrations/versions/2024-01-05_16-20_65a167d4c62e_add_indexes_to_report_models.py +++ b/superset/migrations/versions/2024-01-05_16-20_65a167d4c62e_add_indexes_to_report_models.py @@ -26,37 +26,41 @@ Create Date: 2024-01-05 16:20:31.598995 revision = "65a167d4c62e" down_revision = "06dd9ff00fe8" -from alembic import op # noqa: E402 + +from superset.migrations.shared.utils import create_index, drop_index # noqa: E402 def upgrade(): - op.create_index( - "ix_report_execution_log_report_schedule_id", + create_index( "report_execution_log", + "ix_report_execution_log_report_schedule_id", ["report_schedule_id"], unique=False, ) - op.create_index( - "ix_report_execution_log_start_dttm", + create_index( "report_execution_log", + "ix_report_execution_log_start_dttm", ["start_dttm"], unique=False, ) - op.create_index( - "ix_report_recipient_report_schedule_id", + create_index( "report_recipient", + "ix_report_recipient_report_schedule_id", ["report_schedule_id"], unique=False, ) def downgrade(): - op.drop_index( - "ix_report_recipient_report_schedule_id", table_name="report_recipient" + drop_index( + index_name="ix_report_recipient_report_schedule_id", + table_name="report_recipient", ) - op.drop_index( - "ix_report_execution_log_start_dttm", table_name="report_execution_log" + drop_index( + index_name="ix_report_execution_log_start_dttm", + table_name="report_execution_log", ) - op.drop_index( - "ix_report_execution_log_report_schedule_id", table_name="report_execution_log" + drop_index( + index_name="ix_report_execution_log_report_schedule_id", + table_name="report_execution_log", ) 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 7d9c851285..3c13803ad1 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 @@ -27,6 +27,7 @@ from alembic import op from sqlalchemy_utils import EncryptedType from superset.migrations.shared.utils import ( + create_index, create_table, drop_fks_for_table, ) @@ -77,9 +78,9 @@ def upgrade(): ), sa.PrimaryKeyConstraint("id"), ) - op.create_index( - "idx_user_id_database_id", + create_index( "database_user_oauth2_tokens", + "idx_user_id_database_id", ["user_id", "database_id"], ) diff --git a/superset/migrations/versions/2024-05-02_13-40_3dfd0e78650e_add_query_sql_editor_id_index.py b/superset/migrations/versions/2024-05-02_13-40_3dfd0e78650e_add_query_sql_editor_id_index.py index d4eacaf5a2..4689d91180 100644 --- a/superset/migrations/versions/2024-05-02_13-40_3dfd0e78650e_add_query_sql_editor_id_index.py +++ b/superset/migrations/versions/2024-05-02_13-40_3dfd0e78650e_add_query_sql_editor_id_index.py @@ -22,28 +22,26 @@ Create Date: 2024-05-02 13:40:23.126659 """ +from alembic import op + +from superset.migrations.shared.utils import create_index, drop_index + # revision identifiers, used by Alembic. revision = "3dfd0e78650e" down_revision = "5f57af97bc3f" -from alembic import op # noqa: E402 - -from superset.migrations.shared.utils import table_has_index # noqa: E402 - table = "query" index = "ix_sql_editor_id" def upgrade(): - if not table_has_index(table, index): - op.create_index( - op.f(index), - table, - ["sql_editor_id"], - unique=False, - ) + create_index( + table, + op.f(index), + ["sql_editor_id"], + unique=False, + ) def downgrade(): - if table_has_index(table, index): - op.drop_index(op.f(index), table_name=table) + drop_index(index_name=op.f(index), table_name=table)
