disks: use builder pattern for querying disk usage

Signed-off-by: Hannes Laimer <h.laimer@proxmox.com>

src/api2/node/disks/directory.rs

@@ -13,8 +13,8 @@ use pbs_api_types::{
 };
 
 use crate::tools::disks::{
-    create_file_system, create_single_linux_partition, get_disk_usage_info, get_fs_uuid,
-    DiskManage, DiskUsageType, FileSystemType,
+    create_file_system, create_single_linux_partition, get_fs_uuid, DiskManage, DiskUsageQuery,
+    DiskUsageType, FileSystemType,
 };
 use crate::tools::systemd::{self, types::*};
 
@@ -147,7 +147,7 @@ pub fn create_datastore_disk(
 
     let auth_id = rpcenv.get_auth_id().unwrap();
 
-    let info = get_disk_usage_info(&disk, true, false)?;
+    let info = DiskUsageQuery::new().smart(false).find(&disk)?;
 
     if info.used != DiskUsageType::Unused {
         bail!("disk '{}' is already in use.", disk);

src/api2/node/disks/mod.rs

@@ -12,8 +12,8 @@ use pbs_api_types::{
 };
 
 use crate::tools::disks::{
-    get_disk_usage_info, get_disks, get_smart_data, inititialize_gpt_disk, DiskManage,
-    DiskUsageInfo, DiskUsageType, SmartData,
+    get_smart_data, inititialize_gpt_disk, DiskManage, DiskUsageInfo, DiskUsageQuery,
+    DiskUsageType, SmartData,
 };
 use proxmox_rest_server::WorkerTask;
 
@@ -64,7 +64,11 @@ pub fn list_disks(
 ) -> Result<Vec<DiskUsageInfo>, Error> {
     let mut list = Vec::new();
 
-    for (_, info) in get_disks(None, skipsmart, include_partitions)? {
+    for (_, info) in DiskUsageQuery::new()
+        .smart(!skipsmart)
+        .partitions(include_partitions)
+        .query()?
+    {
         if let Some(ref usage_type) = usage_type {
             if info.used == *usage_type {
                 list.push(info);
@@ -147,7 +151,7 @@ pub fn initialize_disk(
 
     let auth_id = rpcenv.get_auth_id().unwrap();
 
-    let info = get_disk_usage_info(&disk, true, false)?;
+    let info = DiskUsageQuery::new().find(&disk)?;
 
     if info.used != DiskUsageType::Unused {
         bail!("disk '{}' is already in use.", disk);

src/api2/node/disks/zfs.rs

@@ -174,7 +174,7 @@ pub fn create_zpool(
         .map(|v| v.as_str().unwrap().to_string())
         .collect();
 
-    let disk_map = crate::tools::disks::get_disks(None, true, false)?;
+    let disk_map = crate::tools::disks::DiskUsageQuery::new().query()?;
     for disk in devices.iter() {
         match disk_map.get(disk) {
             Some(info) => {

src/tools/disks/mod.rs

@@ -763,19 +763,44 @@ fn scan_partitions(
     Ok(used)
 }
 
-/// Get disk usage information for a single disk
-pub fn get_disk_usage_info(
-    disk: &str,
-    no_smart: bool,
-    include_partitions: bool,
-) -> Result<DiskUsageInfo, Error> {
-    let mut filter = Vec::new();
-    filter.push(disk.to_string());
-    let mut map = get_disks(Some(filter), no_smart, include_partitions)?;
-    if let Some(info) = map.remove(disk) {
-        Ok(info)
-    } else {
-        bail!("failed to get disk usage info - internal error"); // should not happen
+pub struct DiskUsageQuery {
+    smart: bool,
+    partitions: bool,
+}
+
+impl DiskUsageQuery {
+    pub fn new() -> Self {
+        Self {
+            smart: true,
+            partitions: false,
+        }
+    }
+
+    pub fn smart(&mut self, smart: bool) -> &mut Self {
+        self.smart = smart;
+        self
+    }
+
+    pub fn partitions(&mut self, partitions: bool) -> &mut Self {
+        self.partitions = partitions;
+        self
+    }
+
+    pub fn query(&self) -> Result<HashMap<String, DiskUsageInfo>, Error> {
+        get_disks(None, !self.smart, self.partitions)
+    }
+
+    pub fn find(&self, disk: &str) -> Result<DiskUsageInfo, Error> {
+        let mut map = get_disks(Some(vec![disk.to_string()]), !self.smart, self.partitions)?;
+        if let Some(info) = map.remove(disk) {
+            Ok(info)
+        } else {
+            bail!("failed to get disk usage info - internal error"); // should not happen
+        }
+    }
+
+    pub fn find_all(&self, disks: Vec<String>) -> Result<HashMap<String, DiskUsageInfo>, Error> {
+        get_disks(Some(disks), !self.smart, self.partitions)
     }
 }
 
@@ -838,7 +863,7 @@ fn get_partitions_info(
 }
 
 /// Get disk usage information for multiple disks
-pub fn get_disks(
+fn get_disks(
     // filter - list of device names (without leading /dev)
     disks: Option<Vec<String>>,
     // do no include data from smartctl