git.delta.rocks / jrsonnet / refs/commits / a1a72ceb9060

1#![recursion_limit = "512"]23pub(crate) mod cmds;4// pub(crate) mod command;5pub(crate) mod extra_args;67use std::{ffi::OsString, process::ExitCode};89use anyhow::{bail, Result};10use clap::{CommandFactory, Parser};11use cmds::{12	build_systems::{BuildSystems, Deploy},13	complete::Complete,14	info::Info,15	secrets::Secret,16	tf::Tf,17};18use fleet_base::{host::Config, opts::FleetOpts};19use futures::{future::LocalBoxFuture, stream::FuturesUnordered, TryStreamExt};20// use host::Config;21#[cfg(feature = "indicatif")]22use human_repr::HumanCount;23#[cfg(feature = "indicatif")]24use indicatif::{ProgressState, ProgressStyle};25use tracing::{error, info, info_span, Instrument};26#[cfg(feature = "indicatif")]27use tracing_indicatif::IndicatifLayer;28use tracing_subscriber::{prelude::*, EnvFilter};2930#[derive(Parser)]31struct Prefetch {}32impl Prefetch {33	async fn run(&self, config: &Config) -> Result<()> {34		let mut prefetch_dir = config.directory.to_path_buf();35		prefetch_dir.push("prefetch");36		if !prefetch_dir.is_dir() {37			info!("nothing to prefetch: no prefetch directory");38			return Ok(());39		}40		let tasks = <FuturesUnordered<LocalBoxFuture<Result<()>>>>::new();41		for entry in std::fs::read_dir(&prefetch_dir)? {42			tasks.push(Box::pin(async {43				let entry = entry?;44				if !entry.metadata()?.is_file() {45					bail!("only files should exist in prefetch directory");46				}47				let span = info_span!(48					"prefetching",49					name = entry.file_name().to_string_lossy().as_ref()50				);51				let mut path = OsString::new();52				path.push("file://");53				path.push(entry.path());5455				let mut status = config.local_host().cmd("nix").await?;56				status.args(&config.nix_args);57				status.arg("store").arg("prefetch-file").arg(path);58				status.run_nix_string().instrument(span).await?;59				Ok(())60			}));61		}62		tasks.try_collect::<Vec<()>>().await?;63		Ok(())64	}65}6667#[derive(Parser)]68enum Opts {69	/// Prepare systems for deployments70	BuildSystems(BuildSystems),7172	Deploy(Deploy),73	/// Secret management74	#[clap(subcommand)]75	Secret(Secret),76	/// Upload prefetch directory to the nix store77	Prefetch(Prefetch),78	/// Config parsing79	Info(Info),80	/// Command completions81	#[clap(hide(true))]82	Complete(Complete),83	/// Compile and evaluate terranix configuration84	Tf(Tf),85}8687#[derive(Parser)]88#[clap(version, author)]89struct RootOpts {90	#[clap(flatten)]91	fleet_opts: FleetOpts,92	#[clap(subcommand)]93	command: Opts,94}9596async fn run_command(config: &Config, opts: FleetOpts, command: Opts) -> Result<()> {97	match command {98		Opts::BuildSystems(c) => c.run(config, &opts).await?,99		Opts::Deploy(d) => d.run(config, &opts).await?,100		Opts::Secret(s) => s.run(config, &opts).await?,101		Opts::Info(i) => i.run(config).await?,102		Opts::Prefetch(p) => p.run(config).await?,103		Opts::Tf(t) => t.run(config).await?,104		// TODO: actually parse commands before starting the async runtime105		Opts::Complete(c) => {106			tokio::task::spawn_blocking(move || c.run(RootOpts::command())).await?107		}108	};109	Ok(())110}111112fn setup_logging() {113	#[cfg(feature = "indicatif")]114	let indicatif_layer = {115		use std::time::Duration;116117		IndicatifLayer::new().with_progress_style(118			ProgressStyle::with_template(119				"{color_start}{span_child_prefix} {span_name}{{{span_fields}}}{color_end} {wide_msg} {color_start}{download_progress} {elapsed}{color_end}",120			)121				.unwrap()122				.with_key("download_progress", |state: &ProgressState, writer: &mut dyn std::fmt::Write| {123					let Some(len) = state.len() else {124						return;125					};126					let pos = state.pos();127					if pos > len {128						let _ = write!(writer, "{}", pos.human_count_bare());129					} else {130						let _ = write!(writer, "{} / {}", pos.human_count_bare(), len.human_count_bare());131					}132				})133				.with_key(134					"color_start",135					|state: &ProgressState, writer: &mut dyn std::fmt::Write| {136						let elapsed = state.elapsed();137138						if elapsed > Duration::from_secs(60) {139							// Red140							let _ = write!(writer, "\x1b[{}m", 1 + 30);141						} else if elapsed > Duration::from_secs(30) {142							// Yellow143							let _ = write!(writer, "\x1b[{}m", 3 + 30);144						}145					},146				)147				.with_key(148					"color_end",149					|state: &ProgressState, writer: &mut dyn std::fmt::Write| {150						if state.elapsed() > Duration::from_secs(30) {151							let _ = write!(writer, "\x1b[0m");152						}153					},154				),155		)156	};157158	let filter = EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("info"));159160	let reg = tracing_subscriber::registry().with({161		let sub = tracing_subscriber::fmt::layer()162			.without_time()163			.with_target(false);164		#[cfg(feature = "indicatif")]165		let sub = sub.with_writer(indicatif_layer.get_stdout_writer());166		sub.with_filter(filter) // .without,167	});168	// #[cfg(feature = "indicatif")]169	#[cfg(feature = "indicatif")]170	let reg = reg.with(indicatif_layer);171	reg.init();172}173174fn main() -> ExitCode {175	let opts = RootOpts::parse();176	if let Opts::Complete(c) = &opts.command {177		c.run(RootOpts::command());178		return ExitCode::SUCCESS;179	}180181	setup_logging();182	async_main(opts)183}184185#[tokio::main]186async fn async_main(opts: RootOpts) -> ExitCode {187	if let Err(e) = main_real(opts).await {188		// If I remove this line, the next error!() line gets eaten.189		// This is a bug in indicatif, it needs to be fixed190		#[cfg(feature = "indicatif")]191		info!("fixme: this line gets eaten by tracing-indicatif on levels info+");192		error!("{e:#}");193		return ExitCode::FAILURE;194	}195	ExitCode::SUCCESS196}197198async fn main_real(opts: RootOpts) -> Result<()> {199	nix_eval::init_tokio();200201	let nix_args = std::env::var_os("NIX_ARGS")202		.map(|a| extra_args::parse_os(&a))203		.transpose()?204		.unwrap_or_default();205	let config = opts206		.fleet_opts207		.build(208			nix_args,209			matches!(opts.command, Opts::Deploy(_) | Opts::BuildSystems(_)),210		)211		.await?;212213	match run_command(&config, opts.fleet_opts, opts.command).await {214		Ok(()) => {215			config.save()?;216			Ok(())217		}218		Err(e) => {219			let _ = config.save();220			Err(e)221		}222	}223}224225#[cfg(test)]226mod tests {227	use super::*;228229	#[test]230	fn verify_command() {231		use clap::CommandFactory;232		RootOpts::command().debug_assert();233	}234}