git.delta.rocks / jrsonnet / refs/commits / 3f73827e390b

1#![feature(try_blocks)]23pub(crate) mod cmds;4pub(crate) mod command;5pub(crate) mod host;6pub(crate) mod keys;78pub(crate) mod extra_args;9pub(crate) mod better_nix_eval;1011mod fleetdata;1213use std::ffi::OsString;14use std::time::Duration;1516use anyhow::{bail, Result};17use clap::Parser;1819use cmds::{build_systems::BuildSystems, info::Info, secrets::Secrets};20use futures::future::LocalBoxFuture;21use futures::stream::FuturesUnordered;22use futures::TryStreamExt;23use host::{Config, FleetOpts};24use indicatif::{ProgressState, ProgressStyle};25use tracing::{info, metadata::LevelFilter};26use tracing::{info_span, Instrument};27use tracing_indicatif::IndicatifLayer;28use tracing_subscriber::{prelude::*, EnvFilter};2930use crate::command::MyCommand;3132#[derive(Parser)]33struct Prefetch {}34impl Prefetch {35	async fn run(&self, config: &Config) -> Result<()> {36		let mut prefetch_dir = config.directory.to_path_buf();37		prefetch_dir.push("prefetch");38		if !prefetch_dir.is_dir() {39			info!("nothing to prefetch: no prefetch directory");40			return Ok(());41		}42		let tasks = <FuturesUnordered<LocalBoxFuture<Result<()>>>>::new();43		for entry in std::fs::read_dir(&prefetch_dir)? {44			tasks.push(Box::pin(async {45				let entry = entry?;46				if !entry.metadata()?.is_file() {47					bail!("only files should exist in prefetch directory");48				}49				let span = info_span!(50					"prefetching",51					name = entry.file_name().to_string_lossy().as_ref()52				);53				let mut path = OsString::new();54				path.push("file://");55				path.push(entry.path());5657				let mut status = MyCommand::new("nix");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),72	/// Secret management73	#[clap(subcommand)]74	Secrets(Secrets),75	/// Upload prefetch directory to the nix store76	Prefetch(Prefetch),77	/// Config parsing78	Info(Info),79}8081#[derive(Parser)]82#[clap(version = "1.0", author)]83struct RootOpts {84	#[clap(flatten)]85	fleet_opts: FleetOpts,86	#[clap(subcommand)]87	command: Opts,88}8990async fn run_command(config: &Config, command: Opts) -> Result<()> {91	match command {92		Opts::BuildSystems(c) => c.run(config).await?,93		Opts::Secrets(s) => s.run(config).await?,94		Opts::Info(i) => i.run(config).await?,95		Opts::Prefetch(p) => p.run(config).await?,96	};97	Ok(())98}99100// fn main() -> Result<()> {101// 	let pool = r2d2::Builder::<NixSessionPool>::new()102// 		.min_idle(Some(1))103// 		.max_lifetime(Some(Duration::from_secs(10)))104// 		.build(NixSessionPool {105// 			flake: ".".to_owned(),106// 			nix_args: vec![],107// 		})?;108// 	let conn = pool.get()?;109// 	let field = Field::root(conn);110// 	// let builtins = field.get_field("builtins")?;111// 	let cur_sys: String = field.get_field("builtins")?.as_json()?;112// 	eprintln!("current system = {cur_sys}");113// 	let v = field.get_field("fleetConfigurations")?;114// 	eprintln!("configs = {:?}", v.list_fields()?);115// 	let d = v.get_field("default")?;116// 	dbg!(d.list_fields());117// 	Ok(())118// }119//120121fn setup_logging() {122	let indicatif_layer = IndicatifLayer::new().with_progress_style(123		ProgressStyle::with_template(124			"{color_start}{span_child_prefix} {span_name}{{{span_fields}}}{color_end} {wide_msg} {color_start}{pos:>7}/{len:7}{elapsed}{color_end}",125		)126		.unwrap()127		.with_key(128			"color_start",129			|state: &ProgressState, writer: &mut dyn std::fmt::Write| {130				let elapsed = state.elapsed();131132				if elapsed > Duration::from_secs(60) {133					// Red134					let _ = write!(writer, "\x1b[{}m", 1 + 30);135				} else if elapsed > Duration::from_secs(30) {136					// Yellow137					let _ = write!(writer, "\x1b[{}m", 3 + 30);138				}139			},140		)141		.with_key(142			"color_end",143			|state: &ProgressState, writer: &mut dyn std::fmt::Write| {144				if state.elapsed() > Duration::from_secs(30) {145					let _ = write!(writer, "\x1b[0m");146				}147			},148		),149	);150151	let filter = EnvFilter::from_default_env().add_directive(LevelFilter::INFO.into());152153	tracing_subscriber::registry()154		.with(155			tracing_subscriber::fmt::layer()156				.without_time()157				.with_target(false)158				.with_writer(indicatif_layer.get_stderr_writer())159				.with_filter(filter), // .withou,160		)161		.with(indicatif_layer)162		.init();163}164165#[tokio::main]166async fn main() -> Result<()> {167	setup_logging();168	let _ = better_nix_eval::TOKIO_RUNTIME.set(tokio::runtime::Handle::current());169170	let nix_args = std::env::var_os("NIX_ARGS")171		.map(|a| extra_args::parse_os(&a))172		.transpose()?173		.unwrap_or_default();174	let opts = RootOpts::parse();175	let config = opts.fleet_opts.build(nix_args).await?;176177	match run_command(&config, opts.command).await {178		Ok(()) => {179			config.save()?;180			Ok(())181		}182		Err(e) => {183			let _ = config.save();184			Err(e)185		}186	}187}