Add redbot-setup restore cli command (#6709)

redbot/core/_downloader/__init__.py

@@ -12,6 +12,7 @@ from __future__ import annotations
 
 import contextlib
 import dataclasses
+import json
 import os
 import shutil
 import sys
@@ -37,6 +38,7 @@ import discord
 from redbot.core import commands, Config, version_info as red_version_info
 from redbot.core._cog_manager import CogManager
 from redbot.core.data_manager import cog_data_path
+from redbot.core.utils._internal_utils import detailed_progress
 
 from . import errors
 from .log import log
@@ -864,3 +866,157 @@ class CogUnavailableError(Exception):
         self.repo_name = repo_name
         self.cog_name = cog_name
         super().__init__(f"Couldn't find cog {cog_name!r} in {repo_name!r}")
+
+
+async def _restore_from_backup() -> None:
+    """Restore cogs using `repos.json` in cog's data path.
+
+    Used by `redbot-setup restore` cli command.
+    """
+    with _repo_manager.repos_folder.with_name("repos.json").open(encoding="utf-8") as fp:
+        raw_repos = json.load(fp)
+
+    with detailed_progress(unit="repos") as progress:
+        task_id = progress.add_task("Adding repos")
+        for repo_data in progress.track(raw_repos, task_id=task_id):
+            repo_url = repo_data["url"]
+            repo_name = repo_data["name"]
+            repo_branch = repo_data["branch"]
+            progress.update(task_id, description=f"Adding {repo_name!r} repo")
+            try:
+                await _repo_manager.add_repo(repo_url, repo_name, repo_branch)
+            except errors.ExistingGitRepo:
+                # this should not be possible
+                log.error(
+                    "Failed to add repo %r (url: %r branch: %r) as it seems that"
+                    " one with this name already exists.",
+                    repo_name,
+                    repo_url,
+                    repo_branch,
+                )
+            except errors.AuthenticationError:
+                log.error(
+                    "Failed to add repo %r (url: %r branch: %r) due to authentication failure."
+                    " This may also mean that repository no longer exists at that URL.",
+                    repo_name,
+                    repo_url,
+                    repo_branch,
+                )
+            except (errors.CloningError, OSError):
+                log.error(
+                    "Failed to add repo %r (url: %r branch: %r) due to a cloning error.",
+                    repo_name,
+                    repo_url,
+                    repo_branch,
+                )
+
+    _installed_cogs = await installed_cogs()
+    await _config.installed_cogs.clear()
+    await _config.installed_libraries.clear()
+
+    cogs_to_reinstall = []
+    cogs_without_repo: Dict[str, List[str]] = defaultdict(list)
+    for cog in _installed_cogs:
+        if cog.repo is None:
+            cogs_without_repo[cog._json_repo_name].append(cog.name)
+        else:
+            cogs_to_reinstall.append(cog)
+    for repo_name, cogs in cogs_without_repo.items():
+        log.error(
+            "The backup does not contain metadata about repo %r,"
+            " the following cogs cannot be reinstalled automatically: %s",
+            repo_name,
+            ", ".join(cogs),
+        )
+
+    with detailed_progress(unit="cogs") as progress:
+        task_id = progress.add_task("Installing cogs")
+        cog: Installable
+        for cog in progress.track(cogs_to_reinstall, task_id=task_id):
+            progress.update(task_id, description=f"Installing {cog.name!r} cog")
+            if not cog.commit:
+                last_cog_occurrence = await cog.repo.get_last_module_occurrence(cog.name)
+                if last_cog_occurrence is not None and not last_cog_occurrence.disabled:
+                    cog = last_cog_occurrence
+                    log.warning(
+                        "The commit that %r cog was installed from is unknown"
+                        " - will try to reinstall from latest commit where it's still available."
+                    )
+                else:
+                    log.error(
+                        "The commit that %r cog was installed from is unknown"
+                        " and it could not be found in the repo."
+                    )
+                    continue
+
+            try:
+                install_result = await install_cogs(cog.repo, cog.commit, [cog.name])
+            except errors.UnknownRevision:
+                log.error(
+                    "The commit that %r cog was installed from (%s)"
+                    " could not be found in the repo.",
+                    cog.name,
+                    cog.commit,
+                )
+                continue
+
+            # we know we've only tried to install one cog so we can be specific in error messages
+
+            # below 3 should technically never happen with valid config but just in case...
+            if install_result.unavailable_cogs:
+                log.error("Could not find %r cog in %r repo", cog.name, cog.repo.name)
+            if install_result.already_installed:
+                log.error(
+                    "Failed to reinstall %r cog from %r repo as it seems that"
+                    " this cog is already installed.",
+                    cog.name,
+                    cog.repo.name,
+                )
+            if install_result.name_already_used:
+                log.error(
+                    "Failed to reinstall %r cog from %r repo as it seems that"
+                    " a cog with the same name is already installed from a different repo.",
+                    cog.name,
+                    cog.repo.name,
+                )
+
+            if install_result.incompatible_python_version:
+                log.error(
+                    "Failed to reinstall %r cog from %r repo because the instance is"
+                    " being restored into Red running with a lower Python version."
+                    " The minimum Python version required by this cog is %s",
+                    cog.name,
+                    cog.repo.name,
+                    ".".join(map(str, cog.min_python_version)),
+                )
+            if install_result.incompatible_bot_version:
+                log.error(
+                    "Failed to reinstall %r cog from %r repo because the instance is"
+                    " being restored into a Red version that the cog does not support."
+                    " The minimum version required by this cog is %s%s.",
+                    cog.name,
+                    cog.repo.name,
+                    cog.min_bot_version,
+                    (
+                        ""
+                        if cog.min_bot_version > cog.max_bot_version
+                        else f"and maximum allowed is: {cog.max_bot_version}"
+                    ),
+                )
+
+            if install_result.failed_cogs:
+                log.info("Failed to reinstall %r cog", cog.name)
+            if install_result.failed_reqs:
+                log.error(
+                    "Failed to reinstall %r cog from %r repo"
+                    " because the following requirements could not be reinstalled: %s",
+                    cog.name,
+                    cog.repo.name,
+                    ", ".join(install_result.failed_reqs),
+                )
+            if install_result.failed_libs:
+                log.error(
+                    "Failed to reinstall shared libraries for %r cog from %r repo: %s",
+                    cog.repo.name,
+                    ", ".join([lib.name for lib in install_result.failed_libs]),
+                )

redbot/core/utils/_internal_utils.py

@@ -9,9 +9,12 @@ import os
 import re
 import shutil
 import tarfile
+import time
 import warnings
 from datetime import datetime
+from io import BytesIO
 from pathlib import Path
+from tarfile import TarInfo
 from typing import (
     AsyncIterable,
     AsyncIterator,
@@ -24,6 +27,7 @@ from typing import (
     Optional,
     Union,
     TypeVar,
+    TypedDict,
     TYPE_CHECKING,
     Tuple,
     cast,
@@ -33,8 +37,9 @@ import aiohttp
 import discord
 from packaging.requirements import Requirement
 import rapidfuzz
-from rich.progress import ProgressColumn
-from rich.progress_bar import ProgressBar
+import rich.progress
+from rich.console import Console
+from rich.text import Text
 from red_commons.logging import VERBOSE, TRACE
 
 from redbot import VersionInfo
@@ -58,6 +63,8 @@ __all__ = (
     "fetch_latest_red_version_info",
     "deprecated_removed",
     "RichIndefiniteBarColumn",
+    "RichSpeedColumn",
+    "detailed_progress",
     "cli_level_to_log_level",
 )
 
@@ -216,7 +223,27 @@ async def format_fuzzy_results(
         return "Perhaps you wanted one of these? " + box("\n".join(lines), lang="vhdl")
 
 
+def _tar_addfile_from_string(tar: tarfile.TarFile, name: str, string: str) -> None:
+    encoded = string.encode("utf-8")
+    fp = BytesIO(encoded)
+
+    # TarInfo needs `mtime` and `size`
+    # https://stackoverflow.com/q/53306000
+    tar_info = tarfile.TarInfo(name)
+    tar_info.mtime = time.time()
+    tar_info.size = len(encoded)
+
+    tar.addfile(tar_info, fp)
+
+
+class BackupDetails(TypedDict):
+    backup_version: int
+
+
 async def create_backup(dest: Path = Path.home()) -> Optional[Path]:
+    # version of backup
+    BACKUP_VERSION = 2
+
     data_path = Path(data_manager.core_data_path().parent)
     if not data_path.exists():
         return None
@@ -226,13 +253,20 @@ async def create_backup(dest: Path = Path.home()) -> Optional[Path]:
     backup_fpath = dest / f"redv3_{data_manager.instance_name()}_{timestr}.tar.gz"
 
     to_backup = []
+    # we need trailing separator to not exclude files and folders that only start with these names
     exclusions = [
         "__pycache__",
+        # Lavalink will be downloaded on Audio load
         "Lavalink.jar",
-        os.path.join("Downloader", "lib"),
-        os.path.join("CogManager", "cogs"),
-        os.path.join("RepoManager", "repos"),
-        os.path.join("Audio", "logs"),
+        # cogs and repos installed through Downloader can be reinstalled using restore command
+        os.path.join("Downloader", "lib", ""),
+        os.path.join("CogManager", "cogs", ""),
+        os.path.join("RepoManager", "repos", ""),
+        os.path.join("Audio", "logs", ""),
+        # these files are created during backup so we exclude them from data path backup
+        os.path.join("RepoManager", "repos.json"),
+        "instance.json",
+        "backup_details.json",
     ]
 
     # Avoiding circular imports
@@ -243,19 +277,42 @@ async def create_backup(dest: Path = Path.home()) -> Optional[Path]:
     repo_output = []
     for repo in repo_mgr.repos:
         repo_output.append({"url": repo.url, "name": repo.name, "branch": repo.branch})
-    repos_file = data_path / "cogs" / "RepoManager" / "repos.json"
-    with repos_file.open("w") as fs:
-        json.dump(repo_output, fs, indent=4)
-    instance_file = data_path / "instance.json"
-    with instance_file.open("w") as fs:
-        json.dump({data_manager.instance_name(): data_manager.basic_config}, fs, indent=4)
-    for f in data_path.glob("**/*"):
-        if not any(ex in str(f) for ex in exclusions) and f.is_file():
-            to_backup.append(f)
 
-    with tarfile.open(str(backup_fpath), "w:gz") as tar:
-        for f in to_backup:
-            tar.add(str(f), arcname=str(f.relative_to(data_path)), recursive=False)
+    with rich.progress.Progress(
+        rich.progress.SpinnerColumn(),
+        rich.progress.TextColumn("[progress.description]{task.description}"),
+        RichIndefiniteBarColumn(),
+        rich.progress.TextColumn("{task.completed} files processed"),
+        rich.progress.TimeElapsedColumn(),
+    ) as progress:
+        for f in progress.track(
+            data_path.glob("**/*"), description="Preparing files for backup..."
+        ):
+            if not any(ex in str(f) for ex in exclusions) and f.is_file():
+                to_backup.append(f)
+
+    backup_details: BackupDetails = {
+        "backup_version": BACKUP_VERSION,
+    }
+
+    with tarfile.open(str(backup_fpath), "w:gz", dereference=True) as tar:
+        with detailed_progress(unit="files") as progress:
+            progress_tracker = progress.track(to_backup, description="Compressing data")
+            for f in progress_tracker:
+                tar.add(str(f), arcname=str(f.relative_to(data_path)), recursive=False)
+
+        # add repos backup
+        repos_data = json.dumps(repo_output, indent=4)
+        _tar_addfile_from_string(tar, "cogs/RepoManager/repos.json", repos_data)
+
+        # add instance's original data
+        instance_data = json.dumps(
+            {data_manager.instance_name(): data_manager.basic_config}, indent=4
+        )
+        _tar_addfile_from_string(tar, "instance.json", instance_data)
+
+        # add info about backup version
+        _tar_addfile_from_string(tar, "backup_details.json", json.dumps(backup_details))
     return backup_fpath
 
 
@@ -367,10 +424,10 @@ def deprecated_removed(
     )
 
 
-class RichIndefiniteBarColumn(ProgressColumn):
-    def render(self, task):
-        return ProgressBar(
-            pulse=task.completed < task.total,
+class RichIndefiniteBarColumn(rich.progress.ProgressColumn):
+    def render(self, task: rich.progress.Task) -> rich.progress.ProgressBar:
+        return rich.progress.ProgressBar(
+            pulse=task.completed < task.total if task.total is not None else True,
             animation_time=task.get_time(),
             width=40,
             total=task.total,
@@ -378,6 +435,33 @@ class RichIndefiniteBarColumn(ProgressColumn):
         )
 
 
+class RichSpeedColumn(rich.progress.ProgressColumn):
+    def __init__(self, *, unit: str) -> None:
+        self.unit = unit
+        super().__init__()
+
+    def render(self, task: rich.progress.Task) -> Text:
+        speed = task.finished_speed or task.speed
+        if speed is None:
+            return Text("?", style="progress.data.speed")
+        return Text(f"{int(speed)} {self.unit}/s", style="progress.data.speed")
+
+
+def detailed_progress(*, unit: str, console: Optional[Console] = None) -> rich.progress.Progress:
+    return rich.progress.Progress(
+        rich.progress.SpinnerColumn(),
+        rich.progress.TextColumn("[progress.description]{task.description}"),
+        rich.progress.BarColumn(bar_width=None),
+        RichSpeedColumn(unit=unit),
+        rich.progress.TaskProgressColumn(),
+        rich.progress.TextColumn("eta"),
+        rich.progress.TimeRemainingColumn(),
+        rich.progress.TextColumn("elapsed"),
+        rich.progress.TimeElapsedColumn(),
+        console=console,
+    )
+
+
 def cli_level_to_log_level(level: int) -> int:
     if level == 0:
         log_level = logging.INFO