git.delta.rocks / jrsonnet / refs/commits / 213ad7de4b85

1#![recursion_limit = "512"]2#![feature(try_blocks)]34pub(crate) mod cmds;5pub(crate) mod command;6pub(crate) mod host;7pub(crate) mod keys;89pub(crate) mod extra_args;1011mod fleetdata;1213use std::{ffi::OsString, process::ExitCode};1415use anyhow::{bail, Result};16use clap::{CommandFactory, Parser};17use cmds::{18	build_systems::{BuildSystems, Deploy},19	complete::Complete,20	info::Info,21	secrets::Secret,22	tf::Tf,23};24use futures::{future::LocalBoxFuture, stream::FuturesUnordered, TryStreamExt};25use host::{Config, FleetOpts};26#[cfg(feature = "indicatif")]27use human_repr::HumanCount;28#[cfg(feature = "indicatif")]29use indicatif::{ProgressState, ProgressStyle};30use tracing::{error, info, info_span, Instrument};31#[cfg(feature = "indicatif")]32use tracing_indicatif::IndicatifLayer;33use tracing_subscriber::{prelude::*, EnvFilter};3435use crate::command::MyCommand;3637#[derive(Parser)]38struct Prefetch {}39impl Prefetch {40	async fn run(&self, config: &Config) -> Result<()> {41		let mut prefetch_dir = config.directory.to_path_buf();42		prefetch_dir.push("prefetch");43		if !prefetch_dir.is_dir() {44			info!("nothing to prefetch: no prefetch directory");45			return Ok(());46		}47		let tasks = <FuturesUnordered<LocalBoxFuture<Result<()>>>>::new();48		for entry in std::fs::read_dir(&prefetch_dir)? {49			tasks.push(Box::pin(async {50				let entry = entry?;51				if !entry.metadata()?.is_file() {52					bail!("only files should exist in prefetch directory");53				}54				let span = info_span!(55					"prefetching",56					name = entry.file_name().to_string_lossy().as_ref()57				);58				let mut path = OsString::new();59				path.push("file://");60				path.push(entry.path());6162				let mut status = config.local_host().cmd("nix").await?;63				status.args(&config.nix_args);64				status.arg("store").arg("prefetch-file").arg(path);65				status.run_nix_string().instrument(span).await?;66				Ok(())67			}));68		}69		tasks.try_collect::<Vec<()>>().await?;70		Ok(())71	}72}7374#[derive(Parser)]75enum Opts {76	/// Prepare systems for deployments77	BuildSystems(BuildSystems),7879	Deploy(Deploy),80	/// Secret management81	#[clap(subcommand)]82	Secret(Secret),83	/// Upload prefetch directory to the nix store84	Prefetch(Prefetch),85	/// Config parsing86	Info(Info),87	/// Command completions88	#[clap(hide(true))]89	Complete(Complete),90	/// Compile and evaluate terranix configuration91	Tf(Tf),92}9394#[derive(Parser)]95#[clap(version, author)]96struct RootOpts {97	#[clap(flatten)]98	fleet_opts: FleetOpts,99	#[clap(subcommand)]100	command: Opts,101}102103async fn run_command(config: &Config, command: Opts) -> Result<()> {104	match command {105		Opts::BuildSystems(c) => c.run(config).await?,106		Opts::Deploy(d) => d.run(config).await?,107		Opts::Secret(s) => s.run(config).await?,108		Opts::Info(i) => i.run(config).await?,109		Opts::Prefetch(p) => p.run(config).await?,110		Opts::Tf(t) => t.run(config).await?,111		// TODO: actually parse commands before starting the async runtime112		Opts::Complete(c) => {113			tokio::task::spawn_blocking(move || c.run(RootOpts::command())).await?114		}115	};116	Ok(())117}118119fn setup_logging() {120	#[cfg(feature = "indicatif")]121	let indicatif_layer = {122		use std::time::Duration;123124		IndicatifLayer::new().with_progress_style(125			ProgressStyle::with_template(126				"{color_start}{span_child_prefix} {span_name}{{{span_fields}}}{color_end} {wide_msg} {color_start}{download_progress} {elapsed}{color_end}",127			)128				.unwrap()129				.with_key("download_progress", |state: &ProgressState, writer: &mut dyn std::fmt::Write| {130					let Some(len) = state.len() else {131						return;132					};133					let pos = state.pos();134					if pos > len {135						let _ = write!(writer, "{}", pos.human_count_bare());136					} else {137						let _ = write!(writer, "{} / {}", pos.human_count_bare(), len.human_count_bare());138					}139				})140				.with_key(141					"color_start",142					|state: &ProgressState, writer: &mut dyn std::fmt::Write| {143						let elapsed = state.elapsed();144145						if elapsed > Duration::from_secs(60) {146							// Red147							let _ = write!(writer, "\x1b[{}m", 1 + 30);148						} else if elapsed > Duration::from_secs(30) {149							// Yellow150							let _ = write!(writer, "\x1b[{}m", 3 + 30);151						}152					},153				)154				.with_key(155					"color_end",156					|state: &ProgressState, writer: &mut dyn std::fmt::Write| {157						if state.elapsed() > Duration::from_secs(30) {158							let _ = write!(writer, "\x1b[0m");159						}160					},161				),162		)163	};164165	let filter = EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("info"));166167	let reg = tracing_subscriber::registry().with({168		let sub = tracing_subscriber::fmt::layer()169			.without_time()170			.with_target(false);171		#[cfg(feature = "indicatif")]172		let sub = sub.with_writer(indicatif_layer.get_stdout_writer());173		sub.with_filter(filter) // .without,174	});175	// #[cfg(feature = "indicatif")]176	#[cfg(feature = "indicatif")]177	let reg = reg.with(indicatif_layer);178	reg.init();179}180181fn main() -> ExitCode {182	let opts = RootOpts::parse();183	if let Opts::Complete(c) = &opts.command {184		c.run(RootOpts::command());185		return ExitCode::SUCCESS;186	}187188	setup_logging();189	async_main(opts)190}191192#[tokio::main]193async fn async_main(opts: RootOpts) -> ExitCode {194	if let Err(e) = main_real(opts).await {195		// If I remove this line, the next error!() line gets eaten.196		// This is a bug in indicatif, it needs to be fixed197		#[cfg(feature = "indicatif")]198		info!("fixme: this line gets eaten by tracing-indicatif on levels info+");199		error!("{e:#}");200		return ExitCode::FAILURE;201	}202	ExitCode::SUCCESS203}204205async fn main_real(opts: RootOpts) -> Result<()> {206	nix_eval::init_tokio();207208	let nix_args = std::env::var_os("NIX_ARGS")209		.map(|a| extra_args::parse_os(&a))210		.transpose()?211		.unwrap_or_default();212	let config = opts.fleet_opts.build(nix_args).await?;213214	match run_command(&config, opts.command).await {215		Ok(()) => {216			config.save()?;217			Ok(())218		}219		Err(e) => {220			let _ = config.save();221			Err(e)222		}223	}224}225226#[cfg(test)]227mod tests {228	use super::*;229230	#[test]231	fn verify_command() {232		use clap::CommandFactory;233		RootOpts::command().debug_assert();234	}235}