From d4054e3d8554f03ee036aa716b8d9fe88b10bdd6 Mon Sep 17 00:00:00 2001 From: "Michael S. Molina" <70410625+michael-s-molina@users.noreply.github.com> Date: Wed, 26 Jun 2024 09:20:29 -0300 Subject: [PATCH] feat: Adds chart IDs option to migrate-viz (#29361) --- superset-frontend/package-lock.json | 12 +- superset/cli/viz_migrations.py | 154 ++++++++++++------ .../migrations/shared/migrate_viz/base.py | 2 - 3 files changed, 112 insertions(+), 56 deletions(-) diff --git a/superset-frontend/package-lock.json b/superset-frontend/package-lock.json index b8ec2f16d1..66d71dad93 100644 --- a/superset-frontend/package-lock.json +++ b/superset-frontend/package-lock.json @@ -31422,9 +31422,9 @@ } }, "node_modules/dayjs": { - "version": "1.11.10", - "resolved": "https://registry.npmjs.org/dayjs/-/dayjs-1.11.10.tgz", - "integrity": "sha512-vjAczensTgRcqDERK0SR2XMwsF/tSvnvlv6VcF2GIhg6Sx4yOIt/irsr1RDJsKiIyBzJDpCoXiWWq28MqH2cnQ==" + "version": "1.11.11", + "resolved": "https://registry.npmjs.org/dayjs/-/dayjs-1.11.11.tgz", + "integrity": "sha512-okzr3f11N6WuqYtZSvm+F776mB41wRZMhKP+hc34YdW+KmtYYK9iqvHSwo2k9FEH3fhGXvOPV6yz2IcSrfRUDg==" }, "node_modules/deasync": { "version": "0.1.29", @@ -96697,9 +96697,9 @@ "dev": true }, "dayjs": { - "version": "1.11.10", - "resolved": "https://registry.npmjs.org/dayjs/-/dayjs-1.11.10.tgz", - "integrity": "sha512-vjAczensTgRcqDERK0SR2XMwsF/tSvnvlv6VcF2GIhg6Sx4yOIt/irsr1RDJsKiIyBzJDpCoXiWWq28MqH2cnQ==" + "version": "1.11.11", + "resolved": "https://registry.npmjs.org/dayjs/-/dayjs-1.11.11.tgz", + "integrity": "sha512-okzr3f11N6WuqYtZSvm+F776mB41wRZMhKP+hc34YdW+KmtYYK9iqvHSwo2k9FEH3fhGXvOPV6yz2IcSrfRUDg==" }, "deasync": { "version": "0.1.29", diff --git a/superset/cli/viz_migrations.py b/superset/cli/viz_migrations.py index 6b6ba1a704..34550ac2da 100644 --- a/superset/cli/viz_migrations.py +++ b/superset/cli/viz_migrations.py @@ -14,13 +14,35 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. +from __future__ import annotations + from enum import Enum +from typing import Type import click -from click_option_group import optgroup, RequiredMutuallyExclusiveOptionGroup +from click_option_group import optgroup, RequiredAnyOptionGroup from flask.cli import with_appcontext from superset import db +from superset.migrations.shared.migrate_viz.base import ( + MigrateViz, + Slice, +) +from superset.migrations.shared.migrate_viz.processors import ( + MigrateAreaChart, + MigrateBarChart, + MigrateBubbleChart, + MigrateDistBarChart, + MigrateDualLine, + MigrateHeatmapChart, + MigrateHistogramChart, + MigrateLineChart, + MigratePivotTable, + MigrateSankey, + MigrateSunburst, + MigrateTreeMap, +) +from superset.migrations.shared.utils import paginated_update class VizType(str, Enum): @@ -38,6 +60,26 @@ class VizType(str, Enum): TREEMAP = "treemap" +MIGRATIONS: dict[VizType, Type[MigrateViz]] = { + VizType.AREA: MigrateAreaChart, + VizType.BAR: MigrateBarChart, + VizType.BUBBLE: MigrateBubbleChart, + VizType.DIST_BAR: MigrateDistBarChart, + VizType.DUAL_LINE: MigrateDualLine, + VizType.HEATMAP: MigrateHeatmapChart, + VizType.HISTOGRAM: MigrateHistogramChart, + VizType.LINE: MigrateLineChart, + VizType.PIVOT_TABLE: MigratePivotTable, + VizType.SANKEY: MigrateSankey, + VizType.SUNBURST: MigrateSunburst, + VizType.TREEMAP: MigrateTreeMap, +} + +PREVIOUS_VERSION = { + migration.target_viz_type: migration for migration in MIGRATIONS.values() +} + + @click.group() def migrate_viz() -> None: """ @@ -48,68 +90,84 @@ def migrate_viz() -> None: @migrate_viz.command() @with_appcontext @optgroup.group( - "Grouped options", - cls=RequiredMutuallyExclusiveOptionGroup, + cls=RequiredAnyOptionGroup, ) @optgroup.option( "--viz_type", "-t", - help=f"The viz type to migrate: {', '.join(list(VizType))}", + help=f"The viz type to upgrade: {', '.join(list(VizType))}", + type=str, ) -def upgrade(viz_type: str) -> None: +@optgroup.option( + "--id", + "ids", + help="The chart ID to upgrade. It can set set multiple times.", + type=int, + multiple=True, +) +def upgrade(viz_type: str, ids: tuple[int, ...] | None = None) -> None: """Upgrade a viz to the latest version.""" - migrate(VizType(viz_type)) + if ids is None: + migrate_by_viz_type(VizType(viz_type)) + else: + migrate_by_id(ids) @migrate_viz.command() @with_appcontext @optgroup.group( - "Grouped options", - cls=RequiredMutuallyExclusiveOptionGroup, + cls=RequiredAnyOptionGroup, ) @optgroup.option( "--viz_type", "-t", - help=f"The viz type to migrate: {', '.join(list(VizType))}", + help=f"The viz type to downgrade: {', '.join(list(VizType))}", + type=str, ) -def downgrade(viz_type: str) -> None: +@optgroup.option( + "--id", + "ids", + help="The chart ID to downgrade. It can set set multiple times.", + type=int, + multiple=True, +) +def downgrade(viz_type: str, ids: tuple[int, ...] | None = None) -> 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, - MigrateBarChart, - MigrateBubbleChart, - MigrateDistBarChart, - MigrateDualLine, - MigrateHeatmapChart, - MigrateHistogramChart, - MigrateLineChart, - MigratePivotTable, - MigrateSankey, - MigrateSunburst, - MigrateTreeMap, - ) - - migrations = { - VizType.AREA: MigrateAreaChart, - VizType.BAR: MigrateBarChart, - VizType.BUBBLE: MigrateBubbleChart, - VizType.DIST_BAR: MigrateDistBarChart, - VizType.DUAL_LINE: MigrateDualLine, - VizType.HEATMAP: MigrateHeatmapChart, - VizType.HISTOGRAM: MigrateHistogramChart, - VizType.LINE: MigrateLineChart, - VizType.PIVOT_TABLE: MigratePivotTable, - VizType.SANKEY: MigrateSankey, - VizType.SUNBURST: MigrateSunburst, - VizType.TREEMAP: MigrateTreeMap, - } - if is_downgrade: - migrations[viz_type].downgrade(db.session) + if ids is None: + migrate_by_viz_type(VizType(viz_type), is_downgrade=True) else: - migrations[viz_type].upgrade(db.session) + migrate_by_id(ids, is_downgrade=True) + + +def migrate_by_viz_type(viz_type: VizType, is_downgrade: bool = False) -> None: + """ + Migrate all charts of a viz type. + + :param viz_type: The viz type to migrate + :param is_downgrade: Whether to downgrade the charts. Default is upgrade. + """ + migration: Type[MigrateViz] = MIGRATIONS[viz_type] + if is_downgrade: + migration.downgrade(db.session) + else: + migration.upgrade(db.session) + + +def migrate_by_id(ids: tuple[int, ...], is_downgrade: bool = False) -> None: + """ + Migrate a subset of charts by IDs. + + :param id: Tuple of chart IDs to migrate + :param is_downgrade: Whether to downgrade the charts. Default is upgrade. + """ + slices = db.session.query(Slice).filter(Slice.id.in_(ids)) + for slc in paginated_update( + slices, + lambda current, total: print( + f"{('Downgraded' if is_downgrade else 'Upgraded')} {current}/{total} charts" + ), + ): + if is_downgrade: + PREVIOUS_VERSION[slc.viz_type].downgrade_slice(slc) + elif slc.viz_type in MIGRATIONS: + MIGRATIONS[slc.viz_type].upgrade_slice(slc) diff --git a/superset/migrations/shared/migrate_viz/base.py b/superset/migrations/shared/migrate_viz/base.py index 1be0391659..b013aa0be2 100644 --- a/superset/migrations/shared/migrate_viz/base.py +++ b/superset/migrations/shared/migrate_viz/base.py @@ -14,8 +14,6 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -from __future__ import annotations - import copy from typing import Any