From 6b73b69b415ec6b6fcbac80a358f4e31c4ed91b9 Mon Sep 17 00:00:00 2001 From: mknadh Date: Wed, 3 Jul 2024 21:50:05 +0530 Subject: [PATCH] feat(CLI command): Apache Superset "Factory Reset" CLI command #27207 (#27221) --- superset/cli/reset.py | 74 +++++++++++++++++++++++ superset/commands/security/reset.py | 94 +++++++++++++++++++++++++++++ superset/config.py | 2 + 3 files changed, 170 insertions(+) create mode 100644 superset/cli/reset.py create mode 100644 superset/commands/security/reset.py diff --git a/superset/cli/reset.py b/superset/cli/reset.py new file mode 100644 index 000000000000..fd5e7260754c --- /dev/null +++ b/superset/cli/reset.py @@ -0,0 +1,74 @@ +# 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. + +import sys + +import click +from flask.cli import with_appcontext +from werkzeug.security import check_password_hash + +from superset.cli.lib import feature_flags + +if feature_flags.get("ENABLE_FACTORY_RESET_COMMAND"): + + @click.command() + @with_appcontext + @click.option("--username", prompt="Admin Username", help="Admin Username") + @click.option( + "--silent", + is_flag=True, + prompt=( + "Are you sure you want to reset Superset? " + "This action cannot be undone. Continue?" + ), + help="Confirmation flag", + ) + @click.option( + "--exclude-users", + default=None, + help="Comma separated list of users to exclude from reset", + ) + @click.option( + "--exclude-roles", + default=None, + help="Comma separated list of roles to exclude from reset", + ) + def factory_reset( + username: str, silent: bool, exclude_users: str, exclude_roles: str + ) -> None: + """Factory Reset Apache Superset""" + + # pylint: disable=import-outside-toplevel + from superset import security_manager + from superset.commands.security.reset import ResetSupersetCommand + + # Validate the user + password = click.prompt("Admin Password", hide_input=True) + user = security_manager.find_user(username) + if not user or not check_password_hash(user.password, password): + click.secho("Invalid credentials", fg="red") + sys.exit(1) + if not any(role.name == "Admin" for role in user.roles): + click.secho("Permission Denied", fg="red") + sys.exit(1) + + try: + ResetSupersetCommand(silent, user, exclude_users, exclude_roles).run() + click.secho("Factory reset complete", fg="green") + except Exception as ex: # pylint: disable=broad-except + click.secho(f"Factory reset failed: {ex}", fg="red") + sys.exit(1) diff --git a/superset/commands/security/reset.py b/superset/commands/security/reset.py new file mode 100644 index 000000000000..5c93bb46461f --- /dev/null +++ b/superset/commands/security/reset.py @@ -0,0 +1,94 @@ +# 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. + +import logging +from typing import Any, Optional + +from superset import db, security_manager +from superset.commands.base import BaseCommand +from superset.connectors.sqla.models import SqlaTable +from superset.key_value.models import KeyValueEntry +from superset.models.core import Database, FavStar, Log +from superset.models.dashboard import Dashboard +from superset.models.slice import Slice + +logger = logging.getLogger(__name__) + + +class ResetSupersetCommand(BaseCommand): + def __init__( + self, + confirm: bool, + user: Any, + exclude_users: Optional[str] = None, + exclude_roles: Optional[str] = None, + ) -> None: + self._user = user + self._confirm = confirm + self._users_to_exclude = ["admin"] + if exclude_users: + self._users_to_exclude.extend(exclude_users.split(",")) + self._roles_to_exclude = ["Admin", "Public", "Gamma", "Alpha", "sql_lab"] + if exclude_roles: + self._roles_to_exclude.extend(exclude_roles.split(",")) + + def validate(self) -> None: + if not self._confirm: + raise Exception("Reset aborted.") # pylint: disable=broad-exception-raised + if not self._user or not self._user.is_active: + raise Exception("User not found.") # pylint: disable=broad-exception-raised + + def run(self) -> None: + self.validate() + logger.debug("Resetting Superset Started") + db.session.query(SqlaTable).delete() + databases = db.session.query(Database) + for database in databases: + db.session.delete(database) + db.session.query(Dashboard).delete() + db.session.query(Slice).delete() + db.session.query(KeyValueEntry).delete() + db.session.query(Log).delete() + db.session.query(FavStar).delete() + + logger.debug("Ignoring Users: %s", self._users_to_exclude) + users_to_delete = ( + db.session.query(security_manager.user_model) + .filter(security_manager.user_model.username.not_in(self._users_to_exclude)) + .all() + ) + for user in users_to_delete: + if not any(role.name == "Admin" for role in user.roles): + db.session.delete(user) + + logger.debug("Ignoring Roles: %s", self._roles_to_exclude) + roles_to_delete = ( + db.session.query(security_manager.role_model) + .filter(security_manager.role_model.name.not_in(self._roles_to_exclude)) + .all() + ) + for role in roles_to_delete: + db.session.delete(role) + + # Insert new record into Log table + log = Log( + action="Factory Reset", json="{}", user_id=self._user.id, user=self._user + ) + db.session.add(log) + + db.session.commit() # pylint: disable=consider-using-transaction + logger.debug("Resetting Superset Completed") diff --git a/superset/config.py b/superset/config.py index e4dc202537ac..fa31fd069a92 100644 --- a/superset/config.py +++ b/superset/config.py @@ -539,6 +539,8 @@ class D3TimeFormat(TypedDict, total=False): "CHART_PLUGINS_EXPERIMENTAL": False, # Regardless of database configuration settings, force SQLLAB to run async using Celery "SQLLAB_FORCE_RUN_ASYNC": False, + # Set to True to to enable factory resent CLI command + "ENABLE_FACTORY_RESET_COMMAND": False, } # ------------------------------