Merge pull request #126 from ThinkChaos/cli/nix-option

cli: Add --nix-option to allow passing arbitrary nix options

src/cli.rs

@@ -3,8 +3,8 @@
 use std::env;
 
 use clap::{
-    builder::PossibleValue, value_parser, Arg, ArgMatches, ColorChoice, Command as ClapCommand,
+    builder::PossibleValue, value_parser, Arg, ArgAction, ArgMatches, ColorChoice,
-    ValueEnum,
+    Command as ClapCommand, ValueEnum,
 };
 use clap_complete::Shell;
 use const_format::concatcp;
@@ -160,6 +160,17 @@ pub fn build_cli(include_internal: bool) -> ClapCommand {
             .long_help("Passes --impure to Nix commands")
             .global(true)
             .num_args(0))
+        .arg(Arg::new("nix-option")
+            .long("nix-option")
+            .help("Passes an arbitrary option to Nix commands")
+            .long_help(r#"Passes arbitrary options to Nix commands
+
+This only works when building locally.
+"#)
+            .global(true)
+            .num_args(2)
+            .value_names(["NAME", "VALUE"])
+            .action(ArgAction::Append))
         .arg(Arg::new("color")
             .long("color")
             .help("When to colorize the output")

src/command/apply_local.rs

@@ -46,6 +46,7 @@ By default, Colmena will deploy keys set in `deployment.keys` before activating
             .num_args(0))
         .arg(Arg::new("node")
             .long("node")
+            .value_name("NODE")
             .help("Override the node name to use")
             .num_args(1))
 
@@ -104,7 +105,7 @@ pub async fn run(_global_args: &ArgMatches, local_args: &ArgMatches) -> Result<(
 
     let target = {
         if let Some(info) = hive.deployment_info_single(&hostname).await.unwrap() {
-            let nix_options = hive.nix_options_with_builders().await.unwrap();
+            let nix_options = hive.nix_flags_with_builders().await.unwrap();
             if !info.allows_local_deployment() {
                 log::error!(
                     "Local deployment is not enabled for host {}.",

src/nix/deployment/mod.rs

@@ -18,7 +18,7 @@ use futures::future::join_all;
 use itertools::Itertools;
 use tokio_stream::StreamExt;
 
-use super::NixOptions;
+use super::NixFlags;
 use crate::job::{JobHandle, JobMonitor, JobState, JobType};
 use crate::progress::Sender as ProgressSender;
 use crate::util;
@@ -50,7 +50,7 @@ pub struct Deployment {
     options: Options,
 
     /// Options passed to Nix invocations.
-    nix_options: NixOptions,
+    nix_options: NixFlags,
 
     /// Handle to send messages to the ProgressOutput.
     progress: Option<ProgressSender>,
@@ -103,7 +103,7 @@ impl Deployment {
             hive,
             goal,
             options: Options::default(),
-            nix_options: NixOptions::default(),
+            nix_options: NixFlags::default(),
             progress,
             targets,
             parallelism_limit: ParallelismLimit::default(),
@@ -129,7 +129,7 @@ impl Deployment {
             monitor.set_label_width(width);
         }
 
-        let nix_options = self.hive.nix_options_with_builders().await?;
+        let nix_options = self.hive.nix_flags_with_builders().await?;
         self.nix_options = nix_options;
 
         if self.goal == Goal::UploadKeys {
@@ -250,7 +250,7 @@ impl Deployment {
                 evaluator.set_job(job.clone());
 
                 // FIXME: nix-eval-jobs currently does not support IFD with builders
-                let options = self.hive.nix_options();
+                let options = self.hive.nix_flags();
                 let mut stream = evaluator.evaluate(&expr, options).await?;
 
                 let mut futures: Vec<tokio::task::JoinHandle<ColmenaResult<()>>> = Vec::new();

src/nix/evaluator/mod.rs

@@ -17,7 +17,7 @@ use std::result::Result as StdResult;
 use async_trait::async_trait;
 use futures::Stream;
 
-use super::{BuildResult, NixExpression, NixOptions, StoreDerivation, StorePath};
+use super::{BuildResult, NixExpression, NixFlags, StoreDerivation, StorePath};
 use crate::error::{ColmenaError, ColmenaResult};
 use crate::job::JobHandle;
 
@@ -61,7 +61,7 @@ pub trait DrvSetEvaluator {
     async fn evaluate(
         &self,
         expression: &dyn NixExpression,
-        options: NixOptions,
+        flags: NixFlags,
     ) -> ColmenaResult<Pin<Box<dyn Stream<Item = EvalResult>>>>;
 
     /// Sets the maximum number of attributes to evaluate at the same time.

src/nix/evaluator/nix_eval_jobs.rs

@@ -19,7 +19,7 @@ use tokio::process::Command;
 use super::{AttributeError, AttributeOutput, DrvSetEvaluator, EvalError, EvalResult};
 use crate::error::{ColmenaError, ColmenaResult};
 use crate::job::{null_job_handle, JobHandle};
-use crate::nix::{NixExpression, NixOptions, StorePath};
+use crate::nix::{NixExpression, NixFlags, StorePath};
 use crate::util::capture_stream;
 
 /// The pinned nix-eval-jobs binary.
@@ -74,7 +74,7 @@ impl DrvSetEvaluator for NixEvalJobs {
     async fn evaluate(
         &self,
         expression: &dyn NixExpression,
-        options: NixOptions,
+        flags: NixFlags,
     ) -> ColmenaResult<Pin<Box<dyn Stream<Item = EvalResult>>>> {
         let mut command = Command::new(&self.executable);
         command
@@ -82,7 +82,7 @@ impl DrvSetEvaluator for NixEvalJobs {
             .arg(self.workers.to_string())
             .args(&["--expr", &expression.expression()]);
 
-        command.args(options.to_args());
+        command.args(flags.to_args());
 
         if expression.requires_flakes() {
             command.args(&["--extra-experimental-features", "flakes"]);
@@ -235,7 +235,7 @@ mod tests {
 
         block_on(async move {
             let mut stream = evaluator
-                .evaluate(&expr, NixOptions::default())
+                .evaluate(&expr, NixFlags::default())
                 .await
                 .unwrap();
             let mut count = 0;
@@ -259,7 +259,7 @@ mod tests {
 
         block_on(async move {
             let mut stream = evaluator
-                .evaluate(&expr, NixOptions::default())
+                .evaluate(&expr, NixFlags::default())
                 .await
                 .unwrap();
             let mut count = 0;
@@ -283,7 +283,7 @@ mod tests {
 
         block_on(async move {
             let mut stream = evaluator
-                .evaluate(&expr, NixOptions::default())
+                .evaluate(&expr, NixFlags::default())
                 .await
                 .unwrap();
             let mut count = 0;
@@ -323,7 +323,7 @@ mod tests {
 
         block_on(async move {
             let mut stream = evaluator
-                .evaluate(&expr, NixOptions::default())
+                .evaluate(&expr, NixFlags::default())
                 .await
                 .unwrap();
             let mut count = 0;

src/nix/hive/mod.rs

@@ -13,7 +13,7 @@ use validator::Validate;
 
 use super::deployment::TargetNode;
 use super::{
-    Flake, MetaConfig, NixExpression, NixOptions, NodeConfig, NodeFilter, NodeName,
+    Flake, MetaConfig, NixExpression, NixFlags, NodeConfig, NodeFilter, NodeName,
     ProfileDerivation, SerializedNixExpression, StorePath,
 };
 use crate::error::ColmenaResult;
@@ -52,6 +52,9 @@ pub struct Hive {
     /// Whether to pass --impure in Nix commands.
     impure: bool,
 
+    /// Options to pass as --option name value.
+    nix_options: HashMap<String, String>,
+
     meta_config: OnceCell<MetaConfig>,
 }
 
@@ -104,6 +107,7 @@ impl Hive {
             assets,
             show_trace: false,
             impure: false,
+            nix_options: HashMap::new(),
             meta_config: OnceCell::new(),
         })
     }
@@ -131,25 +135,30 @@ impl Hive {
         self.impure = impure;
     }
 
-    /// Returns Nix options to set for this Hive.
+    pub fn add_nix_option(&mut self, name: String, value: String) {
-    pub fn nix_options(&self) -> NixOptions {
+        self.nix_options.insert(name, value);
-        let mut options = NixOptions::default();
-        options.set_show_trace(self.show_trace);
-        options.set_pure_eval(self.path.is_flake());
-        options.set_impure(self.impure);
-        options
     }
 
-    /// Returns Nix options to set for this Hive, with configured remote builders.
+    /// Returns Nix options to set for this Hive.
-    pub async fn nix_options_with_builders(&self) -> ColmenaResult<NixOptions> {
+    pub fn nix_flags(&self) -> NixFlags {
-        let mut options = NixOptions::default();
+        let mut flags = NixFlags::default();
-        options.set_show_trace(self.show_trace);
+        flags.set_show_trace(self.show_trace);
+        flags.set_pure_eval(self.path.is_flake());
+        flags.set_impure(self.impure);
+        flags.set_options(self.nix_options.clone());
+        flags
+    }
+
+    /// Returns Nix flags to set for this Hive, with configured remote builders.
+    pub async fn nix_flags_with_builders(&self) -> ColmenaResult<NixFlags> {
+        let mut flags = NixFlags::default();
+        flags.set_show_trace(self.show_trace);
 
         if let Some(machines_file) = &self.get_meta_config().await?.machines_file {
-            options.set_builders(Some(format!("@{}", machines_file)));
+            flags.set_builders(Some(format!("@{}", machines_file)));
         }
 
-        Ok(options)
+        Ok(flags)
     }
 
     /// Convenience wrapper to filter nodes for CLI actions.
@@ -423,7 +432,7 @@ impl<'hive> NixInstantiate<'hive> {
 
     fn eval(self) -> Command {
         let mut command = self.instantiate();
-        let options = self.hive.nix_options();
+        let flags = self.hive.nix_flags();
         command
             .arg("--eval")
             .arg("--json")
@@ -431,24 +440,24 @@ impl<'hive> NixInstantiate<'hive> {
             // Ensures the derivations are instantiated
             // Required for system profile evaluation and IFD
             .arg("--read-write-mode")
-            .args(options.to_args());
+            .args(flags.to_args());
         command
     }
 
     async fn instantiate_with_builders(self) -> ColmenaResult<Command> {
-        let options = self.hive.nix_options_with_builders().await?;
+        let flags = self.hive.nix_flags_with_builders().await?;
         let mut command = self.instantiate();
 
-        command.args(options.to_args());
+        command.args(flags.to_args());
 
         Ok(command)
     }
 
     async fn eval_with_builders(self) -> ColmenaResult<Command> {
-        let options = self.hive.nix_options_with_builders().await?;
+        let flags = self.hive.nix_flags_with_builders().await?;
         let mut command = self.eval();
 
-        command.args(options.to_args());
+        command.args(flags.to_args());
 
         Ok(command)
     }

src/nix/host/local.rs

@@ -8,7 +8,7 @@ use tokio::process::Command;
 use super::{key_uploader, CopyDirection, CopyOptions, Host};
 use crate::error::{ColmenaError, ColmenaResult};
 use crate::job::JobHandle;
-use crate::nix::{Goal, Key, NixOptions, Profile, StorePath, CURRENT_PROFILE, SYSTEM_PROFILE};
+use crate::nix::{Goal, Key, NixFlags, Profile, StorePath, CURRENT_PROFILE, SYSTEM_PROFILE};
 use crate::util::{CommandExecution, CommandExt};
 
 /// The local machine running Colmena.
@@ -18,12 +18,12 @@ use crate::util::{CommandExecution, CommandExt};
 #[derive(Debug)]
 pub struct Local {
     job: Option<JobHandle>,
-    nix_options: NixOptions,
+    nix_options: NixFlags,
     privilege_escalation_command: Option<Vec<String>>,
 }
 
 impl Local {
-    pub fn new(nix_options: NixOptions) -> Self {
+    pub fn new(nix_options: NixFlags) -> Self {
         Self {
             job: None,
             nix_options,

src/nix/mod.rs

@@ -91,9 +91,9 @@ pub struct MetaConfig {
     pub machines_file: Option<String>,
 }
 
-/// Nix options.
+/// Nix CLI flags.
 #[derive(Debug, Clone, Default)]
-pub struct NixOptions {
+pub struct NixFlags {
     /// Whether to pass --show-trace.
     show_trace: bool,
 
@@ -111,6 +111,9 @@ pub struct NixOptions {
     /// - `@/path/to/machines`
     /// - `builder@host.tld riscv64-linux /home/nix/.ssh/keys/builder.key 8 1 kvm`
     builders: Option<String>,
+
+    /// Options to pass as --option name value.
+    options: HashMap<String, String>,
 }
 
 impl NodeName {
@@ -188,7 +191,7 @@ impl NodeConfig {
     }
 }
 
-impl NixOptions {
+impl NixFlags {
     pub fn set_show_trace(&mut self, show_trace: bool) {
         self.show_trace = show_trace;
     }
@@ -205,11 +208,15 @@ impl NixOptions {
         self.builders = builders;
     }
 
+    pub fn set_options(&mut self, options: HashMap<String, String>) {
+        self.options = options;
+    }
+
     pub fn to_args(&self) -> Vec<String> {
-        let mut options = Vec::new();
+        let mut args = Vec::new();
 
         if let Some(builders) = &self.builders {
-            options.append(&mut vec![
+            args.append(&mut vec![
                 "--option".to_string(),
                 "builders".to_string(),
                 builders.clone(),
@@ -217,18 +224,24 @@ impl NixOptions {
         }
 
         if self.show_trace {
-            options.push("--show-trace".to_string());
+            args.push("--show-trace".to_string());
         }
 
         if self.pure_eval {
-            options.push("--pure-eval".to_string());
+            args.push("--pure-eval".to_string());
         }
 
         if self.impure {
-            options.push("--impure".to_string());
+            args.push("--impure".to_string());
         }
 
-        options
+        for (name, value) in self.options.iter() {
+            args.push("--option".to_string());
+            args.push(name.to_string());
+            args.push(value.to_string());
+        }
+
+        args
     }
 }
 

src/util.rs

@@ -244,17 +244,8 @@ pub async fn hive_from_args(args: &ArgMatches) -> ColmenaResult<Hive> {
                 log::info!("Using flake: {}", flake.uri());
 
                 let hive_path = HivePath::Flake(flake);
-                let mut hive = Hive::new(hive_path).await?;
 
-                if args.get_flag("show-trace") {
+                return hive_from_path(hive_path, args).await;
-                    hive.set_show_trace(true);
-                }
-
-                if args.get_flag("impure") {
-                    hive.set_impure(true);
-                }
-
-                return Ok(hive);
             }
 
             fpath
@@ -263,6 +254,11 @@ pub async fn hive_from_args(args: &ArgMatches) -> ColmenaResult<Hive> {
     };
 
     let hive_path = HivePath::from_path(path).await?;
+
+    hive_from_path(hive_path, args).await
+}
+
+pub async fn hive_from_path(hive_path: HivePath, args: &ArgMatches) -> ColmenaResult<Hive> {
     match &hive_path {
         HivePath::Legacy(p) => {
             log::info!("Using configuration: {}", p.to_string_lossy());
@@ -282,6 +278,17 @@ pub async fn hive_from_args(args: &ArgMatches) -> ColmenaResult<Hive> {
         hive.set_impure(true);
     }
 
+    if let Some(opts) = args.get_many::<String>("nix-option") {
+        let iter = opts.into_iter();
+
+        let names = iter.clone().step_by(2);
+        let values = iter.clone().skip(1).step_by(2);
+
+        for (name, value) in names.zip(values) {
+            hive.add_nix_option(name.to_owned(), value.to_owned());
+        }
+    }
+
     Ok(hive)
 }