feat: Adds CLI commands to execute viz migrations (#25304)

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)

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,

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)

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)

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)

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)

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)