dotfiles/python/tools/zfs_manager.py

"""zfs_manager.

Reconciles the live zfs datasets against a declaration generated by
common/optional/zfs_manager.nix. Datasets are created and properties are
corrected, but nothing is ever destroyed or renamed.
"""

from __future__ import annotations

import json
import logging
import sys
from pathlib import Path  # noqa: TC003 This is required for the typer CLI

import typer

from python.common import configure_logger
from python.signal_alert import signal_alert
from python.zfs import create_dataset, get_properties, list_dataset_names, set_property

logger = logging.getLogger(__name__)

# Properties that can only be chosen at creation time. Attempting to zfs set
# these fails on every run, so a mismatch is reported instead of retried.
CREATE_ONLY_PROPERTIES = frozenset(
    {
        "casesensitivity",
        "encryption",
        "keyformat",
        "normalization",
        "utf8only",
        "volblocksize",
    },
)

# Properties whose values zfs reports in bytes but which are conventionally
# declared with a size suffix, so "16k" and "16384" mean the same thing.
SIZE_PROPERTIES = frozenset(
    {
        "quota",
        "recordsize",
        "refquota",
        "refreservation",
        "reservation",
        "special_small_blocks",
        "volblocksize",
        "volsize",
    },
)

SIZE_SUFFIXES = {"b": 1, "k": 1024, "m": 1024**2, "g": 1024**3, "t": 1024**4, "p": 1024**5}

# Sources that mean the value was deliberately put on this dataset rather than
# inherited from a parent or left at the zfs default.
LOCAL_SOURCES = ("local", "received")


def main(config_file: Path, *, dry_run: bool = False) -> None:
    """Main.

    Args:
        config_file (Path): The path to the generated dataset declaration.
        dry_run (bool): Log the changes that would be made without making them.
    """
    configure_logger(level="DEBUG")
    logger.info(f"Starting zfs_manager {dry_run=}")

    try:
        declared = json.loads(config_file.read_text())["datasets"]
        existing = set(list_dataset_names())
        unusable: set[str] = set()

        # Parents before children so a newly created parent exists by the time
        # its children are reconciled.
        for name in sorted(declared, key=lambda name: (name.count("/"), name)):
            if has_unusable_parent(name, unusable):
                logger.error(f"skipping {name}, its parent could not be created")
                continue

            # Declared purely to record retention, its properties belong to
            # whoever set them.
            if not declared[name].get("manageProperties", True):
                logger.debug(f"{name} is declared but its properties are not managed")
                continue

            properties = declared[name]["properties"]
            if name in existing:
                reconcile_dataset(name, properties, dry_run=dry_run)
            elif create_missing_dataset(name, properties, dry_run=dry_run):
                existing.add(name)
            elif not dry_run:
                unusable.add(name)

        report_undeclared_datasets(existing, declared)
    except Exception:
        logger.exception("zfs_manager failed")
        signal_alert("zfs_manager failed")
        sys.exit(1)
    else:
        logger.info("zfs_manager completed")


def has_unusable_parent(name: str, unusable: set[str]) -> bool:
    """Check whether an ancestor of a dataset could not be created.

    Args:
        name (str): The name of the dataset.
        unusable (set[str]): The datasets that do not exist and could not be created.

    Returns:
        bool: True if any ancestor is unusable.
    """
    parts = name.split("/")
    return any("/".join(parts[:depth]) in unusable for depth in range(1, len(parts)))


def create_missing_dataset(name: str, properties: dict[str, str], *, dry_run: bool) -> bool:
    """Create a dataset that is declared but does not exist yet.

    Pool root filesystems are never created, creating a pool is out of scope.

    Args:
        name (str): The name of the dataset.
        properties (dict[str, str]): The properties to create the dataset with.
        dry_run (bool): Log the change without making it.

    Returns:
        bool: True if the dataset now exists.
    """
    if "/" not in name:
        message = f"pool {name} is declared but does not exist, zfs_manager does not create pools"
        logger.error(message)
        signal_alert(message)
        return False

    if dry_run:
        logger.info(f"would create {name} with {properties}")
        return False

    logger.info(f"creating {name} with {properties}")
    if error := create_dataset(name, properties):
        logger.error(error)
        signal_alert(error)
        return False

    return True


def reconcile_dataset(name: str, properties: dict[str, str], *, dry_run: bool) -> None:
    """Bring an existing dataset in line with its declared properties.

    Args:
        name (str): The name of the dataset.
        properties (dict[str, str]): The declared properties.
        dry_run (bool): Log the changes without making them.
    """
    current = get_properties(name)

    for key, wanted in sorted(properties.items()):
        current_value, _ = current.get(key, ("-", "-"))
        if values_match(key, wanted, current_value):
            continue

        if key in CREATE_ONLY_PROPERTIES:
            message = (
                f"{name} {key} is {current_value} but {wanted} is declared, "
                f"{key} can only be set when the dataset is created"
            )
            logger.error(message)
            signal_alert(message)
            continue

        if dry_run:
            logger.info(f"would set {key}={wanted} on {name}, currently {current_value}")
            continue

        logger.info(f"setting {key}={wanted} on {name}, was {current_value}")
        if error := set_property(name, key, wanted):
            logger.error(error)
            signal_alert(error)

    report_undeclared_properties(name, properties, current)


def report_undeclared_properties(name: str, properties: dict[str, str], current: dict[str, tuple[str, str]]) -> None:
    """Warn about properties set on the dataset but absent from the declaration.

    Inherited and default values are silent, they are not drift. A locally set
    value that nix does not know about was changed outside of this tool and
    will be lost the next time the dataset is recreated, so it is worth saying.

    Args:
        name (str): The name of the dataset.
        properties (dict[str, str]): The declared properties.
        current (dict[str, tuple[str, str]]): The live properties keyed to (value, source).
    """
    for key, (value, source) in sorted(current.items()):
        # User properties such as nixos:shutdown-time are written by other
        # tools and are not something a dataset declaration should own.
        if key in properties or ":" in key or source not in LOCAL_SOURCES:
            continue

        logger.warning(f"{name} has {key}={value} set outside of nix")
        signal_alert(f"{name} has {key}={value} set outside of nix")


def report_undeclared_datasets(existing: set[str], declared: dict[str, dict]) -> None:
    """Warn about datasets that exist but are not declared.

    These are left completely alone. They still get snapshots through the
    default retention table.

    Args:
        existing (set[str]): The names of every live dataset.
        declared (dict[str, dict]): The declaration.
    """
    for name in sorted(existing - set(declared)):
        logger.warning(f"{name} exists but is not declared in nix")


def values_match(key: str, wanted: str, current: str) -> bool:
    """Compare a declared property value against the live one.

    Args:
        key (str): The property name.
        wanted (str): The declared value.
        current (str): The live value.

    Returns:
        bool: True if the two values mean the same thing.
    """
    if key in SIZE_PROPERTIES:
        wanted_size = parse_size(wanted)
        current_size = parse_size(current)
        if wanted_size is not None and current_size is not None:
            return wanted_size == current_size

    return wanted == current


def parse_size(value: str) -> int | None:
    """Convert a zfs size such as 16k or 1M into bytes.

    Args:
        value (str): The size to convert.

    Returns:
        int | None: The size in bytes, or None if it is not a size.
    """
    value = value.strip()
    if value.isdigit():
        return int(value)

    number, suffix = value[:-1], value[-1:].lower()
    if suffix in SIZE_SUFFIXES and number.isdigit():
        return int(number) * SIZE_SUFFIXES[suffix]

    return None


def cli() -> None:
    """CLI."""
    typer.run(main)


if __name__ == "__main__":
    cli()