Merge pull request #584 from Nukesor/edit-task-envs

add(task): Allow editing environment variables

CHANGELOG.md

@@ -82,6 +82,7 @@ TLDR: The new task state representation is more verbose but significantly cleane
 - Allow `pueue status` to order tasks by `enqueue_at`. [#554](https://github.com/Nukesor/pueue/issues/554)
 - Added Windows service on Windows to allow a true daemon experience. [#344](https://github.com/Nukesor/pueue/issues/344) [#567](https://github.com/Nukesor/pueue/pull/567)
 - Add `queued_count` and `stashed_count` to callback template variables. This allows users to fire callbacks when whole groups are finished. [#578](https://github.com/Nukesor/pueue/issues/578)
+- Add new subcommand to set or unset environment variables for tasks. [#503](https://github.com/Nukesor/pueue/issues/503)
 
 ### Fixed
 

pueue/src/client/cli.rs

@@ -10,7 +10,7 @@ use pueue_lib::network::message::Signal;
 
 use super::commands::WaitTargetStatus;
 
-#[derive(Parser, Debug)]
+#[derive(Parser, Debug, Clone)]
 pub enum SubCommand {
     #[command(
         about = "Enqueue a task for execution.\n\
@@ -278,6 +278,12 @@ pub enum SubCommand {
         task_ids: Vec<usize>,
     },
 
+    #[command(about = "Use this to add or remove environment variables from tasks.")]
+    Env {
+        #[command(subcommand)]
+        cmd: EnvCommand,
+    },
+
     #[command(about = "Use this to add or remove groups.\n\
         By default, this will simply display all known groups.")]
     Group {
@@ -492,7 +498,31 @@ https://github.com/Nukesor/pueue/issues/350#issue-1359083118"
     },
 }
 
-#[derive(Parser, Debug)]
+#[derive(Parser, Debug, Clone)]
+pub enum EnvCommand {
+    /// Set a variable for a specific task's environment.
+    Set {
+        /// The id of the task for which the variable should be set.
+        task_id: usize,
+
+        /// The name of the environment variable to set.
+        key: String,
+
+        /// The value of the environment variable to set.
+        value: String,
+    },
+
+    /// Remove a specific variable from a task's environment.
+    Unset {
+        /// The id of the task for which the variable should be set.
+        task_id: usize,
+
+        /// The name of the environment variable to set.
+        key: String,
+    },
+}
+
+#[derive(Parser, Debug, Clone)]
 pub enum GroupCommand {
     /// Add a group by name.
     Add {

pueue/src/client/client.rs

@@ -17,6 +17,8 @@ use crate::client::cli::{CliArguments, ColorChoice, GroupCommand, SubCommand};
 use crate::client::commands::*;
 use crate::client::display::*;
 
+use super::cli::EnvCommand;
+
 /// This struct contains the base logic for the client.
 /// The client is responsible for connecting to the daemon, sending instructions
 /// and interpreting their responses.
@@ -383,7 +385,7 @@ impl Client {
     /// This function is pretty large, but it consists mostly of simple conversions
     /// of [SubCommand] variant to a [Message] variant.
     fn get_message_from_opt(&self) -> Result<Message> {
-        Ok(match &self.subcommand {
+        Ok(match self.subcommand.clone() {
             SubCommand::Add {
                 command,
                 working_directory,
@@ -406,7 +408,7 @@ impl Client {
                 let mut command = command.clone();
                 // The user can request to escape any special shell characters in all parameter strings before
                 // we concatenated them to a single string.
-                if *escape {
+                if escape {
                     command = command
                         .iter()
                         .map(|parameter| shell_escape::escape(Cow::from(parameter)).into_owned())
@@ -418,20 +420,20 @@ impl Client {
                     path,
                     // Catch the current environment for later injection into the task's process.
                     envs: HashMap::from_iter(vars()),
-                    start_immediately: *start_immediately,
-                    stashed: *stashed,
-                    group: group_or_default(group),
-                    enqueue_at: *delay_until,
-                    dependencies: dependencies.to_vec(),
-                    priority: priority.to_owned(),
-                    label: label.clone(),
-                    print_task_id: *print_task_id,
+                    start_immediately,
+                    stashed,
+                    group: group_or_default(&group),
+                    enqueue_at: delay_until,
+                    dependencies,
+                    priority,
+                    label,
+                    print_task_id,
                 }
                 .into()
             }
             SubCommand::Remove { task_ids } => {
                 if self.settings.client.show_confirmation_questions {
-                    self.handle_user_confirmation("remove", task_ids)?;
+                    self.handle_user_confirmation("remove", &task_ids)?;
                 }
                 Message::Remove(task_ids.clone())
             }
@@ -441,19 +443,19 @@ impl Client {
                 all,
                 delay_until,
             } => {
-                let selection = selection_from_params(*all, group, task_ids);
+                let selection = selection_from_params(all, &group, &task_ids);
                 StashMessage {
                     tasks: selection,
-                    enqueue_at: *delay_until,
+                    enqueue_at: delay_until,
                 }
                 .into()
             }
             SubCommand::Switch {
                 task_id_1,
                 task_id_2,
             } => SwitchMessage {
-                task_id_1: *task_id_1,
-                task_id_2: *task_id_2,
+                task_id_1,
+                task_id_2,
             }
             .into(),
             SubCommand::Enqueue {
@@ -462,10 +464,10 @@ impl Client {
                 all,
                 delay_until,
             } => {
-                let selection = selection_from_params(*all, group, task_ids);
+                let selection = selection_from_params(all, &group, &task_ids);
                 EnqueueMessage {
                     tasks: selection,
-                    enqueue_at: *delay_until,
+                    enqueue_at: delay_until,
                 }
             }
             .into(),
@@ -475,7 +477,7 @@ impl Client {
                 all,
                 ..
             } => StartMessage {
-                tasks: selection_from_params(*all, group, task_ids),
+                tasks: selection_from_params(all, &group, &task_ids),
             }
             .into(),
             SubCommand::Pause {
@@ -485,8 +487,8 @@ impl Client {
                 all,
                 ..
             } => PauseMessage {
-                tasks: selection_from_params(*all, group, task_ids),
-                wait: *wait,
+                tasks: selection_from_params(all, &group, &task_ids),
+                wait,
             }
             .into(),
             SubCommand::Kill {
@@ -497,19 +499,32 @@ impl Client {
                 ..
             } => {
                 if self.settings.client.show_confirmation_questions {
-                    self.handle_user_confirmation("kill", task_ids)?;
+                    self.handle_user_confirmation("kill", &task_ids)?;
                 }
                 KillMessage {
-                    tasks: selection_from_params(*all, group, task_ids),
-                    signal: signal.clone(),
+                    tasks: selection_from_params(all, &group, &task_ids),
+                    signal,
                 }
                 .into()
             }
             SubCommand::Send { task_id, input } => SendMessage {
-                task_id: *task_id,
+                task_id,
                 input: input.clone(),
             }
             .into(),
+            SubCommand::Env { cmd } => Message::from(match cmd {
+                EnvCommand::Set {
+                    task_id,
+                    key,
+                    value,
+                } => EnvMessage::Set {
+                    task_id,
+                    key,
+                    value,
+                },
+                EnvCommand::Unset { task_id, key } => EnvMessage::Unset { task_id, key },
+            }),
+
             SubCommand::Group { cmd, .. } => match cmd {
                 Some(GroupCommand::Add { name, parallel }) => GroupMessage::Add {
                     name: name.to_owned(),
@@ -528,8 +543,8 @@ impl Client {
                 all,
                 ..
             } => {
-                let lines = determine_log_line_amount(*full, lines);
-                let selection = selection_from_params(*all, group, task_ids);
+                let lines = determine_log_line_amount(full, &lines);
+                let selection = selection_from_params(all, &group, &task_ids);
 
                 let message = LogRequestMessage {
                     tasks: selection,
@@ -538,17 +553,13 @@ impl Client {
                 };
                 Message::Log(message)
             }
-            SubCommand::Follow { task_id, lines } => StreamRequestMessage {
-                task_id: *task_id,
-                lines: *lines,
-            }
-            .into(),
+            SubCommand::Follow { task_id, lines } => StreamRequestMessage { task_id, lines }.into(),
             SubCommand::Clean {
                 successful_only,
                 group,
             } => CleanMessage {
-                successful_only: *successful_only,
-                group: group.clone(),
+                successful_only,
+                group,
             }
             .into(),
             SubCommand::Reset { force, groups, .. } => {
@@ -570,9 +581,9 @@ impl Client {
                 group,
             } => match parallel_tasks {
                 Some(parallel_tasks) => {
-                    let group = group_or_default(group);
+                    let group = group_or_default(&group);
                     ParallelMessage {
-                        parallel_tasks: *parallel_tasks,
+                        parallel_tasks,
                         group,
                     }
                     .into()

pueue/src/daemon/network/message_handler/env.rs

@@ -0,0 +1,57 @@
+use pueue_lib::{network::message::*, settings::Settings, state::SharedState};
+
+use crate::{
+    daemon::{network::message_handler::ok_or_failure_message, state_helper::save_state},
+    ok_or_save_state_failure,
+};
+
+/// Invoked on `pueue env`.
+/// Manage environment variables for tasks.
+/// - Set environment variables
+/// - Unset environment variables
+pub fn env(settings: &Settings, state: &SharedState, message: EnvMessage) -> Message {
+    let mut state = state.lock().unwrap();
+
+    let message = match message {
+        EnvMessage::Set {
+            task_id,
+            key,
+            value,
+        } => {
+            let Some(task) = state.tasks.get_mut(&task_id) else {
+                return create_failure_message(format!("No task with id {task_id}"));
+            };
+
+            if !(task.is_queued() || task.is_stashed()) {
+                return create_failure_message("You can only edit stashed or queued tasks");
+            }
+
+            task.envs.insert(key, value);
+
+            create_success_message("Environment variable set.")
+        }
+        EnvMessage::Unset { task_id, key } => {
+            let Some(task) = state.tasks.get_mut(&task_id) else {
+                return create_failure_message(format!("No task with id {task_id}"));
+            };
+
+            if !(task.is_queued() || task.is_stashed()) {
+                return create_failure_message("You can only edit stashed or queued tasks");
+            }
+
+            match task.envs.remove(&key) {
+                Some(_) => create_success_message("Environment variable unset."),
+                None => create_failure_message(format!(
+                    "No environment variable with key '{key}' found."
+                )),
+            }
+        }
+    };
+
+    // Save the state if there were any changes.
+    if let Message::Success(_) = message {
+        ok_or_save_state_failure!(save_state(&state, settings));
+    }
+
+    message
+}

pueue/src/daemon/network/message_handler/group.rs

@@ -1,15 +1,21 @@
 use std::collections::BTreeMap;
 
-use pueue_lib::network::message::*;
-use pueue_lib::settings::Settings;
-use pueue_lib::state::{SharedState, PUEUE_DEFAULT_GROUP};
-use pueue_lib::{failure_msg, success_msg};
+use pueue_lib::{
+    failure_msg,
+    network::message::*,
+    settings::Settings,
+    state::{SharedState, PUEUE_DEFAULT_GROUP},
+    success_msg,
+};
 
 use crate::daemon::network::message_handler::ok_or_failure_message;
-use crate::daemon::network::response_helper::ensure_group_exists;
-use crate::daemon::process_handler::initiate_shutdown;
-use crate::daemon::state_helper::save_state;
-use crate::ok_or_save_state_failure;
+use crate::{
+    daemon::{
+        network::response_helper::ensure_group_exists, process_handler::initiate_shutdown,
+        state_helper::save_state,
+    },
+    ok_or_save_state_failure,
+};
 
 /// Invoked on `pueue groups`.
 /// Manage groups.

pueue/src/daemon/network/message_handler/mod.rs

@@ -12,6 +12,7 @@ mod add;
 mod clean;
 mod edit;
 mod enqueue;
+mod env;
 mod group;
 mod kill;
 mod log;
@@ -34,6 +35,7 @@ pub fn handle_message(message: Message, state: &SharedState, settings: &Settings
         Message::Edit(editable_tasks) => edit::edit(settings, state, editable_tasks),
         Message::EditRequest(task_ids) => edit::edit_request(state, task_ids),
         Message::EditRestore(task_ids) => edit::edit_restore(state, task_ids),
+        Message::Env(message) => env::env(settings, state, message),
         Message::Enqueue(message) => enqueue::enqueue(settings, state, message),
         Message::Group(message) => group::group(settings, state, message),
         Message::Kill(message) => kill::kill(settings, state, message),

pueue/src/daemon/process_handler/finish.rs

@@ -115,6 +115,7 @@ pub fn handle_finished_tasks(settings: &Settings, state: &mut LockedState) {
 
             task.clone()
         };
+        info!("WTF");
         spawn_callback(settings, state, &task);
 
         if let TaskResult::Failed(_) = result {