From 16ed9d8c664f1b33bcb88c38881967e225bc4e70 Mon Sep 17 00:00:00 2001 From: Zhaofeng Li Date: Sat, 8 Jan 2022 01:20:36 -0800 Subject: [PATCH] Move nix::NixCommand to util::CommandExt --- src/nix/hive/mod.rs | 3 +- src/nix/host/local.rs | 4 +- src/nix/host/ssh.rs | 4 +- src/nix/mod.rs | 97 +--------------------------------------- src/nix/store.rs | 3 +- src/util.rs | 101 +++++++++++++++++++++++++++++++++++++++++- 6 files changed, 109 insertions(+), 103 deletions(-) diff --git a/src/nix/hive/mod.rs b/src/nix/hive/mod.rs index 6f66e85..46aa805 100644 --- a/src/nix/hive/mod.rs +++ b/src/nix/hive/mod.rs @@ -19,8 +19,7 @@ use super::{ StorePath, }; use super::deployment::TargetNode; -use super::NixCommand; -use crate::util::CommandExecution; +use crate::util::{CommandExecution, CommandExt}; use crate::job::JobHandle; #[cfg(test)] diff --git a/src/nix/host/local.rs b/src/nix/host/local.rs index dfc445e..b4a8dfe 100644 --- a/src/nix/host/local.rs +++ b/src/nix/host/local.rs @@ -6,8 +6,8 @@ use async_trait::async_trait; use tokio::process::Command; use super::{CopyDirection, CopyOptions, Host, key_uploader}; -use crate::nix::{StorePath, Profile, Goal, NixError, NixResult, NixCommand, Key, SYSTEM_PROFILE, CURRENT_PROFILE}; -use crate::util::CommandExecution; +use crate::nix::{StorePath, Profile, Goal, NixError, NixResult, Key, SYSTEM_PROFILE, CURRENT_PROFILE}; +use crate::util::{CommandExecution, CommandExt}; use crate::job::JobHandle; /// The local machine running Colmena. diff --git a/src/nix/host/ssh.rs b/src/nix/host/ssh.rs index bd8e097..0e9d921 100644 --- a/src/nix/host/ssh.rs +++ b/src/nix/host/ssh.rs @@ -7,8 +7,8 @@ use async_trait::async_trait; use tokio::process::Command; use super::{CopyDirection, CopyOptions, Host, key_uploader}; -use crate::nix::{StorePath, Profile, Goal, NixResult, NixCommand, NixError, Key, SYSTEM_PROFILE, CURRENT_PROFILE}; -use crate::util::CommandExecution; +use crate::nix::{StorePath, Profile, Goal, NixResult, NixError, Key, SYSTEM_PROFILE, CURRENT_PROFILE}; +use crate::util::{CommandExecution, CommandExt}; use crate::job::JobHandle; /// A remote machine connected over SSH. diff --git a/src/nix/mod.rs b/src/nix/mod.rs index 853cb5f..a732281 100644 --- a/src/nix/mod.rs +++ b/src/nix/mod.rs @@ -1,21 +1,16 @@ use std::collections::HashMap; -use std::convert::TryFrom; use std::hash::Hash; use std::ops::Deref; use std::os::unix::process::ExitStatusExt; use std::path::Path; -use std::process::{ExitStatus, Stdio}; +use std::process::ExitStatus; -use async_trait::async_trait; -use serde::de::{self, DeserializeOwned}; +use serde::de; use serde::{Deserialize, Deserializer, Serialize}; use snafu::Snafu; -use tokio::process::Command; use users::get_current_username; use validator::{Validate, ValidationErrors, ValidationError as ValidationErrorType}; -use crate::util::CommandExecution; - pub mod host; pub use host::{Host, CopyDirection, CopyOptions}; use host::Ssh; @@ -181,14 +176,6 @@ pub struct NodeConfig { keys: HashMap, } -#[async_trait] -trait NixCommand { - async fn passthrough(&mut self) -> NixResult<()>; - async fn capture_output(&mut self) -> NixResult; - async fn capture_json(&mut self) -> NixResult where T: DeserializeOwned; - async fn capture_store_path(&mut self) -> NixResult; -} - impl NodeName { /// Returns the string. pub fn as_str(&self) -> &str { @@ -258,86 +245,6 @@ impl NodeConfig { } } -#[async_trait] -impl NixCommand 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 { - // 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(&mut self) -> NixResult 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 { - let output = self.capture_output().await?; - let path = output.trim_end().to_owned(); - StorePath::try_from(path) - } -} - -#[async_trait] -impl NixCommand for CommandExecution { - async fn passthrough(&mut self) -> NixResult<()> { - self.run().await - } - - /// Captures output as a String. - async fn capture_output(&mut self) -> NixResult { - self.run().await?; - let (stdout, _) = self.get_logs(); - - Ok(stdout.unwrap().to_owned()) - } - - /// Captures deserialized output from JSON. - async fn capture_json(&mut self) -> NixResult 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 { - let output = self.capture_output().await?; - let path = output.trim_end().to_owned(); - StorePath::try_from(path) - } -} - fn validate_keys(keys: &HashMap) -> Result<(), ValidationErrorType> { // Bad secret names: // - /etc/passwd diff --git a/src/nix/store.rs b/src/nix/store.rs index f1b3d5e..280efd2 100644 --- a/src/nix/store.rs +++ b/src/nix/store.rs @@ -7,7 +7,8 @@ use std::fmt; use serde::{Serialize, Deserialize}; use tokio::process::Command; -use super::{Host, NixCommand, NixResult, NixError}; +use crate::util::CommandExt; +use super::{Host, NixResult, NixError}; /// A Nix store path. #[derive(Debug, Clone, Serialize, Deserialize)] diff --git a/src/util.rs b/src/util.rs index 249557a..2edbd03 100644 --- a/src/util.rs +++ b/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, } +/// 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; + + /// Runs the command, capturing deserialized output from JSON. + async fn capture_json(&mut self) -> NixResult where T: DeserializeOwned; + + /// Runs the command, capturing a single store path. + async fn capture_store_path(&mut self) -> NixResult; +} + 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 { + // 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(&mut self) -> NixResult 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 { + 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 { + self.run().await?; + let (stdout, _) = self.get_logs(); + + Ok(stdout.unwrap().to_owned()) + } + + /// Captures deserialized output from JSON. + async fn capture_json(&mut self) -> NixResult 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 { + 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 { let path = match args.occurrences_of("config") { 0 => {