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

1#![feature(try_blocks)]23pub(crate) mod cmds;4pub(crate) mod command;5pub(crate) mod host;6pub(crate) mod keys;78pub(crate) mod better_nix_eval;9pub(crate) mod extra_args;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 human_repr::HumanCount;25use indicatif::{ProgressState, ProgressStyle};26use tracing::{info, metadata::LevelFilter};27use tracing::{info_span, Instrument};28use tracing_indicatif::IndicatifLayer;29use tracing_subscriber::{prelude::*, EnvFilter};3031use crate::command::MyCommand;3233#[derive(Parser)]34struct Prefetch {}35impl Prefetch {36	async fn run(&self, config: &Config) -> Result<()> {37		let mut prefetch_dir = config.directory.to_path_buf();38		prefetch_dir.push("prefetch");39		if !prefetch_dir.is_dir() {40			info!("nothing to prefetch: no prefetch directory");41			return Ok(());42		}43		let tasks = <FuturesUnordered<LocalBoxFuture<Result<()>>>>::new();44		for entry in std::fs::read_dir(&prefetch_dir)? {45			tasks.push(Box::pin(async {46				let entry = entry?;47				if !entry.metadata()?.is_file() {48					bail!("only files should exist in prefetch directory");49				}50				let span = info_span!(51					"prefetching",52					name = entry.file_name().to_string_lossy().as_ref()53				);54				let mut path = OsString::new();55				path.push("file://");56				path.push(entry.path());5758				let mut status = MyCommand::new("nix");59				status.arg("store").arg("prefetch-file").arg(path);60				status.run_nix_string().instrument(span).await?;61				Ok(())62			}));63		}64		tasks.try_collect::<Vec<()>>().await?;65		Ok(())66	}67}6869#[derive(Parser)]70enum Opts {71	/// Prepare systems for deployments72	BuildSystems(BuildSystems),73	/// Secret management74	#[clap(subcommand)]75	Secrets(Secrets),76	/// Upload prefetch directory to the nix store77	Prefetch(Prefetch),78	/// Config parsing79	Info(Info),80}8182#[derive(Parser)]83#[clap(version = "1.0", author)]84struct RootOpts {85	#[clap(flatten)]86	fleet_opts: FleetOpts,87	#[clap(subcommand)]88	command: Opts,89}9091async fn run_command(config: &Config, command: Opts) -> Result<()> {92	match command {93		Opts::BuildSystems(c) => c.run(config).await?,94		Opts::Secrets(s) => s.run(config).await?,95		Opts::Info(i) => i.run(config).await?,96		Opts::Prefetch(p) => p.run(config).await?,97	};98	Ok(())99}100101// fn main() -> Result<()> {102// 	let pool = r2d2::Builder::<NixSessionPool>::new()103// 		.min_idle(Some(1))104// 		.max_lifetime(Some(Duration::from_secs(10)))105// 		.build(NixSessionPool {106// 			flake: ".".to_owned(),107// 			nix_args: vec![],108// 		})?;109// 	let conn = pool.get()?;110// 	let field = Field::root(conn);111// 	// let builtins = field.get_field("builtins")?;112// 	let cur_sys: String = field.get_field("builtins")?.as_json()?;113// 	eprintln!("current system = {cur_sys}");114// 	let v = field.get_field("fleetConfigurations")?;115// 	eprintln!("configs = {:?}", v.list_fields()?);116// 	let d = v.get_field("default")?;117// 	dbg!(d.list_fields());118// 	Ok(())119// }120//121122fn setup_logging() {123	let indicatif_layer = IndicatifLayer::new().with_progress_style(124		ProgressStyle::with_template(125			"{color_start}{span_child_prefix} {span_name}{{{span_fields}}}{color_end} {wide_msg} {color_start}{download_progress} {elapsed}{color_end}",126		)127		.unwrap()128		.with_key("download_progress", |state: &ProgressState, writer: &mut dyn std::fmt::Write| {129			let Some(len) = state.len() else {130				return;131			};132			let pos = state.pos();133			let _ = write!(writer, "{} / {}", pos.human_count_bare(), len.human_count_bare());134		})135		.with_key(136			"color_start",137			|state: &ProgressState, writer: &mut dyn std::fmt::Write| {138				let elapsed = state.elapsed();139140				if elapsed > Duration::from_secs(60) {141					// Red142					let _ = write!(writer, "\x1b[{}m", 1 + 30);143				} else if elapsed > Duration::from_secs(30) {144					// Yellow145					let _ = write!(writer, "\x1b[{}m", 3 + 30);146				}147			},148		)149		.with_key(150			"color_end",151			|state: &ProgressState, writer: &mut dyn std::fmt::Write| {152				if state.elapsed() > Duration::from_secs(30) {153					let _ = write!(writer, "\x1b[0m");154				}155			},156		),157	);158159	let filter = EnvFilter::from_default_env().add_directive(LevelFilter::INFO.into());160161	tracing_subscriber::registry()162		.with(163			tracing_subscriber::fmt::layer()164				.without_time()165				.with_target(false)166				.with_writer(indicatif_layer.get_stderr_writer())167				.with_filter(filter), // .withou,168		)169		.with(indicatif_layer)170		.init();171}172173#[tokio::main]174async fn main() -> Result<()> {175	setup_logging();176	let _ = better_nix_eval::TOKIO_RUNTIME.set(tokio::runtime::Handle::current());177178	let nix_args = std::env::var_os("NIX_ARGS")179		.map(|a| extra_args::parse_os(&a))180		.transpose()?181		.unwrap_or_default();182	let opts = RootOpts::parse();183	let config = opts.fleet_opts.build(nix_args).await?;184185	match run_command(&config, opts.command).await {186		Ok(()) => {187			config.save()?;188			Ok(())189		}190		Err(e) => {191			let _ = config.save();192			Err(e)193		}194	}195}