Attempt to improve zfs setup a bit

nix/modules/nixos/fs/zfs/default.nix

@@ -120,8 +120,32 @@ in
         "${cfg.mainPoolName}" = { };
       };
     };
+    services = {
+      postgresql = {
+        enable = lib.option {
+          description = "Enables storing postgresql data on a zfs zpool";
+          type = lib.types.bool;
+          default = cfg.enable && config.services.postgresql.enable;
         };
-  config = lib.mkIf cfg.enable {
+        pool = lib.mkOption {
+          type = lib.types.str;
+          default = cfg.mainPoolName;
+        };
+        datasetName = lib.mkOption {
+          type = lib.types.str;
+          default = "database/postgresql";
+        };
+        datasetConfig = lib.mkOption {
+          type = datasetModule;
+          default = {
+            mountpoint = config.services.postgresql.dataDir;
+          };
+        };
+      };
+    };
+  };
+  config = lib.mkMerge [
+    (lib.mkIf cfg.enable {
       # TODO: Verify that each member disk is uniquely named, and exists somewhere?
       assertions = lib.lists.map (
         { name, value }:
@@ -200,5 +224,22 @@ in
       ];
       services.prometheus.exporters.zfs.enable = true;
       khscodes.infrastructure.vault-prometheus-sender.exporters.enabled = [ "zfs" ];
+    })
+    (lib.mkIf (cfg.enable && cfg.services.postgresql.enable) {
+      khscodes.fs.zfs.zpools."${cfg.services.postgresql.pool
+      }".datasets."${cfg.services.postgresql.datasetName}" =
+        cfg.services.postgresql.datasetConfig;
+      systemd.services.postgresql = {
+        after = [ "khscodes-zpool-setup.service" ];
+        unitConfig = {
+          RequiresMountsFor = cfg.services.postgresql.datasetConfig.mountpoint;
         };
+      };
+      systemd.services.khscodes-zpool-setup = {
+        ExecStartPost = ''
+          chown ${config.services.postgresql.user}:${config.services.postgresql.group} ${lib.escapeShellArg cfg.services.postgresql.datasetConfig.mountpoint}
+        '';
+      };
+    })
+  ];
 }

rust/program/zpool-setup/src/main.rs

@@ -1,10 +1,13 @@
 use serde::Deserialize;
-use std::{collections::BTreeMap, path::PathBuf};
+use std::{
+    collections::{BTreeMap, HashMap},
+    path::{Path, PathBuf},
+};
 
 use anyhow::Context as _;
 use clap::{Parser, Subcommand};
 
-use crate::cli::Dataset;
+use crate::cli::{Dataset, VdevMode};
 
 mod cli;
 mod disk_mapping;
@@ -158,23 +161,52 @@ enum ZpoolState {
 #[serde(tag = "vdev_type")]
 enum ZpoolStatusVdev {
     #[serde(rename = "root")]
-    Root {
-        name: String,
-        state: ZpoolState,
-        vdevs: HashMap<String, ZpoolStatusVdev>,
-    },
+    Root(ZpoolStatusVdevRoot),
     #[serde(rename = "disk")]
-    Disk {
-        name: String,
-        state: ZpoolState,
-        path: PathBuf,
-    },
+    Disk(ZpoolStatusVdevDisk),
     #[serde(rename = "mirror")]
-    Mirror {
+    #[allow(dead_code)]
+    Mirror(ZpoolStatusVdevMirror),
+}
+
+impl ZpoolStatusVdev {
+    pub fn as_root(&self) -> anyhow::Result<&ZpoolStatusVdevRoot> {
+        match self {
+            Self::Root(root) => Ok(root),
+            _ => Err(anyhow::format_err!("VDev was not a root vdev")),
+        }
+    }
+    pub fn is_vdev_for_disk(&self, disk_path: &Path) -> bool {
+        matches!(self, Self::Disk(disk) if disk.path == disk_path)
+    }
+}
+
+#[derive(Deserialize)]
+struct ZpoolStatusVdevRoot {
+    #[allow(dead_code)]
     name: String,
+    #[allow(dead_code)]
     state: ZpoolState,
     vdevs: HashMap<String, ZpoolStatusVdev>,
 }
+
+#[derive(Deserialize)]
+struct ZpoolStatusVdevDisk {
+    #[allow(dead_code)]
+    name: String,
+    #[allow(dead_code)]
+    state: ZpoolState,
+    path: PathBuf,
+}
+
+#[derive(Deserialize)]
+struct ZpoolStatusVdevMirror {
+    #[allow(dead_code)]
+    name: String,
+    #[allow(dead_code)]
+    state: ZpoolState,
+    #[allow(dead_code)]
+    vdevs: HashMap<String, ZpoolStatusVdev>,
 }
 
 fn setup_zpool(p: SetupZpool) -> anyhow::Result<()> {
@@ -209,14 +241,33 @@ fn setup_zpool(p: SetupZpool) -> anyhow::Result<()> {
         return Err(anyhow::format_err!("Zpool {} is not online", p.pool_name));
     }
 
-    // TODO: Run through the existing VDevs and add any missing vdevs, and add any missing disks
-    // as needed to any vdevs. Not exactly sure how this should be coded, but I guess we can utilize
-    // the fact we cannot really change vdev type beyond turning a disk vdev into a mirror vdev,
-    // and any single disk can only belong to one vdev. So we can simply not support moving disks between vdevs.
-    // Also, to begin with, we can simply not support any vdev other than disk vdevs, as it doesn't make much
-    // sense for my use case.
+    let root_vdev = pool
+        .vdevs
+        .get(&p.pool_name)
+        .ok_or_else(|| anyhow::format_err!("Root vdev of pool not found"))?;
+    let root_vdev = root_vdev.as_root()?;
 
     for vdev in p.vdevs.iter() {
+        if vdev.mode != VdevMode::Mirror {
+            return Err(anyhow::format_err!(
+                "Vdev contains non mirror mode vdev, this is currently not supported"
+            ));
+        }
+        if vdev.members.len() != 1 {
+            return Err(anyhow::format_err!(
+                "Vdev contains more than one member, this is currently not supported"
+            ));
+        }
+        let main_member = &vdev.members[0];
+        let resolved_main_member = disk_mapping.resolve(main_member)?;
+        let resolved_main_member = PathBuf::from(resolved_main_member);
+        if !root_vdev
+            .vdevs
+            .iter()
+            .any(|(_, vdev)| vdev.is_vdev_for_disk(&resolved_main_member))
+        {
+            zfs::add_vdev_to_pool(&p.pool_name, &disk_mapping, vdev)?;
+        }
         for member in vdev.members.iter() {
             let resolved = disk_mapping.resolve(member)?;
             zfs::resize_disk(&p.pool_name, &resolved)?;

rust/program/zpool-setup/src/zfs.rs

@@ -4,7 +4,11 @@ use anyhow::Context as _;
 use common::proc::Command;
 use serde::Deserialize;
 
-use crate::{SetupZpool, cli::Dataset, disk_mapping::DiskMapping};
+use crate::{
+    SetupZpool,
+    cli::{Dataset, Vdev},
+    disk_mapping::DiskMapping,
+};
 
 #[derive(Debug, Deserialize, PartialEq)]
 enum ZpoolState {
@@ -88,6 +92,19 @@ pub fn create_pool(
     Ok(())
 }
 
+pub fn add_vdev_to_pool(
+    pool_name: &str,
+    disk_mapping: &DiskMapping,
+    vdev: &Vdev,
+) -> anyhow::Result<()> {
+    let mut proc = Command::new("zpool");
+    proc.args(["add", pool_name]);
+    proc.args(vdev.cli_args(disk_mapping)?);
+
+    proc.try_spawn_to_bytes()?;
+    Ok(())
+}
+
 pub fn create_dataset_recursive(
     pool_name: &str,
     dataset_name: &str,