git.delta.rocks / jrsonnet / refs/commits / 89d35672dcfd

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 better_nix_eval;10pub(crate) mod extra_args;1112mod fleetdata;1314use std::ffi::OsString;15use std::time::Duration;1617use anyhow::{bail, Result};18use clap::Parser;1920use cmds::{build_systems::BuildSystems, info::Info, secrets::Secrets};21use futures::future::LocalBoxFuture;22use futures::stream::FuturesUnordered;23use futures::TryStreamExt;24use host::{Config, FleetOpts};25use human_repr::HumanCount;26use indicatif::{ProgressState, ProgressStyle};27use tracing::{info, metadata::LevelFilter};28use tracing::{info_span, Instrument};29use tracing_indicatif::IndicatifLayer;30use tracing_subscriber::{prelude::*, EnvFilter};3132use crate::command::MyCommand;3334#[derive(Parser)]35struct Prefetch {}36impl Prefetch {37	async fn run(&self, config: &Config) -> Result<()> {38		let mut prefetch_dir = config.directory.to_path_buf();39		prefetch_dir.push("prefetch");40		if !prefetch_dir.is_dir() {41			info!("nothing to prefetch: no prefetch directory");42			return Ok(());43		}44		let tasks = <FuturesUnordered<LocalBoxFuture<Result<()>>>>::new();45		for entry in std::fs::read_dir(&prefetch_dir)? {46			tasks.push(Box::pin(async {47				let entry = entry?;48				if !entry.metadata()?.is_file() {49					bail!("only files should exist in prefetch directory");50				}51				let span = info_span!(52					"prefetching",53					name = entry.file_name().to_string_lossy().as_ref()54				);55				let mut path = OsString::new();56				path.push("file://");57				path.push(entry.path());5859				let mut status = MyCommand::new("nix");60				status.arg("store").arg("prefetch-file").arg(path);61				status.run_nix_string().instrument(span).await?;62				Ok(())63			}));64		}65		tasks.try_collect::<Vec<()>>().await?;66		Ok(())67	}68}6970#[derive(Parser)]71enum Opts {72	/// Prepare systems for deployments73	BuildSystems(BuildSystems),74	/// Secret management75	#[clap(subcommand)]76	Secrets(Secrets),77	/// Upload prefetch directory to the nix store78	Prefetch(Prefetch),79	/// Config parsing80	Info(Info),81}8283#[derive(Parser)]84#[clap(version = "1.0", author)]85struct RootOpts {86	#[clap(flatten)]87	fleet_opts: FleetOpts,88	#[clap(subcommand)]89	command: Opts,90}9192async fn run_command(config: &Config, command: Opts) -> Result<()> {93	match command {94		Opts::BuildSystems(c) => c.run(config).await?,95		Opts::Secrets(s) => s.run(config).await?,96		Opts::Info(i) => i.run(config).await?,97		Opts::Prefetch(p) => p.run(config).await?,98	};99	Ok(())100}101102// fn main() -> Result<()> {103// 	let pool = r2d2::Builder::<NixSessionPool>::new()104// 		.min_idle(Some(1))105// 		.max_lifetime(Some(Duration::from_secs(10)))106// 		.build(NixSessionPool {107// 			flake: ".".to_owned(),108// 			nix_args: vec![],109// 		})?;110// 	let conn = pool.get()?;111// 	let field = Field::root(conn);112// 	// let builtins = field.get_field("builtins")?;113// 	let cur_sys: String = field.get_field("builtins")?.as_json()?;114// 	eprintln!("current system = {cur_sys}");115// 	let v = field.get_field("fleetConfigurations")?;116// 	eprintln!("configs = {:?}", v.list_fields()?);117// 	let d = v.get_field("default")?;118// 	dbg!(d.list_fields());119// 	Ok(())120// }121//122123fn setup_logging() {124	let indicatif_layer = 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			let _ = write!(writer, "{} / {}", pos.human_count_bare(), len.human_count_bare());135		})136		.with_key(137			"color_start",138			|state: &ProgressState, writer: &mut dyn std::fmt::Write| {139				let elapsed = state.elapsed();140141				if elapsed > Duration::from_secs(60) {142					// Red143					let _ = write!(writer, "\x1b[{}m", 1 + 30);144				} else if elapsed > Duration::from_secs(30) {145					// Yellow146					let _ = write!(writer, "\x1b[{}m", 3 + 30);147				}148			},149		)150		.with_key(151			"color_end",152			|state: &ProgressState, writer: &mut dyn std::fmt::Write| {153				if state.elapsed() > Duration::from_secs(30) {154					let _ = write!(writer, "\x1b[0m");155				}156			},157		),158	);159160	let filter = EnvFilter::from_default_env().add_directive(LevelFilter::INFO.into());161162	tracing_subscriber::registry()163		.with(164			tracing_subscriber::fmt::layer()165				.without_time()166				.with_target(false)167				.with_writer(indicatif_layer.get_stderr_writer())168				.with_filter(filter), // .withou,169		)170		.with(indicatif_layer)171		.init();172}173174#[tokio::main]175async fn main() -> Result<()> {176	setup_logging();177	let _ = better_nix_eval::TOKIO_RUNTIME.set(tokio::runtime::Handle::current());178179	let nix_args = std::env::var_os("NIX_ARGS")180		.map(|a| extra_args::parse_os(&a))181		.transpose()?182		.unwrap_or_default();183	let opts = RootOpts::parse();184	let config = opts.fleet_opts.build(nix_args).await?;185186	match run_command(&config, opts.command).await {187		Ok(()) => {188			config.save()?;189			Ok(())190		}191		Err(e) => {192			let _ = config.save();193			Err(e)194		}195	}196}