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

1#![recursion_limit = "512"]2#![feature(try_blocks)]34pub(crate) mod cmds;5// pub(crate) mod command;6pub(crate) mod extra_args;78use std::{ffi::OsString, process::ExitCode};910use anyhow::{bail, Result};11use clap::{CommandFactory, Parser};12use cmds::{13	build_systems::{BuildSystems, Deploy},14	complete::Complete,15	info::Info,16	secrets::Secret,17	tf::Tf,18};19use fleet_base::{host::Config, opts::FleetOpts};20use futures::{future::LocalBoxFuture, stream::FuturesUnordered, TryStreamExt};21// use host::Config;22#[cfg(feature = "indicatif")]23use human_repr::HumanCount;24#[cfg(feature = "indicatif")]25use indicatif::{ProgressState, ProgressStyle};26use tracing::{error, info, info_span, Instrument};27#[cfg(feature = "indicatif")]28use tracing_indicatif::IndicatifLayer;29use tracing_subscriber::{prelude::*, EnvFilter};3031#[derive(Parser)]32struct Prefetch {}33impl Prefetch {34	async fn run(&self, config: &Config) -> Result<()> {35		let mut prefetch_dir = config.directory.to_path_buf();36		prefetch_dir.push("prefetch");37		if !prefetch_dir.is_dir() {38			info!("nothing to prefetch: no prefetch directory");39			return Ok(());40		}41		let tasks = <FuturesUnordered<LocalBoxFuture<Result<()>>>>::new();42		for entry in std::fs::read_dir(&prefetch_dir)? {43			tasks.push(Box::pin(async {44				let entry = entry?;45				if !entry.metadata()?.is_file() {46					bail!("only files should exist in prefetch directory");47				}48				let span = info_span!(49					"prefetching",50					name = entry.file_name().to_string_lossy().as_ref()51				);52				let mut path = OsString::new();53				path.push("file://");54				path.push(entry.path());5556				let mut status = config.local_host().cmd("nix").await?;57				status.args(&config.nix_args);58				status.arg("store").arg("prefetch-file").arg(path);59				status.run_nix_string().instrument(span).await?;60				Ok(())61			}));62		}63		tasks.try_collect::<Vec<()>>().await?;64		Ok(())65	}66}6768#[derive(Parser)]69enum Opts {70	/// Prepare systems for deployments71	BuildSystems(BuildSystems),7273	Deploy(Deploy),74	/// Secret management75	#[clap(subcommand)]76	Secret(Secret),77	/// Upload prefetch directory to the nix store78	Prefetch(Prefetch),79	/// Config parsing80	Info(Info),81	/// Command completions82	#[clap(hide(true))]83	Complete(Complete),84	/// Compile and evaluate terranix configuration85	Tf(Tf),86}8788#[derive(Parser)]89#[clap(version, author)]90struct RootOpts {91	#[clap(flatten)]92	fleet_opts: FleetOpts,93	#[clap(subcommand)]94	command: Opts,95}9697async fn run_command(config: &Config, opts: FleetOpts, command: Opts) -> Result<()> {98	match command {99		Opts::BuildSystems(c) => c.run(config, &opts).await?,100		Opts::Deploy(d) => d.run(config, &opts).await?,101		Opts::Secret(s) => s.run(config, &opts).await?,102		Opts::Info(i) => i.run(config).await?,103		Opts::Prefetch(p) => p.run(config).await?,104		Opts::Tf(t) => t.run(config).await?,105		// TODO: actually parse commands before starting the async runtime106		Opts::Complete(c) => {107			tokio::task::spawn_blocking(move || c.run(RootOpts::command())).await?108		}109	};110	Ok(())111}112113fn setup_logging() {114	#[cfg(feature = "indicatif")]115	let indicatif_layer = {116		use std::time::Duration;117118		IndicatifLayer::new().with_progress_style(119			ProgressStyle::with_template(120				"{color_start}{span_child_prefix} {span_name}{{{span_fields}}}{color_end} {wide_msg} {color_start}{download_progress} {elapsed}{color_end}",121			)122				.unwrap()123				.with_key("download_progress", |state: &ProgressState, writer: &mut dyn std::fmt::Write| {124					let Some(len) = state.len() else {125						return;126					};127					let pos = state.pos();128					if pos > len {129						let _ = write!(writer, "{}", pos.human_count_bare());130					} else {131						let _ = write!(writer, "{} / {}", pos.human_count_bare(), len.human_count_bare());132					}133				})134				.with_key(135					"color_start",136					|state: &ProgressState, writer: &mut dyn std::fmt::Write| {137						let elapsed = state.elapsed();138139						if elapsed > Duration::from_secs(60) {140							// Red141							let _ = write!(writer, "\x1b[{}m", 1 + 30);142						} else if elapsed > Duration::from_secs(30) {143							// Yellow144							let _ = write!(writer, "\x1b[{}m", 3 + 30);145						}146					},147				)148				.with_key(149					"color_end",150					|state: &ProgressState, writer: &mut dyn std::fmt::Write| {151						if state.elapsed() > Duration::from_secs(30) {152							let _ = write!(writer, "\x1b[0m");153						}154					},155				),156		)157	};158159	let filter = EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("info"));160161	let reg = tracing_subscriber::registry().with({162		let sub = tracing_subscriber::fmt::layer()163			.without_time()164			.with_target(false);165		#[cfg(feature = "indicatif")]166		let sub = sub.with_writer(indicatif_layer.get_stdout_writer());167		sub.with_filter(filter) // .without,168	});169	// #[cfg(feature = "indicatif")]170	#[cfg(feature = "indicatif")]171	let reg = reg.with(indicatif_layer);172	reg.init();173}174175fn main() -> ExitCode {176	let opts = RootOpts::parse();177	if let Opts::Complete(c) = &opts.command {178		c.run(RootOpts::command());179		return ExitCode::SUCCESS;180	}181182	setup_logging();183	async_main(opts)184}185186#[tokio::main]187async fn async_main(opts: RootOpts) -> ExitCode {188	if let Err(e) = main_real(opts).await {189		// If I remove this line, the next error!() line gets eaten.190		// This is a bug in indicatif, it needs to be fixed191		#[cfg(feature = "indicatif")]192		info!("fixme: this line gets eaten by tracing-indicatif on levels info+");193		error!("{e:#}");194		return ExitCode::FAILURE;195	}196	ExitCode::SUCCESS197}198199async fn main_real(opts: RootOpts) -> Result<()> {200	nix_eval::init_tokio();201202	let nix_args = std::env::var_os("NIX_ARGS")203		.map(|a| extra_args::parse_os(&a))204		.transpose()?205		.unwrap_or_default();206	let config = opts207		.fleet_opts208		.build(209			nix_args,210			matches!(opts.command, Opts::Deploy(_) | Opts::BuildSystems(_)),211		)212		.await?;213214	match run_command(&config, opts.fleet_opts, 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}