feat(CLI command): Apache Superset "Factory Reset" CLI command #27207 (#27221)

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)

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")

superset/config.py

@@ -539,6 +539,8 @@ DEFAULT_FEATURE_FLAGS: dict[str, bool] = {
     "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,
 }
 
 # ------------------------------