More fixes to key deployment and logging

This commit is contained in:
Zhaofeng Li 2021-02-09 14:57:11 -08:00
parent a703a8b031
commit 8934726664
4 changed files with 64 additions and 33 deletions

View file

@ -175,11 +175,16 @@ pub async fn run(_global_args: &ArgMatches<'_>, local_args: &ArgMatches<'_>) {
let mut options = DeploymentOptions::default();
options.set_substituters_push(!local_args.is_present("no-substitutes"));
options.set_upload_keys(!local_args.is_present("no-upload-keys"));
options.set_gzip(!local_args.is_present("no-gzip"));
options.set_progress_bar(!local_args.is_present("verbose"));
options.set_upload_keys(!local_args.is_present("no-keys"));
deployment.set_options(options);
if local_args.is_present("no-keys") && goal_arg == "keys" {
log::error!("--no-keys cannot be used when the goal is to upload keys");
quit::with_code(1);
}
let mut parallelism_limit = ParallelismLimit::default();
parallelism_limit.set_apply_limit({
let limit = local_args.value_of("parallel").unwrap().parse::<usize>().unwrap();

View file

@ -130,7 +130,13 @@ impl Local {
.arg(&temp_path);
let mut execution = CommandExecution::new("local", command);
execution.run().await?;
let exit = execution.run().await;
let (stdout, stderr) = execution.get_logs();
self.logs += stdout.unwrap();
self.logs += stderr.unwrap();
exit?;
}
{
let mut command = Command::new("chown");
@ -139,7 +145,13 @@ impl Local {
.arg(&temp_path);
let mut execution = CommandExecution::new("local", command);
execution.run().await?;
let exit = execution.run().await;
let (stdout, stderr) = execution.get_logs();
self.logs += stdout.unwrap();
self.logs += stderr.unwrap();
exit?;
}
let parent_dir = dest_path.parent().unwrap();

View file

@ -3,13 +3,14 @@ use std::convert::TryInto;
use std::process::Stdio;
use async_trait::async_trait;
use futures::future::join3;
use indicatif::ProgressBar;
use tokio::process::Command;
use tokio::io::AsyncWriteExt;
use tokio::io::{AsyncWriteExt, BufReader};
use super::{CopyDirection, CopyOptions, Host};
use crate::nix::{StorePath, Profile, DeploymentGoal, NixResult, NixCommand, NixError, Key, SYSTEM_PROFILE};
use crate::util::CommandExecution;
use crate::util::{CommandExecution, capture_stream};
const DEPLOY_KEY_TEMPLATE: &'static str = include_str!("./deploy-key.template");
@ -168,8 +169,8 @@ impl Ssh {
let mut command = self.ssh(&["sh", "-c", &remote_command]);
command.stdin(Stdio::piped());
command.stderr(Stdio::null());
command.stdout(Stdio::null());
command.stderr(Stdio::piped());
command.stdout(Stdio::piped());
let mut child = command.spawn()?;
@ -178,7 +179,20 @@ impl Ssh {
stdin.flush().await?;
drop(stdin);
let exit = child.wait().await?;
let stdout = BufReader::new(child.stdout.take().unwrap());
let stderr = BufReader::new(child.stderr.take().unwrap());
let futures = join3(
capture_stream(stdout, &self.friendly_name, self.progress_bar.clone()),
capture_stream(stderr, &self.friendly_name, self.progress_bar.clone()),
child.wait(),
);
let (stdout_str, stderr_str, exit) = futures.await;
self.logs += &stdout_str;
self.logs += &stderr_str;
let exit = exit?;
if exit.success() {
Ok(())
} else {

View file

@ -63,31 +63,6 @@ impl CommandExecution {
let stdout = BufReader::new(child.stdout.take().unwrap());
let stderr = BufReader::new(child.stderr.take().unwrap());
async fn capture_stream<R: AsyncRead + Unpin>(mut stream: BufReader<R>, label: &str, mut progress_bar: Option<ProgressBar>) -> String {
let mut log = String::new();
loop {
let mut line = String::new();
let len = stream.read_line(&mut line).await.unwrap();
if len == 0 {
break;
}
let trimmed = line.trim_end();
if let Some(progress_bar) = progress_bar.as_mut() {
progress_bar.set_message(trimmed);
} else {
eprintln!("{} | {}", style(label).cyan(), trimmed);
}
log += trimmed;
log += "\n";
}
log
}
let futures = join3(
capture_stream(stdout, &self.label, self.progress_bar.clone()),
capture_stream(stderr, &self.label, self.progress_bar.clone()),
@ -243,3 +218,28 @@ fn canonicalize_cli_path(path: String) -> PathBuf {
path.into()
}
}
pub async fn capture_stream<R: AsyncRead + Unpin>(mut stream: BufReader<R>, label: &str, mut progress_bar: Option<ProgressBar>) -> String {
let mut log = String::new();
loop {
let mut line = String::new();
let len = stream.read_line(&mut line).await.unwrap();
if len == 0 {
break;
}
let trimmed = line.trim_end();
if let Some(progress_bar) = progress_bar.as_mut() {
progress_bar.set_message(trimmed);
} else {
eprintln!("{} | {}", style(label).cyan(), trimmed);
}
log += trimmed;
log += "\n";
}
log
}