From b722a95baba33cd901b078d63026ad8baea27f6e Mon Sep 17 00:00:00 2001 From: "Michael S. Molina" <70410625+michael-s-molina@users.noreply.github.com> Date: Tue, 19 Sep 2023 14:34:08 -0300 Subject: [PATCH] feat: Adds CLI commands to execute viz migrations (#25304) --- superset/cli/viz_migrations.py | 91 +++++++++++++++++++ .../migrations/shared/migrate_viz/base.py | 12 +-- ...0_c747c78868b6_migrating_legacy_treemap.py | 12 ++- ...4-00_06e1e70058c7_migrating_legacy_area.py | 11 ++- ...a2ce3086e5_migrate_pivot_table_v1_to_v2.py | 11 ++- ...0-22_4c5da39be729_migrate_treemap_chart.py | 12 ++- ...e58e5c_migrate_dual_line_to_mixed_chart.py | 11 ++- 7 files changed, 140 insertions(+), 20 deletions(-) create mode 100644 superset/cli/viz_migrations.py diff --git a/superset/cli/viz_migrations.py b/superset/cli/viz_migrations.py new file mode 100644 index 0000000000..5d2ee4ae5d --- /dev/null +++ b/superset/cli/viz_migrations.py @@ -0,0 +1,91 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from enum import Enum + +import click +from click_option_group import optgroup, RequiredMutuallyExclusiveOptionGroup +from flask.cli import with_appcontext + +from superset import db + + +class VizType(str, Enum): + TREEMAP = "treemap" + DUAL_LINE = "dual_line" + AREA = "area" + PIVOT_TABLE = "pivot_table" + + +@click.group() +def migrate_viz() -> None: + """ + Migrate a viz from one type to another. + """ + + +@migrate_viz.command() +@with_appcontext +@optgroup.group( + "Grouped options", + cls=RequiredMutuallyExclusiveOptionGroup, +) +@optgroup.option( + "--viz_type", + "-t", + help=f"The viz type to migrate: {', '.join(list(VizType))}", +) +def upgrade(viz_type: str) -> None: + """Upgrade a viz to the latest version.""" + migrate(VizType(viz_type)) + + +@migrate_viz.command() +@with_appcontext +@optgroup.group( + "Grouped options", + cls=RequiredMutuallyExclusiveOptionGroup, +) +@optgroup.option( + "--viz_type", + "-t", + help=f"The viz type to migrate: {', '.join(list(VizType))}", +) +def downgrade(viz_type: str) -> None: + """Downgrade a viz to the previous version.""" + migrate(VizType(viz_type), is_downgrade=True) + + +def migrate(viz_type: VizType, is_downgrade: bool = False) -> None: + """Migrate a viz from one type to another.""" + # pylint: disable=import-outside-toplevel + from superset.migrations.shared.migrate_viz.processors import ( + MigrateAreaChart, + MigrateDualLine, + MigratePivotTable, + MigrateTreeMap, + ) + + migrations = { + VizType.TREEMAP: MigrateTreeMap, + VizType.DUAL_LINE: MigrateDualLine, + VizType.AREA: MigrateAreaChart, + VizType.PIVOT_TABLE: MigratePivotTable, + } + if is_downgrade: + migrations[viz_type].downgrade(db.session) + else: + migrations[viz_type].upgrade(db.session) diff --git a/superset/migrations/shared/migrate_viz/base.py b/superset/migrations/shared/migrate_viz/base.py index 09b77ae271..b9826fee34 100644 --- a/superset/migrations/shared/migrate_viz/base.py +++ b/superset/migrations/shared/migrate_viz/base.py @@ -20,11 +20,11 @@ import copy import json from typing import Any -from alembic import op from sqlalchemy import and_, Column, Integer, String, Text from sqlalchemy.ext.declarative import declarative_base +from sqlalchemy.orm import Session -from superset import conf, db, is_feature_enabled +from superset import conf, is_feature_enabled from superset.constants import TimeGrain from superset.migrations.shared.utils import paginated_update, try_load_json @@ -156,9 +156,7 @@ class MigrateViz: return slc @classmethod - def upgrade(cls) -> None: - bind = op.get_bind() - session = db.Session(bind=bind) + def upgrade(cls, session: Session) -> None: slices = session.query(Slice).filter(Slice.viz_type == cls.source_viz_type) for slc in paginated_update( slices, @@ -170,9 +168,7 @@ class MigrateViz: session.merge(new_viz) @classmethod - def downgrade(cls) -> None: - bind = op.get_bind() - session = db.Session(bind=bind) + def downgrade(cls, session: Session) -> None: slices = session.query(Slice).filter( and_( Slice.viz_type == cls.target_viz_type, diff --git a/superset/migrations/versions/2022-07-07_13-00_c747c78868b6_migrating_legacy_treemap.py b/superset/migrations/versions/2022-07-07_13-00_c747c78868b6_migrating_legacy_treemap.py index 18c9acaf57..be222cdfda 100644 --- a/superset/migrations/versions/2022-07-07_13-00_c747c78868b6_migrating_legacy_treemap.py +++ b/superset/migrations/versions/2022-07-07_13-00_c747c78868b6_migrating_legacy_treemap.py @@ -24,6 +24,7 @@ Create Date: 2022-06-30 22:04:17.686635 from alembic import op from sqlalchemy.dialects.mysql.base import MySQLDialect +from superset import db from superset.migrations.shared.migrate_viz import MigrateTreeMap # revision identifiers, used by Alembic. @@ -32,16 +33,21 @@ down_revision = "cdcf3d64daf4" def upgrade(): + bind = op.get_bind() + # Ensure `slice.params` and `slice.query_context`` in MySQL is MEDIUMTEXT # before migration, as the migration will save a duplicate form_data backup # which may significantly increase the size of these fields. - if isinstance(op.get_bind().dialect, MySQLDialect): + if isinstance(bind.dialect, MySQLDialect): # If the columns are already MEDIUMTEXT, this is a no-op op.execute("ALTER TABLE slices MODIFY params MEDIUMTEXT") op.execute("ALTER TABLE slices MODIFY query_context MEDIUMTEXT") - MigrateTreeMap.upgrade() + session = db.Session(bind=bind) + MigrateTreeMap.upgrade(session) def downgrade(): - MigrateTreeMap.downgrade() + bind = op.get_bind() + session = db.Session(bind=bind) + MigrateTreeMap.downgrade(session) diff --git a/superset/migrations/versions/2022-07-07_14-00_06e1e70058c7_migrating_legacy_area.py b/superset/migrations/versions/2022-07-07_14-00_06e1e70058c7_migrating_legacy_area.py index de40780991..a43f027e2c 100644 --- a/superset/migrations/versions/2022-07-07_14-00_06e1e70058c7_migrating_legacy_area.py +++ b/superset/migrations/versions/2022-07-07_14-00_06e1e70058c7_migrating_legacy_area.py @@ -21,6 +21,9 @@ Revises: c747c78868b6 Create Date: 2022-06-13 14:17:51.872706 """ +from alembic import op + +from superset import db from superset.migrations.shared.migrate_viz import MigrateAreaChart # revision identifiers, used by Alembic. @@ -29,8 +32,12 @@ down_revision = "c747c78868b6" def upgrade(): - MigrateAreaChart.upgrade() + bind = op.get_bind() + session = db.Session(bind=bind) + MigrateAreaChart.upgrade(session) def downgrade(): - MigrateAreaChart.downgrade() + bind = op.get_bind() + session = db.Session(bind=bind) + MigrateAreaChart.downgrade(session) diff --git a/superset/migrations/versions/2023-06-08_09-02_9ba2ce3086e5_migrate_pivot_table_v1_to_v2.py b/superset/migrations/versions/2023-06-08_09-02_9ba2ce3086e5_migrate_pivot_table_v1_to_v2.py index f1df097254..917408a5c0 100644 --- a/superset/migrations/versions/2023-06-08_09-02_9ba2ce3086e5_migrate_pivot_table_v1_to_v2.py +++ b/superset/migrations/versions/2023-06-08_09-02_9ba2ce3086e5_migrate_pivot_table_v1_to_v2.py @@ -21,6 +21,9 @@ Revises: 4ea966691069 Create Date: 2023-08-06 09:02:10.148992 """ +from alembic import op + +from superset import db from superset.migrations.shared.migrate_viz import MigratePivotTable # revision identifiers, used by Alembic. @@ -29,8 +32,12 @@ down_revision = "4ea966691069" def upgrade(): - MigratePivotTable.upgrade() + bind = op.get_bind() + session = db.Session(bind=bind) + MigratePivotTable.upgrade(session) def downgrade(): - MigratePivotTable.downgrade() + bind = op.get_bind() + session = db.Session(bind=bind) + MigratePivotTable.downgrade(session) diff --git a/superset/migrations/versions/2023-06-08_10-22_4c5da39be729_migrate_treemap_chart.py b/superset/migrations/versions/2023-06-08_10-22_4c5da39be729_migrate_treemap_chart.py index 6c1e5753eb..2f36ed8508 100644 --- a/superset/migrations/versions/2023-06-08_10-22_4c5da39be729_migrate_treemap_chart.py +++ b/superset/migrations/versions/2023-06-08_10-22_4c5da39be729_migrate_treemap_chart.py @@ -24,6 +24,7 @@ Create Date: 2023-06-08 10:22:23.192064 from alembic import op from sqlalchemy.dialects.mysql.base import MySQLDialect +from superset import db from superset.migrations.shared.migrate_viz import MigrateTreeMap # revision identifiers, used by Alembic. @@ -32,16 +33,21 @@ down_revision = "9ba2ce3086e5" def upgrade(): + bind = op.get_bind() + # Ensure `slice.params` and `slice.query_context`` in MySQL is MEDIUMTEXT # before migration, as the migration will save a duplicate form_data backup # which may significantly increase the size of these fields. - if isinstance(op.get_bind().dialect, MySQLDialect): + if isinstance(bind.dialect, MySQLDialect): # If the columns are already MEDIUMTEXT, this is a no-op op.execute("ALTER TABLE slices MODIFY params MEDIUMTEXT") op.execute("ALTER TABLE slices MODIFY query_context MEDIUMTEXT") - MigrateTreeMap.upgrade() + session = db.Session(bind=bind) + MigrateTreeMap.upgrade(session) def downgrade(): - MigrateTreeMap.downgrade() + bind = op.get_bind() + session = db.Session(bind=bind) + MigrateTreeMap.downgrade(session) diff --git a/superset/migrations/versions/2023-06-08_11-34_ae58e1e58e5c_migrate_dual_line_to_mixed_chart.py b/superset/migrations/versions/2023-06-08_11-34_ae58e1e58e5c_migrate_dual_line_to_mixed_chart.py index 5d707dc601..5b371e7272 100644 --- a/superset/migrations/versions/2023-06-08_11-34_ae58e1e58e5c_migrate_dual_line_to_mixed_chart.py +++ b/superset/migrations/versions/2023-06-08_11-34_ae58e1e58e5c_migrate_dual_line_to_mixed_chart.py @@ -21,6 +21,9 @@ Revises: 4c5da39be729 Create Date: 2023-06-08 11:34:36.241939 """ +from alembic import op + +from superset import db # revision identifiers, used by Alembic. revision = "ae58e1e58e5c" @@ -30,8 +33,12 @@ from superset.migrations.shared.migrate_viz.processors import MigrateDualLine def upgrade(): - MigrateDualLine.upgrade() + bind = op.get_bind() + session = db.Session(bind=bind) + MigrateDualLine.upgrade(session) def downgrade(): - MigrateDualLine.downgrade() + bind = op.get_bind() + session = db.Session(bind=bind) + MigrateDualLine.downgrade(session)