ss-tools/backend/src/plugins/backup.py

# [DEF:BackupPlugin:Module]
# @SEMANTICS: backup, superset, automation, dashboard, plugin
# @PURPOSE: A plugin that provides functionality to back up Superset dashboards.
# @LAYER: App
# @RELATION: IMPLEMENTS -> PluginBase
# @RELATION: DEPENDS_ON -> superset_tool.client
# @RELATION: DEPENDS_ON -> superset_tool.utils

from typing import Dict, Any
from pathlib import Path
from requests.exceptions import RequestException

from ..core.plugin_base import PluginBase
from ..core.logger import belief_scope
from ..core.superset_client import SupersetClient
from ..core.utils.network import SupersetAPIError
from ..core.utils.fileio import (
    save_and_unpack_dashboard,
    archive_exports,
    sanitize_filename,
    consolidate_archive_folders,
    remove_empty_directories,
    RetentionPolicy
)
from ..dependencies import get_config_manager

# [DEF:BackupPlugin:Class]
# @PURPOSE: Implementation of the backup plugin logic.
class BackupPlugin(PluginBase):
    """
    A plugin to back up Superset dashboards.
    """

    @property
    # [DEF:id:Function]
    # @PURPOSE: Returns the unique identifier for the backup plugin.
    # @PRE: Plugin instance exists.
    # @POST: Returns string ID.
    # @RETURN: str - "superset-backup"
    def id(self) -> str:
        with belief_scope("id"):
            return "superset-backup"
    # [/DEF:id:Function]

    @property
    # [DEF:name:Function]
    # @PURPOSE: Returns the human-readable name of the backup plugin.
    # @PRE: Plugin instance exists.
    # @POST: Returns string name.
    # @RETURN: str - Plugin name.
    def name(self) -> str:
        with belief_scope("name"):
            return "Superset Dashboard Backup"
    # [/DEF:name:Function]

    @property
    # [DEF:description:Function]
    # @PURPOSE: Returns a description of the backup plugin.
    # @PRE: Plugin instance exists.
    # @POST: Returns string description.
    # @RETURN: str - Plugin description.
    def description(self) -> str:
        with belief_scope("description"):
            return "Backs up all dashboards from a Superset instance."
    # [/DEF:description:Function]

    @property
    # [DEF:version:Function]
    # @PURPOSE: Returns the version of the backup plugin.
    # @PRE: Plugin instance exists.
    # @POST: Returns string version.
    # @RETURN: str - "1.0.0"
    def version(self) -> str:
        with belief_scope("version"):
            return "1.0.0"
    # [/DEF:version:Function]

    @property
    # [DEF:ui_route:Function]
    # @PURPOSE: Returns the frontend route for the backup plugin.
    # @RETURN: str - "/tools/backups"
    def ui_route(self) -> str:
        with belief_scope("ui_route"):
            return "/tools/backups"
    # [/DEF:ui_route:Function]

    # [DEF:get_schema:Function]
    # @PURPOSE: Returns the JSON schema for backup plugin parameters.
    # @PRE: Plugin instance exists.
    # @POST: Returns dictionary schema.
    # @RETURN: Dict[str, Any] - JSON schema.
    def get_schema(self) -> Dict[str, Any]:
        with belief_scope("get_schema"):
            config_manager = get_config_manager()
            envs = [e.name for e in config_manager.get_environments()]
        default_path = config_manager.get_config().settings.storage.root_path
        
        return {
            "type": "object",
            "properties": {
                "env": {
                    "type": "string",
                    "title": "Environment",
                    "description": "The Superset environment to back up.",
                    "enum": envs if envs else [],
                },
            },
            "required": ["env"],
        }
    # [/DEF:get_schema:Function]

    # [DEF:execute:Function]
    # @PURPOSE: Executes the dashboard backup logic.
    # @PARAM: params (Dict[str, Any]) - Backup parameters (env, backup_path).
    # @PRE: Target environment must be configured. params must be a dictionary.
    # @POST: All dashboards are exported and archived.
    async def execute(self, params: Dict[str, Any]):
        with belief_scope("execute"):
            config_manager = get_config_manager()
            env_id = params.get("environment_id")
            
            # Resolve environment name if environment_id is provided
            if env_id:
                env_config = next((e for e in config_manager.get_environments() if e.id == env_id), None)
                if env_config:
                    params["env"] = env_config.name
            
            env = params.get("env")
            if not env:
                raise KeyError("env")

            storage_settings = config_manager.get_config().settings.storage
            # Use 'backups' subfolder within the storage root
            backup_path = Path(storage_settings.root_path) / "backups"
            
            from ..core.logger import logger as app_logger
            app_logger.info(f"[BackupPlugin][Entry] Starting backup for {env}.")

            try:
                config_manager = get_config_manager()
                if not config_manager.has_environments():
                    raise ValueError("No Superset environments configured. Please add an environment in Settings.")
                    
                env_config = config_manager.get_environment(env)
                if not env_config:
                    raise ValueError(f"Environment '{env}' not found in configuration.")
                
                client = SupersetClient(env_config)
                
                dashboard_count, dashboard_meta = client.get_dashboards()
                app_logger.info(f"[BackupPlugin][Progress] Found {dashboard_count} dashboards to export in {env}.")

                if dashboard_count == 0:
                    app_logger.info("[BackupPlugin][Exit] No dashboards to back up.")
                    return

                for db in dashboard_meta:
                    dashboard_id = db.get('id')
                    dashboard_title = db.get('dashboard_title', 'Unknown Dashboard')
                    if not dashboard_id:
                        continue

                    try:
                        dashboard_base_dir_name = sanitize_filename(f"{dashboard_title}")
                        dashboard_dir = backup_path / env.upper() / dashboard_base_dir_name
                        dashboard_dir.mkdir(parents=True, exist_ok=True)

                        zip_content, filename = client.export_dashboard(dashboard_id)

                        save_and_unpack_dashboard(
                            zip_content=zip_content,
                            original_filename=filename,
                            output_dir=dashboard_dir,
                            unpack=False
                        )

                        archive_exports(str(dashboard_dir), policy=RetentionPolicy())

                    except (SupersetAPIError, RequestException, IOError, OSError) as db_error:
                        app_logger.error(f"[BackupPlugin][Failure] Failed to export dashboard {dashboard_title} (ID: {dashboard_id}): {db_error}", exc_info=True)
                        continue
                
                consolidate_archive_folders(backup_path / env.upper())
                remove_empty_directories(str(backup_path / env.upper()))

                app_logger.info(f"[BackupPlugin][CoherenceCheck:Passed] Backup logic completed for {env}.")

            except (RequestException, IOError, KeyError) as e:
                app_logger.critical(f"[BackupPlugin][Failure] Fatal error during backup for {env}: {e}", exc_info=True)
                raise e
    # [/DEF:execute:Function]
# [/DEF:BackupPlugin:Class]
# [/DEF:BackupPlugin:Module]