Skip to content

Commit

Permalink
feat: Adds chart ID option to migrate-viz
Browse files Browse the repository at this point in the history
  • Loading branch information
michael-s-molina committed Jun 25, 2024
1 parent fe3ba12 commit 09e5d2f
Show file tree
Hide file tree
Showing 2 changed files with 32 additions and 22 deletions.
42 changes: 23 additions & 19 deletions superset/cli/viz_migrations.py
Original file line number Diff line number Diff line change
Expand Up @@ -46,37 +46,41 @@ def migrate_viz() -> None:

@migrate_viz.command()
@with_appcontext
@optgroup.group(
"Grouped options",
cls=RequiredMutuallyExclusiveOptionGroup,
)
@optgroup.option(
@click.option(
"--viz_type",
"-t",
help=f"The viz type to migrate: {', '.join(list(VizType))}",
help=f"The viz type to upgrade: {', '.join(list(VizType))}",
required=True,
)
@click.option(
"-id",
help="The chart ID to upgrade",
type=int,
)
def upgrade(viz_type: str) -> None:
def upgrade(viz_type: str, id: int | None) -> None:
"""Upgrade a viz to the latest version."""
migrate(VizType(viz_type))
migrate(VizType(viz_type), id)


@migrate_viz.command()
@with_appcontext
@optgroup.group(
"Grouped options",
cls=RequiredMutuallyExclusiveOptionGroup,
)
@optgroup.option(
@click.option(
"--viz_type",
"-t",
help=f"The viz type to migrate: {', '.join(list(VizType))}",
help=f"The viz type to downgrade: {', '.join(list(VizType))}",
required=True,
)
@click.option(
"-id",
help="The chart ID to downgrade",
type=int,
)
def downgrade(viz_type: str) -> None:
def downgrade(viz_type: str, id: int | None) -> None:
"""Downgrade a viz to the previous version."""
migrate(VizType(viz_type), is_downgrade=True)
migrate(VizType(viz_type), id, is_downgrade=True)


def migrate(viz_type: VizType, is_downgrade: bool = False) -> None:
def migrate(viz_type: VizType, id: int | None, 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 (
Expand Down Expand Up @@ -107,6 +111,6 @@ def migrate(viz_type: VizType, is_downgrade: bool = False) -> None:
VizType.TREEMAP: MigrateTreeMap,
}
if is_downgrade:
migrations[viz_type].downgrade(db.session)
migrations[viz_type].downgrade(db.session, id)
else:
migrations[viz_type].upgrade(db.session)
migrations[viz_type].upgrade(db.session, id)
12 changes: 9 additions & 3 deletions superset/migrations/shared/migrate_viz/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -153,20 +153,26 @@ def downgrade_slice(cls, slc: Slice) -> None:
slc.query_context = json.dumps(query_context)

@classmethod
def upgrade(cls, session: Session) -> None:
slices = session.query(Slice).filter(Slice.viz_type == cls.source_viz_type)
def upgrade(cls, session: Session, id: int | None) -> None:
slices = session.query(Slice).filter(
and_(
Slice.viz_type == cls.source_viz_type,
Slice.id == id if id is not None else True,
)
)
for slc in paginated_update(
slices,
lambda current, total: print(f"Upgraded {current}/{total} charts"),
):
cls.upgrade_slice(slc)

@classmethod
def downgrade(cls, session: Session) -> None:
def downgrade(cls, session: Session, id: int | None) -> None:
slices = session.query(Slice).filter(
and_(
Slice.viz_type == cls.target_viz_type,
Slice.params.like(f"%{FORM_DATA_BAK_FIELD_NAME}%"),
Slice.id == id if id is not None else True,
)
)
for slc in paginated_update(
Expand Down

0 comments on commit 09e5d2f

Please sign in to comment.