dotfiles/common/optional/zfs_manager.nix

{
  pkgs,
  lib,
  config,
  ...
}:
let
  cfg = config.services.zfs_manager;

  snapshotOptions = {
    options = {
      "15_min" = lib.mkOption {
        type = lib.types.ints.unsigned;
        default = 0;
        description = "How many 15 minute snapshots to keep.";
      };
      hourly = lib.mkOption {
        type = lib.types.ints.unsigned;
        default = 0;
        description = "How many hourly snapshots to keep.";
      };
      daily = lib.mkOption {
        type = lib.types.ints.unsigned;
        default = 0;
        description = "How many daily snapshots to keep.";
      };
      monthly = lib.mkOption {
        type = lib.types.ints.unsigned;
        default = 0;
        description = "How many monthly snapshots to keep.";
      };
    };
  };

  datasetOptions = {
    options = {
      manageProperties = lib.mkOption {
        type = lib.types.bool;
        default = true;
        description = ''
          Whether zfs_manager owns this dataset's properties. When false the
          dataset only contributes its snapshot retention, which is how
          root_pool datasets are declared.
        '';
      };
      createIfMissing = lib.mkOption {
        type = lib.types.bool;
        default = true;
        description = ''
          Whether zfs_manager may create this dataset when it is absent.

          Set it false for a dataset that has to be provisioned by hand, such
          as an encryption root: encryption is fixed at creation time and
          cannot be expressed here, so creating it automatically would silently
          produce an unencrypted dataset where an encrypted one was intended.
          The dataset is still property checked, and its absence is reported as
          a failure rather than quietly fixed.
        '';
      };
      properties = lib.mkOption {
        type = lib.types.attrsOf lib.types.str;
        default = { };
        description = ''
          The zfs properties this dataset should have. Values are compared
          against the live dataset and corrected when they differ.
        '';
      };
      snapshots = lib.mkOption {
        type = lib.types.submodule snapshotOptions;
        default = cfg.defaultSnapshots;
        description = ''
          Snapshot retention for this dataset. Defaults to defaultSnapshots.
        '';
      };
    };
  };

  # snapshot_manager.py only ever walks datasets below a pool root, so pool
  # roots are left out of the retention table. It also indexes the table
  # directly, which is why every entry carries all four keys.
  snapshotTable = lib.mapAttrs (_: dataset: dataset.snapshots) (
    lib.filterAttrs (name: _: lib.hasInfix "/" name) cfg.datasets
  );

  snapshotConfig = (pkgs.formats.toml { }).generate "snapshot_config.toml" (
    snapshotTable // { default = cfg.defaultSnapshots; }
  );

  # Every declared dataset is emitted, including the ones whose properties are
  # not managed, so the tool can tell "deliberately hands off" apart from
  # "nobody has written this down yet".
  datasetConfig = (pkgs.formats.json { }).generate "zfs_datasets.json" {
    datasets = lib.mapAttrs (_: dataset: {
      inherit (dataset) manageProperties createIfMissing properties;
    }) cfg.datasets;
  };
in
{
  options = {
    services.zfs_manager = {
      enable = lib.mkEnableOption "declarative ZFS dataset management";
      datasets = lib.mkOption {
        type = lib.types.attrsOf (lib.types.submodule datasetOptions);
        default = { };
        example = lib.literalExpression ''
          {
            "media/temp".properties = {
              sync = "disabled";
              redundant_metadata = "none";
            };
          }
        '';
        description = ''
          The datasets to manage, keyed by full dataset name. Missing datasets
          are created and drifted properties are corrected. Nothing is ever
          destroyed, and datasets that are not declared are left alone.

          A name without a "/" is a pool root filesystem. Its properties are
          managed but it is never created, pool creation stays manual.
        '';
      };
      defaultSnapshots = lib.mkOption {
        type = lib.types.submodule snapshotOptions;
        default = { };
        description = ''
          Retention for undeclared datasets and for declared datasets that do
          not override their snapshots. Emitted as the "default" table of the
          snapshot config.
        '';
      };
      dryRun = lib.mkOption {
        type = lib.types.bool;
        default = false;
        description = ''
          Log every change that would be made without touching zfs. Use this to
          validate a new or heavily edited declaration before applying it.
        '';
      };
      PYTHONPATH = lib.mkOption {
        type = lib.types.str;
        description = ''
          the PYTHONPATH to use for the zfs_manager service.
        '';
      };
      package = lib.mkOption {
        type = lib.types.package;
        default = pkgs.my_python;
        description = "Python environment used to run zfs_manager.";
      };
      EnvironmentFile = lib.mkOption {
        type = lib.types.nullOr (lib.types.coercedTo lib.types.path toString lib.types.str);
        default = null;

        description = ''
          Single environment file for the service (e.g. /etc/zfs-manager/env).
          Use a leading "-" to ignore if missing (systemd feature).
        '';
      };
    };
  };

  config = lib.mkIf cfg.enable {
    services.snapshot_manager.path = snapshotConfig;

    systemd = {
      services.zfs_manager = {
        description = "ZFS Dataset Manager";
        requires = [ "zfs-import.target" ];
        after = [
          "zfs-import.target"
          "zfs-mount.service"
        ];
        wantedBy = [ "multi-user.target" ];
        path = [ pkgs.zfs ];
        # Re-run on nixos-rebuild switch whenever the declaration changes.
        restartTriggers = [ datasetConfig ];
        environment = {
          PYTHONPATH = cfg.PYTHONPATH;
        };
        serviceConfig = {
          Type = "oneshot";
          RemainAfterExit = true;
          ExecStart = "${cfg.package}/bin/python -m python.tools.zfs_manager ${lib.escapeShellArg datasetConfig}${lib.optionalString cfg.dryRun " --dry-run"}";
        }
        // lib.optionalAttrs (cfg.EnvironmentFile != null) {
          EnvironmentFile = cfg.EnvironmentFile;
        };
      };
    };
  };
}