From 161dc109e15a14240253f1e3490538ea9b1226db Mon Sep 17 00:00:00 2001 From: Arne Beer Date: Sun, 1 Dec 2024 17:02:30 +0100 Subject: [PATCH] add(task): Allow editing environment variables --- CHANGELOG.md | 1 + pueue/src/client/cli.rs | 34 +++++++- pueue/src/client/client.rs | 81 +++++++++++-------- .../src/daemon/network/message_handler/env.rs | 57 +++++++++++++ .../daemon/network/message_handler/group.rs | 22 +++-- .../src/daemon/network/message_handler/mod.rs | 2 + pueue/src/daemon/process_handler/finish.rs | 1 + pueue/tests/client/integration/env.rs | 60 ++++++++++++++ pueue/tests/client/integration/mod.rs | 1 + pueue_lib/src/network/message.rs | 17 ++++ 10 files changed, 231 insertions(+), 45 deletions(-) create mode 100644 pueue/src/daemon/network/message_handler/env.rs create mode 100644 pueue/tests/client/integration/env.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index bd2722c3..66e63a0d 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/pueue/src/client/cli.rs b/pueue/src/client/cli.rs index d530352b..5f88942b 100644 --- a/pueue/src/client/cli.rs +++ b/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, }, + #[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 { diff --git a/pueue/src/client/client.rs b/pueue/src/client/client.rs index febdf2e4..bef14f04 100644 --- a/pueue/src/client/client.rs +++ b/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 { - 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,10 +443,10 @@ 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() } @@ -452,8 +454,8 @@ impl Client { 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() diff --git a/pueue/src/daemon/network/message_handler/env.rs b/pueue/src/daemon/network/message_handler/env.rs new file mode 100644 index 00000000..31ddac05 --- /dev/null +++ b/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 +} diff --git a/pueue/src/daemon/network/message_handler/group.rs b/pueue/src/daemon/network/message_handler/group.rs index 5bbfa123..e1eac39f 100644 --- a/pueue/src/daemon/network/message_handler/group.rs +++ b/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. diff --git a/pueue/src/daemon/network/message_handler/mod.rs b/pueue/src/daemon/network/message_handler/mod.rs index 24d29c67..0bf39eee 100644 --- a/pueue/src/daemon/network/message_handler/mod.rs +++ b/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), diff --git a/pueue/src/daemon/process_handler/finish.rs b/pueue/src/daemon/process_handler/finish.rs index 292a18a7..3c2716a0 100644 --- a/pueue/src/daemon/process_handler/finish.rs +++ b/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 { diff --git a/pueue/tests/client/integration/env.rs b/pueue/tests/client/integration/env.rs new file mode 100644 index 00000000..9d0ae0e7 --- /dev/null +++ b/pueue/tests/client/integration/env.rs @@ -0,0 +1,60 @@ +use anyhow::Result; + +use crate::client::helper::*; + +/// Set an environment variable an make sure it's there afterwards. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn set_environment() -> Result<()> { + let daemon = daemon().await?; + let shared = &daemon.settings.shared; + + // Add a stashed task so we can edit it. + run_client_command(shared, &["add", "--stashed", "echo $TEST_VARIABLE"])?; + + // Set the environment variable + run_client_command(shared, &["env", "set", "0", "TEST_VARIABLE", "thisisatest"])?; + + // Now start the command and wait for it to finish + run_client_command(shared, &["enqueue", "0"])?; + wait_for_task_condition(shared, 0, |task| task.is_done()).await?; + + let state = get_state(shared).await?; + println!("{:#?}", state.tasks[&0].envs); + + // Make sure the environment variable has been set. + let output = run_client_command(shared, &["follow", "0"])?; + let stdout = String::from_utf8_lossy(&output.stdout); + assert_eq!("thisisatest", stdout.trim()); + + Ok(()) +} + +/// Set an environment variable an make sure it's there afterwards. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn unset_environment() -> Result<()> { + let daemon = daemon().await?; + let shared = &daemon.settings.shared; + + // Add a stashed task so we can edit it. + run_client_command(shared, &["add", "--stashed", "echo $TEST_VARIABLE"])?; + + // Set the environment variable + run_client_command(shared, &["env", "set", "0", "TEST_VARIABLE", "thisisatest"])?; + + // Unset the environment variable again. + run_client_command(shared, &["env", "unset", "0", "TEST_VARIABLE"])?; + + // Now start the command and wait for it to finish + run_client_command(shared, &["enqueue", "0"])?; + wait_for_task_condition(shared, 0, |task| task.is_done()).await?; + + let state = get_state(shared).await?; + println!("{:#?}", state.tasks[&0].envs); + + // Make sure the environment variable has been set. + let output = run_client_command(shared, &["follow", "0"])?; + let stdout = String::from_utf8_lossy(&output.stdout); + assert_eq!("", stdout.trim()); + + Ok(()) +} diff --git a/pueue/tests/client/integration/mod.rs b/pueue/tests/client/integration/mod.rs index eedfb7a8..8dfd57dc 100644 --- a/pueue/tests/client/integration/mod.rs +++ b/pueue/tests/client/integration/mod.rs @@ -1,6 +1,7 @@ mod completions; mod configuration; mod edit; +mod env; mod follow; mod group; mod log; diff --git a/pueue_lib/src/network/message.rs b/pueue_lib/src/network/message.rs index 9da08da1..73146380 100644 --- a/pueue_lib/src/network/message.rs +++ b/pueue_lib/src/network/message.rs @@ -66,6 +66,8 @@ pub enum Message { /// The client sends the edited details to the daemon. Edit(Vec), + Env(EnvMessage), + Group(GroupMessage), GroupResponse(GroupResponseMessage), @@ -273,6 +275,21 @@ impl EditableTask { } } +#[derive(PartialEq, Eq, Clone, Debug, Deserialize, Serialize)] +pub enum EnvMessage { + Set { + task_id: usize, + key: String, + value: String, + }, + Unset { + task_id: usize, + key: String, + }, +} + +impl_into_message!(EnvMessage, Message::Env); + #[derive(PartialEq, Eq, Clone, Debug, Deserialize, Serialize)] pub enum GroupMessage { Add {