Move nix::NixCommand to util::CommandExt

src/util.rs

@@ -1,12 +1,15 @@
+use std::convert::TryFrom;
 use std::path::PathBuf;
 use std::process::Stdio;
 
+use async_trait::async_trait;
 use clap::{App, Arg, ArgMatches};
 use futures::future::join3;
+use serde::de::DeserializeOwned;
 use tokio::io::{AsyncRead, AsyncBufReadExt, BufReader};
 use tokio::process::Command;
 
-use super::nix::{Flake, Hive, HivePath, NixResult};
+use super::nix::{Flake, Hive, HivePath, NixResult, NixError, StorePath};
 use super::nix::deployment::TargetNodeMap;
 use super::job::JobHandle;
 
@@ -19,6 +22,22 @@ pub struct CommandExecution {
     stderr: Option<String>,
 }
 
+/// Helper extensions for Commands.
+#[async_trait]
+pub trait CommandExt {
+    /// Runs the command with stdout and stderr passed through to the user.
+    async fn passthrough(&mut self) -> NixResult<()>;
+
+    /// Runs the command, capturing the output as a String.
+    async fn capture_output(&mut self) -> NixResult<String>;
+
+    /// Runs the command, capturing deserialized output from JSON.
+    async fn capture_json<T>(&mut self) -> NixResult<T> where T: DeserializeOwned;
+
+    /// Runs the command, capturing a single store path.
+    async fn capture_store_path(&mut self) -> NixResult<StorePath>;
+}
+
 impl CommandExecution {
     pub fn new(command: Command) -> Self {
         Self {
@@ -81,6 +100,86 @@ impl CommandExecution {
     }
 }
 
+#[async_trait]
+impl CommandExt for Command {
+    /// Runs the command with stdout and stderr passed through to the user.
+    async fn passthrough(&mut self) -> NixResult<()> {
+        let exit = self
+            .spawn()?
+            .wait()
+            .await?;
+
+        if exit.success() {
+            Ok(())
+        } else {
+            Err(exit.into())
+        }
+    }
+
+    /// Captures output as a String.
+    async fn capture_output(&mut self) -> NixResult<String> {
+        // We want the user to see the raw errors
+        let output = self
+            .stdout(Stdio::piped())
+            .stderr(Stdio::inherit())
+            .spawn()?
+            .wait_with_output()
+            .await?;
+
+        if output.status.success() {
+            // FIXME: unwrap
+            Ok(String::from_utf8(output.stdout).unwrap())
+        } else {
+            Err(output.status.into())
+        }
+    }
+
+    /// Captures deserialized output from JSON.
+    async fn capture_json<T>(&mut self) -> NixResult<T> where T: DeserializeOwned {
+        let output = self.capture_output().await?;
+        serde_json::from_str(&output).map_err(|_| NixError::BadOutput {
+            output: output.clone()
+        })
+    }
+
+    /// Captures a single store path.
+    async fn capture_store_path(&mut self) -> NixResult<StorePath> {
+        let output = self.capture_output().await?;
+        let path = output.trim_end().to_owned();
+        StorePath::try_from(path)
+    }
+}
+
+#[async_trait]
+impl CommandExt for CommandExecution {
+    async fn passthrough(&mut self) -> NixResult<()> {
+        self.run().await
+    }
+
+    /// Captures output as a String.
+    async fn capture_output(&mut self) -> NixResult<String> {
+        self.run().await?;
+        let (stdout, _) = self.get_logs();
+
+        Ok(stdout.unwrap().to_owned())
+    }
+
+    /// Captures deserialized output from JSON.
+    async fn capture_json<T>(&mut self) -> NixResult<T> where T: DeserializeOwned {
+        let output = self.capture_output().await?;
+        serde_json::from_str(&output).map_err(|_| NixError::BadOutput {
+            output: output.clone()
+        })
+    }
+
+    /// Captures a single store path.
+    async fn capture_store_path(&mut self) -> NixResult<StorePath> {
+        let output = self.capture_output().await?;
+        let path = output.trim_end().to_owned();
+        StorePath::try_from(path)
+    }
+}
+
 pub async fn hive_from_args(args: &ArgMatches) -> NixResult<Hive> {
     let path = match args.occurrences_of("config") {
         0 => {