zfs-backup/zfs_backup/core.py

from __future__ import annotations

from subprocess import Popen

import config
from dataset import Dataset
from pool import Pool, ExternalPool
from snapshot import Snapshot
from zfs import ZFS


class Manager:
    """Controls the overall backup process."""

    def __init__(self, local_pool: Pool, external_pool: ExternalPool):
        self._local_pool = local_pool
        self._external_pool = external_pool

    def backup(self) -> None:
        """Backup workflow."""
        self._external_pool.import_()
        try:
            self._external_pool.scrub(skip_if_recently_scrubbed=True)
            self._backup_all_datasets()
            self._external_pool.clean_old_snapshots()
            self._external_pool.scrub(skip_if_recently_scrubbed=False)
        finally:
            self._external_pool.export()

    def _backup_all_datasets(self) -> None:
        """Send snapshots of local datasets to the backup pool."""
        print("Backing up datasets")
        for (
            local_dataset,
            remote_dataset,
        ) in self._get_user_filtered_dataset_backup_map().items():
            if remote_dataset is None:
                stream = self._send_absolute(local_dataset)
            else:
                stream = self._send_incremental(local_dataset, remote_dataset)
                if stream is None:
                    continue
            ZFS.receive(stream, self._external_pool)

    def _get_user_filtered_dataset_backup_map(self) -> dict[Dataset, Dataset | None]:
        """Get the dataset map, let user confirm new datasets."""
        return self._user_confirm_new_datasets(self._get_raw_dataset_map())

    def _user_confirm_new_datasets(
        self, dataset_map: dict[Dataset, Dataset | None]
    ) -> dict[Dataset, Dataset | None]:
        """For datasets that exist only locally, ask the user if he wants to back it up.

        If he doesn't, remove it from the dict.
        """
        return {
            local_dataset: remote_dataset
            for local_dataset, remote_dataset in dataset_map.items()
            if remote_dataset is not None or self._ask_user_confirmation(local_dataset)
        }

    def _get_raw_dataset_map(self) -> dict[Dataset, Dataset | None]:
        """Return map from local to remote datasets.

        If the local dataset has no remote match, set the remote dataset to None.
        """
        return {
            local_dataset: self._search_matching_dataset_in_remote_pool(
                local_dataset, self._external_pool
            )
            for local_dataset in self._datasets_to_backup
        }

    def _send_incremental(
        self, local_dataset: Dataset, remote_dataset: Dataset
    ) -> Popen | None:
        """Send incremental stream between two datasets to remote, return the pipe. If local and remote snapshot are identical, perform nothing and return None."""
        start_snapshot = self._get_newest_common_snapshot_with_backup_tags(
            local_dataset, remote_dataset
        )
        end_snapshot = self._get_newest_snapshot_with_backup_tag(local_dataset)
        if start_snapshot == end_snapshot:
            return None
        stream = ZFS.send_incremental(start_snapshot, end_snapshot)
        return stream

    def _send_absolute(self, local_dataset: Dataset) -> Popen:
        """Send absolute stream to remote."""
        snapshot = self._get_newest_snapshot_with_backup_tag(local_dataset)
        return ZFS.send_absolute(snapshot)

    def _get_newest_snapshot_with_backup_tag(self, dataset: Dataset) -> Snapshot:
        """From all snapshots in the dataset, get the newest one that matches the backup tags in the config."""
        snapshots = self._find_snapshots_with_backup_tag(dataset)
        snapshots.sort()
        return snapshots[-1]

    @property
    def _datasets_to_backup(self) -> list[Dataset]:
        """Return local datasets to consider for a backup."""
        return [
            dataset
            for dataset in self._local_pool.datasets
            if dataset.qualified_name not in config.do_not_backup
        ]

    @staticmethod
    def _find_snapshots_with_backup_tag(dataset: Dataset) -> list[Snapshot]:
        """Find snapshots of a dataset that carry the backup tags."""
        regex = _get_regex_matching_snapshots_with_tags(
            [config.snapshot_tag, config.snapshot_interval]
        )
        return [
            snapshot for snapshot in dataset.snapshots if snapshot.matches_regex(regex)
        ]

    @staticmethod
    def _search_matching_dataset_in_remote_pool(
        dataset: Dataset, pool: Pool
    ) -> Dataset | None:
        """For a dataset in the local pool, get the matching dataset in the remote pool.

        If there is no match, return None.
        """
        dataset_to_search = dataset.replace_pool(pool.name)
        for dataset in pool.datasets:
            if dataset == dataset_to_search:
                return dataset

    @classmethod
    def _get_newest_common_snapshot_with_backup_tags(
        cls, local_dataset: Dataset, remote_dataset: Dataset
    ) -> Snapshot:
        """In a local and a remote dataset, find the newest snapshot with backup tags which exists in both datasets.

        Return the local snapshot."""
        common_snapshots = [
            local_snap
            for local_snap in list(cls._find_snapshots_with_backup_tag(local_dataset))
            for remote_snap in list(cls._find_snapshots_with_backup_tag(remote_dataset))
            if local_snap.snapshot_name == remote_snap.snapshot_name
        ]
        common_snapshots.sort()
        return common_snapshots[-1]

    @staticmethod
    def _ask_user_confirmation(dataset: Dataset) -> bool:  # type: ignore
        msg = f"Dataset {dataset.qualified_name} does not exist in backup pool. Backup? [y/n] "
        reply = ""
        while not reply:
            reply = input(msg).lower()
            if reply == "y":
                return True
            elif reply == "n":
                return False
            else:
                reply = ""


def _get_regex_matching_snapshots_with_tags(tags: list[str]):
    return "@" + ".*".join(tags)